289 lines
10 KiB
Python
289 lines
10 KiB
Python
import threading
|
|
import time
|
|
from typing import List
|
|
|
|
|
|
class _ThreadBuckets:
|
|
"""Per-thread bucket storage for rolling window.
|
|
|
|
Each thread gets its own instance to avoid lock contention on the hot path.
|
|
"""
|
|
|
|
# This is a performance optimization to avoid creating a dictionary for the instance.
|
|
__slots__ = ("buckets", "current_bucket_idx", "last_rotation_time")
|
|
|
|
def __init__(self, num_buckets: int):
|
|
self.buckets = [0.0] * num_buckets
|
|
self.current_bucket_idx = 0
|
|
self.last_rotation_time = time.time()
|
|
|
|
|
|
class _ThreadLocalRef(threading.local):
|
|
"""Thread-local reference to the thread's _ThreadBuckets instance."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
# by using threading.local, each thread gets its own instance of _ThreadBuckets.
|
|
self.data: _ThreadBuckets = None
|
|
|
|
|
|
class _RollingWindowBase:
|
|
"""Base class for rolling window trackers.
|
|
|
|
Provides the shared infrastructure: bucketing, rotation, thread-local
|
|
storage, and thread registration. Subclasses define how values are
|
|
recorded into buckets and how buckets are aggregated.
|
|
|
|
Uses bucketing for memory efficiency - divides the window into N buckets
|
|
and rotates them as time passes. This allows efficient tracking of values
|
|
over a sliding window without storing individual data points.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
window_duration_s: float,
|
|
num_buckets: int = 60,
|
|
):
|
|
if window_duration_s <= 0:
|
|
raise ValueError(
|
|
f"window_duration_s must be positive, got {window_duration_s}"
|
|
)
|
|
if num_buckets <= 0:
|
|
raise ValueError(f"num_buckets must be positive, got {num_buckets}")
|
|
|
|
self._window_duration_s = window_duration_s
|
|
self._num_buckets = num_buckets
|
|
self._bucket_duration_s = window_duration_s / num_buckets
|
|
|
|
# Thread-local reference to per-thread bucket data
|
|
self._local = _ThreadLocalRef()
|
|
|
|
# Track all per-thread bucket instances for aggregation
|
|
self._all_thread_data: List[_ThreadBuckets] = []
|
|
self._registry_lock = threading.Lock()
|
|
|
|
@property
|
|
def window_duration_s(self) -> float:
|
|
"""The total duration of the rolling window in seconds."""
|
|
return self._window_duration_s
|
|
|
|
@property
|
|
def num_buckets(self) -> int:
|
|
"""The number of buckets in the rolling window."""
|
|
return self._num_buckets
|
|
|
|
@property
|
|
def bucket_duration_s(self) -> float:
|
|
"""The duration of each bucket in seconds."""
|
|
return self._bucket_duration_s
|
|
|
|
def _ensure_initialized(self) -> _ThreadBuckets:
|
|
"""Ensure thread-local storage is initialized for the current thread.
|
|
|
|
This is called on every add() but the fast path (already initialized)
|
|
is just a single attribute check with no locking.
|
|
|
|
Returns:
|
|
The _ThreadBuckets instance for the current thread.
|
|
"""
|
|
data = self._local.data
|
|
if data is not None:
|
|
return data
|
|
|
|
# Slow path: first call from this thread
|
|
data = _ThreadBuckets(self._num_buckets)
|
|
self._local.data = data
|
|
|
|
# Register for aggregation (only happens once per thread)
|
|
with self._registry_lock:
|
|
self._all_thread_data.append(data)
|
|
|
|
return data
|
|
|
|
def _rotate_buckets_if_needed(self, data: _ThreadBuckets) -> None:
|
|
"""Rotate buckets for the given thread's storage.
|
|
|
|
Advances the current bucket index and clears old buckets as time passes.
|
|
"""
|
|
now = time.time()
|
|
elapsed = now - data.last_rotation_time
|
|
buckets_to_advance = int(elapsed / self._bucket_duration_s)
|
|
|
|
if buckets_to_advance > 0:
|
|
if buckets_to_advance >= self._num_buckets:
|
|
# All buckets have expired, reset everything
|
|
data.buckets = [0.0] * self._num_buckets
|
|
data.current_bucket_idx = 0
|
|
else:
|
|
# Clear old buckets as we advance
|
|
for _ in range(buckets_to_advance):
|
|
data.current_bucket_idx = (
|
|
data.current_bucket_idx + 1
|
|
) % self._num_buckets
|
|
data.buckets[data.current_bucket_idx] = 0.0
|
|
|
|
data.last_rotation_time = now
|
|
|
|
def get_num_registered_threads(self) -> int:
|
|
"""Get the number of threads that have called add().
|
|
|
|
Useful for debugging and testing.
|
|
|
|
Returns:
|
|
The number of threads registered with this accumulator.
|
|
"""
|
|
with self._registry_lock:
|
|
return len(self._all_thread_data)
|
|
|
|
|
|
class RollingWindowAccumulator(_RollingWindowBase):
|
|
"""Tracks cumulative values over a rolling time window.
|
|
|
|
Uses thread-local storage for lock-free writes on the hot path (add()).
|
|
Only get_total() requires synchronization to aggregate across threads.
|
|
|
|
Example:
|
|
# Create a 10-minute rolling window with 60 buckets (10s each)
|
|
accumulator = RollingWindowAccumulator(
|
|
window_duration_s=600.0,
|
|
num_buckets=60,
|
|
)
|
|
|
|
# Add values (lock-free, safe from multiple threads)
|
|
accumulator.add(100.0)
|
|
accumulator.add(50.0)
|
|
|
|
# Get total (aggregates across all threads)
|
|
total = accumulator.get_total()
|
|
|
|
Thread Safety:
|
|
- add() is lock-free after the first call from each thread
|
|
- get_total() acquires a lock to aggregate across threads
|
|
- Safe to call from multiple threads concurrently
|
|
"""
|
|
|
|
def add(self, value: float) -> None:
|
|
"""Add a value to the current bucket.
|
|
|
|
This operation is lock-free for the calling thread after the first call.
|
|
Safe to call from multiple threads concurrently.
|
|
|
|
Args:
|
|
value: The value to add to the accumulator.
|
|
"""
|
|
# Fast path: just check if initialized (no lock)
|
|
data = self._ensure_initialized()
|
|
|
|
# Lock-free: only touches thread-local data
|
|
self._rotate_buckets_if_needed(data)
|
|
data.buckets[data.current_bucket_idx] += value
|
|
|
|
def get_total(self) -> float:
|
|
"""Get total value across all buckets in the window.
|
|
|
|
This aggregates values from all threads that have called add().
|
|
Expired buckets (older than window_duration_s) are not included.
|
|
|
|
Note: We are accepting some inaccuracy in the total value to avoid the overhead of a lock.
|
|
This is acceptable because we are only using this for utilization metrics, which are not
|
|
critical for the overall system. Given that the default window duration is 600s and the
|
|
default report interval is 10s, the inaccuracy is less than 0.16%.
|
|
|
|
Returns:
|
|
The sum of all non-expired values in the rolling window.
|
|
"""
|
|
total = 0.0
|
|
now = time.time()
|
|
|
|
with self._registry_lock:
|
|
for data in self._all_thread_data:
|
|
# Calculate which buckets are still valid for this thread's data
|
|
elapsed = now - data.last_rotation_time
|
|
buckets_expired = int(elapsed / self._bucket_duration_s)
|
|
|
|
if buckets_expired >= self._num_buckets:
|
|
# All buckets have expired for this thread
|
|
continue
|
|
|
|
# Sum buckets that haven't expired
|
|
# Buckets are arranged in a circular buffer, with current_bucket_idx
|
|
# being the most recent. We need to skip buckets that have expired.
|
|
for i in range(self._num_buckets - buckets_expired):
|
|
# Go backwards from current bucket
|
|
idx = (data.current_bucket_idx - i) % self._num_buckets
|
|
total += data.buckets[idx]
|
|
|
|
return total
|
|
|
|
|
|
class RollingWindowMax(_RollingWindowBase):
|
|
"""Tracks the maximum value over a rolling time window.
|
|
|
|
Uses the same bucketed rolling window approach as RollingWindowAccumulator,
|
|
but each bucket stores the maximum observed value instead of a cumulative
|
|
sum. Querying returns the max across all non-expired buckets.
|
|
|
|
Example:
|
|
# Create a 30-second rolling window with 6 buckets (5s each)
|
|
tracker = RollingWindowMax(
|
|
window_duration_s=30.0,
|
|
num_buckets=6,
|
|
)
|
|
|
|
# Record values (lock-free, safe from multiple threads)
|
|
tracker.add(100.0)
|
|
tracker.add(500.0)
|
|
tracker.add(50.0)
|
|
|
|
# Get max in the window (aggregates across all threads)
|
|
maximum = tracker.get_max() # returns 500.0
|
|
|
|
Thread Safety:
|
|
- add() is lock-free after the first call from each thread
|
|
- get_max() acquires a lock to aggregate across threads
|
|
- Safe to call from multiple threads concurrently
|
|
"""
|
|
|
|
def add(self, value: float) -> None:
|
|
"""Record a value, updating the current bucket's max if exceeded.
|
|
|
|
This operation is lock-free for the calling thread after the first call.
|
|
Safe to call from multiple threads concurrently.
|
|
|
|
Args:
|
|
value: The value to record.
|
|
"""
|
|
data = self._ensure_initialized()
|
|
|
|
self._rotate_buckets_if_needed(data)
|
|
if value > data.buckets[data.current_bucket_idx]:
|
|
data.buckets[data.current_bucket_idx] = value
|
|
|
|
def get_max(self) -> float:
|
|
"""Get max value across all non-expired buckets in the window.
|
|
|
|
This aggregates values from all threads that have called add().
|
|
Expired buckets (older than window_duration_s) are not included.
|
|
|
|
Returns:
|
|
The maximum value observed in the rolling window, or 0.0
|
|
if no values have been recorded.
|
|
"""
|
|
result = 0.0
|
|
now = time.time()
|
|
|
|
with self._registry_lock:
|
|
for data in self._all_thread_data:
|
|
elapsed = now - data.last_rotation_time
|
|
buckets_expired = int(elapsed / self._bucket_duration_s)
|
|
|
|
if buckets_expired >= self._num_buckets:
|
|
continue
|
|
|
|
for i in range(self._num_buckets - buckets_expired):
|
|
idx = (data.current_bucket_idx - i) % self._num_buckets
|
|
if data.buckets[idx] > result:
|
|
result = data.buckets[idx]
|
|
|
|
return result
|