import json import os import time from logging import getLogger from typing import NamedTuple from peewee import OperationalError from defence360agent.internals.feature_flags import ( MESSAGE_LOSS_OBSERVABILITY_FLAG, is_enabled, ) from defence360agent.internals.message_status_publisher import Gen, publisher from defence360agent.model.instance import db from defence360agent.model.messages_to_send import MessageToSend from defence360agent.utils import is_disk_error from defence360agent.utils.common import rate_limit logger = getLogger(__name__) _reporter_gen_evicted = Gen() ROWS_ENV = "IMUNIFY360_SEND_QUEUE_MAX_ROWS" BYTES_ENV = "IMUNIFY360_SEND_QUEUE_MAX_BYTES" DEFAULT_STORAGE_LIMIT = 10_000 # ponytail: the sender loads the whole stored queue into memory every round # (_read_round_batch), so this is a per-round RAM budget as much as a disk # one; page the round before raising it DEFAULT_STORAGE_BYTES_LIMIT = 32 * 1024 * 1024 def _is_storage_full(exc: Exception) -> bool: if isinstance(exc, OperationalError): return "database or disk is full" in str(exc).lower() return is_disk_error(exc) class QueueMetadata(NamedTuple): method: str | None group_id: str | None class PersistentMessagesQueue: """ The queue to store messages sent to the server if it is unavailable. - stores more recent data; if a limit is exceeded, older messages are deleted. - no duplicate messages are sent NOTE: it is worth remembering that when writing a large number of messages, the amount of memory used may increase by the size of the sqlite cache (this may not be immediately obvious). https://www.sqlite.org/pragma.html#pragma_cache_size """ #: evicted first: a rescan re-derives these, unlike any other bulk method SHED_FIRST_METHODS = frozenset({"MALWARE_SCAN_RESULT"}) #: evicted last: Accumulate stamps it before its own chunks, so plain #: oldest-first sheds it first RETAINED_METHODS = frozenset({"MALWARE_SCAN_SUMMARY"}) #: share of storage_limit, in rows; bounds how far retained rows can #: displace ordinary ones, and is not what keeps eviction progressing RETAINED_SHARE = 0.5 def __init__( self, buffer_limit=20, storage_limit=None, storage_bytes_limit=None, model=None, ): if storage_limit is None: storage_limit = int( os.environ.get(ROWS_ENV, DEFAULT_STORAGE_LIMIT) ) if storage_bytes_limit is None: storage_bytes_limit = int( os.environ.get(BYTES_ENV, DEFAULT_STORAGE_BYTES_LIMIT) ) # a cap below 1 can never be met, so every flush would purge the queue if storage_limit < 1: raise ValueError( f"{ROWS_ENV} must be a positive integer, got {storage_limit}" ) if storage_bytes_limit < 1: raise ValueError( f"{BYTES_ENV} must be a positive integer," f" got {storage_bytes_limit}" ) self._buffer_limit = buffer_limit self._storage_limit = storage_limit self._storage_bytes_limit = storage_bytes_limit self._buffer = [] # [(timestamp, message, *QueueMetadata, size),...] #: {id(row): (row, QueueMetadata)} for the current drain only. Holding #: the row is what keeps its id() from being recycled under us. self._inflight = {} self._model = model or MessageToSend self.dropped_total = 0 self._evicted = 0 self._warn_full = rate_limit(period=10 * 60)(logger.warning) def pop_evicted(self) -> int: """Evictions since the last call, then reset (delta for metrics).""" evicted, self._evicted = self._evicted, 0 return evicted def push_buffer_to_storage(self) -> None: try: self._push_buffer_to_storage() except (OperationalError, OSError) as e: if not _is_storage_full(e): raise self._warn_full("Can't store queued messages: %s", e) # keep the newest messages, the buffer must not grow unbounded overflow = len(self._buffer) - self._buffer_limit if overflow > 0: del self._buffer[:overflow] self.dropped_total += overflow if is_enabled(MESSAGE_LOSS_OBSERVABILITY_FLAG): self._evicted += overflow def _push_buffer_to_storage(self) -> None: if self._buffer: evicted = [] removed = freed = 0 with db.atomic(): # buffer may contain older messages than db, # so remove oldest items after insert self._model.insert_many(self._buffer) over_rows = self.storage_size - self._storage_limit over_bytes = self.storage_bytes - self._storage_bytes_limit if over_rows > 0 or over_bytes > 0: removed, freed, evicted = self._evict( over_rows, over_bytes ) # after the commit: a failed commit rolls everything back, so the # buffer is still the only copy and nothing was evicted self._buffer = [] if removed: # This is the last point at which the messages exist, so it # is the only place their loss can be reported. self.dropped_total += removed if is_enabled(MESSAGE_LOSS_OBSERVABILITY_FLAG): self._evicted += removed logger.warning( "Persistent message queue overflow: dropped %d" " message(s) (%d bytes), storage_limit=%d," " storage_bytes_limit=%d, dropped_total=%d", removed, freed, self._storage_limit, self._storage_bytes_limit, self.dropped_total, ) # A send round holds a snapshot of the store, so a row evicted here # can still be delivered by that round: this stage marks the # deletion, a later agent-sent means delivery, and the server # counts a loss only when no later stage follows. for message in evicted: publisher.report( message, _reporter_gen_evicted, stage="agent-evicted" ) def _evict( self, over_rows: int, over_bytes: int ) -> tuple[int, int, list[dict]]: victims, freed = [], 0 for row_id, size in self._model.eviction_order( shed_first=self.SHED_FIRST_METHODS, retained=self.RETAINED_METHODS, retained_keep=int(self._storage_limit * self.RETAINED_SHARE), ): if len(victims) >= over_rows and freed >= over_bytes: break victims.append(row_id) freed += size removed, evicted = 0, [] # sqlite may have internal limit of variables-per-query for i in range(0, len(victims), 100): chunk = victims[i : i + 100] evicted += self._decode_victims(chunk) removed += self._model.delete_in(chunk) return removed, freed, evicted def _decode_victims(self, ids: list) -> list[dict]: if not is_enabled("mqtt_tracking"): return [] blobs = ( self._model.select(self._model.message) .where(self._model.id.in_(ids)) .tuples() .iterator() ) messages = [] for (blob,) in blobs: try: message = json.loads(blob) except ValueError: continue if isinstance(message, dict): # only what the report reads; a scan result carries its hits messages.append( { key: message.get(key) for key in ( "method", "message_reporter_id", "message_reporter_increment", ) } ) return messages def peek_stored(self) -> list: """Return stored rows as (id, timestamp, message) oldest-first without deleting (buffer is neither flushed nor included).""" return list(self._model.get_all_ordered().tuples()) def drain_buffer(self) -> list: """Return and clear the buffer as round rows (None, timestamp, message). Each row's metadata is retained, keyed by that row object, until the next drain replaces the map.""" items, self._buffer = self._buffer, [] rows = [(None, ts, msg) for ts, msg, _method, _group, _size in items] self._inflight = { id(row): (row, QueueMetadata(method, group_id)) for row, (_ts, _msg, method, group_id, _size) in zip(rows, items) } return rows def metadata_for(self, row) -> QueueMetadata: entry = self._inflight.get(id(row)) if entry is None or entry[0] is not row: logger.warning("No retained metadata for a re-queued row") return QueueMetadata(None, None) return entry[1] def requeue(self, rows) -> None: self.put_many( [ (row[1], row[2], *self.metadata_for(row), len(row[2])) for row in rows ] ) self.push_buffer_to_storage() def delete(self, ids: list) -> None: if ids: try: with db.atomic(): self._model.delete_in(ids) except (OperationalError, OSError) as e: if not _is_storage_full(e): raise self._warn_full("Can't delete sent messages: %s", e) def update_message(self, message_id: int, message: bytes) -> None: try: with db.atomic(): self._model.set_message(message_id, message) except (OperationalError, OSError) as e: if not _is_storage_full(e): raise # the stored row keeps its old retry count self._warn_full("Can't update a stored message: %s", e) def empty(self) -> bool: return self.qsize() == 0 def qsize(self) -> int: return self.storage_size + len(self._buffer) @property def buffer_size(self) -> int: return len(self._buffer) @property def storage_size(self) -> int: return self._model.select().count() @property def storage_bytes(self) -> int: return self._model.total_bytes() def put( self, message: bytes, timestamp=None, meta: QueueMetadata | None = None ): if timestamp is None: timestamp = time.time() if meta is None: meta = QueueMetadata(None, None) self._buffer.append((timestamp, message, *meta, len(message))) if self.buffer_size >= self._buffer_limit: self.push_buffer_to_storage() def put_many(self, messages: list[tuple]) -> None: self._buffer.extend(messages) if self.buffer_size >= self._buffer_limit: self.push_buffer_to_storage()