diff --git a/CHANGELOG.md b/CHANGELOG.md index 0b5978dfb..2048ee46d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,7 @@ and this project adheres to - ♿️(frontend) use semantic `
` structure in document info card #2379 - 💄(frontend) use the same highlight color for cells and moves #2575 +- ⚡️(backend) optimize media_auth endpoint ### Fixed diff --git a/src/backend/core/api/utils.py b/src/backend/core/api/utils.py index 19cb03f3e..b44f486cf 100644 --- a/src/backend/core/api/utils.py +++ b/src/backend/core/api/utils.py @@ -13,6 +13,8 @@ import botocore from lasuite.oidc_login.decorators import refresh_oidc_access_token from rest_framework.throttling import BaseThrottle +from core.utils.s3 import get_s3_client, get_unsigned_s3_client + def nest_tree(flat_list, steplen): """ @@ -76,14 +78,14 @@ def generate_s3_authorization_headers(key): - access control is truly realtime - the object storage service does not need to be exposed on internet """ - url = default_storage.unsigned_connection.meta.client.generate_presigned_url( + url = get_unsigned_s3_client().generate_presigned_url( "get_object", ExpiresIn=0, Params={"Bucket": default_storage.bucket_name, "Key": key}, ) request = botocore.awsrequest.AWSRequest(method="get", url=url) - s3_client = default_storage.connection.meta.client + s3_client = get_s3_client() # pylint: disable=protected-access credentials = s3_client._request_signer._credentials # noqa: SLF001 frozen_credentials = credentials.get_frozen_credentials() diff --git a/src/backend/core/api/viewsets.py b/src/backend/core/api/viewsets.py index 28b2463d9..6efec4586 100644 --- a/src/backend/core/api/viewsets.py +++ b/src/backend/core/api/viewsets.py @@ -73,6 +73,7 @@ from core.tasks.mail import send_ask_for_access_mail from core.utils.analytics import PosthogEventName, posthog_capture from core.utils.dicts import lowercase_keys from core.utils.paths import filter_descendants +from core.utils.s3 import get_s3_client from core.utils.s3_response_stream import content_stream from core.utils.treebeard import create_tree_node_with_retry from core.utils.users import users_sharing_documents_with @@ -2043,8 +2044,12 @@ class DocumentViewSet( logger.debug("User '%s' lacks permission for attachment", user) raise drf.exceptions.PermissionDenied() - # Check if the attachment is ready - s3_client = default_storage.connection.meta.client + # Check if the attachment is ready. Use the process-global S3 client + # (see core.utils.s3): django-storages caches its client per thread, so + # relying on default_storage.connection here rebuilds the boto3 client + # on every fresh thread -- profiling showed that client construction, + # not the DB, dominated this endpoint's CPU under load. + s3_client = get_s3_client() bucket_name = default_storage.bucket_name try: head_resp = s3_client.head_object(Bucket=bucket_name, Key=key) diff --git a/src/backend/core/tests/documents/test_api_documents_media_auth.py b/src/backend/core/tests/documents/test_api_documents_media_auth.py index d3cccfcb6..ee06ed000 100644 --- a/src/backend/core/tests/documents/test_api_documents_media_auth.py +++ b/src/backend/core/tests/documents/test_api_documents_media_auth.py @@ -12,6 +12,7 @@ from django.utils import timezone import pytest import requests +from botocore.client import BaseClient from freezegun import freeze_time from rest_framework.test import APIClient diff --git a/src/backend/core/utils/s3.py b/src/backend/core/utils/s3.py new file mode 100644 index 000000000..411f92aa5 --- /dev/null +++ b/src/backend/core/utils/s3.py @@ -0,0 +1,52 @@ +"""Process-global boto3 S3 client accessors. + +django-storages caches its S3 connection on **thread-local** storage +(``S3Storage.connection`` / ``unsigned_connection``). Under a multi-threaded or +async-threadpool server, every thread that has not served this storage yet +rebuilds the boto3 client from scratch — and building a client makes botocore +load its service model off disk (millions of ``stat``/``listdir`` syscalls plus +a JSON decode of the model). Profiling the ``media-auth`` hot path showed that +this client construction, not PostgreSQL, dominated its CPU: ~2.6 fresh clients +were built per call, one for ``head_object`` and two more inside +``generate_s3_authorization_headers`` (signed + unsigned). + +boto3 *clients* are thread-safe once built (only *resources* are not, and we use +``.meta.client`` exclusively), and their attached credential provider still +refreshes rotating credentials on demand. So we can build each variant **once +per process** and share it across every thread and request, which removes the +per-request/per-thread rebuild entirely. + +The clients are captured from ``default_storage`` so they inherit its exact +configuration (endpoint, region, signature version, credentials, TLS). +""" + +import threading + +from django.core.files.storage import default_storage + +_LOCK = threading.Lock() +_CLIENTS = {} + + +def _cached(name, factory): + """Return the process-global client ``name``, building it once via factory.""" + client = _CLIENTS.get(name) + if client is None: + with _LOCK: + # Re-check inside the lock: another thread may have built it while we + # waited (double-checked locking). + client = _CLIENTS.get(name) + if client is None: + client = factory() + _CLIENTS[name] = client + return client + + +def get_s3_client(): + """Return a process-global signed S3 client (thread-safe, built once).""" + return _cached("signed", lambda: default_storage.connection.meta.client) + + +def get_unsigned_s3_client(): + """Return a process-global unsigned S3 client, for presigning URLs.""" + return _cached("unsigned", lambda: default_storage.unsigned_connection.meta.client)