Improve caching (HTTP sig and thumbnails)

app/httpsig.py

@@ -10,6 +10,7 @@ from dataclasses import dataclass
 from datetime import datetime
 from typing import Any
 from typing import Dict
+from typing import MutableMapping
 from typing import Optional
 
 import fastapi
@@ -27,7 +28,7 @@ from app.database import get_db_session
 from app.key import Key
 from app.key import get_key
 
-_KEY_CACHE = LFUCache(256)
+_KEY_CACHE: MutableMapping[str, Key] = LFUCache(256)
 
 
 def _build_signed_string(
@@ -73,6 +74,7 @@ async def _get_public_key(db_session: AsyncSession, key_id: str) -> Key:
 
     # Check if the key belongs to an actor already in DB
     from app import models
+
     existing_actor = (
         await db_session.scalars(
             select(models.Actor).where(models.Actor.ap_id == key_id.split("#")[0])

app/main.py

@@ -5,6 +5,7 @@ import time
 from datetime import timezone
 from io import BytesIO
 from typing import Any
+from typing import MutableMapping
 from typing import Type
 
 import httpx
@@ -57,7 +58,7 @@ from app.utils import pagination
 from app.utils.emoji import EMOJIS_BY_NAME
 from app.webfinger import get_remote_follow_template
 
-_RESIZED_CACHE = LFUCache(32)
+_RESIZED_CACHE: MutableMapping[tuple[str, int], tuple[bytes, str, Any]] = LFUCache(32)
 
 
 # TODO(ts):
@@ -743,17 +744,13 @@ async def serve_proxy_media_resized(
     # Decode the base64-encoded URL
     url = base64.urlsafe_b64decode(encoded_url).decode()
 
-    is_cached = False
-    is_resized = False
     if cached_resp := _RESIZED_CACHE.get((url, size)):
-        is_resized, resized_content, resized_mimetype, resp_headers = cached_resp
-        if is_resized:
-            return PlainTextResponse(
-                resized_content,
-                media_type=resized_mimetype,
-                headers=resp_headers,
-            )
-        is_cached = True
+        resized_content, resized_mimetype, resp_headers = cached_resp
+        return PlainTextResponse(
+            resized_content,
+            media_type=resized_mimetype,
+            headers=resp_headers,
+        )
 
     # Request the URL (and filter request headers)
     async with httpx.AsyncClient() as client:
@@ -773,7 +770,7 @@ async def serve_proxy_media_resized(
             ]
             + [(b"user-agent", USER_AGENT.encode())],
         )
-    if proxy_resp.status_code != 200 or (is_cached and not is_resized):
+    if proxy_resp.status_code != 200:
         return PlainTextResponse(
             proxy_resp.content,
             status_code=proxy_resp.status_code,
@@ -804,12 +801,14 @@ async def serve_proxy_media_resized(
         resized_buf.seek(0)
         resized_content = resized_buf.read()
         resized_mimetype = i.get_format_mimetype()  # type: ignore
-        _RESIZED_CACHE[(url, size)] = (
-            True,
-            resized_content,
-            resized_mimetype,
-            proxy_resp_headers,
-        )
+
+        # Only cache images < 1MB
+        if len(resized_content) < 2**20:
+            _RESIZED_CACHE[(url, size)] = (
+                resized_content,
+                resized_mimetype,
+                proxy_resp_headers,
+            )
         return PlainTextResponse(
             resized_content,
             media_type=resized_mimetype,

poetry.lock

@@ -1014,6 +1014,14 @@ category = "dev"
 optional = false
 python-versions = "*"
 
+[[package]]
+name = "types-cachetools"
+version = "5.2.1"
+description = "Typing stubs for cachetools"
+category = "dev"
+optional = false
+python-versions = "*"
+
 [[package]]
 name = "types-emoji"
 version = "1.7.2"
@@ -1143,7 +1151,7 @@ dev = ["pytest (>=4.6.2)", "black (>=19.3b0)"]
 [metadata]
 lock-version = "1.1"
 python-versions = "^3.10"
-content-hash = "91e35a13d21bb5fd3e8916aee95c0a8019bec3cf4f0c677bb86641f1d88dcfe3"
+content-hash = "cbfda21eb816f33407cd124db1ed037e2b1da7cdd72247de5793ba90e52ee648"
 
 [metadata.files]
 aiosqlite = [
@@ -1919,6 +1927,10 @@ types-bleach = [
     {file = "types-bleach-5.0.2.tar.gz", hash = "sha256:e1498c512a62117496cf82be3d129972bb89fd1d6482b001cdeb2759ab3c82f5"},
     {file = "types_bleach-5.0.2-py3-none-any.whl", hash = "sha256:6fcb75ee4b69190fe60340147b66442cecddaefe3c0629433a4240da1ec2dcf6"},
 ]
+types-cachetools = [
+    {file = "types-cachetools-5.2.1.tar.gz", hash = "sha256:069cfc825697cd51445c1feabbe4edc1fae2b2315870e7a9a179a7c4a5851bee"},
+    {file = "types_cachetools-5.2.1-py3-none-any.whl", hash = "sha256:b496b7e364ba050c4eaadcc6582f2c9fbb04f8ee7141eb3b311a8589dbd4506a"},
+]
 types-emoji = [
     {file = "types-emoji-1.7.2.tar.gz", hash = "sha256:a7660fb507b30cb80bcec2d01417d828f1258b9b2cd9fa80918e8e5470c5e037"},
     {file = "types_emoji-1.7.2-py3-none-any.whl", hash = "sha256:f4c18bb43e33dc267c650b73d7ae0cd71708c75c79063706d0b91fa9416190c8"},

pyproject.toml

@@ -59,6 +59,7 @@ factory-boy = "^3.2.1"
 pytest-asyncio = "^0.18.3"
 types-Pillow = "^9.0.20"
 types-emoji = "^1.7.2"
+types-cachetools = "^5.2.1"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

tests/test_httpsig.py

@@ -8,6 +8,8 @@ from fastapi.testclient import TestClient
 
 from app import activitypub as ap
 from app import httpsig
+from app.database import AsyncSession
+from app.httpsig import _KEY_CACHE
 from app.httpsig import HTTPSigInfo
 from app.key import Key
 from tests import factories
@@ -56,6 +58,7 @@ def test_enforce_httpsig__no_signature(
 @pytest.mark.asyncio
 async def test_enforce_httpsig__with_valid_signature(
     respx_mock: respx.MockRouter,
+    async_db_session: AsyncSession,
 ) -> None:
     # Given a remote actor
     privkey, pubkey = factories.generate_key()
@@ -69,7 +72,7 @@ async def test_enforce_httpsig__with_valid_signature(
     auth = httpsig.HTTPXSigAuth(k)
     respx_mock.get(ra.ap_id).mock(return_value=httpx.Response(200, json=ra.ap_actor))
 
-    httpsig._get_public_key.cache_clear()
+    _KEY_CACHE.clear()
 
     async with httpx.AsyncClient(app=_test_app, base_url="http://test") as client:
         response = await client.post(
@@ -105,6 +108,7 @@ def test_httpsig_checker__no_signature(
 @pytest.mark.asyncio
 async def test_httpsig_checker__with_valid_signature(
     respx_mock: respx.MockRouter,
+    async_db_session: AsyncSession,
 ) -> None:
     # Given a remote actor
     privkey, pubkey = factories.generate_key()
@@ -118,7 +122,7 @@ async def test_httpsig_checker__with_valid_signature(
     k.load(privkey)
     auth = httpsig.HTTPXSigAuth(k)
 
-    httpsig._get_public_key.cache_clear()
+    _KEY_CACHE.clear()
 
     async with httpx.AsyncClient(app=_test_app, base_url="http://test") as client:
         response = await client.get(
@@ -137,6 +141,7 @@ async def test_httpsig_checker__with_valid_signature(
 @pytest.mark.asyncio
 async def test_httpsig_checker__with_invvalid_signature(
     respx_mock: respx.MockRouter,
+    async_db_session: AsyncSession,
 ) -> None:
     # Given a remote actor
     privkey, pubkey = factories.generate_key()
@@ -158,7 +163,7 @@ async def test_httpsig_checker__with_invvalid_signature(
     assert ra.ap_id == ra2.ap_id
     respx_mock.get(ra.ap_id).mock(return_value=httpx.Response(200, json=ra2.ap_actor))
 
-    httpsig._get_public_key.cache_clear()
+    _KEY_CACHE.clear()
 
     async with httpx.AsyncClient(app=_test_app, base_url="http://test") as client:
         response = await client.get(