ndb global cache: limit to users and profile objects, not other objec…

…ts/activities

common.py

@@ -353,3 +353,16 @@ def report_exception(**kwargs):
             logger.warning(f'Failed to report error! {kwargs}', exc_info=True)
 
 
+PROFILE_ID_RE = re.compile(
+    fr"""
+      /users?/[^/]+$ |
+      /app.bsky.actor.profile/self$ |
+      ^did:[a-z0-9:.]+$ |
+      ^https://{DOMAIN_RE[1:-1]}/?$
+    """, re.VERBOSE)
+
+def global_cache_policy(key):
+    """Cache users and profile objects, not other objects or activities."""
+    return (key and
+            (key.kind in ('ActivityPub', 'ATProto', 'MagicKey')
+             or key.kind == 'Object' and PROFILE_ID_RE.search(key.name)))

flask_app.py

@@ -17,6 +17,8 @@
     flask_util,
 )
 
+from common import global_cache_policy
+
 logger = logging.getLogger(__name__)
 # logging.getLogger('lexrpc').setLevel(logging.INFO)
 logging.getLogger('negotiator').setLevel(logging.WARNING)
@@ -44,6 +46,7 @@
 app.wsgi_app = flask_util.ndb_context_middleware(
     app.wsgi_app, client=appengine_config.ndb_client,
     global_cache=_InProcessGlobalCache(),
+    global_cache_policy=global_cache_policy,
     global_cache_timeout_policy=lambda key: 3600,  # 1 hour
     # disable context-local cache since we're using a global in memory cache
     cache_policy=lambda key: False)

router.py

@@ -11,7 +11,7 @@
 
 # all protocols
 import activitypub, atproto, web
-from common import USER_AGENT
+from common import global_cache_policy, USER_AGENT
 import models
 import protocol
 
@@ -27,6 +27,7 @@
 app.wsgi_app = flask_util.ndb_context_middleware(
     app.wsgi_app, client=appengine_config.ndb_client,
     global_cache=_InProcessGlobalCache(),
+    global_cache_policy=global_cache_policy,
     global_cache_timeout_policy=lambda key: 3600,  # 1 hour
     # disable context-local cache since we're using a global in memory cache
     cache_policy=lambda key: False)

tests/test_common.py

@@ -5,7 +5,9 @@
 from activitypub import ActivityPub
 from atproto import ATProto
 import common
+from arroba.datastore_storage import AtpBlock
 from flask_app import app
+from models import Follower, Object
 from ui import UIProtocol
 from web import Web
 
@@ -103,3 +105,26 @@ def test_host_url(self):
 
         with app.test_request_context(base_url='https://bsky.brid.gy', path='/foo'):
             self.assertEqual('https://bsky.brid.gy/asdf', common.host_url('asdf'))
+
+    def test_global_cache_policy(self):
+        for good in (
+            ATProto(id='alice'),
+            ActivityPub(id='alice'),
+            Web(id='alice'),
+            Object(id='https://mastodon.social/users/alice'),
+            Object(id='https://mastodon.social/users/alice#main-key'),
+            Object(id='did:plc:foo'),
+            Object(id='did:web:foo.com'),
+            Object(id='at://did:plc:user/app.bsky.actor.profile/self'),
+        ):
+            self.assertTrue(common.global_cache_policy(good.key._key))
+
+        for bad in (
+            Follower(id='abc'),
+            Object(id='abc'),
+            Object(id='https://mastodon.social/users/alice/statuses/123'),
+            Object(id='at://did:plc:user/app.bsky.feed.post/abc'),
+            Object(id='https://web.site/post'),
+            AtpBlock(id='abc123'),
+        ):
+            self.assertFalse(common.global_cache_policy(bad.key._key))

tests/testutil.py

@@ -34,7 +34,7 @@
 
 # other modules are imported _after_ Fake etc classes is defined so that it's in
 # PROTOCOLS when URL routes are registered.
-from common import add, long_to_base64, TASKS_LOCATION
+from common import add, global_cache_policy, long_to_base64, TASKS_LOCATION
 import ids
 import models
 from models import KEY_BITS, Object, PROTOCOLS, Target, User
@@ -256,6 +256,7 @@ def setUp(self):
         requests.post(f'http://{ndb_client.host}/reset')
         self.ndb_context = ndb_client.context(
             global_cache=_InProcessGlobalCache(),
+            global_cache_policy=global_cache_policy,
             global_cache_timeout_policy=lambda key: 3600,  # 1 hour
             cache_policy=lambda key: False)
         self.ndb_context.__enter__()