From 1f10c3367f2e0177d995ccc11498fc5f8ced0a0c Mon Sep 17 00:00:00 2001 From: Thomas Sileo Date: Thu, 30 Jun 2022 09:43:28 +0200 Subject: [PATCH] Improve caching (HTTP sig and thumbnails) --- app/httpsig.py | 4 +++- app/main.py | 35 +++++++++++++++++------------------ poetry.lock | 14 +++++++++++++- pyproject.toml | 1 + tests/test_httpsig.py | 11 ++++++++--- 5 files changed, 42 insertions(+), 23 deletions(-) diff --git a/app/httpsig.py b/app/httpsig.py index 713087a..3a10c88 100644 --- a/app/httpsig.py +++ b/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]) diff --git a/app/main.py b/app/main.py index 930ef8d..1467d56 100644 --- a/app/main.py +++ b/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, diff --git a/poetry.lock b/poetry.lock index 5768f8d..1457be5 100644 --- a/poetry.lock +++ b/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"}, diff --git a/pyproject.toml b/pyproject.toml index 1f7be3f..62212af 100644 --- a/pyproject.toml +++ b/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"] diff --git a/tests/test_httpsig.py b/tests/test_httpsig.py index 1538ce1..a36775a 100644 --- a/tests/test_httpsig.py +++ b/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(