Stop retrying rejected callbacks, and cut the retry budget to one

tests/unit/test_finish_query_callback_logs.py

@@ -13,11 +13,13 @@
 import pytest
 
 from workers.finish_query.worker import (
+    CALLBACK_ATTEMPTS,
     CALLBACK_ERROR_BODY_BYTES,
-    CALLBACK_RETRIES,
     _append_log_entry,
     _describe_callback_failure,
+    _is_retryable,
     finish_query,
+    send_callback,
 )
 
 logger = logging.getLogger(__name__)
@@ -83,7 +85,7 @@ async def test_successful_callback_logs_duration_and_size(redis_mock, mocker, ca
     assert "http://callback" in sent[0]
     # "<n>.<mmm>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]
+    assert f"attempt 1/{CALLBACK_ATTEMPTS}" in sent[0]
 
 
 @pytest.mark.asyncio
@@ -103,10 +105,10 @@ async def test_failed_callback_logs_status_and_body(redis_mock, mocker, caplog):
     failures = [
         r.message for r in caplog.records if "Failed to send callback" in r.message
     ]
-    assert len(failures) == CALLBACK_RETRIES
+    assert len(failures) == CALLBACK_ATTEMPTS
     assert "HTTP 502" in failures[0]
     assert "upstream exploded" in failures[0]
-    assert f"attempt 1/{CALLBACK_RETRIES}" in failures[0]
+    assert f"attempt 1/{CALLBACK_ATTEMPTS}" 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
@@ -128,7 +130,7 @@ async def test_failure_is_spliced_into_the_retry_payload(redis_mock, mocker):
 
     async def record(*args, **kwargs):
         payloads.append(kwargs["content"])
-        if len(payloads) < 3:
+        if len(payloads) < CALLBACK_ATTEMPTS:
             return _http_error_response(503, b"try again later")
         return _http_error_response(200, b"ok")
 
@@ -137,16 +139,16 @@ async def record(*args, **kwargs):
 
     await finish_query(TASK, logger)
 
-    assert len(payloads) == 3
-    first, second, third = (orjson.loads(p) for p in payloads)
+    assert len(payloads) == CALLBACK_ATTEMPTS
+    first, second = (orjson.loads(p) for p in payloads)
     assert [entry["message"] for entry in first["logs"]] == ["earlier"]
+    # The retry adds only the one failure that preceded it, and the message
+    # itself survives the splice intact.
     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"] == {}
+    assert second["message"] == {}
 
 
 @pytest.mark.asyncio
@@ -165,7 +167,7 @@ async def record(*args, **kwargs):
 
     await finish_query(TASK, logger)
 
-    assert len(payloads) == CALLBACK_RETRIES
+    assert len(payloads) == CALLBACK_ATTEMPTS
     assert len(set(payloads)) == 1
 
 
@@ -244,3 +246,109 @@ 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
+
+
+@pytest.mark.asyncio
+async def test_client_error_is_not_retried(redis_mock, mocker, caplog):
+    """A 4xx is the receiver's verdict on these bytes, so we stop after one POST.
+
+    Retrying costs a full extra round trip -- up to CALLBACK_TIMEOUT of it --
+    with the whole payload pinned in memory, for a response that cannot change.
+    """
+    _patch_async_query(mocker)
+    mock_post = mocker.patch(
+        "httpx.AsyncClient.post",
+        new_callable=mocker.AsyncMock,
+        return_value=_http_error_response(400, b"malformed TRAPI"),
+    )
+    mocker.patch("asyncio.sleep", new_callable=mocker.AsyncMock)
+
+    with caplog.at_level(logging.INFO):
+        await finish_query(TASK, logger)
+
+    assert mock_post.call_count == 1
+    messages = [r.message for r in caplog.records]
+    assert any("Not retrying the callback" in m for m in messages)
+    assert any("HTTP 400" in m and "malformed TRAPI" in m for m in messages)
+    # The give-up line reports the attempt actually spent, not the budget.
+    gave_up = [m for m in messages if "Gave up sending callback" in m]
+    assert len(gave_up) == 1
+    assert "1 attempt(s)" in gave_up[0]
+
+
+@pytest.mark.asyncio
+async def test_server_error_is_still_retried(redis_mock, mocker):
+    """A 5xx may well be transient, so the retry budget still applies to it."""
+    _patch_async_query(mocker)
+    mock_post = mocker.patch(
+        "httpx.AsyncClient.post",
+        new_callable=mocker.AsyncMock,
+        return_value=_http_error_response(502, b"bad gateway"),
+    )
+    mocker.patch("asyncio.sleep", new_callable=mocker.AsyncMock)
+
+    await finish_query(TASK, logger)
+
+    assert mock_post.call_count == CALLBACK_ATTEMPTS
+
+
+@pytest.mark.asyncio
+async def test_rate_limit_is_retried_despite_being_4xx(redis_mock, mocker):
+    """429 asks us to come back later -- the one 4xx where a retry is the point."""
+    _patch_async_query(mocker)
+    mock_post = mocker.patch(
+        "httpx.AsyncClient.post",
+        new_callable=mocker.AsyncMock,
+        return_value=_http_error_response(429, b"slow down"),
+    )
+    mocker.patch("asyncio.sleep", new_callable=mocker.AsyncMock)
+
+    await finish_query(TASK, logger)
+
+    assert mock_post.call_count == CALLBACK_ATTEMPTS
+
+
+@pytest.mark.asyncio
+async def test_callback_reports_undelivered_on_a_client_error(redis_mock, mocker):
+    """Bailing out early reports undelivered, and skips the backoff entirely.
+
+    Driven through ``send_callback`` rather than ``finish_query`` so the retry
+    loop is the only thing that could reach ``asyncio.sleep`` -- the wrap-up's
+    own db retries have a backoff of their own.
+    """
+    mocker.patch(
+        "httpx.AsyncClient.post",
+        new_callable=mocker.AsyncMock,
+        return_value=_http_error_response(404, b"no such message"),
+    )
+    mock_sleep = mocker.patch("asyncio.sleep", new_callable=mocker.AsyncMock)
+
+    assert await send_callback("http://callback", b'{"logs":[]}', logger) is False
+    # The backoff exists only to space out a retry we are no longer making.
+    mock_sleep.assert_not_awaited()
+
+
+def test_is_retryable_splits_client_from_server_errors():
+    """4xx stops the loop; 5xx and the transient 4xx codes keep it going."""
+
+    def status_error(code: int) -> httpx.HTTPStatusError:
+        return httpx.HTTPStatusError(
+            "boom",
+            request=httpx.Request("POST", "http://callback"),
+            response=_http_error_response(code, b""),
+        )
+
+    for code in (400, 401, 403, 404, 413, 422):
+        assert _is_retryable(status_error(code)) is False, code
+    for code in (408, 429, 500, 502, 503, 504):
+        assert _is_retryable(status_error(code)) is True, code
+
+
+def test_is_retryable_keeps_transport_and_unknown_failures():
+    """Narrowing the loop must not silently stop retrying real transients."""
+    assert _is_retryable(httpx.ConnectError("")) is True
+    assert _is_retryable(httpx.ReadTimeout("")) is True
+    assert _is_retryable(httpx.RemoteProtocolError("")) is True
+    # An exception we don't recognize keeps the old behavior rather than
+    # quietly becoming terminal.
+    assert _is_retryable(Exception("simulated network error")) is True

tests/unit/test_finish_query_edge_cases.py

@@ -8,7 +8,7 @@
 import orjson
 import pytest
 
-from workers.finish_query.worker import finish_query
+from workers.finish_query.worker import CALLBACK_ATTEMPTS, finish_query
 
 logger = logging.getLogger(__name__)
 
@@ -94,8 +94,10 @@ async def test_finish_query_propagates_status_to_set_query_completed(
 
 @pytest.mark.asyncio
 async def test_finish_async_query_retries_callback_on_failure(redis_mock, mocker):
-    """If the first POST raises, finish_query should retry up to CALLBACK_RETRIES
-    times with backoff before giving up and still mark the query completed."""
+    """A transport-level failure is retried up to the callback budget.
+
+    ``finish_query`` should spend CALLBACK_ATTEMPTS POSTs with backoff before
+    giving up, and still mark the query completed either way."""
     mocker.patch(
         "workers.finish_query.worker.get_query_state",
         new_callable=mocker.AsyncMock,
@@ -136,8 +138,7 @@ async def test_finish_async_query_retries_callback_on_failure(redis_mock, mocker
         ],
         logger,
     )
-    # 3 retries baked into the worker.
-    assert mock_post.call_count == 3
+    assert mock_post.call_count == CALLBACK_ATTEMPTS
     assert mock_set_query_completed.called
 
 

workers/finish_query/worker.py

@@ -32,8 +32,18 @@
 TASK_LIMIT = 10
 tracer = setup_tracer(STREAM)
 LOGGER = get_worker_logger(STREAM)
-CALLBACK_RETRIES = 3
+# Retries *after* the first attempt, so the total number of POSTs is
+# CALLBACK_ATTEMPTS. Kept deliberately small: this worker holds the entire
+# (potentially very large) decompressed response in memory for every second a
+# callback is in flight, and each attempt can burn CALLBACK_TIMEOUT seconds
+# before it even fails. A long retry budget therefore multiplies the worker's
+# peak memory residency far more than it improves delivery odds.
+CALLBACK_RETRIES = 1
+CALLBACK_ATTEMPTS = CALLBACK_RETRIES + 1
 CALLBACK_TIMEOUT = 120
+# 4xx codes that describe a *transient* condition and explicitly invite another
+# attempt, unlike the rest of the 4xx range. See ``_is_retryable``.
+RETRYABLE_CLIENT_ERROR_STATUS = frozenset({408, 429})
 # 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
@@ -86,6 +96,27 @@ def _describe_callback_failure(e: Exception) -> str:
     return f"{type(e).__name__}: {e}"
 
 
+def _is_retryable(e: Exception) -> bool:
+    """Whether another attempt at this callback could plausibly succeed.
+
+    A 4xx means the receiver understood the request and rejected it, so sending
+    the same bytes again gets the same answer. Retrying is not merely useless
+    here: this worker keeps the whole payload resident for every second of
+    every attempt, so a doomed retry costs real memory on a worker whose peak
+    memory is what gets it OOM-killed. The exceptions are the 4xx codes that
+    signal a transient condition rather than a bad request.
+
+    Everything else -- 5xx, timeouts, connect/protocol failures, and any
+    exception we don't recognize -- stays retryable, so this narrows the retry
+    loop only where a retry is known to be pointless.
+    """
+    if isinstance(e, httpx.HTTPStatusError):
+        status = e.response.status_code
+        if 400 <= status < 500:
+            return status in RETRYABLE_CLIENT_ERROR_STATUS
+    return True
+
+
 def _append_log_entry(payload: bytes, entry: dict) -> bytes:
     """Return ``payload`` with ``entry`` appended to its trailing logs array.
 
@@ -117,6 +148,10 @@ async def send_callback(
     receiver that eventually gets the response can see the attempts that didn't
     make it.
 
+    Retries stop early on a failure ``_is_retryable`` rules out -- a payload the
+    receiver has rejected outright is not worth holding in memory for another
+    round trip.
+
     Returns True if the response was delivered.
     """
     headers = {"Content-Type": "application/json"}
@@ -129,10 +164,11 @@ async def send_callback(
     started = time.time()
     payload_size = len(message_bytes)
     delivered = False
+    retryable = True
     attempts = 0
     wait = 0.0
     backoff = 0.0
-    for attempt in range(1, CALLBACK_RETRIES + 1):
+    for attempt in range(1, CALLBACK_ATTEMPTS + 1):
         attempts = attempt
         attempt_start = time.time()
         try:
@@ -148,17 +184,18 @@ async def send_callback(
                 logger.info(
                     f"Sent response back to {callback_url} in {elapsed:.3f}s "
                     f"({len(message_bytes)} bytes, "
-                    f"attempt {attempt}/{CALLBACK_RETRIES})"
+                    f"attempt {attempt}/{CALLBACK_ATTEMPTS})"
                 )
                 delivered = True
                 break
         except Exception as e:
             elapsed = time.time() - attempt_start
             wait += elapsed
+            reason = _describe_callback_failure(e)
             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)}"
+                f"(attempt {attempt}/{CALLBACK_ATTEMPTS}, "
+                f"{len(message_bytes)} bytes): {reason}"
             )
             logger.error(failure)
             span.add_event(
@@ -168,7 +205,15 @@ async def send_callback(
                     "callback.attempt_duration_ms": int(elapsed * 1000),
                 },
             )
-            if attempt < CALLBACK_RETRIES:
+            if not _is_retryable(e):
+                retryable = False
+                logger.error(
+                    f"Not retrying the callback to {callback_url}: {reason} is a "
+                    "client error, so an identical retry would be rejected the "
+                    "same way."
+                )
+                break
+            if attempt < CALLBACK_ATTEMPTS:
                 if len(message_bytes) <= RETRY_LOG_SPLICE_MAX_BYTES:
                     message_bytes = _append_log_entry(
                         message_bytes, _log_entry(failure)
@@ -181,7 +226,7 @@ async def send_callback(
     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 "
+            f"{attempts} attempt(s) and {total:.3f}s. The response was "
             "not delivered."
         )
     elif attempts > 1:
@@ -195,6 +240,9 @@ async def send_callback(
     span.set_attribute("callback.wait_ms", int(wait * 1000))
     span.set_attribute("callback.backoff_ms", int(backoff * 1000))
     span.set_attribute("callback.attempts", attempts)
+    # False means we stopped before spending the budget because the receiver
+    # rejected the payload outright -- distinguishes "gave up" from "ran out".
+    span.set_attribute("callback.retryable", retryable)
     span.set_attribute("callback.payload_bytes", payload_size)
     span.set_attribute("callback.delivered", delivered)
     return delivered