feat: warm caches during app startup

src/auslib/AUS.py

@@ -4,7 +4,7 @@
 from random import randint
 from urllib.parse import urlparse
 
-from auslib.blobs.base import ServeUpdate, createBlob
+from auslib.blobs.base import ServeUpdate
 from auslib.global_state import cache, dbo
 from auslib.services import releases
 from auslib.util.versions import PinVersion
@@ -135,18 +135,7 @@ def evaluateRules(self, updateQuery, transaction=None):
         # 3) Incoming release is older than the one in the mapping, defined as one of:
         #    * version decreases
         #    * version is the same and buildID doesn't increase
-        def get_blob(mapping):
-            release = releases.get_release(mapping, transaction, include_sc=False)
-            blob = None
-            if release:
-                blob = createBlob(release["blob"])
-            # TODO: remove me when old releases table dies
-            else:
-                release = dbo.releases.getReleases(name=mapping, limit=1, transaction=transaction)[0]
-                blob = release["data"]
-            return blob
-
-        blob = get_blob(mapping)
+        blob = releases.get_release_blob(mapping)
         if not blob:
             return None, None, eval_metadata
         candidate = blob.shouldServeUpdate(updateQuery)
@@ -179,7 +168,7 @@ def get_blob(mapping):
             # installations vulnerable.
             if pin_mapping is not None:
                 mapping = pin_mapping
-                blob = get_blob(mapping)
+                blob = releases.get_release_blob(mapping)
                 if not blob or not blob.shouldServeUpdate(updateQuery):
                     return None, None, eval_metadata
 

src/auslib/blobs/superblob.py

@@ -37,3 +37,9 @@ def getInnerHeaderXML(self, updateQuery, update_type, allowlistedDomains, specia
 
     def getInnerFooterXML(self, updateQuery, update_type, allowlistedDomains, specialForceHosts):
         return "    </addons>"
+
+    def getReferencedReleases(self):
+        if responseBlobs := self.getResponseBlobs():
+            return set(responseBlobs)
+        else:
+            return set()

src/auslib/services/releases.py

@@ -927,3 +927,17 @@ def get_version(mapping, trans):
             data_version=old_row["data_version"],
         )
         return {".": {"sc_id": sc_id, "change_type": "update", "data_version": old_row["data_version"] + 1, "signoffs": {}, "when": when}}
+
+
+def get_release_blob(name, trans=None):
+    """Get a release blob regardless of whether it's in the old releases or new releases."""
+    release = releases.get_release(name, trans, include_sc=False)
+    blob = None
+    if release:
+        blob = createBlob(release["blob"])
+    # TODO: remove me when old releases table dies
+    else:
+        release = dbo.releases.getReleases(name=name, limit=1, transaction=trans)
+        if release:
+            blob = release[0]["data"]
+    return blob

src/auslib/web/public/helpers.py

@@ -5,6 +5,7 @@
 
 from auslib.AUS import AUS
 from auslib.global_state import cache, dbo
+from auslib.services.releases import get_release_blob
 from auslib.util.autograph import make_hash, sign_hash
 
 log = logging.getLogger(__name__)
@@ -52,4 +53,24 @@ def sign():
     return headers
 
 
+def warm_caches():
+    """Fetch all releases pointed at by a rule. This is intended to run at app
+    startup, to ensure that it has all necessary information is in memory before
+    it begins serving requests."""
+    mapped_releases = set()
+    for rule in dbo.rules.select(columns=[dbo.rules.mapping, dbo.rules.fallbackMapping], distinct=True):
+        if rule["mapping"]:
+            mapped_releases.add(rule["mapping"])
+        if rule["fallbackMapping"]:
+            mapped_releases.add(rule["fallbackMapping"])
+
+    for release in mapped_releases:
+        blob = get_release_blob(release)
+        if not blob:
+            continue
+
+        for referenced_release in blob.getReferencedReleases():
+            get_release_blob(referenced_release)
+
+
 AUS = AUS()

tests/web/test_helpers.py

@@ -1,7 +1,8 @@
 from unittest.mock import MagicMock
 
 import auslib.web.public.helpers
-import auslib.web.public.json
+from auslib.blobs.base import createBlob
+from auslib.global_state import cache, dbo
 
 
 def test_get_content_signature_headers(monkeypatch):
@@ -32,3 +33,79 @@ def mock_sign_hash(_, key, *__):
     assert auslib.web.public.helpers.get_content_signature_headers(content, product) == {"Content-Signature": f"x5u={x5u}; p384ecdsa={ecdsa}"}
 
     assert mocksign.call_count == 1
+
+
+def test_warm_caches(db_schema, insert_release, firefox_54_0_1_build1, firefox_67_0_build1, superblob_e8f4a19, hotfix_bug_1548973_1_1_4, timecop_1_0):
+    dbo.setDb("sqlite:///:memory:")
+    db_schema.create_all(dbo.engine)
+    dbo.rules.t.insert().execute(
+        rule_id=5,
+        priority=90,
+        backgroundRate=100,
+        mapping="Firefox-67.0-build1",
+        fallbackMapping="Firefox-54.0.1-build1",
+        update_type="minor",
+        product="Firefox",
+        mig64=True,
+        data_version=1,
+    )
+    dbo.rules.t.insert().execute(
+        priority=300,
+        product="SystemAddons",
+        channel="releasesjson",
+        mapping="Superblob-e8f4a19cfd695bf0eb66a2115313c31cc23a2369c0dc7b736d2f66d9075d7c66",
+        backgroundRate=100,
+        update_type="minor",
+        data_version=1,
+    )
+    dbo.releases.t.insert().execute(
+        name="Firefox-54.0.1-build1",
+        product="Firefox",
+        data_version=1,
+        data=createBlob(firefox_54_0_1_build1),
+    )
+    insert_release(firefox_67_0_build1, "Firefox", history=False)
+    insert_release(superblob_e8f4a19, "SystemAddons", history=False)
+    insert_release(hotfix_bug_1548973_1_1_4, "SystemAddons", history=False)
+    insert_release(timecop_1_0, "SystemAddons", history=False)
+    cache.reset()
+    cache.make_cache("blob", 50, 3600)
+    cache.make_cache("releases", 50, 3600)
+    cache.make_cache("release_assets", 50, 3600)
+
+    for cache_name in ("blob", "releases", "release_assets"):
+        c = cache.caches[cache_name]
+        assert c.lookups == 0, cache_name
+        assert c.hits == 0, cache_name
+        assert c.misses == 0, cache_name
+
+    auslib.web.public.helpers.warm_caches()
+
+    # one lookup per release, all misses
+    assert cache.caches["blob"].lookups == 1
+    assert cache.caches["blob"].hits == 0
+    assert cache.caches["blob"].misses == 1
+    for cache_name in ("releases", "release_assets"):
+        c = cache.caches[cache_name]
+        # There are 13 lookups here, which cover the following:
+        # - an attempt to look up 54.0.1 in this table                (1)
+        # - an attempt to look up the 3 releases referenced by 54.0.1 (4)
+        # - the superblob                                             (5)
+        # - the two releases referenced by the superblob              (7)
+        # - 67.0                                                      (8)
+        # - the 5 releases referenced by 67.0                         (13)
+        assert c.lookups == 13, cache_name
+        assert c.hits == 0, cache_name
+        assert c.misses == 13, cache_name
+
+    auslib.web.public.helpers.warm_caches()
+
+    # another lookup per release, all hits this time
+    assert cache.caches["blob"].lookups == 2
+    assert cache.caches["blob"].hits == 1
+    assert cache.caches["blob"].misses == 1
+    for cache_name in ("releases", "release_assets"):
+        c = cache.caches[cache_name]
+        assert c.lookups == 26, cache_name
+        assert c.hits == 13, cache_name
+        assert c.misses == 13, cache_name

uwsgi/public.wsgi

@@ -70,6 +70,7 @@ statsd.defaults.PREFIX = "balrog.public"
 
 from auslib.global_state import cache, dbo  # noqa
 from auslib.web.public.base import create_app
+from auslib.web.public.helpers import warm_caches
 
 application = create_app().app
 
@@ -133,3 +134,5 @@ if os.environ.get("CACHE_CONTROL"):
 
 if STAGING:
     application.config["SWAGGER_DEBUG"] = True
+
+warm_caches()