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 datetime import datetime
from typing import Any from typing import Any
from typing import Dict from typing import Dict
from typing import MutableMapping
from typing import Optional from typing import Optional
import fastapi import fastapi
@ -27,7 +28,7 @@ from app.database import get_db_session
from app.key import Key from app.key import Key
from app.key import get_key from app.key import get_key
_KEY_CACHE = LFUCache(256) _KEY_CACHE: MutableMapping[str, Key] = LFUCache(256)
def _build_signed_string( 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 # Check if the key belongs to an actor already in DB
from app import models from app import models
existing_actor = ( existing_actor = (
await db_session.scalars( await db_session.scalars(
select(models.Actor).where(models.Actor.ap_id == key_id.split("#")[0]) 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 datetime import timezone
from io import BytesIO from io import BytesIO
from typing import Any from typing import Any
from typing import MutableMapping
from typing import Type from typing import Type
import httpx import httpx
@ -57,7 +58,7 @@ from app.utils import pagination
from app.utils.emoji import EMOJIS_BY_NAME from app.utils.emoji import EMOJIS_BY_NAME
from app.webfinger import get_remote_follow_template 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): # TODO(ts):
@ -743,17 +744,13 @@ async def serve_proxy_media_resized(
# Decode the base64-encoded URL # Decode the base64-encoded URL
url = base64.urlsafe_b64decode(encoded_url).decode() url = base64.urlsafe_b64decode(encoded_url).decode()
is_cached = False
is_resized = False
if cached_resp := _RESIZED_CACHE.get((url, size)): if cached_resp := _RESIZED_CACHE.get((url, size)):
is_resized, resized_content, resized_mimetype, resp_headers = cached_resp resized_content, resized_mimetype, resp_headers = cached_resp
if is_resized:
return PlainTextResponse( return PlainTextResponse(
resized_content, resized_content,
media_type=resized_mimetype, media_type=resized_mimetype,
headers=resp_headers, headers=resp_headers,
) )
is_cached = True
# Request the URL (and filter request headers) # Request the URL (and filter request headers)
async with httpx.AsyncClient() as client: async with httpx.AsyncClient() as client:
@ -773,7 +770,7 @@ async def serve_proxy_media_resized(
] ]
+ [(b"user-agent", USER_AGENT.encode())], + [(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( return PlainTextResponse(
proxy_resp.content, proxy_resp.content,
status_code=proxy_resp.status_code, status_code=proxy_resp.status_code,
@ -804,8 +801,10 @@ async def serve_proxy_media_resized(
resized_buf.seek(0) resized_buf.seek(0)
resized_content = resized_buf.read() resized_content = resized_buf.read()
resized_mimetype = i.get_format_mimetype() # type: ignore resized_mimetype = i.get_format_mimetype() # type: ignore
# Only cache images < 1MB
if len(resized_content) < 2**20:
_RESIZED_CACHE[(url, size)] = ( _RESIZED_CACHE[(url, size)] = (
True,
resized_content, resized_content,
resized_mimetype, resized_mimetype,
proxy_resp_headers, proxy_resp_headers,

14
poetry.lock generated
View file

@ -1014,6 +1014,14 @@ category = "dev"
optional = false optional = false
python-versions = "*" python-versions = "*"
[[package]]
name = "types-cachetools"
version = "5.2.1"
description = "Typing stubs for cachetools"
category = "dev"
optional = false
python-versions = "*"
[[package]] [[package]]
name = "types-emoji" name = "types-emoji"
version = "1.7.2" version = "1.7.2"
@ -1143,7 +1151,7 @@ dev = ["pytest (>=4.6.2)", "black (>=19.3b0)"]
[metadata] [metadata]
lock-version = "1.1" lock-version = "1.1"
python-versions = "^3.10" python-versions = "^3.10"
content-hash = "91e35a13d21bb5fd3e8916aee95c0a8019bec3cf4f0c677bb86641f1d88dcfe3" content-hash = "cbfda21eb816f33407cd124db1ed037e2b1da7cdd72247de5793ba90e52ee648"
[metadata.files] [metadata.files]
aiosqlite = [ aiosqlite = [
@ -1919,6 +1927,10 @@ types-bleach = [
{file = "types-bleach-5.0.2.tar.gz", hash = "sha256:e1498c512a62117496cf82be3d129972bb89fd1d6482b001cdeb2759ab3c82f5"}, {file = "types-bleach-5.0.2.tar.gz", hash = "sha256:e1498c512a62117496cf82be3d129972bb89fd1d6482b001cdeb2759ab3c82f5"},
{file = "types_bleach-5.0.2-py3-none-any.whl", hash = "sha256:6fcb75ee4b69190fe60340147b66442cecddaefe3c0629433a4240da1ec2dcf6"}, {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 = [ types-emoji = [
{file = "types-emoji-1.7.2.tar.gz", hash = "sha256:a7660fb507b30cb80bcec2d01417d828f1258b9b2cd9fa80918e8e5470c5e037"}, {file = "types-emoji-1.7.2.tar.gz", hash = "sha256:a7660fb507b30cb80bcec2d01417d828f1258b9b2cd9fa80918e8e5470c5e037"},
{file = "types_emoji-1.7.2-py3-none-any.whl", hash = "sha256:f4c18bb43e33dc267c650b73d7ae0cd71708c75c79063706d0b91fa9416190c8"}, {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" pytest-asyncio = "^0.18.3"
types-Pillow = "^9.0.20" types-Pillow = "^9.0.20"
types-emoji = "^1.7.2" types-emoji = "^1.7.2"
types-cachetools = "^5.2.1"
[build-system] [build-system]
requires = ["poetry-core>=1.0.0"] 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 activitypub as ap
from app import httpsig from app import httpsig
from app.database import AsyncSession
from app.httpsig import _KEY_CACHE
from app.httpsig import HTTPSigInfo from app.httpsig import HTTPSigInfo
from app.key import Key from app.key import Key
from tests import factories from tests import factories
@ -56,6 +58,7 @@ def test_enforce_httpsig__no_signature(
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_enforce_httpsig__with_valid_signature( async def test_enforce_httpsig__with_valid_signature(
respx_mock: respx.MockRouter, respx_mock: respx.MockRouter,
async_db_session: AsyncSession,
) -> None: ) -> None:
# Given a remote actor # Given a remote actor
privkey, pubkey = factories.generate_key() privkey, pubkey = factories.generate_key()
@ -69,7 +72,7 @@ async def test_enforce_httpsig__with_valid_signature(
auth = httpsig.HTTPXSigAuth(k) auth = httpsig.HTTPXSigAuth(k)
respx_mock.get(ra.ap_id).mock(return_value=httpx.Response(200, json=ra.ap_actor)) 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: async with httpx.AsyncClient(app=_test_app, base_url="http://test") as client:
response = await client.post( response = await client.post(
@ -105,6 +108,7 @@ def test_httpsig_checker__no_signature(
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_httpsig_checker__with_valid_signature( async def test_httpsig_checker__with_valid_signature(
respx_mock: respx.MockRouter, respx_mock: respx.MockRouter,
async_db_session: AsyncSession,
) -> None: ) -> None:
# Given a remote actor # Given a remote actor
privkey, pubkey = factories.generate_key() privkey, pubkey = factories.generate_key()
@ -118,7 +122,7 @@ async def test_httpsig_checker__with_valid_signature(
k.load(privkey) k.load(privkey)
auth = httpsig.HTTPXSigAuth(k) 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: async with httpx.AsyncClient(app=_test_app, base_url="http://test") as client:
response = await client.get( response = await client.get(
@ -137,6 +141,7 @@ async def test_httpsig_checker__with_valid_signature(
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_httpsig_checker__with_invvalid_signature( async def test_httpsig_checker__with_invvalid_signature(
respx_mock: respx.MockRouter, respx_mock: respx.MockRouter,
async_db_session: AsyncSession,
) -> None: ) -> None:
# Given a remote actor # Given a remote actor
privkey, pubkey = factories.generate_key() privkey, pubkey = factories.generate_key()
@ -158,7 +163,7 @@ async def test_httpsig_checker__with_invvalid_signature(
assert ra.ap_id == ra2.ap_id assert ra.ap_id == ra2.ap_id
respx_mock.get(ra.ap_id).mock(return_value=httpx.Response(200, json=ra2.ap_actor)) 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: async with httpx.AsyncClient(app=_test_app, base_url="http://test") as client:
response = await client.get( response = await client.get(