Better benchmark

2 files changed, 28 insertions, 14 deletions

diff --git a/synapse/storage/databases/main/events_worker.py b/synapse/storage/databases/main/events_worker.py
index 2d6bc53f20..4b3f958b19 100644
--- a/synapse/storage/databases/main/events_worker.py
+++ b/synapse/storage/databases/main/events_worker.py
@@ -1480,7 +1480,7 @@ class EventsWorkerStore(SQLBaseStore):
         # the batches as big as possible.
 
         remaining_event_ids: Set[str] = set()
-        for chunk in batch_iter(event_ids, 500):
+        for chunk in batch_iter(event_ids, 1000):
             remaining_event_ids_from_chunk = await self._have_seen_events_dict(chunk)
             remaining_event_ids.update(remaining_event_ids_from_chunk)
 
@@ -1495,8 +1495,8 @@ class EventsWorkerStore(SQLBaseStore):
         """
 
         # if the event cache contains the event, obviously we've seen it.
-        event_entry_map = self._get_events_from_local_cache(event_ids)
-        event_ids_in_cache = event_entry_map.keys()
+        event_cache_entry_map = self._get_events_from_local_cache(event_ids)
+        event_ids_in_cache = event_cache_entry_map.keys()
         remaining_event_ids = {
             event_id for event_id in event_ids if event_id not in event_ids_in_cache
         }
diff --git a/tests/storage/databases/main/test_events_worker.py b/tests/storage/databases/main/test_events_worker.py
index ca65e77d3e..8a64af082a 100644
--- a/tests/storage/databases/main/test_events_worker.py
+++ b/tests/storage/databases/main/test_events_worker.py
@@ -94,6 +94,7 @@ class HaveSeenEventsTestCase(unittest.HomeserverTestCase):
             )
             self.event_ids.append(event_id)
 
+    # TODO: Remove me before mergin
     def test_benchmark(self):
         import time
 
@@ -130,17 +131,30 @@ class HaveSeenEventsTestCase(unittest.HomeserverTestCase):
         )
 
         with LoggingContext(name="test") as ctx:
-            benchmark_start_time = time.time()
-            remaining_event_ids = self.get_success(
-                self.store.have_seen_events(room_id, event_ids)
-            )
-            benchmark_end_time = time.time()
-            logger.info("afewewf %s %s", benchmark_start_time, benchmark_end_time)
-            logger.info(
-                "Benchmark time: %s",
-                (benchmark_end_time - benchmark_start_time),
-            )
-            # self.assertEqual(remaining_event_ids, set())
+
+            def time_have_seen_events(test_index: int, event_ids):
+                benchmark_start_time = time.time()
+                remaining_event_ids = self.get_success(
+                    self.store.have_seen_events(room_id, event_ids)
+                )
+                benchmark_end_time = time.time()
+                logger.info(
+                    "Benchmark time%s: %s",
+                    test_index,
+                    (benchmark_end_time - benchmark_start_time),
+                )
+                self.assertIsNotNone(remaining_event_ids)
+
+            event_ids_odd = event_ids[::2]
+            event_ids_even = event_ids[1::2]
+
+            time_have_seen_events(1, event_ids)
+            time_have_seen_events(2, event_ids)
+            time_have_seen_events(3, event_ids)
+            time_have_seen_events(4, event_ids_odd)
+            time_have_seen_events(5, event_ids_odd)
+            time_have_seen_events(6, event_ids_even)
+            time_have_seen_events(7, event_ids_even)
 
             # that should result in a many db queries
             self.assertEqual(ctx.get_resource_usage().db_txn_count, 1)