Improve caching (HTTP sig and thumbnails)

This commit is contained in:
Thomas Sileo 2022-06-30 09:43:28 +02:00
parent 6458d2a6c7
commit 1f10c3367f
5 changed files with 42 additions and 23 deletions

View file

@ -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])

View file

@ -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:
resized_content, resized_mimetype, resp_headers = cached_resp
return PlainTextResponse(
resized_content,
media_type=resized_mimetype,
headers=resp_headers,
)
is_cached = True
# 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,8 +801,10 @@ async def serve_proxy_media_resized(
resized_buf.seek(0)
resized_content = resized_buf.read()
resized_mimetype = i.get_format_mimetype() # type: ignore
# Only cache images < 1MB
if len(resized_content) < 2**20:
_RESIZED_CACHE[(url, size)] = (
True,
resized_content,
resized_mimetype,
proxy_resp_headers,

14
poetry.lock generated
View file

@ -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"},

View file

@ -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"]

View file

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