Viewing File: /opt/imunify360/venv/versions/imunify-core-8.14.0-1/defence360agent/internals/persistent_message.py
import time
from logging import getLogger
from typing import NamedTuple
from defence360agent.internals.feature_flags import (
MESSAGE_LOSS_OBSERVABILITY_FLAG,
is_enabled,
)
from defence360agent.model.instance import db
from defence360agent.model.messages_to_send import MessageToSend
logger = getLogger(__name__)
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
"""
def __init__(self, buffer_limit=20, storage_limit=1000, model=None):
self._buffer_limit = buffer_limit
self._storage_limit = storage_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
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:
if self._buffer:
with db.atomic():
# buffer may contain older messages than db,
# so remove oldest items after insert
self._model.insert_many(self._buffer)
need_to_remove = self.storage_size - self._storage_limit
if need_to_remove > 0:
# keep only the most recent messages
removed = self._model.delete_old(need_to_remove)
# 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 oldest"
" message(s), storage_limit=%d, dropped_total=%d",
removed,
self._storage_limit,
self.dropped_total,
)
self._buffer = []
def pop_all(self) -> list:
"""Empty the queue into (timestamp, message) pairs, oldest first.
Deleting before the send is confirmed is what made an interrupted
round lose the batch, so this exists only for firewall releases whose
im360 still drives the queue this way; in-tree callers use
peek_stored/drain_buffer/requeue/delete.
TODO: remove once the firewall using _read_round_batch is stable.
"""
with db.atomic():
stored = [
(timestamp, message)
for _id, timestamp, message in self.peek_stored()
]
self._model.delete().execute()
buffered = [
(timestamp, message)
for timestamp, message, *_metadata in self._buffer
]
self._buffer = []
self._inflight = {}
# by timestamp only: a stable sort keeps the flushed rows ahead of
# the buffered ones at an equal timestamp, where comparing the whole
# pair would order them by payload bytes instead
return sorted(stored + buffered, key=lambda row: row[0])
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:
with db.atomic():
self._model.delete_in(ids)
def update_message(self, message_id: int, message: bytes) -> None:
with db.atomic():
self._model.set_message(message_id, message)
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()
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(
self._with_metadata(message) for message in messages
)
if self.buffer_size >= self._buffer_limit:
self.push_buffer_to_storage()
@staticmethod
def _with_metadata(item: tuple) -> tuple:
"""Pad a bare pair from pop_all to a stored row; see pop_all."""
if len(item) == 2:
timestamp, message = item
return (timestamp, message, None, None, len(message))
return item
Back to Directory