From 0d86a1bfffd3041bf1d2a2c657cc1e59a3b30b2a Mon Sep 17 00:00:00 2001 From: Brianna Smart Date: Thu, 9 Jul 2026 16:10:56 -0700 Subject: [PATCH] Update ingester so that it can close out when pods scale down Move handle_kafka_message to async so that it doesn't slow things down Move the shutdown to warning so it appears in logs Update ingester to batch process messages and update unit tests Claude assisted Remove excess code and allow for proper multhreading Add arguments --- alertingest/bin/alertdb_ingester.py | 70 +++++++- alertingest/ingester.py | 263 ++++++++++++++-------------- tests/test_ingester.py | 133 +++++++++++++- 3 files changed, 335 insertions(+), 131 deletions(-) diff --git a/alertingest/bin/alertdb_ingester.py b/alertingest/bin/alertdb_ingester.py index 1d5acff..6a051fb 100644 --- a/alertingest/bin/alertdb_ingester.py +++ b/alertingest/bin/alertdb_ingester.py @@ -2,6 +2,7 @@ import asyncio import logging import os +import signal from alertingest.ingester import IngestWorker, KafkaConnectionParams from alertingest.schema_registry import SchemaRegistryClient @@ -138,6 +139,40 @@ def main(): default=30, help="Maximum number of idle-prefix summaries to remember (default: 30)", ) + parser.add_argument( + "--batch-size", + type=int, + default=20, + help="Maximum messages to fetch and process concurrently per loop iteration (default: 20)", + ) + parser.add_argument( + "--commit-timeout", + type=int, + default=600, + help=( + "Maximum seconds to hold uncommitted offsets before forcing a commit, " + "regardless of commit interval (default: 600)" + ), + ) + parser.add_argument( + "--commit-interval", + type=int, + default=100, + help="Minimum number of messages between offset commits (default: 100)", + ) + parser.add_argument( + "--limit", + type=int, + default=-1, + help="Maximum number of messages to copy; -1 means no limit (default: -1)", + ) + parser.add_argument( + "--auto-offset-reset", + type=str, + choices=("latest", "earliest"), + default="latest", + help="Where to start reading when joining a new topic (default: latest)", + ) args = parser.parse_args() @@ -184,4 +219,37 @@ def main(): prefix_idle_timeout=args.prefix_idle_timeout, max_logged_prefixes=args.max_logged_prefixes, ) - asyncio.get_event_loop().run_until_complete(worker.run()) + asyncio.run( + _run_worker( + worker, + batch_size=args.batch_size, + commit_timeout=args.commit_timeout, + commit_interval=args.commit_interval, + limit=args.limit, + auto_offset_reset=args.auto_offset_reset, + ) + ) + + +async def _run_worker( + worker, + batch_size=20, + commit_timeout=600, + commit_interval=100, + limit=-1, + auto_offset_reset="latest", +): + loop = asyncio.get_running_loop() + task = asyncio.current_task() + for sig in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(sig, task.cancel) + try: + await worker.run( + batch_size=batch_size, + commit_timeout=commit_timeout, + commit_interval=commit_interval, + limit=limit, + auto_offset_reset=auto_offset_reset, + ) + except asyncio.CancelledError: + pass diff --git a/alertingest/ingester.py b/alertingest/ingester.py index 5e40e33..3cf098d 100644 --- a/alertingest/ingester.py +++ b/alertingest/ingester.py @@ -3,6 +3,7 @@ """ import asyncio +import concurrent.futures import datetime import io import logging @@ -120,12 +121,17 @@ def __init__( self.log_check_timeout = log_check_timeout self.prefix_idle_timeout = prefix_idle_timeout self.max_logged_prefixes = max_logged_prefixes + self._executor: concurrent.futures.ThreadPoolExecutor | None = None + # We have to use this because S#/boto3 and the schema registry + # are synchronous and we need to use them in an async context async def run( self, limit: int = -1, commit_interval: int = 100, auto_offset_reset: str = "latest", + batch_size: int = 20, + commit_timeout: int = 600, ): """Run the consumer, copying messages from Kafka to the IngestWorker's backend. @@ -136,14 +142,26 @@ async def run( Maximum number of messages to copy. If this value is less than 1, no limit is used. The default is -1. commit_interval : int - Interval (measured in messages) between committing the offset of - the worker. Higher values will require more repeated work if the - IngestWorker crashes or backends are unavailable, while lower - values will cost more overhead communicating with Kafka. + Minimum number of messages between offset commits. Commits are + aligned to batch boundaries, so the actual interval is between + commit_interval and commit_interval + batch_size messages. Higher + values will require more repeated work if the IngestWorker crashes + or backends are unavailable, while lower values will cost more + overhead communicating with Kafka. auto_offset_reset : str When reading from a new topic, where should the worker start? Options are "latest" and "earliest". + batch_size : int + Maximum number of messages to fetch and process concurrently per + loop iteration. Higher values increase throughput at the cost of + more memory and thread-pool workers. The default is 20. + commit_timeout : int + Maximum seconds to hold uncommitted offsets. If this many seconds + pass since the last commit and there are pending messages, a commit + is forced regardless of commit_interval. The default is 600 (10 + minutes). """ + self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=batch_size) consumer = self._create_consumer(auto_offset_reset) await consumer.start() @@ -161,53 +179,22 @@ async def run( "logged_prefixes": deque( maxlen=self.max_logged_prefixes ), # Prefixes whose idle summaries have been logged (max 30) + "last_commit_time": asyncio.get_event_loop().time(), # Time of the last commit } logger.info("Ingest worker run loop start.") while True: - try: - logger.debug("Waiting for message.") - msg = await asyncio.wait_for( - consumer.__anext__(), timeout=self.message_timeout - ) - - # Process messages and update the state tracker. Will set - # new_messages to True if new messages have been read. - state_tracker.update( - await self.process_message( - msg, consumer, **state_tracker, worker=self, limit=limit - ) - ) - - # Check if the commit_interval has been reached and submit - # messages if it has - if state_tracker["commit_interval_counter"] == commit_interval: - state_tracker["commit_interval_counter"] = ( - await self.handle_commit( - consumer, state_tracker["commit_interval_counter"] - ) - ) - logger.info( - "Alerts stored today: %s", state_tracker["daily_stored"] - ) - self._check_daily_reset(datetime.datetime.now(), state_tracker) - - # Check message limit - if limit > 0 and state_tracker["limit_n"] >= limit: - logger.info("limit reached - returning") - await self.handle_commit( - consumer, state_tracker["commit_interval_counter"] - ) - self._log_final_summary(state_tracker) - return + logger.debug("Waiting for messages.") + batch = await consumer.getmany( + max_records=batch_size, + timeout_ms=self.message_timeout * 1000, + ) + msgs = [msg for msgs in batch.values() for msg in msgs] - except asyncio.TimeoutError: + if not msgs: logger.info("Waiting timed out, checking for new messages...") - # Check if we are reading new messages. If new messages - # is set to false, don't try and read the partitions. - # If new messages is set to true, we will try and read - # any remaining messages from the partitions. current_time = asyncio.get_event_loop().time() + prev_counter = state_tracker["commit_interval_counter"] ( state_tracker["new_messages"], state_tracker["commit_interval_counter"], @@ -218,6 +205,11 @@ async def run( state_tracker["commit_interval_counter"], state_tracker["new_messages"], ) + if ( + prev_counter > 0 + and state_tracker["commit_interval_counter"] == 0 + ): + state_tracker["last_commit_time"] = current_time self._check_idle_prefixes( current_time, state_tracker["prefix_counts"], @@ -225,94 +217,94 @@ async def run( state_tracker["logged_prefixes"], ) self._check_daily_reset(datetime.datetime.now(), state_tracker) + continue - except Exception as e: - logger.error("Error during message processing: %s", e) - raise - - finally: - await consumer.stop() - - async def process_message( - self, - msg, - consumer, - last_message_time, - commit_interval_counter, - limit_n, - worker, - new_messages, - limit=-1, - daily_stored=0, - day_start_time=None, - prefix_counts=None, - prefix_last_write=None, - logged_prefixes=None, - ): - """Process a single Kafka message. - - The function reads a single kafka message and updates the state - tracker. - - Parameters - ---------- - commit_interval_counter: int - The number of messages since the last commit. - - limit_n : int - Counts the number of messages which have been processed - since the last commit. Will commit once the required number - of messages has been reached and end the loop. Not tracked if - limit is less than 1. + # Process all messages in the batch concurrently; each fires + # its S3 write into the thread pool via run_in_executor. + results = await asyncio.gather( + *[self.handle_kafka_message(msg) for msg in msgs], + return_exceptions=True, + ) - worker : IngestWorker - The ingester worker which is handling the message. + # Log per-message errors before updating state, then re-raise. + first_exc = None + for msg, result in zip(msgs, results): + if isinstance(result, BaseException): + logger.error( + "Error processing message at offset %s: %s", + msg.offset, + result, + ) + if first_exc is None: + first_exc = result + if first_exc is not None: + raise first_exc + + # Update state for the whole batch. + now = asyncio.get_event_loop().time() + n = len(msgs) + state_tracker["last_message_time"] = now + state_tracker["new_messages"] = True + state_tracker["commit_interval_counter"] += n + state_tracker["daily_stored"] += n + if limit > 0: + state_tracker["limit_n"] += n + + for alert_id in results: + alert_prefix = str(alert_id)[:6] + state_tracker["prefix_counts"][alert_prefix] = ( + state_tracker["prefix_counts"].get(alert_prefix, 0) + 1 + ) + state_tracker["prefix_last_write"][alert_prefix] = now - new_messages : bool - Track if we have received new messages, but keep the current state - (to be updated upon timeout) if we have not. + # Commit when interval threshold is reached or exceeded. + if state_tracker["commit_interval_counter"] >= commit_interval: + state_tracker["commit_interval_counter"] = await self.handle_commit( + consumer, state_tracker["commit_interval_counter"] + ) + state_tracker["last_commit_time"] = now + logger.info( + "Alerts stored today: %s", state_tracker["daily_stored"] + ) + self._check_daily_reset(datetime.datetime.now(), state_tracker) + elif ( + state_tracker["commit_interval_counter"] > 0 + and now - state_tracker["last_commit_time"] >= commit_timeout + ): + logger.info( + "No commit in %s seconds, committing %s pending messages.", + commit_timeout, + state_tracker["commit_interval_counter"], + ) + state_tracker["commit_interval_counter"] = await self.handle_commit( + consumer, state_tracker["commit_interval_counter"] + ) + state_tracker["last_commit_time"] = now - limit : int - The maximum number of messages to process. If this value is less - than 1, we do not track the number of messages processed. + # Check message limit. + if limit > 0 and state_tracker["limit_n"] >= limit: + logger.info("limit reached - returning") + await self.handle_commit( + consumer, state_tracker["commit_interval_counter"] + ) + self._log_final_summary(state_tracker) + return - """ - if prefix_counts is None: - prefix_counts = {} - if prefix_last_write is None: - prefix_last_write = {} - if logged_prefixes is None: - logged_prefixes = deque(maxlen=self.max_logged_prefixes) + except asyncio.CancelledError: + logger.warning( + "Shutdown signal received, committing pending offsets before exit." + ) + await self.handle_commit(consumer, state_tracker["commit_interval_counter"]) + raise - try: - alert_id = worker.handle_kafka_message(msg) except Exception as e: - logger.error("Error processing message at offset %s: %s", msg.offset, e) - logger.exception("full traceback") + logger.error("Error during message processing: %s", e) raise - logger.debug("handle complete") - if limit > 0: - limit_n += 1 - will_return = True if msg else new_messages - - now = asyncio.get_event_loop().time() - daily_stored += 1 - alert_prefix = str(alert_id)[:6] - prefix_counts[alert_prefix] = prefix_counts.get(alert_prefix, 0) + 1 - prefix_last_write[alert_prefix] = now - - return { - "last_message_time": now, - "commit_interval_counter": commit_interval_counter + 1, - "limit_n": limit_n, - "new_messages": will_return, - "daily_stored": daily_stored, - "day_start_time": day_start_time, - "prefix_counts": prefix_counts, - "prefix_last_write": prefix_last_write, - "logged_prefixes": logged_prefixes, - } + finally: + await consumer.stop() + self._executor.shutdown(wait=True) + self._executor = None async def handle_commit(self, consumer, commit_interval_counter): """Handle committing of consumer offsets. @@ -545,7 +537,7 @@ def _create_mtls_consumer(self, auto_offset_reset): consumer.subscribe(topics=self.kafka_params.topics) return consumer - def handle_kafka_message(self, msg: ConsumerRecord): + async def handle_kafka_message(self, msg: ConsumerRecord): """Handle a single Kafka message. Parses out the schema ID and alert ID from the message. Stores the @@ -564,19 +556,32 @@ def handle_kafka_message(self, msg: ConsumerRecord): msg.partition, msg.offset, ) + loop = asyncio.get_running_loop() raw_msg = msg.value - schema_id, alert_id = self._parse_alert_msg(raw_msg) + + # _parse_alert_msg may fetch from the schema registry on first call + schema_id, alert_id = await loop.run_in_executor( + self._executor, self._parse_alert_msg, raw_msg + ) logger.debug("Parsed message: schema_id=%s, alert_id=%s", schema_id, alert_id) - if not self.backend.schema_exists(schema_id): + if not await loop.run_in_executor( + self._executor, self.backend.schema_exists, schema_id + ): logger.info("%s is a new schema ID - storing it", schema_id) - encoded_schema = self.schema_registry.get_raw_schema(schema_id) - self.backend.store_schema(schema_id, encoded_schema) + encoded_schema = await loop.run_in_executor( + self._executor, self.schema_registry.get_raw_schema, schema_id + ) + await loop.run_in_executor( + self._executor, self.backend.store_schema, schema_id, encoded_schema + ) else: logger.debug("Schema %s already exists, skipping storage", schema_id) logger.debug("Storing alert %s to backend.", alert_id) - self.backend.store_alert(alert_id, raw_msg) + await loop.run_in_executor( + self._executor, self.backend.store_alert, alert_id, raw_msg + ) logger.debug("Alert %s stored successfully.", alert_id) return alert_id diff --git a/tests/test_ingester.py b/tests/test_ingester.py index 884596e..802ed9e 100644 --- a/tests/test_ingester.py +++ b/tests/test_ingester.py @@ -1,7 +1,8 @@ +import asyncio import datetime import logging from collections import deque -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -180,3 +181,133 @@ def test_confluent_wire_format_parsing(): data = b"\x00\x00\x00\x00\x04\xd3" have = _read_confluent_wire_format_header(data) assert have == 4 + + +def _make_msg(offset=0): + """Return a minimal mock ConsumerRecord.""" + msg = MagicMock() + msg.offset = offset + msg.topic = "test-topic" + msg.partition = 0 + return msg + + +def _make_consumer(batches): + """Return a mock aiokafka consumer whose getmany yields successive batches. + + Each element of `batches` is a dict that will be returned by one + getmany() call. An empty dict simulates a timeout (no messages). + """ + consumer = MagicMock() + consumer.getmany = AsyncMock(side_effect=batches) + consumer.start = AsyncMock() + consumer.stop = AsyncMock() + consumer.commit = AsyncMock() + consumer.assignment.return_value = [] # skip per-partition logging in + # handle_commit + return consumer + + +def test_run_processes_full_batch(): + """All messages in a batch are passed to handle_kafka_message.""" + worker = _make_worker() + msgs = [_make_msg(offset=i) for i in range(3)] + consumer = _make_consumer(batches=[{"tp": msgs}]) + worker._create_consumer = MagicMock(return_value=consumer) + worker.handle_kafka_message = AsyncMock(side_effect=[100000, 200000, 300000]) + + asyncio.run(worker.run(limit=3, commit_interval=100)) + + assert worker.handle_kafka_message.call_count == 3 + consumer.start.assert_awaited_once() + consumer.stop.assert_awaited_once() + + +def test_run_state_updated_for_batch(): + """daily_stored increments by the full batch size.""" + worker = _make_worker() + msgs = [_make_msg(offset=i) for i in range(4)] + consumer = _make_consumer(batches=[{"tp": msgs}]) + worker._create_consumer = MagicMock(return_value=consumer) + # Use alert IDs with distinct 6-char prefixes to produce two prefix + # buckets. + worker.handle_kafka_message = AsyncMock( + side_effect=[111111000, 111111001, 222222000, 222222001] + ) + + # Capture state by hooking _log_final_summary (called when limit is hit). + captured = {} + original = worker._log_final_summary + + def capture(state): + captured.update(state) + original(state) + + worker._log_final_summary = capture + + asyncio.run(worker.run(limit=4, commit_interval=100)) + + assert captured["daily_stored"] == 4 + assert captured["prefix_counts"]["111111"] == 2 + assert captured["prefix_counts"]["222222"] == 2 + + +def test_run_empty_batch_invokes_process_timeout(): + """An empty getmany result (timeout) calls process_timeout.""" + worker = _make_worker() + msg = _make_msg() + consumer = _make_consumer(batches=[{}, {"tp": [msg]}]) + worker._create_consumer = MagicMock(return_value=consumer) + worker.handle_kafka_message = AsyncMock(return_value=123456789) + worker.process_timeout = AsyncMock(return_value=(False, 0)) + + asyncio.run(worker.run(limit=1, commit_interval=100)) + + worker.process_timeout.assert_awaited_once() + + +def test_run_batch_error_is_logged_and_reraised(caplog): + """A failing message is logged with its offset and the exception + propagates.""" + worker = _make_worker() + msgs = [_make_msg(offset=0), _make_msg(offset=7)] + consumer = _make_consumer(batches=[{"tp": msgs}]) + worker._create_consumer = MagicMock(return_value=consumer) + exc = ValueError("bad alert") + worker.handle_kafka_message = AsyncMock(side_effect=[123456789, exc]) + + with caplog.at_level(logging.ERROR, logger="alertingest.ingester"): + with pytest.raises(ValueError, match="bad alert"): + asyncio.run(worker.run(limit=10, commit_interval=100)) + + assert "offset 7" in caplog.text + + +def test_run_commits_when_interval_reached(): + """Commit fires when commit_interval_counter meets or exceeds the + threshold.""" + worker = _make_worker() + msgs = [_make_msg(offset=i) for i in range(5)] + consumer = _make_consumer(batches=[{"tp": msgs}]) + worker._create_consumer = MagicMock(return_value=consumer) + worker.handle_kafka_message = AsyncMock(side_effect=list(range(100000, 100005))) + + asyncio.run(worker.run(limit=5, commit_interval=3)) + + consumer.commit.assert_awaited() + + +def test_run_commits_on_commit_timeout(): + """Commit fires when commit_timeout seconds elapse since the last commit, + even if commit_interval hasn't been reached.""" + worker = _make_worker() + msgs = [_make_msg(offset=i) for i in range(2)] + consumer = _make_consumer(batches=[{"tp": msgs}]) + worker._create_consumer = MagicMock(return_value=consumer) + worker.handle_kafka_message = AsyncMock(side_effect=[100000, 200000]) + + # commit_interval=10 means the count-based commit won't fire for 2 messages + # commit_timeout=0 ensures any elapsed time triggers a time-based commit + asyncio.run(worker.run(limit=2, commit_interval=10, commit_timeout=0)) + + consumer.commit.assert_awaited()