diff --git a/tests/unit/test_finish_query_callback_logs.py b/tests/unit/test_finish_query_callback_logs.py new file mode 100644 index 0000000..a0b47e2 --- /dev/null +++ b/tests/unit/test_finish_query_callback_logs.py @@ -0,0 +1,246 @@ +"""Tests for the callback timing/failure logging in ``finish_query``. + +Covers what the delivery logs have to answer after the fact: how long the +callback POST took, and -- when it didn't land -- why, both in the worker's +logs and in the payload the receiver eventually gets. +""" + +import json +import logging + +import httpx +import orjson +import pytest + +from workers.finish_query.worker import ( + CALLBACK_ERROR_BODY_BYTES, + CALLBACK_RETRIES, + _append_log_entry, + _describe_callback_failure, + finish_query, +) + +logger = logging.getLogger(__name__) + +TASK = [ + "test", + { + "query_id": "test", + "response_id": "rid", + "workflow": json.dumps([]), + "log_level": "20", + }, +] + + +def _patch_async_query(mocker, message=None, logs=None): + """Patch the db reads for an async (callback) query.""" + mocker.patch( + "workers.finish_query.worker.get_query_state", + new_callable=mocker.AsyncMock, + return_value=["", "", "", "", "", "", "", "rid", "http://callback"], + ) + mocker.patch( + "workers.finish_query.worker.set_query_completed", + new_callable=mocker.AsyncMock, + ) + mocker.patch( + "workers.finish_query.worker.get_message", + new_callable=mocker.AsyncMock, + return_value=orjson.dumps(message if message is not None else {"message": {}}), + ) + mocker.patch( + "workers.finish_query.worker.get_logs", + new_callable=mocker.AsyncMock, + return_value=logs if logs is not None else [], + ) + + +def _http_error_response(status_code: int, body: bytes) -> httpx.Response: + """A real httpx response, so ``raise_for_status`` raises the real error.""" + return httpx.Response( + status_code, + content=body, + request=httpx.Request("POST", "http://callback"), + ) + + +@pytest.mark.asyncio +async def test_successful_callback_logs_duration_and_size(redis_mock, mocker, caplog): + """A delivered callback logs how long the POST took and how big it was.""" + _patch_async_query(mocker) + mocker.patch( + "httpx.AsyncClient.post", + new_callable=mocker.AsyncMock, + return_value=_http_error_response(200, b"ok"), + ) + + with caplog.at_level(logging.INFO): + await finish_query(TASK, logger) + + sent = [r.message for r in caplog.records if "Sent response back" in r.message] + assert len(sent) == 1 + assert "http://callback" in sent[0] + # ".s" duration and a byte count, both after the send completed. + assert "s (" in sent[0] and "bytes" in sent[0] + assert f"attempt 1/{CALLBACK_RETRIES}" in sent[0] + + +@pytest.mark.asyncio +async def test_failed_callback_logs_status_and_body(redis_mock, mocker, caplog): + """A rejected callback logs the status code and the server's explanation.""" + _patch_async_query(mocker) + mocker.patch( + "httpx.AsyncClient.post", + new_callable=mocker.AsyncMock, + return_value=_http_error_response(502, b"upstream exploded"), + ) + mocker.patch("asyncio.sleep", new_callable=mocker.AsyncMock) + + with caplog.at_level(logging.INFO): + await finish_query(TASK, logger) + + failures = [ + r.message for r in caplog.records if "Failed to send callback" in r.message + ] + assert len(failures) == CALLBACK_RETRIES + assert "HTTP 502" in failures[0] + assert "upstream exploded" in failures[0] + assert f"attempt 1/{CALLBACK_RETRIES}" in failures[0] + # And a single summary line once the retries are spent. + gave_up = [ + r.message for r in caplog.records if "Gave up sending callback" in r.message + ] + assert len(gave_up) == 1 + assert "http://callback" in gave_up[0] + + +@pytest.mark.asyncio +async def test_failure_is_spliced_into_the_retry_payload(redis_mock, mocker): + """The receiver sees the attempts that didn't make it. + + A callback that fails and is retried can't carry its own failure, but it + can carry the previous attempt's -- so the eventual recipient knows the + response is late and why. + """ + _patch_async_query(mocker, logs=[{"message": "earlier", "level": "INFO"}]) + payloads = [] + + async def record(*args, **kwargs): + payloads.append(kwargs["content"]) + if len(payloads) < 3: + return _http_error_response(503, b"try again later") + return _http_error_response(200, b"ok") + + mocker.patch("httpx.AsyncClient.post", side_effect=record) + mocker.patch("asyncio.sleep", new_callable=mocker.AsyncMock) + + await finish_query(TASK, logger) + + assert len(payloads) == 3 + first, second, third = (orjson.loads(p) for p in payloads) + assert [entry["message"] for entry in first["logs"]] == ["earlier"] + assert len(second["logs"]) == 2 + assert "HTTP 503" in second["logs"][1]["message"] + assert second["logs"][1]["level"] == "ERROR" + assert second["logs"][1]["timestamp"] + # Each attempt adds only its own failure, and the message survives intact. + assert len(third["logs"]) == 3 + assert third["message"] == {} + + +@pytest.mark.asyncio +async def test_oversized_payload_skips_the_inline_retry_note(redis_mock, mocker): + """A huge payload isn't rebuilt just to carry a note that's in the logs.""" + _patch_async_query(mocker) + mocker.patch("workers.finish_query.worker.RETRY_LOG_SPLICE_MAX_BYTES", 1) + payloads = [] + + async def record(*args, **kwargs): + payloads.append(kwargs["content"]) + return _http_error_response(500, b"nope") + + mocker.patch("httpx.AsyncClient.post", side_effect=record) + mocker.patch("asyncio.sleep", new_callable=mocker.AsyncMock) + + await finish_query(TASK, logger) + + assert len(payloads) == CALLBACK_RETRIES + assert len(set(payloads)) == 1 + + +@pytest.mark.asyncio +async def test_callback_logs_are_persisted_for_the_query(redis_mock, mocker): + """The delivery outcome lands in the query's logs, not just stdout. + + ``finish_query`` acks directly instead of going through ``wrap_up_task``, + so it has to flush its own logs -- otherwise a failed callback is invisible + to anyone reading the query back. + """ + _patch_async_query(mocker) + mock_save_logs = mocker.patch( + "workers.finish_query.worker.save_logs", + new_callable=mocker.AsyncMock, + ) + mocker.patch( + "httpx.AsyncClient.post", + new_callable=mocker.AsyncMock, + return_value=_http_error_response(200, b"ok"), + ) + + await finish_query(TASK, logger) + + mock_save_logs.assert_awaited_once_with("rid", logger) + + +@pytest.mark.asyncio +async def test_failing_to_save_logs_does_not_fail_the_task(redis_mock, mocker): + """A log flush that blows up must not take the whole wrap-up with it.""" + _patch_async_query(mocker) + mocker.patch( + "workers.finish_query.worker.save_logs", + new_callable=mocker.AsyncMock, + side_effect=Exception("redis down"), + ) + mocker.patch( + "httpx.AsyncClient.post", + new_callable=mocker.AsyncMock, + return_value=_http_error_response(200, b"ok"), + ) + + await finish_query(TASK, logger) + + +def test_describe_callback_failure_truncates_the_body(): + """A rejecting server can return anything; only the head of it is logged.""" + error = httpx.HTTPStatusError( + "boom", + request=httpx.Request("POST", "http://callback"), + response=_http_error_response(413, b"x" * (CALLBACK_ERROR_BODY_BYTES * 10)), + ) + described = _describe_callback_failure(error) + assert described.startswith("HTTP 413: ") + assert len(described) < CALLBACK_ERROR_BODY_BYTES + 50 + + +def test_describe_callback_failure_names_silent_exceptions(): + """httpx's connect/timeout errors often stringify to nothing useful.""" + assert "ConnectTimeout" in _describe_callback_failure(httpx.ConnectTimeout("")) + assert "ConnectError" in _describe_callback_failure(httpx.ConnectError("")) + + +def test_append_log_entry_handles_both_array_shapes(): + """Empty and populated logs arrays both stay valid JSON.""" + entry = {"message": "late", "level": "ERROR"} + empty = orjson.dumps({"message": {}, "logs": []}) + assert orjson.loads(_append_log_entry(empty, entry))["logs"] == [entry] + + populated = orjson.dumps({"message": {}, "logs": [{"message": "first"}]}) + appended = orjson.loads(_append_log_entry(populated, entry))["logs"] + assert [e["message"] for e in appended] == ["first", "late"] + + +def test_append_log_entry_leaves_an_unexpected_tail_alone(): + """Rather than corrupt a payload we don't recognize, send it as-is.""" + payload = orjson.dumps({"logs": [], "message": {}}) + assert _append_log_entry(payload, {"message": "late"}) == payload diff --git a/workers/finish_query/worker.py b/workers/finish_query/worker.py index 2cbed87..98fb131 100644 --- a/workers/finish_query/worker.py +++ b/workers/finish_query/worker.py @@ -7,8 +7,10 @@ import uuid import orjson +from datetime import datetime, timezone + from opentelemetry.propagate import inject -from opentelemetry.trace import Status, StatusCode +from opentelemetry.trace import Status, StatusCode, get_current_span from shepherd_utils.broker import mark_task_as_complete from shepherd_utils.db import ( @@ -16,6 +18,7 @@ get_logs, get_message, get_query_state, + save_logs, set_query_completed, ) from shepherd_utils.shared import get_tasks @@ -30,6 +33,160 @@ tracer = setup_tracer(STREAM) LOGGER = get_worker_logger(STREAM) CALLBACK_RETRIES = 3 +CALLBACK_TIMEOUT = 120 +# How much of a rejecting server's response body goes into the failure log. The +# body is already in memory (we don't stream the response), but it can be an +# arbitrarily large HTML error page, and this string is copied into the query's +# logs -- and possibly into the retry payload -- so keep only the head of it. +CALLBACK_ERROR_BODY_BYTES = 500 +# Ceiling on the payload size for which we splice a failed attempt's note into +# the *next* attempt's body. Splicing rebuilds the whole buffer, so for a large +# response the transient second copy costs far more than the note is worth. The +# note is in the query's own logs either way, so oversized payloads just skip +# the inline copy. +RETRY_LOG_SPLICE_MAX_BYTES = 64 * 1024 * 1024 + + +def _log_entry(message: str, level: str = "ERROR") -> dict: + """Build a TRAPI LogEntry, matching ReasonerLogEntryFormatter's shape. + + Used for entries we splice straight into an outgoing payload, which never + pass through the logging handler that would otherwise format them. + """ + return { + "message": message, + "timestamp": datetime.now(timezone.utc).isoformat(), + "level": level, + } + + +def _describe_callback_failure(e: Exception) -> str: + """One bounded line explaining why a callback POST failed. + + The reason is the whole point of logging the failure -- "callback failed" + alone doesn't say whether the receiver is down, slow, or rejecting the + payload -- so pull out the status code and the head of the response body + for an HTTP error, and the exception type otherwise (httpx reports connect + failures, TLS errors and timeouts as distinct classes, and several of them + stringify to an empty message). + """ + if isinstance(e, httpx.HTTPStatusError): + detail = "" + try: + body = e.response.content[:CALLBACK_ERROR_BODY_BYTES] + if body: + detail = f": {body.decode('utf-8', 'replace')}" + except Exception: + # Body not readable (streamed/closed response) -- the status code + # is still worth reporting on its own. + pass + return f"HTTP {e.response.status_code}{detail}" + if isinstance(e, httpx.TimeoutException): + return f"{type(e).__name__} (no response within {CALLBACK_TIMEOUT}s)" + return f"{type(e).__name__}: {e}" + + +def _append_log_entry(payload: bytes, entry: dict) -> bytes: + """Return ``payload`` with ``entry`` appended to its trailing logs array. + + Only sound for a payload this worker built, which always ends with the logs + array followed by the closing brace. Rebuilding costs a transient second + copy of the payload, so callers guard on size; the rebind releases the old + buffer immediately. If the payload doesn't have the expected tail, hand it + back untouched rather than risk shipping malformed JSON. + """ + entry_bytes = orjson.dumps(entry) + if payload.endswith(b"[]}"): + return payload[:-3] + b"[" + entry_bytes + b"]}" + if payload.endswith(b"]}"): + return payload[:-2] + b"," + entry_bytes + b"]}" + return payload + + +async def send_callback( + callback_url: str, + message_bytes: bytes, + logger: logging.Logger, +) -> bool: + """POST the finished response to the caller's callback URL. + + Every attempt is timed and logged *after* the send completes -- how long a + callback takes is a property of the receiver we otherwise have no record + of, and a failure is only actionable with the reason attached. Failures are + also spliced into the next attempt's payload (size permitting), so a + receiver that eventually gets the response can see the attempts that didn't + make it. + + Returns True if the response was delivered. + """ + headers = {"Content-Type": "application/json"} + # Propagate the otel trace context through the callback. + # Matches the inject() carrier pattern used by the + # lookup workers; the active span comes from process_task's + # start_as_current_span. + inject(headers) + span = get_current_span() + started = time.time() + payload_size = len(message_bytes) + delivered = False + attempts = 0 + for attempt in range(1, CALLBACK_RETRIES + 1): + attempts = attempt + attempt_start = time.time() + try: + async with httpx.AsyncClient(timeout=CALLBACK_TIMEOUT) as client: + response = await client.post( + callback_url, + content=message_bytes, + headers=headers, + ) + response.raise_for_status() + elapsed = time.time() - attempt_start + logger.info( + f"Sent response back to {callback_url} in {elapsed:.3f}s " + f"({len(message_bytes)} bytes, " + f"attempt {attempt}/{CALLBACK_RETRIES})" + ) + delivered = True + break + except Exception as e: + elapsed = time.time() - attempt_start + failure = ( + f"Failed to send callback to {callback_url} after {elapsed:.3f}s " + f"(attempt {attempt}/{CALLBACK_RETRIES}, " + f"{len(message_bytes)} bytes): {_describe_callback_failure(e)}" + ) + logger.error(failure) + span.add_event( + "callback_attempt_failed", + {"attempt": attempt, "duration_ms": int(elapsed * 1000)}, + ) + if attempt < CALLBACK_RETRIES: + if len(message_bytes) <= RETRY_LOG_SPLICE_MAX_BYTES: + message_bytes = _append_log_entry( + message_bytes, _log_entry(failure) + ) + await asyncio.sleep(1 * (2 ** (attempt - 1))) + + total = time.time() - started + if not delivered: + logger.error( + f"Gave up sending callback to {callback_url} after " + f"{CALLBACK_RETRIES} attempts and {total:.3f}s. The response was " + "not delivered." + ) + elif attempts > 1: + logger.info( + f"Callback to {callback_url} succeeded on attempt {attempts} " + f"after {total:.3f}s total." + ) + # Attributes rather than another log line: same numbers, no per-query log + # storage, and they're queryable alongside the rest of the trace. + span.set_attribute("callback.duration_ms", int(total * 1000)) + span.set_attribute("callback.attempts", attempts) + span.set_attribute("callback.payload_bytes", payload_size) + span.set_attribute("callback.delivered", delivered) + return delivered async def finish_query(task, logger: logging.Logger): @@ -63,29 +220,22 @@ async def finish_query(task, logger: logging.Logger): ) else: message = orjson.loads(message_bytes) + # Re-insert rather than assign in place so "logs" is last in + # the serialized payload -- send_callback appends retry notes + # by rewriting the payload's tail. + message.pop("logs", None) message["logs"] = logs message_bytes = orjson.dumps(message) del message - headers = {"Content-Type": "application/json"} - # Propagate the otel trace context through the callback. - # Matches the inject() carrier pattern used by the - # lookup workers; the active span comes from process_task's - # start_as_current_span. - inject(headers) - for attempt in range(CALLBACK_RETRIES): - try: - async with httpx.AsyncClient(timeout=120) as client: - response = await client.post( - callback_url, - content=message_bytes, - headers=headers, - ) - response.raise_for_status() - logger.info(f"Sent response back to {callback_url}") - break - except Exception as e: - logger.error(f"Failed to send callback to {callback_url}: {e}") - await asyncio.sleep(1 * (2**attempt)) + # The logs list and its serialization are a full second copy of + # every log line the query produced; they're inside the payload + # now, so drop them before the send rather than holding them for + # its duration. + del logs, logs_bytes + + await send_callback(callback_url, message_bytes, logger) + # Release the payload before the remaining db round trips. + del message_bytes await set_query_completed(query_id, status, logger) @@ -98,6 +248,17 @@ async def finish_query(task, logger: logging.Logger): logger.info(f"Finished task {task[0]} in {time.time() - start}") + # This worker acks directly instead of going through wrap_up_task, so + # nothing else flushes what it logged. Persist here so the callback + # outcome -- how long delivery took, or why it failed -- survives in the + # query's logs (GET /response/{query_id}) instead of only in the pod's + # stdout. Draining also clears the handler's queue, which for this + # process-wide logger would otherwise just accumulate. + try: + await save_logs(response_id, logger) + except Exception as e: + logger.error(f"Failed to save logs for {response_id}: {e}") + async def process_task(task, parent_ctx, logger: logging.Logger, limiter): """Process a given task and ACK in redis."""