From bae966518e2ba86f3c8d2525f036d1a5ff261ec0 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Fri, 6 May 2022 15:21:44 +0100 Subject: [PATCH] Use a set for `WheelTimer` to better handle duplicates. --- synapse/util/wheel_timer.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/synapse/util/wheel_timer.py b/synapse/util/wheel_timer.py index e108adc4604f..ab5d5a27de77 100644 --- a/synapse/util/wheel_timer.py +++ b/synapse/util/wheel_timer.py @@ -11,9 +11,9 @@ # 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. -from typing import Generic, List, TypeVar +from typing import Generic, Hashable, List, Set, TypeVar -T = TypeVar("T") +T = TypeVar("T", bound=Hashable) class _Entry(Generic[T]): @@ -21,7 +21,10 @@ class _Entry(Generic[T]): def __init__(self, end_key: int) -> None: self.end_key: int = end_key - self.queue: List[T] = [] + + # We use a set here as otherwise we can end up with a lot of duplicate + # entries. + self.queue: Set[T] = set() class WheelTimer(Generic[T]): @@ -55,7 +58,7 @@ def insert(self, now: int, obj: T, then: int) -> None: 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].queue.append(obj) + self.entries[max(min_key, then_key) - min_key].queue.add(obj) return next_key = int(now / self.bucket_size) + 1 @@ -71,7 +74,7 @@ def insert(self, now: int, obj: T, then: int) -> None: # 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].queue.append(obj) + self.entries[-1].queue.add(obj) def fetch(self, now: int) -> List[T]: """Fetch any objects that have timed out @@ -84,7 +87,7 @@ def fetch(self, now: int) -> List[T]: """ now_key = int(now / self.bucket_size) - ret = [] + ret: List[T] = [] while self.entries and self.entries[0].end_key <= now_key: ret.extend(self.entries.pop(0).queue)