Small improvement with the caching system

Signed-off-by: Aurélien Bompard <aurelien@bompard.org>

fedbadges/cached.py

@@ -12,6 +12,15 @@
 cache = make_region()
 
 
+def _query_has_single_arg(search_kwargs, required_kwargs):
+    query_keys = set(search_kwargs.keys())
+    with suppress(KeyError):
+        query_keys.remove("rows_per_page")
+    if query_keys != set(required_kwargs):
+        return False
+    return all(len(search_kwargs[arg]) == 1 for arg in required_kwargs)
+
+
 class CachedValue:
 
     def __init__(self):
@@ -22,45 +31,62 @@ def _get_key(self, **kwargs):
 
     def get(self, **kwargs):
         key = self._get_key(**kwargs)
-        return cache.get_or_create(key, creator=self.compute, creator_args=((), kwargs))
-
-    def set(self, *, value, **kwargs):
-        key = self._get_key(**kwargs)
-        log.debug("Updating cached value %s with key %s", self.__class__.__name__, key)
+        current_value = cache.get(key)
+        if current_value == NO_VALUE:
+            log.debug(
+                "Updating cached value %s (key is %s)",
+                self.__class__.__name__,
+                key,
+            )
+            value = self.compute(**kwargs)
+        else:
+            # That's the previous value, add this message
+            value = current_value + 1
         cache.set(key, value)
+        return value
 
     def compute(self, **kwargs):
         raise NotImplementedError
 
     def on_message(self, message: Message):
         raise NotImplementedError
 
+    @classmethod
+    def is_applicable(cls, search_kwargs, badge_dict):
+        raise NotImplementedError
+
 
-class TopicUserCount(CachedValue):
+class TopicCount(CachedValue):
 
     def compute(self, *, topic, username):
         total, pages, query = datanommer.models.Message.grep(
             topics=[topic], users=[username], defer=True
         )
         return total
 
-    def on_message(self, message: Message):
-        for username in message.usernames:
-            current_value = self.get(topic=message.topic, username=username)
-            if current_value == NO_VALUE:
-                continue  # Don't update the value if no one has ever requested it
-            self.set(topic=message.topic, username=username, value=current_value + 1)
-
     @classmethod
-    def is_applicable(cls, search_kwargs):
+    def is_applicable(cls, search_kwargs, badge_dict):
         """Return whether we can use this cached value for this datanommer query"""
-        query_keys = set(search_kwargs.keys())
-        with suppress(KeyError):
-            query_keys.remove("rows_per_page")
-        if query_keys != {"topics", "users"}:
+        if badge_dict.get("operation") != "count":
             return False
-        return len(search_kwargs["topics"]) == 1 and len(search_kwargs["users"]) == 1
+        return _query_has_single_arg(search_kwargs, ["topics", "users"])
 
 
-def on_message(message: Message):
-    TopicUserCount().on_message(message)
+class NewBuildsCount(CachedValue):
+
+    def compute(self, *, topic, username):
+        total, pages, messages = datanommer.models.Message.grep(topics=[topic], users=[username])
+        return len(msg for msg in messages if msg["msg"]["new"] == 1)
+
+    @classmethod
+    def is_applicable(cls, search_kwargs, badge_dict):
+        """Return whether we can use this cached value for this datanommer query"""
+        try:
+            topic = search_kwargs["topic"][0]
+        except (KeyError, IndexError):
+            return False
+        if not topic.endswith("buildsys.build.state.change"):
+            return False
+        if badge_dict.get("lambda") != "sum(1 for msg in query.all() if msg.msg['new'] == 1)":
+            return False
+        return _query_has_single_arg(search_kwargs, ["topics", "users"])

fedbadges/consumer.py

@@ -20,7 +20,6 @@
 
 from .aio import Periodic
 from .cached import cache
-from .cached import on_message as update_cache_on_message
 from .rulesrepo import RulesRepo
 from .utils import datanommer_has_message, notification_callback
 
@@ -137,8 +136,6 @@ def __call__(self, message: Message):
         # Award every badge as appropriate.
         log.debug("Received %s, %s", message.topic, message.id)
 
-        update_cache_on_message(message)
-
         tahrir = self._get_tahrir_client()
         for badge_rule in self.badge_rules:
             try:

fedbadges/rules.py

@@ -17,7 +17,7 @@
 from fedora_messaging.api import Message
 from tahrir_api.dbapi import TahrirDatabase
 
-from fedbadges.cached import TopicUserCount
+from fedbadges.cached import NewBuildsCount, TopicCount
 from fedbadges.utils import (
     # These are all in-process utilities
     construct_substitutions,
@@ -470,10 +470,14 @@ def _construct_query(self, msg):
             if users:
                 kwargs["users"] = users
 
-        if TopicUserCount.is_applicable(kwargs) and self._d["operation"] == "count":
-            cached_value = TopicUserCount()
-            total = cached_value.get(topic=kwargs["topics"][0], username=kwargs["users"][0])
-            return total, None, None
+        for CachedClass in (TopicCount, NewBuildsCount):
+            if CachedClass.is_applicable(kwargs, self._d):
+                log.debug(
+                    f"Using the cached datanommer value (or setting it) for {CachedClass.__name__}"
+                )
+                cached_value = CachedClass()
+                total = cached_value.get(topic=kwargs["topics"][0], username=kwargs["users"][0])
+                return total, None, None
 
         log.debug("Making datanommer query: %r", kwargs)
         kwargs["defer"] = True