mirror of
https://mau.dev/maunium/synapse.git
synced 2024-07-04 18:31:48 +00:00
d8cc86eff4
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.
108 lines
3.4 KiB
Python
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)
|