feat(queue): merge queued heartbeats per bucket before dispatch

aw_client/client.py

@@ -3,6 +3,7 @@
 import json
 import logging
 import os
+import re
 import socket
 import sqlite3
 import threading
@@ -552,6 +553,39 @@ def _retry_delay(response: req.Response) -> float:
     return max(RETRY_DELAY_DEFAULT, min(delay, RETRY_DELAY_MAX))
 
 
+# Matches the endpoint the queue stores heartbeats under, capturing the bucket
+# id and the pulsetime needed to merge consecutive queued heartbeats.
+_HEARTBEAT_ENDPOINT_RE = re.compile(
+    r"^buckets/(?P<bucket_id>.+)/heartbeat\?pulsetime=(?P<pulsetime>[0-9]+(?:\.[0-9]+)?)$"
+)
+
+
+def _parse_heartbeat_endpoint(endpoint: str) -> Optional[Tuple[str, float]]:
+    """Return (bucket_id, pulsetime) for a queued heartbeat endpoint, else None.
+
+    A malformed pulsetime must not raise: it just falls back to sending the
+    request verbatim, so one poisoned endpoint cannot wedge the queue.
+    """
+    match = _HEARTBEAT_ENDPOINT_RE.match(endpoint)
+    if match is None:
+        return None
+    try:
+        pulsetime = float(match.group("pulsetime"))
+    except ValueError:
+        return None
+    return match.group("bucket_id"), pulsetime
+
+
+def _try_event(data: Any) -> Optional[Event]:
+    """Parse a queued payload as an Event, or None if it is not one."""
+    if not isinstance(data, dict):
+        return None
+    try:
+        return Event(**data)
+    except TypeError:
+        return None
+
+
 class RequestQueue(threading.Thread):
     """Used to asynchronously send heartbeats.
 
@@ -567,6 +601,10 @@ class RequestQueue(threading.Thread):
     # 503 in particular is sent by aw-server when the heartbeat lock times out.
     RETRY_STATUS_CODES = {429, 500, 502, 503, 504}
 
+    # How many queued requests to pop at once, so consecutive heartbeats for
+    # the same bucket can be merged before dispatch.
+    BATCH_SIZE = 1000
+
     def __init__(
         self,
         client: ActivityWatchClient,
@@ -607,21 +645,49 @@ def __init__(
         self._persistqueue = persistqueue.FIFOSQLiteQueue(
             persistqueue_path, multithreading=True, auto_commit=False
         )
-        self._current = None  # type: Optional[QueuedRequest]
+        # The requests popped from the queue but not yet acknowledged. Held
+        # across dispatches so a transient error retries them instead of
+        # dropping them. A single task_done() deletes every row <= the cursor,
+        # i.e. the whole batch, so no per-item bookkeeping is needed.
+        self._current_batch = []  # type: List[QueuedRequest]
+        # `_current_batch` after coalescing, plus how many of those have been
+        # handled. Dispatching one coalesced request per call means a transient
+        # error retries only the failed request - never replays ones that
+        # already reached the server - and lets the run loop observe stop().
+        self._coalesced_batch = []  # type: List[QueuedRequest]
+        self._coalesced_index = 0
         self._queue_write_failing = False
 
     def _get_next(self) -> Optional[QueuedRequest]:
-        # self._current will always hold the next not-yet-sent event,
-        # until self._task_done() is called.
-        if not self._current:
+        # Returns the head of the in-flight batch, otherwise pops a single
+        # request. Used by tests/callers that only need one item.
+        if not self._current_batch:
             try:
-                self._current = self._persistqueue.get(block=False)
+                self._current_batch = [self._persistqueue.get(block=False)]
             except persistqueue.exceptions.Empty:
                 return None
-        return self._current
+        return self._current_batch[0]
+
+    def _get_next_batch(self) -> List[QueuedRequest]:
+        # Pop up to BATCH_SIZE requests in one go so consecutive heartbeats
+        # for the same bucket can be merged before dispatch. The in-memory
+        # batch is only acknowledged (deleted) by _task_done(), so a crash
+        # mid-batch leaves every popped request on disk for the next run.
+        if self._current_batch:
+            return self._current_batch
+        batch = []  # type: List[QueuedRequest]
+        while len(batch) < self.BATCH_SIZE:
+            try:
+                batch.append(self._persistqueue.get(block=False))
+            except persistqueue.exceptions.Empty:
+                break
+        self._current_batch = batch
+        return batch
 
     def _task_done(self) -> None:
-        self._current = None
+        self._current_batch = []
+        self._coalesced_batch = []
+        self._coalesced_index = 0
         self._persistqueue.task_done()
 
     def _create_buckets(self) -> None:
@@ -658,31 +724,101 @@ def wait_for_queue_empty(self, timeout: Optional[float] = None) -> bool:
         :return: True if the queue became empty, False if the timeout was reached.
         """
         if not self.is_alive():
-            return self._persistqueue.qsize() == 0 and self._current is None
+            return self._persistqueue.qsize() == 0 and not self._current_batch
 
         start_time = monotonic()
-        while self._persistqueue.qsize() > 0 or self._current is not None:
+        while self._persistqueue.qsize() > 0 or self._current_batch:
             if timeout is not None and monotonic() - start_time >= timeout:
                 return False
             if self.wait(0.1):
                 # stop() was called while waiting
                 return False
         return True
 
+    def _coalesce(self, batch: List[QueuedRequest]) -> List[QueuedRequest]:
+        """Merge consecutive queued heartbeats for the same bucket.
+
+        When the server is unreachable the queue accumulates one heartbeat per
+        commit interval. Consecutive heartbeats for a bucket usually carry
+        identical data, so aw_transform.heartbeat_merge collapses them into a
+        handful of long events. Sending the merged events to the heartbeat
+        endpoint produces the same server-side result with far fewer requests.
+
+        Requests are grouped by bucket, not by full endpoint: grouping by
+        endpoint would reorder a bucket's heartbeats when their pulsetimes
+        differ. Within a bucket only a contiguous run sharing one endpoint is
+        merged, so the order of a bucket's requests is preserved exactly.
+        """
+        groups = {}  # type: Dict[str, List[QueuedRequest]]
+        for request in batch:
+            parsed = _parse_heartbeat_endpoint(request.endpoint)
+            key = parsed[0] if parsed is not None else request.endpoint
+            groups.setdefault(key, []).append(request)
+
+        coalesced = []  # type: List[QueuedRequest]
+
+        def flush(endpoint: Optional[str], merged: List[Event]) -> None:
+            coalesced.extend(
+                QueuedRequest(endpoint, event.to_json_dict()) for event in merged
+            )
+
+        for requests in groups.values():
+            merged = []  # type: List[Event]
+            merged_endpoint = None  # type: Optional[str]
+            for request in requests:
+                parsed = _parse_heartbeat_endpoint(request.endpoint)
+                event = _try_event(request.data) if parsed is not None else None
+                if (
+                    parsed is None
+                    or event is None
+                    or request.endpoint != merged_endpoint
+                ):
+                    # Unknown endpoint shape, non-event payload, or a new
+                    # endpoint: close the current run and start a new one.
+                    flush(merged_endpoint, merged)
+                    merged = []
+                    merged_endpoint = None
+                if parsed is None or event is None:
+                    coalesced.append(request)
+                    continue
+                if merged:
+                    merged_event = heartbeat_merge(merged[-1], event, parsed[1])
+                    if merged_event is not None:
+                        merged[-1] = merged_event
+                        continue
+                    # Not mergeable: close the run before starting a new one.
+                    flush(merged_endpoint, merged)
+                    merged = []
+                merged.append(event)
+                merged_endpoint = request.endpoint
+            flush(merged_endpoint, merged)
+        return coalesced
+
     def _dispatch_request(self) -> None:
-        request = self._get_next()
-        if not request:
+        batch = self._get_next_batch()
+        if not batch:
             self.wait(0.2)  # seconds to wait before re-polling the empty queue
             return
 
+        if not self._coalesced_batch:
+            self._coalesced_batch = self._coalesce(batch)
+            if not self._coalesced_batch:
+                self._task_done()
+                return
+
+        # Dispatch one coalesced request per call: a transient error then
+        # retries only the failed request (never replays ones that already
+        # reached the server), and the run loop can observe stop() between
+        # requests instead of blocking on a long batch.
+        request = self._coalesced_batch[self._coalesced_index]
         try:
             self.client._post(request.endpoint, request.data)
         except (req.exceptions.ConnectionError, req.exceptions.Timeout):
             # Triggered by:
             #   - server not running (connection refused)
             #   - server not responding (timeout)
-            # Keep the request in the queue and go back to waiting for the
-            # server to become available (the run loop reconnects).
+            # Keep the batch in memory and go back to waiting for the server
+            # to become available (the run loop reconnects).
             self.connected = False
             logger.warning(
                 "Connection refused or timeout, will queue requests until connection is available."
@@ -701,9 +837,9 @@ def _dispatch_request(self) -> None:
             if response is not None and status_code in self.RETRY_STATUS_CODES:
                 # Transient server-side problem (busy, overloaded, restarting
                 # or behind a flaky proxy) - the request itself is likely
-                # fine, so keep it in the queue and retry. Heartbeats are safe
-                # to replay: a duplicate of an already-processed heartbeat
-                # merges into the last event as a no-op.
+                # fine, so retry it. Heartbeats are safe to replay: a
+                # duplicate of an already-processed heartbeat merges into the
+                # last event as a no-op.
                 delay = _retry_delay(response)
                 logger.warning(
                     f"Server error {status_code}, will retry in {delay}s: {request.endpoint}"
@@ -714,16 +850,19 @@ def _dispatch_request(self) -> None:
             else:
                 # Client errors (e.g. HTTP 400 - bad request, see
                 # https://github.com/ActivityWatch/activitywatch/issues/815)
-                # are likely to fail forever, so drop the request instead of
-                # blocking the queue.
-                logger.exception(
+                # are likely to fail forever, so drop this request and move on
+                # to the next one.
+                logger.error(
                     f"Request failed ({status_code}), not retrying: {request.data}"
                 )
         except Exception:
             logger.exception(f"Unknown error, not retrying: {request.data}")
 
-        # Mark the request as done
-        self._task_done()
+        # Handled (delivered or dropped); advance. Acknowledge the batch only
+        # once every request in it has been handled.
+        self._coalesced_index += 1
+        if self._coalesced_index >= len(self._coalesced_batch):
+            self._task_done()
 
     def run(self) -> None:
         self._stop_event.clear()