diff --git a/shepherd_utils/config.py b/shepherd_utils/config.py index 0c6b07a..d3526a7 100644 --- a/shepherd_utils/config.py +++ b/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 diff --git a/shepherd_utils/shared.py b/shepherd_utils/shared.py index d9675d8..2fb37d6 100644 --- a/shepherd_utils/shared.py +++ b/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,10 +1020,11 @@ 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 @@ -1025,9 +1032,9 @@ def merge_kgraph(og_message, new_message, source, logger: logging.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( diff --git a/tests/unit/test_merge_message_retry.py b/tests/unit/test_merge_message_retry.py new file mode 100644 index 0000000..7f50b11 --- /dev/null +++ b/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" diff --git a/tests/unit/test_shared_utils.py b/tests/unit/test_shared_utils.py index fb14202..8279acb 100644 --- a/tests/unit/test_shared_utils.py +++ b/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 diff --git a/workers/merge_message/worker.py b/workers/merge_message/worker.py index 7d4d186..f08eca9 100644 --- a/workers/merge_message/worker.py +++ b/workers/merge_message/worker.py @@ -46,6 +46,9 @@ GROUP = "consumer" CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 10 +# Task field carrying how many times this wake task's merge has failed in a +# row. See _reenqueue_wake_task. +MERGE_ATTEMPT_FIELD = "merge_attempt" tracer = setup_tracer(STREAM) LOGGER = get_worker_logger(STREAM) @@ -828,15 +831,81 @@ def merge_messages_by_id( return bool(merged) -async def _reenqueue_wake_task(task, logger): +async def _clear_batch(response_id, callback_ids, logger): + """Drop a processed batch from the ready index and callbacks table. + + Every callback we attempted is removed (not just the ones that merged) + so a callback whose payload has vanished can't wedge the drain loop. + """ + await asyncio.gather( + *(clear_ready_callback(response_id, cb, logger) for cb in callback_ids) + ) + await asyncio.gather(*(remove_callback_id(cb, logger) for cb in callback_ids)) + + +async def _handle_merge_failure(task, response_id, ready, logger): + """Back off and retry a failed merge, giving up on a batch that can't merge. + + A merge failure re-enqueues the wake task and leaves the batch in the + ready index so nothing is dropped on a transient error. But a batch that + fails *deterministically* (a malformed callback the merge chokes on) + then fails again on every retry, and because the re-enqueue was + immediate and unbounded the query spun in a hot loop -- re-merging, + raising, re-enqueueing many times a second for as long as the query + lived, burning a pool slot and flooding the logs. + + So: back off exponentially between retries, and once the same wake task + has failed ``merge_max_attempts`` times in a row, drop the batch it + keeps failing on (clearing it from the ready index and callbacks table) + and re-enqueue with the counter reset. The query loses those callbacks + -- which it was never going to merge anyway -- and goes on to merge the + rest and finish, instead of spinning until it times out. + """ + try: + attempt = int(task[1].get(MERGE_ATTEMPT_FIELD, 0) or 0) + 1 + except (TypeError, ValueError): + attempt = 1 + max_attempts = settings.merge_max_attempts + if max_attempts > 0 and attempt >= max_attempts and ready: + logger.error( + f"Merge failed {attempt} times in a row for {response_id}; " + f"discarding {len(ready)} unmergeable callback(s): " + f"{', '.join(ready)}" + ) + await _clear_batch(response_id, ready, logger) + # Counter reset: the poison batch is gone, so whatever else is + # ready deserves a clean run of attempts. + await _reenqueue_wake_task(task, logger) + return + backoff = min( + settings.merge_retry_backoff * (2 ** (attempt - 1)), + settings.merge_retry_backoff_max, + ) + if backoff > 0: + await asyncio.sleep(backoff) + await _reenqueue_wake_task(task, logger, attempt) + + +async def _reenqueue_wake_task(task, logger, attempt: int = 0): """Put a fresh merge_message wake task back on the stream. Used when this worker can't make progress on a callback right now (the query's lock is held by someone else, or a merge failed). The callback's entry in the ready index is left intact; the new wake task simply drives another drain attempt later, so callbacks are retried rather than dropped. + + ``attempt`` carries the consecutive-failure count forward on the retry + path. Re-enqueueing mints a brand new stream message, so Redis' + ``times_delivered`` (what the reclaim poison-pill breaker reads) resets to + 1 every time and can never trip here -- the count has to ride in the task + fields instead. It is dropped when 0 so the normal path enqueues exactly + the fields it always did. """ fields = {k: v for k, v in task[1].items() if k != "_started_at"} + if attempt: + fields[MERGE_ATTEMPT_FIELD] = str(attempt) + else: + fields.pop(MERGE_ATTEMPT_FIELD, None) await add_task(STREAM, fields, logger) @@ -860,17 +929,6 @@ async def poll_for_tasks(): task_timeout=settings.pool_task_timeout_sec, ) - async def _clear_batch(response_id, callback_ids, logger): - """Drop a processed batch from the ready index and callbacks table. - - Every callback we attempted is removed (not just the ones that merged) - so a callback whose payload has vanished can't wedge the drain loop. - """ - await asyncio.gather( - *(clear_ready_callback(response_id, cb, logger) for cb in callback_ids) - ) - await asyncio.gather(*(remove_callback_id(cb, logger) for cb in callback_ids)) - def _ingest_merge_logs(logger, entries): """Fold the merge child's log records into this task's query logger. @@ -932,6 +990,9 @@ async def process_query(task, parent_ctx, logger, limiter): return lock_time = time.time() + # Bound before the try: the failure handler reports the batch + # that failed, and get_ready_callbacks itself can raise. + ready: list[str] = [] # Drain the query to empty: one load + one save per iteration. # Re-reading the set each pass sweeps up callbacks that arrived # while we were merging, so the holder does all of this query's @@ -959,11 +1020,14 @@ async def process_query(task, parent_ctx, logger, limiter): await refresh_lock(response_id, CONSUMER, 45000, logger) except BrokenProcessPool: # pool.run already swapped in a fresh executor; here we just - # release the lock and re-enqueue so the callback is retried. + # release the lock and retry the batch. Counted like any + # other failure: a child that is OOM-killed (or timed out) + # by one particular batch is killed by it again on every + # retry, so that batch has to age out too. logger.error(f"[{callback_id}] Process pool broken; re-enqueuing.") await remove_lock(response_id, CONSUMER, logger) - await _reenqueue_wake_task(task, logger) span.set_attribute("merge.drained_callbacks", drained) + await _handle_merge_failure(task, response_id, ready, logger) return except Exception: logger.error( @@ -971,8 +1035,8 @@ async def process_query(task, parent_ctx, logger, limiter): f"{traceback.format_exc()}" ) await remove_lock(response_id, CONSUMER, logger) - await _reenqueue_wake_task(task, logger) span.set_attribute("merge.drained_callbacks", drained) + await _handle_merge_failure(task, response_id, ready, logger) return span.set_attribute("merge.drained_callbacks", drained)