perf: optimize heartbeat writes, queries, bucket listings, and exports

Makefile

@@ -24,7 +24,7 @@ test:
 	@# Note that extensive integration tests are also run in the bundle repo,
 	@# for both aw-server and aw-server-rust, but without code coverage.
 	python -c 'import aw_server'
-	python -m pytest tests/test_server.py tests/test_profile.py tests/test_profile_config.py tests/test_query_cache.py tests/test_import.py tests/test_host_header.py tests/test_database_concurrency.py
+	python -m pytest tests/test_server.py tests/test_profile.py tests/test_profile_config.py tests/test_query_cache.py tests/test_import.py tests/test_host_header.py tests/test_database_concurrency.py tests/test_performance.py
 
 typecheck:
 	python -m mypy aw_server tests --ignore-missing-imports

aw_server/api.py

@@ -1,10 +1,12 @@
+import copy
 import functools
 import json
 import logging
 import threading
 from datetime import datetime
 from pathlib import Path
 from socket import gethostname
+from contextlib import closing
 from typing import (
     Any,
     Dict,
@@ -57,8 +59,14 @@ def get_device_id() -> str:
 def check_bucket_exists(f):
     @functools.wraps(f)
     def g(self, bucket_id, *args, **kwargs):
-        if bucket_id not in self.db.buckets():
-            raise NotFound("NoSuchBucket", f"There's no bucket named {bucket_id}")
+        # Datastore caches bucket handles and invalidates them on deletion.
+        # Reuse that lookup instead of loading every bucket's metadata per call.
+        try:
+            self.db[bucket_id]
+        except KeyError:
+            raise NotFound(
+                "NoSuchBucket", f"There's no bucket named {bucket_id}"
+            ) from None
         return f(self, bucket_id, *args, **kwargs)
 
     return g
@@ -77,6 +85,19 @@ def g(self, *args, **kwargs):
     return g
 
 
+def _forget_last_event(f):
+    """Drop the cached heartbeat state for the bucket before a non-heartbeat
+    write, so a later heartbeat re-reads the stored last event instead of
+    merging into (and replacing by ID) a stale one. Use under _serialized."""
+
+    @functools.wraps(f)
+    def g(self, bucket_id, *args, **kwargs):
+        self.last_event.pop(bucket_id, None)
+        return f(self, bucket_id, *args, **kwargs)
+
+    return g
+
+
 class ServerAPI:
     def __init__(self, db, testing, query_cache: bool = True) -> None:
         self.db = db
@@ -110,15 +131,7 @@ def get_info(self) -> Dict[str, Any]:
     def get_buckets(self) -> Dict[str, Dict]:
         """Get dict {bucket_name: Bucket} of all buckets"""
         logger.debug("Received get request for buckets")
-        buckets = self.db.buckets()
-        for b in buckets:
-            # TODO: Move this code to aw-core?
-            last_events = self.db[b].get(limit=1)
-            if len(last_events) > 0:
-                last_event = last_events[0]
-                last_updated = last_event.timestamp + last_event.duration
-                buckets[b]["last_updated"] = last_updated.isoformat()
-        return buckets
+        return self.db.buckets(include_last_updated=True)
 
     @check_bucket_exists
     def get_bucket_metadata(self, bucket_id: str) -> Dict[str, Any]:
@@ -150,15 +163,59 @@ def export_bucket_csv(
 
     def export_all(self) -> Dict[str, Any]:
         """Exports all buckets and their events to a format consistent across versions"""
-        buckets = self.get_buckets()
+        buckets = self.db.buckets()
         exported_buckets = {}
         for bid in buckets.keys():
             exported_buckets[bid] = self.export_bucket(bid)
         return exported_buckets
 
+    def stream_export(self, bucket_id=None):
+        # Validate before returning the generator so a missing bucket is a 404,
+        # not an exception after the HTTP response has started.
+        if bucket_id is not None:
+            buckets = {bucket_id: self.get_bucket_metadata(bucket_id)}
+        else:
+            buckets = self.db.buckets()
+
+        def generate():
+            yield '{"buckets":{'
+            for index, (bid, metadata) in enumerate(buckets.items()):
+                if index:
+                    yield ","
+                yield json.dumps(bid) + ":"
+                yield json.dumps(metadata)[:-1] + ',"events":['
+                with closing(self.db[bid].iter_events()) as events:
+                    for event_index, event in enumerate(events):
+                        if event_index:
+                            yield ","
+                        payload = event.to_json_dict()
+                        payload.pop("id", None)
+                        yield json.dumps(payload)
+                yield "]}"
+            yield "}}"
+
+        def buffered():
+            # Avoid a socket write for every separator/small event while retaining
+            # only a bounded batch (plus the largest individual event).
+            with closing(generate()) as fragments:
+                batch = []
+                size = 0
+                for fragment in fragments:
+                    batch.append(fragment)
+                    size += len(fragment)
+                    if size >= 64 * 1024:
+                        yield "".join(batch)
+                        batch = []
+                        size = 0
+                if batch:
+                    yield "".join(batch)
+
+        return buffered()
+
     @_serialized
     def import_bucket(self, bucket_data: Any):
         bucket_id = bucket_data["id"]
+        self.last_event.pop(bucket_id, None)
         logger.info(f"Importing bucket {bucket_id}")
 
         if bucket_id in self.db.buckets():
@@ -250,6 +307,7 @@ def import_all(self, buckets: Dict[str, Any]):
             raise
 
     @_serialized
+    @_forget_last_event
     def create_bucket(
         self,
         bucket_id: str,
@@ -269,7 +327,7 @@ def create_bucket(
         """
         if created is None:
             created = datetime.now()
-        if bucket_id in self.db.buckets():
+        if self.db.has_bucket(bucket_id):
             return False
         if hostname == "!local":
             info = self.get_info()
@@ -292,6 +350,7 @@ def create_bucket(
 
     @check_bucket_exists
     @_serialized
+    @_forget_last_event
     def update_bucket(
         self,
         bucket_id: str,
@@ -304,7 +363,7 @@ def update_bucket(
         try:
             self.db.update_bucket(
                 bucket_id,
-                type=event_type,
+                type_id=event_type,
                 client=client,
                 hostname=hostname,
                 data=data,
@@ -316,14 +375,15 @@ def update_bucket(
 
     @check_bucket_exists
     @_serialized
+    @_forget_last_event
     def delete_bucket(self, bucket_id: str) -> None:
         """Delete a bucket"""
         try:
             self.db.delete_bucket(bucket_id)
         finally:
             self._invalidate_all()
             self.last_event.pop(bucket_id, None)
-        logger.debug(f"Deleted bucket '{bucket_id}'")
+        logger.debug("Deleted bucket '%s'", bucket_id)
         return None
 
     @check_bucket_exists
@@ -334,7 +394,7 @@ def get_event(
     ) -> Optional[Event]:
         """Get a single event from a bucket"""
         logger.debug(
-            f"Received get request for event {event_id} in bucket '{bucket_id}'"
+            "Received get request for event %s in bucket '%s'", event_id, bucket_id
         )
         event = self.db[bucket_id].get_by_id(event_id)
         return event.to_json_dict() if event else None
@@ -348,7 +408,7 @@ def get_events(
         end: Optional[datetime] = None,
     ) -> List[Event]:
         """Get events from a bucket"""
-        logger.debug(f"Received get request for events in bucket '{bucket_id}'")
+        logger.debug("Received get request for events in bucket '%s'", bucket_id)
         if limit is None:  # Let limit = None also mean "no limit"
             limit = -1
         events = [
@@ -358,6 +418,7 @@ def get_events(
 
     @check_bucket_exists
     @_serialized
+    @_forget_last_event
     def create_events(self, bucket_id: str, events: List[Event]) -> List[Event]:
         """Create events for a bucket. Can handle both single events and multiple ones.
 
@@ -394,11 +455,12 @@ def get_eventcount(
         end: Optional[datetime] = None,
     ) -> int:
         """Get eventcount from a bucket"""
-        logger.debug(f"Received get request for eventcount in bucket '{bucket_id}'")
+        logger.debug("Received get request for eventcount in bucket '%s'", bucket_id)
         return self.db[bucket_id].get_eventcount(start, end)
 
     @check_bucket_exists
     @_serialized
+    @_forget_last_event
     def delete_event(self, bucket_id: str, event_id) -> bool:
         """Delete a single event from a bucket"""
         old = self.db[bucket_id].get_by_id(event_id)
@@ -463,41 +525,29 @@ def heartbeat(self, bucket_id: str, heartbeat: Event, pulsetime: float) -> Event
         Inspired by: https://wakatime.com/developers#heartbeats
         """
         logger.debug(
-            "Received heartbeat in bucket '{}'\n\ttimestamp: {}, duration: {}, pulsetime: {}\n\tdata: {}".format(
-                bucket_id,
-                heartbeat.timestamp,
-                heartbeat.duration,
-                pulsetime,
-                heartbeat.data,
-            )
+            "Received heartbeat in bucket '%s'\n\ttimestamp: %s, duration: %s, pulsetime: %s\n\tdata: %s",
+            bucket_id,
+            heartbeat.timestamp,
+            heartbeat.duration,
+            pulsetime,
+            heartbeat.data,
         )
 
-        # The endtime here is set such that in the event that the heartbeat is older than an
-        # existing event we should try to merge it with the last event before the heartbeat instead.
-        # FIXME: This (the endtime=heartbeat.timestamp) gets rid of the "heartbeat was older than last event"
-        #        warning and also causes a already existing "newer" event to be overwritten in the
-        #        replace_last call below. This is problematic.
-        # Solution: This could be solved if we were able to replace arbitrary events.
-        #           That way we could double check that the event has been applied
-        #           and if it hasn't we simply replace it with the updated counterpart.
-
         last_event = None
         if bucket_id not in self.last_event:
             last_events = self.db[bucket_id].get(limit=1)
             if len(last_events) > 0:
                 last_event = last_events[0]
         else:
-            last_event = self.last_event[bucket_id]
+            last_event = copy.copy(self.last_event[bucket_id])
 
         if last_event:
             if last_event.data == heartbeat.data:
                 merged = heartbeat_merge(last_event, heartbeat, pulsetime)
                 if merged is not None:
                     # Heartbeat was merged into last_event
                     logger.debug(
-                        "Received valid heartbeat, merging. (bucket: {})".format(
-                            bucket_id
-                        )
+                        "Received valid heartbeat, merging. (bucket: %s)", bucket_id
                     )
                     affected = [event_range(last_event), event_range(merged)]
                     if merged.id is None and self.query_cache:
@@ -506,45 +556,47 @@ def heartbeat(self, bucket_id: str, heartbeat: Event, pulsetime: float) -> Event
                         # out-of-order heartbeat), so invalidate its range too.
                         stored = self.db[bucket_id].get(limit=1)
                         affected += [event_range(e) for e in stored]
-                    self.last_event[bucket_id] = merged
+                    replaced = False
                     try:
                         # Replace the event we merged into by id. replace_last()
                         # targets whichever event sorts last by timestamp, which
                         # after an out-of-order heartbeat is a different event:
                         # it gets overwritten with a copy of this one, leaving
                         # same-timestamp duplicates (aw-watcher-afk#61).
                         if merged.id is not None:
-                            self.db[bucket_id].replace(merged.id, merged)
+                            replaced = self.db[bucket_id].replace(merged.id, merged)
                         else:
                             self.db[bucket_id].replace_last(merged)
+                            replaced = True
                     finally:
                         self._invalidate(affected)
-                    return merged
+                    if replaced:
+                        self.last_event[bucket_id] = merged
+                        return merged
+                    # The row may have been removed by another datastore user.
+                    self.last_event.pop(bucket_id, None)
                 else:
                     logger.info(
-                        "Received heartbeat after pulse window, inserting as new event. (bucket: {})".format(
-                            bucket_id
-                        )
+                        "Received heartbeat after pulse window, inserting as new event. (bucket: %s)",
+                        bucket_id,
                     )
             else:
                 logger.debug(
-                    "Received heartbeat with differing data, inserting as new event. (bucket: {})".format(
-                        bucket_id
-                    )
+                    "Received heartbeat with differing data, inserting as new event. (bucket: %s)",
+                    bucket_id,
                 )
         else:
             logger.info(
-                "Received heartbeat, but bucket was previously empty, inserting as new event. (bucket: {})".format(
-                    bucket_id
-                )
+                "Received heartbeat, but bucket was previously empty, inserting as new event. (bucket: %s)",
+                bucket_id,
             )
 
         try:
-            self.db[bucket_id].insert(heartbeat)
+            inserted = self.db[bucket_id].insert(heartbeat)
         finally:
             self._invalidate([event_range(heartbeat)])
-        self.last_event[bucket_id] = heartbeat
-        return heartbeat
+        self.last_event[bucket_id] = inserted
+        return inserted
 
     def query2(self, name, query, timeperiods, cache: bool = True):
         """Run a query for each timeperiod.