From c02393ce741a49be9960dee0a82cd1009a0c63e7 Mon Sep 17 00:00:00 2001 From: Renn F Date: Wed, 14 Oct 2026 21:27:50 +0200 Subject: [PATCH 1/4] fix(buffer): persist to Redis outside the buffer lock and never recurse into persistence Also accept unpadded base64url encryption keys and include the cause in the startup encryption failure message. --- guard_agent/_buffer_lifecycle.py | 4 + guard_agent/_buffer_queue.py | 44 +++-- guard_agent/_buffer_redis.py | 45 +++++ guard_agent/_transport_lifecycle.py | 3 +- guard_agent/buffer.py | 1 + guard_agent/encryption.py | 3 +- tests/conftest.py | 26 +++ tests/test_buffer_redis_reentrancy.py | 241 ++++++++++++++++++++++++++ 8 files changed, 347 insertions(+), 20 deletions(-) create mode 100644 tests/test_buffer_redis_reentrancy.py diff --git a/guard_agent/_buffer_lifecycle.py b/guard_agent/_buffer_lifecycle.py index e2c835e..45d3e4c 100644 --- a/guard_agent/_buffer_lifecycle.py +++ b/guard_agent/_buffer_lifecycle.py @@ -18,6 +18,7 @@ class BufferLifecycleMixin: _flush_semaphore: asyncio.Semaphore | None _running: bool _inflight_flush_tasks: set[asyncio.Task[None]] + _inflight_persist_tasks: set[asyncio.Task[None]] async def start_auto_flush(self) -> None: if self._flush_task and not self._flush_task.done(): @@ -41,6 +42,9 @@ async def stop_auto_flush(self) -> None: if self._inflight_flush_tasks: await asyncio.gather(*self._inflight_flush_tasks, return_exceptions=True) self._inflight_flush_tasks.clear() + if self._inflight_persist_tasks: + await asyncio.gather(*self._inflight_persist_tasks, return_exceptions=True) + self._inflight_persist_tasks.clear() self._flush_semaphore = None async def stop(self) -> None: diff --git a/guard_agent/_buffer_queue.py b/guard_agent/_buffer_queue.py index 7402b02..adc3ed3 100644 --- a/guard_agent/_buffer_queue.py +++ b/guard_agent/_buffer_queue.py @@ -1,15 +1,33 @@ import asyncio import time +from collections.abc import Callable, Coroutine +from typing import Any, TypeVar from guard_agent._buffer_lifecycle import BufferLifecycleMixin from guard_agent._buffer_overflow import BLOCK_POLICY_POLL_INTERVAL, BufferOverflowMixin +from guard_agent._buffer_redis import _PERSISTING from guard_agent.models import SecurityEvent, SecurityMetric +T = TypeVar("T") + class BufferQueueMixin(BufferOverflowMixin, BufferLifecycleMixin): events_flushed: int metrics_flushed: int + def _spawn_persist( + self, persist: Callable[[T], Coroutine[Any, Any, None]], item: T + ) -> None: + """Run a Redis persist as a tracked background task so the request + path never waits on Redis; stop_auto_flush drains it. Nothing is + spawned without a Redis handler, or while a persist is in flight in + this context (a Redis failure reporting an event must not recurse).""" + if not self.redis_handler or _PERSISTING.get(): + return + task = asyncio.create_task(persist(item)) + self._inflight_persist_tasks.add(task) + task.add_done_callback(self._inflight_persist_tasks.discard) + async def add_event(self, event: SecurityEvent) -> None: """Add security event to buffer honoring the configured overflow policy. @@ -23,15 +41,8 @@ async def add_event(self, event: SecurityEvent) -> None: while True: must_wait = await self._resolve_event_overflow() if not must_wait: - try: - self.event_buffer.append(event) - self.events_buffered += 1 - if self.redis_handler: - key = await self._persist_event_to_redis(event) - if key is not None: - self._event_redis_keys[id(event)] = key - except Exception as e: - self.logger.error(f"Failed to buffer event: {str(e)}") + self.event_buffer.append(event) + self.events_buffered += 1 break try: await asyncio.wait_for( @@ -40,6 +51,8 @@ async def add_event(self, event: SecurityEvent) -> None: except asyncio.TimeoutError: pass + self._spawn_persist(self._persist_and_track_event, event) + if ( len(self.event_buffer) >= self.config.buffer_size * self.config.high_watermark_ratio @@ -58,15 +71,8 @@ async def add_metric(self, metric: SecurityMetric) -> None: while True: must_wait = await self._resolve_metric_overflow() if not must_wait: - try: - self.metric_buffer.append(metric) - self.metrics_buffered += 1 - if self.redis_handler: - key = await self._persist_metric_to_redis(metric) - if key is not None: - self._metric_redis_keys[id(metric)] = key - except Exception as e: - self.logger.error(f"Failed to buffer metric: {str(e)}") + self.metric_buffer.append(metric) + self.metrics_buffered += 1 break try: await asyncio.wait_for( @@ -75,6 +81,8 @@ async def add_metric(self, metric: SecurityMetric) -> None: except asyncio.TimeoutError: pass + self._spawn_persist(self._persist_and_track_metric, metric) + if ( len(self.metric_buffer) >= self.config.buffer_size * self.config.high_watermark_ratio diff --git a/guard_agent/_buffer_redis.py b/guard_agent/_buffer_redis.py index 17de6d5..3e2309b 100644 --- a/guard_agent/_buffer_redis.py +++ b/guard_agent/_buffer_redis.py @@ -1,4 +1,5 @@ import asyncio +import contextvars import logging import time import uuid @@ -8,6 +9,10 @@ from guard_agent.protocols import RedisHandlerProtocol from guard_agent.utils import safe_json_deserialize, safe_json_serialize +_PERSISTING: contextvars.ContextVar[bool] = contextvars.ContextVar( + "guard_agent_redis_persisting", default=False +) + class BufferRedisMixin: logger: logging.Logger @@ -83,6 +88,46 @@ async def _persist_metric_to_redis(self, metric: SecurityMetric) -> str | None: self.logger.warning(f"Failed to persist metric to Redis: {str(e)}") return None + async def _persist_and_track_event(self, event: SecurityEvent) -> None: + """Persist an already-buffered event, outside the buffer lock. + + Skipped while a persist is in flight in this context, so a Redis + failure that emits an event cannot recurse into Redis. If the event + left the buffer meanwhile (flushed), the new key is deleted instead + of tracked, so it neither orphans nor replays after a restart. + """ + token = _PERSISTING.set(True) + try: + key = await self._persist_event_to_redis(event) + except Exception as e: + self.logger.error(f"Failed to buffer event: {str(e)}") + return + finally: + _PERSISTING.reset(token) + if key is None: + return + if any(e is event for e in reversed(self.event_buffer)): + self._event_redis_keys[id(event)] = key + else: + await self.confirm_event_redis_keys([key]) + + async def _persist_and_track_metric(self, metric: SecurityMetric) -> None: + """Metric twin of _persist_and_track_event.""" + token = _PERSISTING.set(True) + try: + key = await self._persist_metric_to_redis(metric) + except Exception as e: + self.logger.error(f"Failed to buffer metric: {str(e)}") + return + finally: + _PERSISTING.reset(token) + if key is None: + return + if any(m is metric for m in reversed(self.metric_buffer)): + self._metric_redis_keys[id(metric)] = key + else: + await self.confirm_metric_redis_keys([key]) + async def _load_from_redis(self) -> None: """Load persisted events/metrics from Redis on startup; track keys. diff --git a/guard_agent/_transport_lifecycle.py b/guard_agent/_transport_lifecycle.py index 91049b2..c3b8249 100644 --- a/guard_agent/_transport_lifecycle.py +++ b/guard_agent/_transport_lifecycle.py @@ -74,7 +74,8 @@ def _init_encryption(self) -> None: raise except Exception as exc: raise EncryptionConfigError( - "Encryption round-trip failed at startup; refusing plaintext fallback" + "Encryption round-trip failed at startup; refusing plaintext " + f"fallback (cause: {exc})" ) from exc async def initialize(self) -> None: diff --git a/guard_agent/buffer.py b/guard_agent/buffer.py index fe5f16a..3570219 100644 --- a/guard_agent/buffer.py +++ b/guard_agent/buffer.py @@ -30,6 +30,7 @@ def __init__( self._flush_semaphore: asyncio.Semaphore | None = None self._running = False self._inflight_flush_tasks: set[asyncio.Task[None]] = set() + self._inflight_persist_tasks: set[asyncio.Task[None]] = set() self._event_redis_keys: dict[int, str] = {} self._metric_redis_keys: dict[int, str] = {} diff --git a/guard_agent/encryption.py b/guard_agent/encryption.py index c8f5080..20cee37 100644 --- a/guard_agent/encryption.py +++ b/guard_agent/encryption.py @@ -69,7 +69,8 @@ def __init__(self, project_key: str) -> None: raise EncryptionError("Project key cannot be empty") try: - key_bytes = base64.urlsafe_b64decode(project_key.encode()) + padded = project_key.strip() + "=" * (-len(project_key.strip()) % 4) + key_bytes = base64.urlsafe_b64decode(padded.encode()) if len(key_bytes) != self.KEY_SIZE: raise EncryptionError( diff --git a/tests/conftest.py b/tests/conftest.py index 6b83d68..e518531 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -24,6 +24,7 @@ django.setup() from datetime import datetime, timezone +from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest @@ -116,3 +117,28 @@ def security_metric() -> SecurityMetric: metric_type="request_count", value=1.0, ) + + +@pytest.fixture(autouse=True) +def settle_redis_persists( + request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch +) -> None: + """Persist is a background task in production; in tests that assert on + Redis state right after an add, wait for it. Tests that exercise the + background behaviour live in test_buffer_redis_reentrancy.py.""" + if request.module.__name__.endswith("test_buffer_redis_reentrancy"): + return + import asyncio + + for name in ("add_event", "add_metric"): + original = getattr(EventBuffer, name) + + def make(orig: Any) -> Any: + async def wrapper(self: EventBuffer, item: Any) -> None: + await orig(self, item) + if self._inflight_persist_tasks: + await asyncio.gather(*self._inflight_persist_tasks) + + return wrapper + + monkeypatch.setattr(EventBuffer, name, make(original)) diff --git a/tests/test_buffer_redis_reentrancy.py b/tests/test_buffer_redis_reentrancy.py new file mode 100644 index 0000000..ebb9fe6 --- /dev/null +++ b/tests/test_buffer_redis_reentrancy.py @@ -0,0 +1,241 @@ +import asyncio +import base64 +from datetime import datetime, timezone +from typing import Any + +import pytest + +from guard_agent.buffer import EventBuffer +from guard_agent.encryption import PayloadEncryptor +from guard_agent.models import AgentConfig, SecurityEvent, SecurityMetric +from guard_agent.transport import HTTPTransport + +FAKE_TIMEOUT = 3.0 + + +def _event(reason: str = "r") -> SecurityEvent: + return SecurityEvent( + timestamp=datetime.now(timezone.utc), + event_type="ip_banned", + ip_address="127.0.0.1", + action_taken="block", + reason=reason, + ) + + +def _metric() -> SecurityMetric: + return SecurityMetric( + timestamp=datetime.now(timezone.utc), metric_type="request_count", value=1.0 + ) + + +def _buffer(**kw: Any) -> EventBuffer: + return EventBuffer( + AgentConfig(api_key="k", endpoint="http://localhost:8000", project_id="p", **kw) + ) + + +class GuardCoreLikeRedis: + """set_key fails and, like guard-core's RedisManager, reports the failure + through the agent's send_event before raising.""" + + def __init__(self, buffer: EventBuffer) -> None: + self.buffer = buffer + self.calls = 0 + + async def set_key(self, *args: Any, **kwargs: Any) -> bool: + self.calls += 1 + await self.buffer.add_event(_event("redis_error")) + await self.buffer.add_metric(_metric()) + raise ConnectionError("redis down") + + async def delete(self, *args: Any, **kwargs: Any) -> int: + return 0 + + async def get_key(self, *args: Any, **kwargs: Any) -> None: + return None + + async def keys(self, *args: Any, **kwargs: Any) -> list[str]: + return [] + + async def initialize(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_redis_failure_reporting_through_agent_does_not_deadlock() -> None: + buf = _buffer() + redis = GuardCoreLikeRedis(buf) + await buf.initialize_redis(redis) # type: ignore[arg-type] + await asyncio.wait_for(buf.add_event(_event()), FAKE_TIMEOUT) + await asyncio.wait_for(buf.add_metric(_metric()), FAKE_TIMEOUT) + await asyncio.wait_for(buf.add_event(_event()), FAKE_TIMEOUT) + await _settle(buf) + assert redis.calls == 3 + assert len(buf.event_buffer) == 5 + assert len(buf.metric_buffer) == 4 + assert buf._event_redis_keys == {} + + +class GatedRedis: + """set_key blocks on an asyncio.Event, so ordering is deterministic.""" + + def __init__(self) -> None: + self.release = asyncio.Event() + self.entered = asyncio.Event() + self.store: dict[str, str] = {} + + async def set_key(self, ns: str, key: str, value: str, ttl: Any = None) -> bool: + self.entered.set() + await self.release.wait() + self.store[key] = value + return True + + async def delete(self, ns: str, key: str) -> int: + return int(self.store.pop(key, None) is not None) + + async def get_key(self, *args: Any, **kwargs: Any) -> None: + return None + + async def keys(self, *args: Any, **kwargs: Any) -> list[str]: + return [] + + async def initialize(self) -> None: + return None + + +async def _settle(buf: EventBuffer) -> None: + await asyncio.wait_for(asyncio.gather(*buf._inflight_persist_tasks), FAKE_TIMEOUT) + + +@pytest.mark.asyncio +async def test_add_returns_before_slow_redis_completes() -> None: + buf = _buffer(buffer_size=100) + redis = GatedRedis() + await buf.initialize_redis(redis) # type: ignore[arg-type] + await asyncio.wait_for( + asyncio.gather( + *[buf.add_event(_event()) for _ in range(10)], + *[buf.add_metric(_metric()) for _ in range(10)], + ), + FAKE_TIMEOUT, + ) + assert len(buf.event_buffer) == 10 and len(buf.metric_buffer) == 10 + assert redis.store == {} + assert buf._event_redis_keys == {} + redis.release.set() + await _settle(buf) + assert len(redis.store) == 20 + assert len(buf._event_redis_keys) == 10 + assert len(buf._metric_redis_keys) == 10 + + +@pytest.mark.asyncio +async def test_stop_drains_pending_persist() -> None: + buf = _buffer() + redis = GatedRedis() + await buf.initialize_redis(redis) # type: ignore[arg-type] + await buf.start_auto_flush() + await buf.add_event(_event()) + await buf.add_metric(_metric()) + await asyncio.wait_for(redis.entered.wait(), FAKE_TIMEOUT) + stopper = asyncio.create_task(buf.stop_auto_flush()) + await asyncio.sleep(0) + assert not stopper.done() + redis.release.set() + await asyncio.wait_for(stopper, FAKE_TIMEOUT) + assert len(redis.store) == 2 + assert not buf._inflight_persist_tasks + + +@pytest.mark.asyncio +async def test_flush_during_persist_deletes_key_instead_of_orphaning() -> None: + buf = _buffer() + redis = GatedRedis() + await buf.initialize_redis(redis) # type: ignore[arg-type] + event, metric = _event(), _metric() + await buf.add_event(event) + await buf.add_metric(metric) + await asyncio.wait_for(redis.entered.wait(), FAKE_TIMEOUT) + events, ekeys = await buf.flush_events_with_keys() + metrics, mkeys = await buf.flush_metrics_with_keys() + assert events == [event] and ekeys == [""] + assert metrics == [metric] and mkeys == [""] + redis.release.set() + await _settle(buf) + assert redis.store == {} + assert buf._event_redis_keys == {} + assert buf._metric_redis_keys == {} + + +@pytest.mark.asyncio +async def test_requeue_during_persist_tracks_key() -> None: + buf = _buffer() + redis = GatedRedis() + await buf.initialize_redis(redis) # type: ignore[arg-type] + await buf.add_event(_event()) + await asyncio.wait_for(redis.entered.wait(), FAKE_TIMEOUT) + events, keys = await buf.flush_events_with_keys() + await buf.requeue_events_in_memory(events, keys) + redis.release.set() + await _settle(buf) + assert len(redis.store) == 1 + assert list(buf._event_redis_keys.values()) == list(redis.store) + + +@pytest.mark.asyncio +async def test_persist_failure_is_not_tracked() -> None: + buf = _buffer() + + class Failing(GatedRedis): + async def set_key(self, *a: Any, **k: Any) -> bool: + raise ConnectionError("down") + + await buf.initialize_redis(Failing()) # type: ignore[arg-type] + await buf.add_event(_event()) + await buf.add_metric(_metric()) + await _settle(buf) + assert buf._event_redis_keys == {} and buf._metric_redis_keys == {} + assert buf.redis_persist_failures == 2 + + +def test_encryptor_accepts_unpadded_key() -> None: + raw = bytes(range(31, 63)) + b"" + padded = base64.urlsafe_b64encode(raw).decode() + assert padded.endswith("=") and len(padded) == 44 + a = PayloadEncryptor(padded) + b = PayloadEncryptor(padded.rstrip("=")) + assert b.verify_key() + assert a.decrypt(b.encrypt({"x": 1})) == {"x": 1} + + +def test_encryptor_still_rejects_wrong_length_key() -> None: + short = base64.urlsafe_b64encode(b"a" * 16).decode().rstrip("=") + with pytest.raises(Exception, match="Invalid key size"): + PayloadEncryptor(short) + + +def test_startup_failure_message_includes_cause() -> None: + from guard_agent.encryption import EncryptionConfigError + + cfg = AgentConfig( + api_key="x", + endpoint="https://example.com", + project_encryption_key="not-a-valid-key", + ) + with pytest.raises(EncryptionConfigError, match=r"refusing plaintext.*cause: .+"): + HTTPTransport(cfg) + + +def test_encryptor_rejects_non_base64_key() -> None: + with pytest.raises(Exception, match="Invalid project key format"): + PayloadEncryptor("A") + + +@pytest.mark.asyncio +async def test_no_persist_task_without_redis_handler() -> None: + buf = _buffer() + await buf.add_event(_event()) + await buf.add_metric(_metric()) + assert not buf._inflight_persist_tasks + assert len(buf.event_buffer) == 1 and len(buf.metric_buffer) == 1 From d32170a0915dd1a1cf1a119139f97abdc3e4f113 Mon Sep 17 00:00:00 2001 From: Renn F Date: Thu, 15 Oct 2026 02:20:00 +0200 Subject: [PATCH 2/4] fix(buffer): bound in-flight Redis persists and skip items already evicted --- guard_agent/_buffer_queue.py | 19 ++++++++++++- guard_agent/_buffer_redis.py | 7 ++++- guard_agent/buffer.py | 1 + tests/test_buffer_redis_reentrancy.py | 40 +++++++++++++++++++++++++++ 4 files changed, 65 insertions(+), 2 deletions(-) diff --git a/guard_agent/_buffer_queue.py b/guard_agent/_buffer_queue.py index adc3ed3..4851ef7 100644 --- a/guard_agent/_buffer_queue.py +++ b/guard_agent/_buffer_queue.py @@ -14,6 +14,8 @@ class BufferQueueMixin(BufferOverflowMixin, BufferLifecycleMixin): events_flushed: int metrics_flushed: int + redis_persist_failures: int + _persist_overflow_logged: bool def _spawn_persist( self, persist: Callable[[T], Coroutine[Any, Any, None]], item: T @@ -24,9 +26,24 @@ def _spawn_persist( this context (a Redis failure reporting an event must not recurse).""" if not self.redis_handler or _PERSISTING.get(): return + if len(self._inflight_persist_tasks) >= self.config.buffer_size: + self.redis_persist_failures += 1 + if not self._persist_overflow_logged: + self._persist_overflow_logged = True + self.logger.warning( + "Redis persist backlog at buffer_size=" + f"{self.config.buffer_size}; buffering in memory only " + "until it drains" + ) + return task = asyncio.create_task(persist(item)) self._inflight_persist_tasks.add(task) - task.add_done_callback(self._inflight_persist_tasks.discard) + task.add_done_callback(self._persist_done) + + def _persist_done(self, task: "asyncio.Task[None]") -> None: + self._inflight_persist_tasks.discard(task) + if not self._inflight_persist_tasks: + self._persist_overflow_logged = False async def add_event(self, event: SecurityEvent) -> None: """Add security event to buffer honoring the configured overflow diff --git a/guard_agent/_buffer_redis.py b/guard_agent/_buffer_redis.py index 3e2309b..35ec64e 100644 --- a/guard_agent/_buffer_redis.py +++ b/guard_agent/_buffer_redis.py @@ -94,8 +94,11 @@ async def _persist_and_track_event(self, event: SecurityEvent) -> None: Skipped while a persist is in flight in this context, so a Redis failure that emits an event cannot recurse into Redis. If the event left the buffer meanwhile (flushed), the new key is deleted instead - of tracked, so it neither orphans nor replays after a restart. + of tracked, so it neither orphans nor replays after a restart. An + event already gone before the task starts is never written. """ + if not any(e is event for e in reversed(self.event_buffer)): + return token = _PERSISTING.set(True) try: key = await self._persist_event_to_redis(event) @@ -113,6 +116,8 @@ async def _persist_and_track_event(self, event: SecurityEvent) -> None: async def _persist_and_track_metric(self, metric: SecurityMetric) -> None: """Metric twin of _persist_and_track_event.""" + if not any(m is metric for m in reversed(self.metric_buffer)): + return token = _PERSISTING.set(True) try: key = await self._persist_metric_to_redis(metric) diff --git a/guard_agent/buffer.py b/guard_agent/buffer.py index 3570219..b478595 100644 --- a/guard_agent/buffer.py +++ b/guard_agent/buffer.py @@ -31,6 +31,7 @@ def __init__( self._running = False self._inflight_flush_tasks: set[asyncio.Task[None]] = set() self._inflight_persist_tasks: set[asyncio.Task[None]] = set() + self._persist_overflow_logged = False self._event_redis_keys: dict[int, str] = {} self._metric_redis_keys: dict[int, str] = {} diff --git a/tests/test_buffer_redis_reentrancy.py b/tests/test_buffer_redis_reentrancy.py index ebb9fe6..8c959f4 100644 --- a/tests/test_buffer_redis_reentrancy.py +++ b/tests/test_buffer_redis_reentrancy.py @@ -239,3 +239,43 @@ async def test_no_persist_task_without_redis_handler() -> None: await buf.add_metric(_metric()) assert not buf._inflight_persist_tasks assert len(buf.event_buffer) == 1 and len(buf.metric_buffer) == 1 + + +@pytest.mark.asyncio +async def test_inflight_persists_capped_and_overflow_counted( + caplog: pytest.LogCaptureFixture, +) -> None: + buf = _buffer(buffer_size=3) + redis = GatedRedis() + await buf.initialize_redis(redis) # type: ignore[arg-type] + for _ in range(4): + await buf.add_event(_event()) + for _ in range(3): + await buf.add_metric(_metric()) + assert len(buf._inflight_persist_tasks) == 3 + assert buf.redis_persist_failures == 4 + assert sum("persist backlog" in r.message for r in caplog.records) == 1 + redis.release.set() + await _settle(buf) + assert not buf._inflight_persist_tasks + redis.release.clear() + for _ in range(4): + await buf.add_event(_event()) + assert sum("persist backlog" in r.message for r in caplog.records) == 2 + redis.release.set() + await _settle(buf) + + +@pytest.mark.asyncio +async def test_item_gone_before_task_starts_is_never_written() -> None: + buf = _buffer() + redis = GatedRedis() + redis.release.set() + await buf.initialize_redis(redis) # type: ignore[arg-type] + await buf.add_event(_event()) + await buf.add_metric(_metric()) + buf.event_buffer.clear() + buf.metric_buffer.clear() + await _settle(buf) + assert not redis.entered.is_set() + assert redis.store == {} From 60c99698c26ca63cbd1a0177217497e661630654 Mon Sep 17 00:00:00 2001 From: Renn F Date: Thu, 15 Oct 2026 02:36:22 +0200 Subject: [PATCH 3/4] fix(buffer): replace persist membership scans with O(1) pending-id tracking --- guard_agent/_buffer_overflow.py | 4 +++ guard_agent/_buffer_queue.py | 28 ++++++++++++++++--- guard_agent/_buffer_redis.py | 12 ++++++--- guard_agent/buffer.py | 2 ++ tests/test_buffer_redis_reentrancy.py | 39 +++++++++++++++++++++++++-- 5 files changed, 76 insertions(+), 9 deletions(-) diff --git a/guard_agent/_buffer_overflow.py b/guard_agent/_buffer_overflow.py index 93159fe..3480c29 100644 --- a/guard_agent/_buffer_overflow.py +++ b/guard_agent/_buffer_overflow.py @@ -85,24 +85,28 @@ def _forget_oldest_event_key(self) -> str | None: if not self.event_buffer: return None oldest = self.event_buffer[0] + self._pending_event_persists.pop(id(oldest), None) return self._event_redis_keys.pop(id(oldest), None) def _forget_oldest_metric_key(self) -> str | None: if not self.metric_buffer: return None oldest = self.metric_buffer[0] + self._pending_metric_persists.pop(id(oldest), None) return self._metric_redis_keys.pop(id(oldest), None) def _forget_newest_event_key(self) -> str | None: if not self.event_buffer: return None newest = self.event_buffer[-1] + self._pending_event_persists.pop(id(newest), None) return self._event_redis_keys.pop(id(newest), None) def _forget_newest_metric_key(self) -> str | None: if not self.metric_buffer: return None newest = self.metric_buffer[-1] + self._pending_metric_persists.pop(id(newest), None) return self._metric_redis_keys.pop(id(newest), None) def _is_event_buffer_full(self) -> bool: diff --git a/guard_agent/_buffer_queue.py b/guard_agent/_buffer_queue.py index 4851ef7..99698c9 100644 --- a/guard_agent/_buffer_queue.py +++ b/guard_agent/_buffer_queue.py @@ -16,9 +16,14 @@ class BufferQueueMixin(BufferOverflowMixin, BufferLifecycleMixin): metrics_flushed: int redis_persist_failures: int _persist_overflow_logged: bool + _pending_event_persists: dict[int, bool] + _pending_metric_persists: dict[int, bool] def _spawn_persist( - self, persist: Callable[[T], Coroutine[Any, Any, None]], item: T + self, + persist: Callable[[T], Coroutine[Any, Any, None]], + item: T, + pending: dict[int, bool], ) -> None: """Run a Redis persist as a tracked background task so the request path never waits on Redis; stop_auto_flush drains it. Nothing is @@ -36,6 +41,7 @@ def _spawn_persist( "until it drains" ) return + pending[id(item)] = True task = asyncio.create_task(persist(item)) self._inflight_persist_tasks.add(task) task.add_done_callback(self._persist_done) @@ -68,7 +74,9 @@ async def add_event(self, event: SecurityEvent) -> None: except asyncio.TimeoutError: pass - self._spawn_persist(self._persist_and_track_event, event) + self._spawn_persist( + self._persist_and_track_event, event, self._pending_event_persists + ) if ( len(self.event_buffer) @@ -98,7 +106,9 @@ async def add_metric(self, metric: SecurityMetric) -> None: except asyncio.TimeoutError: pass - self._spawn_persist(self._persist_and_track_metric, metric) + self._spawn_persist( + self._persist_and_track_metric, metric, self._pending_metric_persists + ) if ( len(self.metric_buffer) @@ -129,6 +139,9 @@ async def flush_events_with_keys( condition = self._get_event_condition() async with condition: events = list(self.event_buffer) + for item in events: + if id(item) in self._pending_event_persists: + self._pending_event_persists[id(item)] = False keys = [self._event_redis_keys.pop(id(event), "") for event in events] self.event_buffer.clear() self.events_flushed += len(events) @@ -147,6 +160,9 @@ async def flush_metrics_with_keys( condition = self._get_metric_condition() async with condition: metrics = list(self.metric_buffer) + for item in metrics: + if id(item) in self._pending_metric_persists: + self._pending_metric_persists[id(item)] = False keys = [self._metric_redis_keys.pop(id(metric), "") for metric in metrics] self.metric_buffer.clear() self.metrics_flushed += len(metrics) @@ -171,6 +187,8 @@ async def requeue_events_in_memory( if evicted_key: evicted_keys.append(evicted_key) self.event_buffer.appendleft(event) + if id(event) in self._pending_event_persists: + self._pending_event_persists[id(event)] = True if key: self._event_redis_keys[id(event)] = key return evicted_keys @@ -191,6 +209,8 @@ async def requeue_metrics_in_memory( if evicted_key: evicted_keys.append(evicted_key) self.metric_buffer.appendleft(metric) + if id(metric) in self._pending_metric_persists: + self._pending_metric_persists[id(metric)] = True if key: self._metric_redis_keys[id(metric)] = key return evicted_keys @@ -202,12 +222,14 @@ async def clear_buffer(self) -> None: async with event_condition: self.event_buffer.clear() self._event_redis_keys.clear() + self._pending_event_persists.clear() event_condition.notify_all() metric_condition = self._get_metric_condition() async with metric_condition: self.metric_buffer.clear() self._metric_redis_keys.clear() + self._pending_metric_persists.clear() metric_condition.notify_all() if self.redis_handler: diff --git a/guard_agent/_buffer_redis.py b/guard_agent/_buffer_redis.py index 35ec64e..9396977 100644 --- a/guard_agent/_buffer_redis.py +++ b/guard_agent/_buffer_redis.py @@ -24,6 +24,8 @@ class BufferRedisMixin: redis_persist_failures: int _event_redis_keys: dict[int, str] _metric_redis_keys: dict[int, str] + _pending_event_persists: dict[int, bool] + _pending_metric_persists: dict[int, bool] def _get_event_condition(self) -> asyncio.Condition: raise NotImplementedError @@ -97,7 +99,8 @@ async def _persist_and_track_event(self, event: SecurityEvent) -> None: of tracked, so it neither orphans nor replays after a restart. An event already gone before the task starts is never written. """ - if not any(e is event for e in reversed(self.event_buffer)): + if not self._pending_event_persists.get(id(event)): + self._pending_event_persists.pop(id(event), None) return token = _PERSISTING.set(True) try: @@ -109,14 +112,15 @@ async def _persist_and_track_event(self, event: SecurityEvent) -> None: _PERSISTING.reset(token) if key is None: return - if any(e is event for e in reversed(self.event_buffer)): + if self._pending_event_persists.pop(id(event), False): self._event_redis_keys[id(event)] = key else: await self.confirm_event_redis_keys([key]) async def _persist_and_track_metric(self, metric: SecurityMetric) -> None: """Metric twin of _persist_and_track_event.""" - if not any(m is metric for m in reversed(self.metric_buffer)): + if not self._pending_metric_persists.get(id(metric)): + self._pending_metric_persists.pop(id(metric), None) return token = _PERSISTING.set(True) try: @@ -128,7 +132,7 @@ async def _persist_and_track_metric(self, metric: SecurityMetric) -> None: _PERSISTING.reset(token) if key is None: return - if any(m is metric for m in reversed(self.metric_buffer)): + if self._pending_metric_persists.pop(id(metric), False): self._metric_redis_keys[id(metric)] = key else: await self.confirm_metric_redis_keys([key]) diff --git a/guard_agent/buffer.py b/guard_agent/buffer.py index b478595..bd58ccd 100644 --- a/guard_agent/buffer.py +++ b/guard_agent/buffer.py @@ -32,6 +32,8 @@ def __init__( self._inflight_flush_tasks: set[asyncio.Task[None]] = set() self._inflight_persist_tasks: set[asyncio.Task[None]] = set() self._persist_overflow_logged = False + self._pending_event_persists: dict[int, bool] = {} + self._pending_metric_persists: dict[int, bool] = {} self._event_redis_keys: dict[int, str] = {} self._metric_redis_keys: dict[int, str] = {} diff --git a/tests/test_buffer_redis_reentrancy.py b/tests/test_buffer_redis_reentrancy.py index 8c959f4..0ad4e89 100644 --- a/tests/test_buffer_redis_reentrancy.py +++ b/tests/test_buffer_redis_reentrancy.py @@ -274,8 +274,43 @@ async def test_item_gone_before_task_starts_is_never_written() -> None: await buf.initialize_redis(redis) # type: ignore[arg-type] await buf.add_event(_event()) await buf.add_metric(_metric()) - buf.event_buffer.clear() - buf.metric_buffer.clear() + await buf.flush_events_with_keys() + await buf.flush_metrics_with_keys() await _settle(buf) assert not redis.entered.is_set() assert redis.store == {} + + +@pytest.mark.asyncio +async def test_overflow_eviction_and_clear_drop_pending_persists() -> None: + buf = _buffer(buffer_size=2) + redis = GatedRedis() + redis.release.set() + await buf.initialize_redis(redis) # type: ignore[arg-type] + first = _event() + await buf.add_event(first) + await buf.add_event(_event()) + await buf.add_event(_event()) + assert id(first) not in buf._pending_event_persists + await buf.add_metric(_metric()) + await buf.clear_buffer() + assert buf._pending_event_persists == {} + assert buf._pending_metric_persists == {} + await _settle(buf) + assert redis.store == {} + assert buf._event_redis_keys == {} and buf._metric_redis_keys == {} + assert buf._pending_event_persists == {} + + +@pytest.mark.asyncio +async def test_metric_requeue_during_persist_tracks_key() -> None: + buf = _buffer() + redis = GatedRedis() + await buf.initialize_redis(redis) # type: ignore[arg-type] + await buf.add_metric(_metric()) + await asyncio.wait_for(redis.entered.wait(), FAKE_TIMEOUT) + metrics, keys = await buf.flush_metrics_with_keys() + await buf.requeue_metrics_in_memory(metrics, keys) + redis.release.set() + await _settle(buf) + assert list(buf._metric_redis_keys.values()) == list(redis.store) From f7f6517af0d25d51b95c7ec5457971b182757e16 Mon Sep 17 00:00:00 2001 From: Renn F Date: Thu, 15 Oct 2026 02:38:29 +0200 Subject: [PATCH 4/4] fix(buffer): always release the pending persist id, including on failed writes --- guard_agent/_buffer_redis.py | 10 ++++++---- tests/test_buffer_redis_reentrancy.py | 18 ++++++++++++++++++ 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/guard_agent/_buffer_redis.py b/guard_agent/_buffer_redis.py index 9396977..abeb2e7 100644 --- a/guard_agent/_buffer_redis.py +++ b/guard_agent/_buffer_redis.py @@ -107,12 +107,13 @@ async def _persist_and_track_event(self, event: SecurityEvent) -> None: key = await self._persist_event_to_redis(event) except Exception as e: self.logger.error(f"Failed to buffer event: {str(e)}") - return + key = None finally: _PERSISTING.reset(token) + still_buffered = self._pending_event_persists.pop(id(event), False) if key is None: return - if self._pending_event_persists.pop(id(event), False): + if still_buffered: self._event_redis_keys[id(event)] = key else: await self.confirm_event_redis_keys([key]) @@ -127,12 +128,13 @@ async def _persist_and_track_metric(self, metric: SecurityMetric) -> None: key = await self._persist_metric_to_redis(metric) except Exception as e: self.logger.error(f"Failed to buffer metric: {str(e)}") - return + key = None finally: _PERSISTING.reset(token) + still_buffered = self._pending_metric_persists.pop(id(metric), False) if key is None: return - if self._pending_metric_persists.pop(id(metric), False): + if still_buffered: self._metric_redis_keys[id(metric)] = key else: await self.confirm_metric_redis_keys([key]) diff --git a/tests/test_buffer_redis_reentrancy.py b/tests/test_buffer_redis_reentrancy.py index 0ad4e89..2bba4bb 100644 --- a/tests/test_buffer_redis_reentrancy.py +++ b/tests/test_buffer_redis_reentrancy.py @@ -314,3 +314,21 @@ async def test_metric_requeue_during_persist_tracks_key() -> None: redis.release.set() await _settle(buf) assert list(buf._metric_redis_keys.values()) == list(redis.store) + + +@pytest.mark.asyncio +async def test_failed_persists_leave_no_pending_ids() -> None: + buf = _buffer() + + class Failing(GatedRedis): + async def set_key(self, *a: Any, **k: Any) -> bool: + raise ConnectionError("down") + + await buf.initialize_redis(Failing()) # type: ignore[arg-type] + for _ in range(3): + await buf.add_event(_event()) + await buf.add_metric(_metric()) + await _settle(buf) + assert buf.redis_persist_failures == 6 + assert buf._pending_event_persists == {} + assert buf._pending_metric_persists == {}