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()