The_Great_Archive/ArchiveBox
1
0
Fork 0
mirror of https://github.com/ArchiveBox/ArchiveBox.git synced 2026-09-12 19:50:57 +05:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

Keep OpenCode SSE startup off ASGI request threads

Browse Source
This commit is contained in:
Nick Sweeting 2026-09-02 22:02:03 -07:00
parent b44ab723d2
commit 75cf3f1261
No known key found for this signature in database
3 changed files with 74 additions and 68 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

59
archivebox/core/asgi.py
View File

@ -13,64 +13,5 @@ from archivebox.config.django import setup_django
setup_django(check_db=True)
def _patch_thread_sensitive_context_shutdown() -> None:
"""Stop ``ThreadSensitiveContext.__aexit__`` from blocking the daphne loop.
Django 6.0's ASGIHandler wraps every request in ``async with
ThreadSensitiveContext():`` (django/core/handlers/asgi.py:169). On exit
asgiref calls ``executor.shutdown()`` with the default ``wait=True``
(asgiref/sync.py:148), which is a *synchronous* ``Thread.join()`` inside
an async function — so it blocks the daphne event loop until the
executor's worker thread exits.
That's normally fine because the request handler has already awaited
every ``sync_to_async`` it submitted, so the worker is idle and dies as
soon as the shutdown sentinel reaches it. The blocking turns into a
problem when a client disconnects mid-request:
* ``SyncToAsync.__call__`` shields the executor work with
``await asyncio.shield(exec_coro)`` (asgiref/sync.py:506) so that the
sync DB call doesn't get torn down halfway through.
* On cancellation it calls ``exec_coro.cancel()`` (line 522) which only
flips the asyncio ``Future`` to cancelled — the underlying thread
keeps running the SQL query.
* Control unwinds to ``__aexit__`` while the orphaned thread is still
mid-query. ``shutdown(wait=True)`` then blocks the event loop until
that orphan finishes.
Under heavy SQLite contention (the load-test scenario that surfaced
this on cabbage) those orphan threads can take 30 seconds each waiting
for write locks, and the daphne loop is single-threaded — so every
such orphan stalls every other in-flight request, healthchecks time
out, and the container goes ``unhealthy``.
Switching to ``shutdown(wait=False)`` queues the sentinel and returns
immediately; the worker thread still exits cleanly once its current
task finishes, and asgiref's ``WeakKeyDictionary`` releases the
executor as soon as the request's context is GC'd. No per-request
teardown guarantee is lost — there was no caller relying on it.
"""
from asgiref import sync as _asgiref_sync
original_aexit = _asgiref_sync.ThreadSensitiveContext.__aexit__
async def __aexit__(self, exc, value, tb): # type: ignore[no-redef]
if not self.token:
return
executor = _asgiref_sync.SyncToAsync.context_to_thread_executor.pop(self, None)
if executor is not None:
executor.shutdown(wait=False)
_asgiref_sync.SyncToAsync.thread_sensitive_context.reset(self.token)
# Idempotent: only patch once even if asgi.py is reloaded.
if getattr(original_aexit, "_archivebox_patched", False):
return
__aexit__._archivebox_patched = True # type: ignore[attr-defined]
_asgiref_sync.ThreadSensitiveContext.__aexit__ = __aexit__
_patch_thread_sensitive_context_shutdown()
# Standard Django ASGI application (no websockets/channels needed)
application = get_asgi_application()

35
archivebox/opencode/views.py
View File

@ -20,6 +20,7 @@ from abx_plugins.plugins import opencode as opencode_plugin
from archivebox.config import CONSTANTS
from archivebox.config.common import get_config
from archivebox.core.routes_util import build_admin_url, get_api_base_url, get_base_url
from asgiref.sync import sync_to_async
from django.http import (
Http404,
HttpRequest,
@ -555,17 +556,27 @@ def _request_params(request: HttpRequest) -> tuple[tuple[str, str], ...]:
return tuple((key, str(value)) for key, value in dict(request.GET).items())
async def _event_chunks(request: HttpRequest, settings: dict, path: str | None):
async def _event_chunks(
settings: dict,
path: str | None,
method: str,
params: tuple[tuple[str, str], ...],
headers: dict[str, str],
):
ok, error = await sync_to_async(_ensure_opencode, thread_sensitive=False)(settings)
if not ok:
_LOGGER.warning("OpenCode event stream unavailable: %s", error)
return
timeout = httpx.Timeout(settings["timeout"], read=None)
url = _proxy_url(settings, path)
method = request.method or "GET"
try:
async with httpx.AsyncClient(timeout=timeout, follow_redirects=False) as client:
async with client.stream(
method,
url,
params=_request_params(request),
headers=_request_headers(request, settings),
params=params,
headers=headers,
) as upstream:
async for chunk in upstream.aiter_raw(chunk_size=512):
yield chunk
@ -662,19 +673,25 @@ def opencode_proxy_view(request: HttpRequest, path: str | None = None):
settings["archivebox_admin_url"] = admin_url
settings["archivebox_api_url"] = api_url
ok, error = _ensure_opencode(settings)
if not ok:
return _proxy_error_response(error)
if request.method == "GET" and (path or "").endswith("/event"):
response = StreamingHttpResponse(
_event_chunks(request, settings, path),
_event_chunks(
settings,
path,
request.method or "GET",
_request_params(request),
_request_headers(request, settings),
),
content_type="text/event-stream",
)
response.headers["Cache-Control"] = "no-store"
response.headers["X-Accel-Buffering"] = "no"
return response
ok, error = _ensure_opencode(settings)
if not ok:
return _proxy_error_response(error)
try:
method = request.method or "GET"
upstream = requests.request(

48
archivebox/tests/test_opencode_agent.py
View File

@ -1,3 +1,4 @@
import asyncio
import os
import socket
import subprocess
@ -8,6 +9,7 @@ from urllib.parse import quote
import pytest
import requests
from asgiref.testing import ApplicationCommunicator
from archivebox.tests.conftest import ADMIN_TEST_HOST, run_archivebox_cmd
@ -268,10 +270,56 @@ def test_opencode_proxy_sse_response_is_unbuffered(admin_client, live_opencode):
assert response.status_code == 200
assert response.streaming
assert response.is_async
assert response.headers["X-Accel-Buffering"] == "no"
assert response.headers["Cache-Control"] == "no-store"
def test_opencode_proxy_sse_returns_headers_before_restart_finishes(admin_client, live_opencode):
from archivebox.core.asgi import application
from archivebox.opencode import views
from django.conf import settings as django_settings
views._stop_owned_process()
session_cookie_name = django_settings.SESSION_COOKIE_NAME
session_cookie = admin_client.cookies[session_cookie_name].value
path = "/admin/agent/opencode/global/event"
async def request_event_stream():
communicator = ApplicationCommunicator(
application,
{
"type": "http",
"asgi": {"version": "3.0"},
"http_version": "1.1",
"method": "GET",
"scheme": "http",
"path": path,
"raw_path": path.encode(),
"query_string": b"",
"headers": [
(b"host", ADMIN_TEST_HOST.encode()),
(b"cookie", f"{session_cookie_name}={session_cookie}".encode()),
(b"sec-fetch-site", b"same-origin"),
],
"client": ("127.0.0.1", 12345),
"server": ("127.0.0.1", 8000),
},
)
views._PROCESS_LOCK.acquire()
try:
await communicator.send_input({"type": "http.request", "body": b"", "more_body": False})
response_start = await communicator.receive_output(timeout=2)
assert response_start["type"] == "http.response.start"
assert response_start["status"] == 200
finally:
views._PROCESS_LOCK.release()
await communicator.send_input({"type": "http.disconnect"})
await communicator.wait(timeout=5)
asyncio.run(request_event_stream())
def test_opencode_starts_with_isolated_state(live_opencode):
workdir = str(live_opencode.config.data_dir.resolve())
state_dir = live_opencode.config.state_dir