synapse-product/synapse/util/wheel_timer.py
Patrick Cloke d8cc86eff4
Remove redundant types from comments. (#14412)
Remove type hints from comments which have been added
as Python type hints. This helps avoid drift between comments
and reality, as well as removing redundant information.

Also adds some missing type hints which were simple to fill in.
2022-11-16 15:25:24 +00:00

108 lines
3.4 KiB
Python

# Copyright 2016 OpenMarket Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
from typing import Generic, Hashable, List, Set, TypeVar
import attr
logger = logging.getLogger(__name__)
T = TypeVar("T", bound=Hashable)
@attr.s(slots=True, frozen=True, auto_attribs=True)
class _Entry(Generic[T]):
end_key: int
elements: Set[T] = attr.Factory(set)
class WheelTimer(Generic[T]):
"""Stores arbitrary objects that will be returned after their timers have
expired.
"""
def __init__(self, bucket_size: int = 5000) -> None:
"""
Args:
bucket_size: Size of buckets in ms. Corresponds roughly to the
accuracy of the timer.
"""
self.bucket_size: int = bucket_size
self.entries: List[_Entry[T]] = []
self.current_tick: int = 0
def insert(self, now: int, obj: T, then: int) -> None:
"""Inserts object into timer.
Args:
now: Current time in msec
obj: Object to be inserted
then: When to return the object strictly after.
"""
then_key = int(then / self.bucket_size) + 1
now_key = int(now / self.bucket_size)
if self.entries:
min_key = self.entries[0].end_key
max_key = self.entries[-1].end_key
if min_key < now_key - 10:
# If we have ten buckets that are due and still nothing has
# called `fetch()` then we likely have a bug that is causing a
# memory leak.
logger.warning(
"Inserting into a wheel timer that hasn't been read from recently. Item: %s",
obj,
)
if then_key <= max_key:
# The max here is to protect against inserts for times in the past
self.entries[max(min_key, then_key) - min_key].elements.add(obj)
return
next_key = now_key + 1
if self.entries:
last_key = self.entries[-1].end_key
else:
last_key = next_key
# Handle the case when `then` is in the past and `entries` is empty.
then_key = max(last_key, then_key)
# Add empty entries between the end of the current list and when we want
# to insert. This ensures there are no gaps.
self.entries.extend(_Entry(key) for key in range(last_key, then_key + 1))
self.entries[-1].elements.add(obj)
def fetch(self, now: int) -> List[T]:
"""Fetch any objects that have timed out
Args:
now: Current time in msec
Returns:
List of objects that have timed out
"""
now_key = int(now / self.bucket_size)
ret: List[T] = []
while self.entries and self.entries[0].end_key <= now_key:
ret.extend(self.entries.pop(0).elements)
return ret
def __len__(self) -> int:
return sum(len(entry.elements) for entry in self.entries)