Fix merge_message crash loop on nodes missing optional TRAPI fields

shepherd_utils/config.py

@@ -281,6 +281,16 @@ class Settings(BaseSettings):
     # under pathological bursts; leftover ready callbacks are swept by the next
     # drain iteration. 0 disables the cap (fold everything ready in one pass).
     merge_max_fold: int = 25
+    # A merge that fails is retried by re-enqueueing its wake task. Space the
+    # retries out (exponential from merge_retry_backoff, capped at
+    # merge_retry_backoff_max) so a deterministic failure can't spin the worker,
+    # and after merge_max_attempts consecutive failures give up on the batch --
+    # discard those callbacks so the query merges the rest and finishes rather
+    # than retrying an unmergeable payload until it times out. 0 attempts
+    # disables the breaker (unbounded retries).
+    merge_max_attempts: int = 3
+    merge_retry_backoff: float = 1.0
+    merge_retry_backoff_max: float = 10.0
     # Cap on how many log entries are lifted out of a single callback message.
     # Subservices report their retrieval work in the TRAPI ``logs`` list they
     # post back, and those entries are folded into the query's log list. A

shepherd_utils/shared.py

@@ -893,10 +893,9 @@ def recursive_get_auxgraph_edges(
 
 def is_support_edge(edge) -> bool:
     """Checks if a given edge is a support edge."""
-    if "attributes" not in edge:
-        return False
-    for attribute in edge["attributes"]:
-        if attribute["attribute_type_id"] == "biolink:support_graphs":
+    # ``attributes`` is optional and may be absent or null.
+    for attribute in edge.get("attributes") or []:
+        if attribute.get("attribute_type_id") == "biolink:support_graphs":
             return True
     return False
 
@@ -980,21 +979,28 @@ def merge_kgraph(og_message, new_message, source, logger: logging.Logger):
         if existing is None:
             og_nodes[key] = value
             continue
-        # Overlapping node: merge fields onto the existing entry.
-        if value["name"]:
-            existing["name"] = value["name"]
-        new_categories = value["categories"]
+        # Overlapping node: merge fields onto the existing entry. ``name``,
+        # ``categories`` and ``attributes`` are all optional in TRAPI, and a
+        # subservice is free to omit one entirely or send it as null. Read
+        # every field with .get() so a node that leaves one out merges as an
+        # absent value instead of raising KeyError -- a single such node used
+        # to abort the whole batch merge and, because the failure path
+        # re-enqueues the wake task, wedged the query in a retry loop.
+        new_name = value.get("name")
+        if new_name:
+            existing["name"] = new_name
+        new_categories = value.get("categories")
         if new_categories:
-            existing_categories = existing["categories"]
+            existing_categories = existing.get("categories")
             if existing_categories:
                 existing["categories"] = list(
                     set(existing_categories) | set(new_categories)
                 )
             else:
                 existing["categories"] = new_categories
-        new_attrs = value["attributes"]
+        new_attrs = value.get("attributes")
         if new_attrs:
-            existing_attrs = existing["attributes"]
+            existing_attrs = existing.get("attributes")
             if existing_attrs:
                 existing["attributes"] = combine_unique_dicts(
                     existing_attrs, new_attrs, logger
@@ -1014,20 +1020,21 @@ def merge_kgraph(og_message, new_message, source, logger: logging.Logger):
                 if aggregator_source not in sources:
                     sources.append(aggregator_source)
             continue
-        # Overlapping edge: merge attributes and sources.
-        new_attrs = value["attributes"]
+        # Overlapping edge: merge attributes and sources. Same as for nodes,
+        # read optional fields with .get() rather than subscripting.
+        new_attrs = value.get("attributes")
         if new_attrs:
-            existing_attrs = existing["attributes"]
+            existing_attrs = existing.get("attributes")
             if existing_attrs:
                 existing["attributes"] = combine_unique_dicts(
                     existing_attrs, new_attrs, logger
                 )
             else:
                 existing["attributes"] = new_attrs
 
-        new_sources = value["sources"]
+        new_sources = value.get("sources")
         if new_sources:
-            existing_sources = existing["sources"]
+            existing_sources = existing.get("sources")
             if existing_sources:
                 # TODO: there might need to be some sort of upstream resource id merging to do past this?
                 existing["sources"] = combine_unique_dicts(

tests/unit/test_merge_message_retry.py

@@ -0,0 +1,147 @@
+"""Regression guards on the merge_message failure path.
+
+A merge that raises used to re-enqueue its wake task immediately and
+unconditionally. That is right for a transient failure, but a batch that fails
+*deterministically* -- e.g. a callback node missing an optional TRAPI field,
+which raised ``KeyError: 'name'`` out of ``merge_kgraph`` -- fails again on
+every retry. Observed in production as one query logging the same traceback
+dozens of times a second for as long as it lived, never finishing.
+
+So the failure path now backs off between attempts and, after
+``merge_max_attempts`` consecutive failures, discards the batch it can't merge
+so the query gets on with the rest.
+"""
+
+import logging
+
+import pytest
+
+from shepherd_utils.config import settings
+from workers.merge_message.worker import (
+    MERGE_ATTEMPT_FIELD,
+    STREAM,
+    _handle_merge_failure,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _task(attempt=None):
+    fields = {
+        "query_id": "q1",
+        "response_id": "rid",
+        "callback_id": "cb1",
+        "target": "aragorn",
+        "_started_at": "123",
+    }
+    if attempt is not None:
+        fields[MERGE_ATTEMPT_FIELD] = str(attempt)
+    return ("1-1", fields)
+
+
+@pytest.fixture(autouse=True)
+def _no_sleep(mocker):
+    """The backoff itself is asserted on; don't actually wait it out."""
+    return mocker.patch(
+        "workers.merge_message.worker.asyncio.sleep", new=mocker.AsyncMock()
+    )
+
+
+async def test_first_failure_backs_off_and_carries_the_attempt_count(mocker):
+    add_task = mocker.patch(
+        "workers.merge_message.worker.add_task", new=mocker.AsyncMock()
+    )
+    sleep = mocker.patch(
+        "workers.merge_message.worker.asyncio.sleep", new=mocker.AsyncMock()
+    )
+
+    await _handle_merge_failure(_task(), "rid", ["cb1"], logger)
+
+    sleep.assert_awaited_once_with(settings.merge_retry_backoff)
+    stream, fields, _ = add_task.await_args.args
+    assert stream == STREAM
+    assert fields[MERGE_ATTEMPT_FIELD] == "1"
+    # The re-enqueued task is otherwise the one we got, minus the bookkeeping
+    # field the broker adds.
+    assert "_started_at" not in fields
+    assert fields["callback_id"] == "cb1"
+
+
+async def test_backoff_grows_and_is_capped(mocker):
+    """Successive failures wait longer, up to merge_retry_backoff_max."""
+    mocker.patch("workers.merge_message.worker.add_task", new=mocker.AsyncMock())
+    sleep = mocker.patch(
+        "workers.merge_message.worker.asyncio.sleep", new=mocker.AsyncMock()
+    )
+    mocker.patch.object(settings, "merge_max_attempts", 0)  # breaker off
+
+    waits = []
+    for attempt in (1, 2, 3, 20):
+        sleep.reset_mock()
+        await _handle_merge_failure(_task(attempt), "rid", ["cb1"], logger)
+        waits.append(sleep.await_args.args[0])
+
+    assert waits[0] < waits[1] < waits[2]
+    assert waits[-1] == settings.merge_retry_backoff_max
+    assert all(w <= settings.merge_retry_backoff_max for w in waits)
+
+
+async def test_batch_is_discarded_after_max_attempts(mocker):
+    add_task = mocker.patch(
+        "workers.merge_message.worker.add_task", new=mocker.AsyncMock()
+    )
+    clear = mocker.patch(
+        "workers.merge_message.worker.clear_ready_callback", new=mocker.AsyncMock()
+    )
+    remove = mocker.patch(
+        "workers.merge_message.worker.remove_callback_id", new=mocker.AsyncMock()
+    )
+    sleep = mocker.patch(
+        "workers.merge_message.worker.asyncio.sleep", new=mocker.AsyncMock()
+    )
+
+    last = settings.merge_max_attempts - 1
+    await _handle_merge_failure(_task(last), "rid", ["cb1", "cb2"], logger)
+
+    # The unmergeable callbacks are dropped from the ready index and the
+    # callbacks table, so the next drain pass can't pick them up again.
+    assert {c.args[1] for c in clear.await_args_list} == {"cb1", "cb2"}
+    assert {c.args[0] for c in remove.await_args_list} == {"cb1", "cb2"}
+    # No point sleeping: this batch isn't being retried.
+    sleep.assert_not_awaited()
+    # A wake task still goes back so anything else ready still gets merged --
+    # with the counter reset, since the poison batch is gone.
+    _, fields, _ = add_task.await_args.args
+    assert MERGE_ATTEMPT_FIELD not in fields
+
+
+async def test_breaker_does_not_fire_with_nothing_to_discard(mocker):
+    """A failure before any batch was read (get_ready_callbacks itself raising)
+    has nothing to drop, so it just retries."""
+    add_task = mocker.patch(
+        "workers.merge_message.worker.add_task", new=mocker.AsyncMock()
+    )
+    clear = mocker.patch(
+        "workers.merge_message.worker.clear_ready_callback", new=mocker.AsyncMock()
+    )
+
+    await _handle_merge_failure(
+        _task(settings.merge_max_attempts + 5), "rid", [], logger
+    )
+
+    clear.assert_not_awaited()
+    _, fields, _ = add_task.await_args.args
+    assert fields[MERGE_ATTEMPT_FIELD] == str(settings.merge_max_attempts + 6)
+
+
+async def test_garbled_attempt_field_does_not_break_the_retry(mocker):
+    add_task = mocker.patch(
+        "workers.merge_message.worker.add_task", new=mocker.AsyncMock()
+    )
+    task = _task()
+    task[1][MERGE_ATTEMPT_FIELD] = "not-a-number"
+
+    await _handle_merge_failure(task, "rid", ["cb1"], logger)
+
+    _, fields, _ = add_task.await_args.args
+    assert fields[MERGE_ATTEMPT_FIELD] == "1"

tests/unit/test_shared_utils.py

@@ -691,3 +691,74 @@ def task(msg_id):
         "callback two retrieval log",
     ]
     first_logger.handlers.clear()
+
+
+def test_merge_kgraph_tolerates_nodes_missing_optional_fields():
+    """A node that omits ``name`` (or ``categories``/``attributes``) must merge.
+
+    ``name``, ``categories`` and ``attributes`` are all optional in TRAPI, and
+    subservices do leave them out. merge_kgraph used to subscript them, so a
+    single such node on an id already in the accumulator raised
+    ``KeyError: 'name'`` out of the process-pool child and aborted the whole
+    batch merge -- which the worker then retried forever.
+    """
+    og = {
+        "nodes": {
+            "MONDO:1": {"name": "Original", "categories": ["biolink:Disease"]},
+            "MONDO:2": {},
+        },
+        "edges": {},
+    }
+    new = {
+        # No 'name', no 'attributes' -- and 'categories' present on only one.
+        "nodes": {
+            "MONDO:1": {"categories": ["biolink:NamedThing"]},
+            "MONDO:2": {"name": "Filled in", "attributes": [{"a": 1}]},
+        },
+        "edges": {},
+    }
+    merged = merge_kgraph(og, new, "infores:test", logger)
+    # The existing name survives a new node that simply doesn't carry one.
+    assert merged["nodes"]["MONDO:1"]["name"] == "Original"
+    assert set(merged["nodes"]["MONDO:1"]["categories"]) == {
+        "biolink:Disease",
+        "biolink:NamedThing",
+    }
+    # ...and fields absent from the existing entry are adopted from the new one.
+    assert merged["nodes"]["MONDO:2"]["name"] == "Filled in"
+    assert merged["nodes"]["MONDO:2"]["attributes"] == [{"a": 1}]
+
+
+def test_merge_kgraph_tolerates_edges_missing_optional_fields():
+    """Same for overlapping edges: absent ``attributes``/``sources`` merge."""
+    og = {
+        "nodes": {},
+        "edges": {
+            "e1": {"subject": "A", "object": "B"},
+            "e2": {
+                "subject": "A",
+                "object": "B",
+                "sources": [{"resource_id": "infores:one"}],
+            },
+        },
+    }
+    new = {
+        "nodes": {},
+        "edges": {
+            "e1": {
+                "subject": "A",
+                "object": "B",
+                "sources": [{"resource_id": "infores:two"}],
+            },
+            # Nothing optional at all on this one.
+            "e2": {"subject": "A", "object": "B"},
+        },
+    }
+    merged = merge_kgraph(og, new, "infores:test", logger)
+    assert merged["edges"]["e1"]["sources"] == [{"resource_id": "infores:two"}]
+    assert merged["edges"]["e2"]["sources"] == [{"resource_id": "infores:one"}]
+
+
+def test_is_support_edge_handles_missing_and_null_attributes():
+    assert is_support_edge({"subject": "A", "object": "B"}) is False
+    assert is_support_edge({"attributes": None}) is False