From 08cfac1fa670c62d35ad85481e12e68845648270 Mon Sep 17 00:00:00 2001 From: Bob Date: Tue, 25 Aug 2026 11:57:44 +0000 Subject: [PATCH 1/2] perf(aw-transform): cache categorize/tag results per event data MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Same fix as ActivityWatch/aw-server-rust#657 — categorize() and tag() were re-evaluating every regex rule against every event individually. Heartbeat-based watchers emit the same app+title payload repeatedly, so the same regex matches were being recomputed thousands of times. Added a function-local cache keyed on json.dumps(e.data, sort_keys=True). Only the first occurrence of each distinct data fingerprint is matched against the rule set; subsequent events with identical data reuse the cached result in O(1). Each event receives its own list copy to prevent mutation aliasing across events. For a month with 50 000 events but ~200 distinct app/title pairs and 20 category rules: 1 000 000 → 4 000 regex evaluations (~250× less). Adds test_categorize_cache_correctness: 101 events (50 + 1 + 50), two distinct data shapes, verifies correct category for each and that mutating one event's category list does not affect others. --- aw_transform/classify.py | 19 +++++++++++++++++-- tests/test_transforms.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/aw_transform/classify.py b/aw_transform/classify.py index 898f416a..f77b9b25 100644 --- a/aw_transform/classify.py +++ b/aw_transform/classify.py @@ -1,3 +1,4 @@ +import json from typing import Pattern, List, Iterable, Tuple, Dict, Optional, Any from functools import reduce import re @@ -43,7 +44,15 @@ def match(self, e: Event) -> bool: def categorize( events: List[Event], classes: List[Tuple[Category, Rule]] ) -> List[Event]: - return [_categorize_one(e, classes) for e in events] + cache: Dict[str, Category] = {} + for e in events: + key = json.dumps(e.data, sort_keys=True) + if key not in cache: + cache[key] = _pick_category( + [_cls for _cls, rule in classes if rule.match(e)] + ) + e.data["$category"] = list(cache[key]) + return events def _categorize_one(e: Event, classes: List[Tuple[Category, Rule]]) -> Event: @@ -54,7 +63,13 @@ def _categorize_one(e: Event, classes: List[Tuple[Category, Rule]]) -> Event: def tag(events: List[Event], classes: List[Tuple[Tag, Rule]]) -> List[Event]: - return [_tag_one(e, classes) for e in events] + cache: Dict[str, List[Tag]] = {} + for e in events: + key = json.dumps(e.data, sort_keys=True) + if key not in cache: + cache[key] = [_cls for _cls, rule in classes if rule.match(e)] + e.data["$tags"] = list(cache[key]) + return events def _tag_one(e: Event, classes: List[Tuple[Tag, Rule]]) -> Event: diff --git a/tests/test_transforms.py b/tests/test_transforms.py index df627dbd..b7bf8837 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -418,6 +418,36 @@ def test_categorize(): assert events[3].data["$category"] == ["Uncategorized"] +def test_categorize_cache_correctness(): + """Cache reuses category for identical data; distinct data gets its own category.""" + now = datetime.now(timezone.utc) + + classes = [ + (["Browser"], Rule({"regex": "Firefox"})), + (["Editor"], Rule({"regex": "vim"})), + ] + firefox_data = {"app": "Firefox", "title": "Home"} + vim_data = {"app": "vim", "title": "classify.py"} + + # 50 Firefox events, 1 vim event, 50 more Firefox events + events = ( + [Event(timestamp=now, duration=0, data=dict(firefox_data)) for _ in range(50)] + + [Event(timestamp=now, duration=0, data=dict(vim_data))] + + [Event(timestamp=now, duration=0, data=dict(firefox_data)) for _ in range(50)] + ) + result = categorize(events, classes) + + for e in result[:50]: + assert e.data["$category"] == ["Browser"] + assert result[50].data["$category"] == ["Editor"] + for e in result[51:]: + assert e.data["$category"] == ["Browser"] + + # Mutating one event's category must not affect others sharing the same data fingerprint + result[0].data["$category"].append("MUTATED") + assert result[1].data["$category"] == ["Browser"] + + def test_tags(): now = datetime.now(timezone.utc) From d41a40640b366d610f73242203b5c93b55a80563 Mon Sep 17 00:00:00 2001 From: Bob Date: Tue, 25 Aug 2026 12:12:42 +0000 Subject: [PATCH 2/2] fix(classify): handle non-JSON-serializable event data in cache key MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit json.dumps raises TypeError when e.data contains non-serializable values (e.g. nested Event objects from chunk_events_by_key). query2.py catches TypeError and maps it to 'invalid amount of arguments', which obscures the real error and breaks test_query2_query_functions. Fall back to str(id(e.data)) for non-serializable data — no caching benefit for those events, but correctness is preserved. --- aw_transform/classify.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/aw_transform/classify.py b/aw_transform/classify.py index f77b9b25..7f3ba4fa 100644 --- a/aw_transform/classify.py +++ b/aw_transform/classify.py @@ -46,7 +46,10 @@ def categorize( ) -> List[Event]: cache: Dict[str, Category] = {} for e in events: - key = json.dumps(e.data, sort_keys=True) + try: + key = json.dumps(e.data, sort_keys=True) + except TypeError: + key = str(id(e.data)) if key not in cache: cache[key] = _pick_category( [_cls for _cls, rule in classes if rule.match(e)] @@ -65,7 +68,10 @@ def _categorize_one(e: Event, classes: List[Tuple[Category, Rule]]) -> Event: def tag(events: List[Event], classes: List[Tuple[Tag, Rule]]) -> List[Event]: cache: Dict[str, List[Tag]] = {} for e in events: - key = json.dumps(e.data, sort_keys=True) + try: + key = json.dumps(e.data, sort_keys=True) + except TypeError: + key = str(id(e.data)) if key not in cache: cache[key] = [_cls for _cls, rule in classes if rule.match(e)] e.data["$tags"] = list(cache[key])