chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,288 @@
|
||||
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
|
||||
Reference in New Issue
Block a user