Prevent ack timeout from dropping callback ids

shepherd_utils/config.py

@@ -81,6 +81,10 @@ class Settings(BaseSettings):
     # byte count; 0 (or unparseable) disables the limit.
     callback_max_request_size: str = "0"
     kg_retrieval_url: str = "http://host.docker.internal:8080/asyncquery"
+    # How long a lookup worker waits for the retrieval service to ACK an
+    # /asyncquery submission. This bounds only the acknowledgement, the
+    # results arrive later via callback.
+    kg_retrieval_submit_timeout: float = 100.0
     sync_kg_retrieval_url: str = "http://host.docker.internal:8080/query"
     kg_rehydrate_url: str = "http://host.docker.internal:8080/rehydrate"
     default_data_tier: int = 0

tests/unit/aragorn/test_aragorn_lookup_submit.py

@@ -0,0 +1,130 @@
+"""Tests for how ``workers.aragorn_lookup.worker`` handles the /asyncquery
+submission ACK.
+
+A creative query fans out to ~20 submissions at once and the retrieval
+service has been seen taking 20s+ to ACK the burst. A submission whose ACK
+times out has still been received -- the service runs it and POSTs the
+callback later -- so its callback id must survive, or the callback handler
+rejects the results with a 500 and the lookup waits out its deadline for
+nothing.
+"""
+
+import json
+import logging
+
+import httpx
+import pytest
+
+from workers.aragorn_lookup import worker as al
+from workers.aragorn_lookup.worker import (
+    AsyncResponse,
+    aragorn_lookup,
+    run_async_lookup,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _make_task():
+    return [
+        "test",
+        {
+            "query_id": "qid",
+            "response_id": "rid",
+            "workflow": json.dumps([{"id": "aragorn.lookup"}]),
+            "log_level": "20",
+            "otel": json.dumps({}),
+        },
+    ]
+
+
+INFERRED_MSG = {
+    "message": {
+        "query_graph": {
+            "nodes": {
+                "a": {"ids": ["X:1"], "categories": ["biolink:Drug"]},
+                "b": {"categories": ["biolink:Disease"]},
+            },
+            "edges": {
+                "e0": {
+                    "subject": "a",
+                    "object": "b",
+                    "knowledge_type": "inferred",
+                    "predicates": ["biolink:treats"],
+                }
+            },
+        }
+    },
+    "parameters": {"timeout": 5},
+}
+
+
+@pytest.mark.asyncio
+async def test_run_async_lookup_read_timeout_is_in_flight(redis_mock, mocker):
+    mocker.patch.object(al, "add_callback_id", new_callable=mocker.AsyncMock)
+
+    client = mocker.Mock()
+    client.timeout.read = 100.0
+    client.post = mocker.AsyncMock(side_effect=httpx.ReadTimeout(""))
+
+    out = await run_async_lookup(client, {"message": {}}, "qid", logger)
+    assert out.success is False
+    assert out.in_flight is True
+    assert "ReadTimeout" in out.error
+    assert "100.0s" in out.error
+
+
+@pytest.mark.asyncio
+async def test_run_async_lookup_connect_error_is_not_in_flight(redis_mock, mocker):
+    mocker.patch.object(al, "add_callback_id", new_callable=mocker.AsyncMock)
+
+    client = mocker.Mock()
+    client.post = mocker.AsyncMock(side_effect=httpx.ConnectError("boom"))
+
+    out = await run_async_lookup(client, {"message": {}}, "qid", logger)
+    assert out.success is False
+    assert out.in_flight is False
+    assert "ConnectError: boom" == out.error
+
+
+@pytest.mark.asyncio
+async def test_aragorn_lookup_keeps_in_flight_callback_ids(redis_mock, mocker):
+    """Only submissions that definitely failed lose their callback id."""
+    mocker.patch.object(
+        al, "get_message", new_callable=mocker.AsyncMock, return_value=INFERRED_MSG
+    )
+    mocker.patch.object(
+        al,
+        "expand_aragorn_query",
+        return_value=[
+            {"message": {"query_graph": {}}, "parameters": {}, "submitter": "t"},
+            {"message": {"query_graph": {}}, "parameters": {}, "submitter": "t"},
+            {"message": {"query_graph": {}}, "parameters": {}, "submitter": "t"},
+        ],
+    )
+    mocker.patch.object(
+        al,
+        "run_async_lookup",
+        new_callable=mocker.AsyncMock,
+        side_effect=[
+            AsyncResponse(status_code=200, success=True, callback_id="ok-cb"),
+            AsyncResponse(
+                status_code=500,
+                success=False,
+                callback_id="slow-ack-cb",
+                error="ReadTimeout",
+                in_flight=True,
+            ),
+            AsyncResponse(
+                status_code=500, success=False, callback_id="failed-cb", error="x"
+            ),
+        ],
+    )
+    mock_remove = mocker.patch.object(
+        al, "remove_callback_id", new_callable=mocker.AsyncMock
+    )
+    mocker.patch.object(
+        al, "get_running_callbacks", new_callable=mocker.AsyncMock, return_value=[]
+    )
+    await aragorn_lookup(_make_task(), logger)
+    mock_remove.assert_awaited_once_with("failed-cb", logger)

tests/unit/test_bte_lookup_branches.py

@@ -90,6 +90,39 @@ async def test_run_async_lookup_returns_500_when_post_raises(redis_mock, mocker)
     assert out.success is False
     assert out.status_code == 500
     assert "boom" in out.error
+    assert out.in_flight is False
+
+
+@pytest.mark.asyncio
+async def test_run_async_lookup_read_timeout_is_in_flight(redis_mock, mocker):
+    """A read timeout means the submission was sent but not ACKed: the service
+    may still deliver the callback, so the response is flagged in_flight."""
+    mocker.patch.object(btel, "add_callback_id", new_callable=mocker.AsyncMock)
+
+    client = mocker.Mock()
+    client.timeout.read = 100.0
+    client.post = mocker.AsyncMock(side_effect=httpx.ReadTimeout(""))
+
+    out = await run_async_lookup(client, {"message": {}}, "qid", logger)
+    assert out.success is False
+    assert out.in_flight is True
+    # httpx timeouts stringify to nothing; the error should still say what happened.
+    assert "ReadTimeout" in out.error
+    assert "100.0s" in out.error
+
+
+@pytest.mark.asyncio
+async def test_run_async_lookup_connect_timeout_is_not_in_flight(redis_mock, mocker):
+    """A connect timeout never reached the service, so nothing is coming back."""
+    mocker.patch.object(btel, "add_callback_id", new_callable=mocker.AsyncMock)
+
+    client = mocker.Mock()
+    client.post = mocker.AsyncMock(side_effect=httpx.ConnectTimeout(""))
+
+    out = await run_async_lookup(client, {"message": {}}, "qid", logger)
+    assert out.success is False
+    assert out.in_flight is False
+    assert "ConnectTimeout" in out.error
 
 
 @pytest.mark.asyncio
@@ -220,6 +253,75 @@ async def test_bte_lookup_inferred_removes_failed_callback_ids(redis_mock, mocke
     mock_remove.assert_awaited_once_with("failed-cb", logger)
 
 
+@pytest.mark.asyncio
+async def test_bte_lookup_inferred_keeps_in_flight_callback_ids(redis_mock, mocker):
+    """A submission that timed out waiting for the ACK keeps its callback id:
+    the retrieval service may still POST the results, and removing the id
+    would make the callback handler reject them with a 500."""
+    inferred_msg = {
+        "message": {
+            "query_graph": {
+                "nodes": {
+                    "a": {"ids": ["X:1"], "categories": ["biolink:Drug"]},
+                    "b": {"categories": ["biolink:Disease"]},
+                },
+                "edges": {
+                    "e0": {
+                        "subject": "a",
+                        "object": "b",
+                        "knowledge_type": "inferred",
+                        "predicates": ["biolink:treats"],
+                    }
+                },
+            }
+        },
+        "parameters": {"timeout": 5},
+    }
+    mocker.patch.object(
+        btel,
+        "get_message",
+        new_callable=mocker.AsyncMock,
+        return_value=inferred_msg,
+    )
+    mocker.patch.object(
+        btel,
+        "expand_bte_query",
+        return_value=[
+            {"message": {"query_graph": {}}, "parameters": {}, "submitter": "t"},
+            {"message": {"query_graph": {}}, "parameters": {}, "submitter": "t"},
+        ],
+    )
+    mocker.patch.object(
+        btel,
+        "run_async_lookup",
+        new_callable=mocker.AsyncMock,
+        side_effect=[
+            AsyncResponse(
+                status_code=500,
+                success=False,
+                callback_id="slow-ack-cb",
+                error="ReadTimeout",
+                in_flight=True,
+            ),
+            AsyncResponse(
+                status_code=500, success=False, callback_id="failed-cb", error="x"
+            ),
+        ],
+    )
+    mock_remove = mocker.patch.object(
+        btel, "remove_callback_id", new_callable=mocker.AsyncMock
+    )
+    mocker.patch.object(
+        btel,
+        "get_running_callbacks",
+        new_callable=mocker.AsyncMock,
+        return_value=[],
+    )
+    await bte_lookup(_make_task(), logger)
+    # Only the genuinely failed submission is dropped.
+    mock_remove.assert_awaited_once_with("failed-cb", logger)
+
+
 @pytest.mark.asyncio
 async def test_bte_lookup_inferred_logs_exception_responses(redis_mock, mocker):
     """An exception in ``asyncio.gather`` (return_exceptions=True) is logged

workers/aragorn_lookup/worker.py

@@ -84,6 +84,11 @@ class AsyncResponse:
     success: bool
     callback_id: str
     error: Optional[str] = None
+    # True when the submission was sent but we never heard the ACK. The
+    # retrieval service may well be running the lookup and will POST the
+    # callback later, so the callback id must be kept or the results will be
+    # rejected on arrival. False means the request never reached the service.
+    in_flight: bool = False
 
 
 async def run_async_lookup(
@@ -119,13 +124,25 @@ async def run_async_lookup(
                 success=response.status_code == 200,
                 callback_id=callback_id,
             )
+        except httpx.ReadTimeout as e:
+            # The request went out and the service just didn't ACK in time.
+            # It has the query and the callback URL, so treat it as still
+            # running rather than failed.
+            span.record_exception(e)
+            return AsyncResponse(
+                status_code=500,
+                success=False,
+                callback_id=callback_id,
+                error=f"{type(e).__name__}: no ACK within {client.timeout.read}s",
+                in_flight=True,
+            )
         except Exception as e:
             span.record_exception(e)
             return AsyncResponse(
                 status_code=500,
                 success=False,
                 callback_id=callback_id,
-                error=str(e),
+                error=f"{type(e).__name__}: {e}",
             )
 
 
@@ -168,7 +185,9 @@ async def aragorn_lookup(task, logger: logging.Logger):
         )
         with tracer.start_as_current_span("aragorn.lookup") as span:
             span.set_attribute("callback.id", callback_id)
-            async with httpx.AsyncClient(timeout=100) as client:
+            async with httpx.AsyncClient(
+                timeout=settings.kg_retrieval_submit_timeout
+            ) as client:
                 await client.post(
                     settings.kg_retrieval_url,
                     json=message,
@@ -180,7 +199,9 @@ async def aragorn_lookup(task, logger: logging.Logger):
 
         requests = []
         # send all messages to lookup service
-        async with httpx.AsyncClient(timeout=20) as client:
+        async with httpx.AsyncClient(
+            timeout=settings.kg_retrieval_submit_timeout
+        ) as client:
             for expanded_message in expanded_messages:
                 requests.append(
                     run_async_lookup(client, expanded_message, query_id, logger)
@@ -196,12 +217,20 @@ async def aragorn_lookup(task, logger: logging.Logger):
                         f"Failed to do lookup and unable to remove callback id: {response}"
                     )
                 elif isinstance(response, AsyncResponse):
-                    if not response.success:
-                        logger.error(
-                            f"[{response.callback_id}] Failed to do lookup, "
-                            f"removing callback id: {response.error}"
+                    if response.success:
+                        continue
+                    if response.in_flight:
+                        logger.warning(
+                            f"[{response.callback_id}] Lookup submitted but not "
+                            f"acknowledged, waiting for its callback anyway: "
+                            f"{response.error}"
                         )
-                        await remove_callback_id(response.callback_id, logger)
+                        continue
+                    logger.error(
+                        f"[{response.callback_id}] Failed to do lookup, "
+                        f"removing callback id: {response.error}"
+                    )
+                    await remove_callback_id(response.callback_id, logger)
                 else:
                     logger.error(
                         f"Failed to do lookup and unable to remove callback id: {response}"