feat: warm caches during app startup (#3500)

* feat: implement getReferencedReleases for SuperBlob

This will be useful when warming caches.

* refactor: add helper to find release regardless of which table it is in

* feat: warm caches during app startup

This code runs before the public app begins accepting requests, which should mean that we don't need to hit the database very much after beginning to serve requests. Doing this at start-up also means we fetch the releases sequentially, which should (slightly) spread the load out more.

I verified locally that the app is incapable of responding to requests until `public.py` finishes executing by adding a `time.sleep` to the end of it. While the sleep was running requests to the app hung. This means that `/__heartbeat__` requests won't complete successfully until caches have been warmed up.

Hopefully this will help with database load during scaling events. See: https://mozilla-hub.atlassian.net/browse/SVCSE-3425

Author: bhearsum

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()