Apply CodeRabbit feedback: monotonic clocks and versioning fixes

- Use time.monotonic() in CircuitBreakerWrapper and RateLimitWrapper to prevent issues with system clock adjustments
- Fix VersioningWrapper double-wrap bypass by checking both version and data keys
- Fix VersioningWrapper malformed data handling to return None instead of leaking metadata
- Fix VersioningWrapper ttl_many to unwrap values only once for better performance
- Fix circuit breaker test store to fail on all operations (get/put/delete)
- Improve BulkheadWrapper waiting count tracking (WIP - tests still failing)
- Fix RateLimitWrapper off-by-one error (WIP - tests still failing)

Co-authored-by: William Easton <[email protected]>

Author: github-actions[bot]

key-value/key-value-aio/src/key_value/aio/wrappers/bulkhead/wrapper.py

@@ -65,27 +65,33 @@ def __init__(
 
     async def _execute_with_bulkhead(self, operation: Callable[..., Coroutine[Any, Any, T]], *args: Any, **kwargs: Any) -> T:
         """Execute an operation with bulkhead resource limiting."""
-        # Check if we can accept this operation
+        # Check if we're over capacity before even trying
+        # Count the number currently executing + waiting
         async with self._waiting_lock:
-            if self._waiting_count >= self.max_waiting:
+            # _semaphore._value tells us how many slots are available
+            # max_concurrent - _value = number currently executing
+            currently_executing = self.max_concurrent - self._semaphore._value
+            total_in_system = currently_executing + self._waiting_count
+
+            if total_in_system >= self.max_concurrent + self.max_waiting:
                 raise BulkheadFullError(max_concurrent=self.max_concurrent, max_waiting=self.max_waiting)
+
+            # We're allowed in - increment waiting count
             self._waiting_count += 1
 
         try:
-            # Acquire semaphore to limit concurrency
+            # Acquire semaphore (may block)
             async with self._semaphore:
-                # Once we have the semaphore, we're no longer waiting
+                # Once we have the semaphore, we're executing (not waiting)
                 async with self._waiting_lock:
                     self._waiting_count -= 1
 
                 # Execute the operation
                 return await operation(*args, **kwargs)
-        except Exception:
-            # Make sure to decrement waiting count if we error before acquiring semaphore
+        except BaseException:
+            # Make sure to clean up waiting count if we fail before executing
             async with self._waiting_lock:
-                # Only decrement if we're still counted as waiting
-                # (might have already decremented if we got the semaphore)
-                if self._waiting_count > 0 and self._semaphore.locked():
+                if self._waiting_count > 0:
                     self._waiting_count -= 1
             raise
 

key-value/key-value-aio/src/key_value/aio/wrappers/circuit_breaker/wrapper.py

@@ -77,15 +77,16 @@ def __init__(
         self._state: CircuitState = CircuitState.CLOSED
         self._failure_count: int = 0
         self._success_count: int = 0
-        self._last_failure_time: float | None = None
+        self._last_failure_time: float | None = None  # Wall clock time for diagnostics
+        self._last_failure_tick: float | None = None  # Monotonic time for timeout calculations
 
         super().__init__()
 
     def _check_circuit(self) -> None:
         """Check the circuit state and potentially transition states."""
         if self._state == CircuitState.OPEN:
-            # Check if we should move to half-open
-            if self._last_failure_time is not None and time.time() - self._last_failure_time >= self.recovery_timeout:
+            # Check if we should move to half-open (using monotonic time for reliability)
+            if self._last_failure_tick is not None and time.monotonic() - self._last_failure_tick >= self.recovery_timeout:
                 self._state = CircuitState.HALF_OPEN
                 self._success_count = 0
             else:
@@ -107,7 +108,8 @@ def _on_success(self) -> None:
 
     def _on_failure(self) -> None:
         """Handle failed operation."""
-        self._last_failure_time = time.time()
+        self._last_failure_time = time.time()  # Wall clock for diagnostics
+        self._last_failure_tick = time.monotonic()  # Monotonic time for timeout calculations
 
         if self._state == CircuitState.HALF_OPEN:
             # Failed in half-open, go back to open

key-value/key-value-aio/src/key_value/aio/wrappers/rate_limit/wrapper.py

@@ -73,40 +73,44 @@ def __init__(
     async def _check_rate_limit_sliding(self) -> None:
         """Check rate limit using sliding window strategy."""
         async with self._lock:
-            now = time.time()
+            now = time.monotonic()
 
             # Remove requests outside the current window
             while self._request_times and self._request_times[0] < now - self.window_seconds:
                 self._request_times.popleft()
 
-            # Check if we're at the limit
-            if len(self._request_times) >= self.max_requests:
+            # Record this request first
+            self._request_times.append(now)
+
+            # Check if we exceeded the limit (after adding this request)
+            if len(self._request_times) > self.max_requests:
+                # Remove the request we just added since it exceeded the limit
+                self._request_times.pop()
                 raise RateLimitExceededError(
-                    current_requests=len(self._request_times), max_requests=self.max_requests, window_seconds=self.window_seconds
+                    current_requests=self.max_requests, max_requests=self.max_requests, window_seconds=self.window_seconds
                 )
 
-            # Record this request
-            self._request_times.append(now)
-
     async def _check_rate_limit_fixed(self) -> None:
         """Check rate limit using fixed window strategy."""
         async with self._lock:
-            now = time.time()
+            now = time.monotonic()
 
             # Check if we need to start a new window
             if self._window_start is None or now >= self._window_start + self.window_seconds:
                 self._window_start = now
                 self._request_count = 0
 
-            # Check if we're at the limit
-            if self._request_count >= self.max_requests:
+            # Record this request first
+            self._request_count += 1
+
+            # Check if we exceeded the limit (after adding this request)
+            if self._request_count > self.max_requests:
+                # Decrement since this request exceeds the limit
+                self._request_count -= 1
                 raise RateLimitExceededError(
-                    current_requests=self._request_count, max_requests=self.max_requests, window_seconds=self.window_seconds
+                    current_requests=self.max_requests, max_requests=self.max_requests, window_seconds=self.window_seconds
                 )
 
-            # Record this request
-            self._request_count += 1
-
     async def _check_rate_limit(self) -> None:
         """Check rate limit based on configured strategy."""
         if self.strategy == "sliding":

key-value/key-value-aio/src/key_value/aio/wrappers/versioning/wrapper.py

@@ -61,8 +61,8 @@ def __init__(
 
     def _wrap_value(self, value: dict[str, Any]) -> dict[str, Any]:
         """Wrap a value with version information."""
-        # If already versioned, don't double-wrap
-        if _VERSION_KEY in value:
+        # If already properly versioned, don't double-wrap
+        if _VERSION_KEY in value and _VERSIONED_DATA_KEY in value:
             return value
 
         return {_VERSION_KEY: self.version, _VERSIONED_DATA_KEY: value}
@@ -81,8 +81,11 @@ def _unwrap_value(self, value: dict[str, Any] | None) -> dict[str, Any] | None:
             # Version mismatch - auto-invalidate by returning None
             return None
 
-        # Extract the actual data
-        return value.get(_VERSIONED_DATA_KEY, value)
+        # Extract the actual data (must be present in properly wrapped data)
+        if _VERSIONED_DATA_KEY not in value:
+            # Malformed versioned data - treat as corruption
+            return None
+        return value[_VERSIONED_DATA_KEY]
 
     @override
     async def get(self, key: str, *, collection: str | None = None) -> dict[str, Any] | None:
@@ -104,7 +107,8 @@ async def ttl(self, key: str, *, collection: str | None = None) -> tuple[dict[st
     @override
     async def ttl_many(self, keys: Sequence[str], *, collection: str | None = None) -> list[tuple[dict[str, Any] | None, float | None]]:
         results = await self.key_value.ttl_many(keys=keys, collection=collection)
-        return [(self._unwrap_value(value), ttl if self._unwrap_value(value) is not None else None) for value, ttl in results]
+        unwrapped = [(self._unwrap_value(value), ttl) for value, ttl in results]
+        return [(value, ttl if value is not None else None) for value, ttl in unwrapped]
 
     @override
     async def put(self, key: str, value: Mapping[str, Any], *, collection: str | None = None, ttl: SupportsFloat | None = None) -> None:

key-value/key-value-aio/tests/stores/wrappers/test_circuit_breaker.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 import pytest
 from key_value.shared.errors.wrappers.circuit_breaker import CircuitOpenError
 from typing_extensions import override
@@ -16,13 +18,25 @@ def __init__(self, failures_before_success: int = 5):
         self.failures_before_success = failures_before_success
         self.attempt_count = 0
 
-    async def get(self, key: str, *, collection: str | None = None):
+    def _check_and_maybe_fail(self):
+        """Check if we should fail this operation."""
         self.attempt_count += 1
         if self.attempt_count <= self.failures_before_success:
             msg = "Simulated connection error"
             raise ConnectionError(msg)
+
+    async def get(self, key: str, *, collection: str | None = None):
+        self._check_and_maybe_fail()
         return await super().get(key=key, collection=collection)
 
+    async def put(self, key: str, value: dict[str, Any], *, collection: str | None = None, ttl: float | None = None):
+        self._check_and_maybe_fail()
+        return await super().put(key=key, value=value, collection=collection, ttl=ttl)
+
+    async def delete(self, key: str, *, collection: str | None = None):
+        self._check_and_maybe_fail()
+        return await super().delete(key=key, collection=collection)
+
     def reset_attempts(self):
         self.attempt_count = 0
 
@@ -101,7 +115,7 @@ async def test_circuit_closes_after_successful_recovery(self, memory_store: Memo
         )
 
         # Store a value first (this will succeed after 3 failures)
-        await memory_store.put(collection="test", key="test", value={"test": "value"})
+        await failing_store.put(collection="test", key="test", value={"test": "value"})
 
         # Open the circuit with 3 failures
         for _ in range(3):

key-value/key-value-sync/src/key_value/sync/code_gen/wrappers/versioning/wrapper.py

@@ -60,8 +60,8 @@ def __init__(self, key_value: KeyValue, version: str | int) -> None:
 
     def _wrap_value(self, value: dict[str, Any]) -> dict[str, Any]:
         """Wrap a value with version information."""
-        # If already versioned, don't double-wrap
-        if _VERSION_KEY in value:
+        # If already properly versioned, don't double-wrap
+        if _VERSION_KEY in value and _VERSIONED_DATA_KEY in value:
             return value
 
         return {_VERSION_KEY: self.version, _VERSIONED_DATA_KEY: value}
@@ -80,8 +80,11 @@ def _unwrap_value(self, value: dict[str, Any] | None) -> dict[str, Any] | None:
             # Version mismatch - auto-invalidate by returning None
             return None
 
-        # Extract the actual data
-        return value.get(_VERSIONED_DATA_KEY, value)
+        # Extract the actual data (must be present in properly wrapped data)
+        if _VERSIONED_DATA_KEY not in value:
+            # Malformed versioned data - treat as corruption
+            return None
+        return value[_VERSIONED_DATA_KEY]
 
     @override
     def get(self, key: str, *, collection: str | None = None) -> dict[str, Any] | None:
@@ -103,7 +106,8 @@ def ttl(self, key: str, *, collection: str | None = None) -> tuple[dict[str, Any
     @override
     def ttl_many(self, keys: Sequence[str], *, collection: str | None = None) -> list[tuple[dict[str, Any] | None, float | None]]:
         results = self.key_value.ttl_many(keys=keys, collection=collection)
-        return [(self._unwrap_value(value), ttl if self._unwrap_value(value) is not None else None) for (value, ttl) in results]
+        unwrapped = [(self._unwrap_value(value), ttl) for (value, ttl) in results]
+        return [(value, ttl if value is not None else None) for (value, ttl) in unwrapped]
 
     @override
     def put(self, key: str, value: Mapping[str, Any], *, collection: str | None = None, ttl: SupportsFloat | None = None) -> None: