vvzvlad/pyrogram-bridge
1
0
Fork 0
Code Issues 6 Pull Requests Actions 8 Packages Projects Releases Wiki Activity

perf(stability): batch SQLite access-time writes out of the /media hot path (#5)

Browse Source 
Stage 5. A /media cache hit no longer touches SQLite: it records the access
timestamp into a module-level accumulator (a dict write on the event loop,
cheap and atomic), and a supervised 60s background task flushes the whole
batch in one executemany UPDATE. This removes both per-hit write sites — the
awaited to_thread in download_media_file and the fire-and-forget create_task
in the pre-semaphore fast path — so under active RSS polling the threadpool is
no longer starved by per-request access-time UPDATEs.

- file_io: add update_media_file_access_bulk_sync (one connection, executemany;
  empty batch is a no-op).
- api_server: _access_updates accumulator + _flush_access_updates (snapshot-
  then-clear atomically before the await so writes during the flush land in the
  fresh dict; re-queue the batch with setdefault on write failure so a fresher
  concurrent write is never clobbered and no access-time is lost) +
  _access_flush_loop under _supervised + a final flush on shutdown, ordered
  after the loop task is cancelled and before the io threadpool is shut down.
- Keys use str(channel) to match the TEXT channel column (a str/int mix would
  make the UPDATE WHERE silently never match, evicting still-used files).
- tests/test_stage5_sqlite.py: 7 tests (no-sync-write hot path, str-key
  discipline, hit->flush->DB, empty no-op, snapshot-then-clear race, re-queue-
  without-clobbering-fresh, bulk SQL).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
agent_coder
2026-07-05 10:04:47 +03:00
parent 8d21390294
commit a04588740b
3 changed files with 323 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api_server.py
+73 -21
View File
@@ -40,7 +40,8 @@ from rss_generator import generate_channel_rss, generate_channel_html
from post_parser import PostParser
from url_signer import verify_media_digest, generate_media_digest
from file_io import (DB_PATH, init_db_sync, get_all_media_file_ids_sync,
update_media_file_access_sync, remove_media_file_ids_sync,
update_media_file_access_sync, update_media_file_access_bulk_sync,
remove_media_file_ids_sync,
get_mime_type_sync, set_mime_type_sync)
# Global python-magic instance for MIME type detection
@@ -139,6 +140,57 @@ async def _supervised(factory, name: str, min_restart_interval: float = 60.0):
await asyncio.sleep(min_restart_interval - elapsed)
# Access-time write accumulator. A /media cache hit used to touch SQLite on the hot path
# (a threadpool hop + connect + UPDATE per request), which starves the threadpool under
# active RSS polling. Instead a cache hit just records the timestamp here — a dict write
# on the single-threaded event loop is cheap and atomic — and a periodic background task
# flushes the whole batch to SQLite in one executemany. Keys use str(channel) to match the
# TEXT channel column; mixing str/int would make the UPDATE's WHERE silently never match,
# so the access-time would stop advancing and the file would fall out of the 20-day cache.
ACCESS_FLUSH_INTERVAL = 60 # seconds between access-time flushes
_access_updates: dict[tuple[str, int, str], float] = {}
async def _flush_access_updates() -> None:
"""Flush the accumulated access timestamps to SQLite in one bulk UPDATE.
Snapshot-then-clear atomically on the loop: capture the current dict reference and
replace the module global with a fresh empty dict in ONE synchronous step (before any
await), so cache-hit writes arriving DURING the flush land in the new dict and are not
lost. The bulk UPDATE runs off-loop via asyncio.to_thread. An empty batch is a no-op.
"""
global _access_updates
if not _access_updates:
return
pending = _access_updates
_access_updates = {}
entries = [(channel, post_id, file_unique_id, added)
for (channel, post_id, file_unique_id), added in pending.items()]
try:
await asyncio.to_thread(update_media_file_access_bulk_sync, DB_PATH, entries)
except Exception:
# Bulk write failed: re-queue this batch so the access-times are not lost (a lost
# timestamp would eventually evict a still-used file from the 20-day cache). Use
# setdefault so any FRESHER write accumulated during the flush is never overwritten
# by our stale snapshot. Runs on the loop with no await before the mutation, so this
# is race-free. Re-raise so the flush loop logs it.
for key, added in pending.items():
_access_updates.setdefault(key, added)
raise
async def _access_flush_loop() -> None:
"""Periodically flush the access-time accumulator (runs under _supervised)."""
while True:
await asyncio.sleep(ACCESS_FLUSH_INTERVAL)
try:
await _flush_access_updates()
except Exception as e:
# Log and keep looping: a transient SQLite error must not drop the batch's
# successors. (_supervised still restarts us if this ever raises out.)
logger.error(f"access_flush_error: {e}")
@asynccontextmanager
async def lifespan(_: FastAPI):
setup_logging(Config["log_level"])
@@ -161,9 +213,11 @@ async def lifespan(_: FastAPI):
# CRITICAL and restarted, so a crash can no longer silently stop cache sweeping or downloads.
background_task = asyncio.create_task(_supervised(cache_media_files, "cache_media_files"))
worker_task = asyncio.create_task(_supervised(background_download_worker, "background_download_worker"))
access_flush_task = asyncio.create_task(_supervised(_access_flush_loop, "access_flush_loop"))
yield
background_task.cancel() # Cleanup
worker_task.cancel()
access_flush_task.cancel()
try:
await background_task
except asyncio.CancelledError:
@@ -172,6 +226,17 @@ async def lifespan(_: FastAPI):
await worker_task
except asyncio.CancelledError:
pass
try:
await access_flush_task
except asyncio.CancelledError:
pass
# Final flush AFTER the loop task is cancelled (no race with a loop-driven flush) and
# BEFORE the threadpool is shut down (to_thread still has its executor), so the last
# <=ACCESS_FLUSH_INTERVAL seconds of access-times are persisted on shutdown.
try:
await _flush_access_updates()
except Exception as e:
logger.error(f"access_flush_shutdown_error: {e}")
await client.stop()
# Shut the io threadpool down so its threads don't linger past a reload/restart.
io_executor.shutdown(wait=False)
@@ -538,16 +603,11 @@ async def download_media_file(channel: Union[str, int], post_id: int, file_uniqu
logger.error(f"cleanup_error: Failed to remove zero-size cached file {cache_path}: {e}")
# Do not raise error here, proceed to download below
else:
# File exists and is not zero size, update access timestamp and return
# File exists and is not zero size, record access timestamp and return.
# Record into the accumulator instead of touching SQLite on the hot path; the
# background flush persists it. Key channel as str(channel) — see _access_updates.
logger.info(f"Found cached media file: {cache_path}")
try:
await asyncio.to_thread(
update_media_file_access_sync,
DB_PATH, str(channel), post_id, file_unique_id,
datetime.now().timestamp()
)
except Exception as e:
logger.error(f"Failed to update timestamp for {channel}/{post_id}/{file_unique_id}: {str(e)}")
_access_updates[(str(channel), post_id, file_unique_id)] = datetime.now().timestamp()
return cache_path, False
file_id = await find_file_id_in_message(message, file_unique_id)
@@ -1074,17 +1134,9 @@ async def get_media(channel: str, post_id: int, file_unique_id: str, request: Re
if os.path.exists(cache_path) and os.path.getsize(cache_path) > 0:
# File is already in cache — skip semaphore and serve directly
logger.info(f"pre_semaphore_cache_hit: {channel}/{post_id}/{file_unique_id}")
# Fire-and-forget timestamp update with error handling to avoid silent failures
async def _update_access(_ch, _pid, _fid):
try:
await asyncio.to_thread(
update_media_file_access_sync,
DB_PATH, str(_ch), _pid, _fid,
datetime.now().timestamp()
)
except Exception as _e:
logger.warning(f"Failed to update access time for {_ch}/{_pid}/{_fid}: {_e}")
asyncio.create_task(_update_access(channel, post_id, file_unique_id))
# Record the access time into the accumulator instead of firing a per-hit
# SQLite write. Key channel as str(channel) — see _access_updates.
_access_updates[(str(channel), post_id, file_unique_id)] = datetime.now().timestamp()
return await prepare_file_response(cache_path, request=request,
media_key=(str(channel), post_id, file_unique_id))
file_io.py
+21
View File
@@ -96,6 +96,27 @@ def update_media_file_access_sync(db_path: str, channel: str, post_id: int, file
)
def update_media_file_access_bulk_sync(db_path: str, entries: List[tuple]) -> None:
"""Update the access timestamp for multiple existing media file ID records.
entries: iterable of (channel, post_id, file_unique_id, added) tuples.
Uses executemany (one connection, one commit) so a batch of cache-hit access
updates costs a single SQLite transaction instead of one connect+UPDATE per hit.
Rows that do not exist are simply not matched by the WHERE clause (no-op), mirroring
the single-row update_media_file_access_sync. An empty batch is a no-op.
"""
if not entries:
return
with _db_connection(db_path) as conn:
conn.executemany(
"UPDATE media_file_ids SET added = ? WHERE channel = ? AND post_id = ? AND file_unique_id = ?",
# Reorder each (channel, post_id, file_unique_id, added) tuple to match the
# UPDATE's placeholder order (added first, then the WHERE key columns).
[(added, channel, post_id, file_unique_id)
for (channel, post_id, file_unique_id, added) in entries],
)
def get_all_media_file_ids_sync(db_path: str) -> List[dict]:
"""Return all rows from media_file_ids as a list of dicts."""
with _db_connection(db_path) as conn:
tests/test_stage5_sqlite.py
+229
View File
@@ -0,0 +1,229 @@
# flake8: noqa
# pylint: disable=protected-access, missing-function-docstring, missing-class-docstring
# pylint: disable=redefined-outer-name, logging-fstring-interpolation, line-too-long
# pylance: disable=reportMissingImports, reportMissingModuleSource
"""
Stage 5 (SQLite access-time batching) tests.
Covers:
- 5.1 accumulator: a /media cache hit records into api_server._access_updates and does NOT
call update_media_file_access_sync (no synchronous SQLite write on the hot path).
- 5.1 flush: seeding _access_updates + running the flush once writes the accumulated
timestamps to a real temp DB (hit -> flush -> value updated in DB), and clears the dict.
- 5.1 bulk fn update_media_file_access_bulk_sync: empty no-op, multi-row, and updating an
EXISTING row's `added` (WHERE matches on the str channel key).
- 5.1 snapshot-then-clear: an update arriving DURING the flush lands in the fresh dict and
is not lost.
- gotcha: str(channel) key discipline — an int-ish channel on the hot path keys the
accumulator (and thus the UPDATE) by the string form.
"""
import os
import sys
import sqlite3
import asyncio
from types import SimpleNamespace
import pytest
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
sys.modules['config'] = __import__('tests.mock_config', fromlist=['get_settings'])
import api_server
from file_io import init_db_sync, update_media_file_access_bulk_sync
def _fake_plain_message():
"""A non-poll, non-video message so download_media_file takes the normal cache flow."""
return SimpleNamespace(media=None, video=None, empty=False)
@pytest.fixture(autouse=True)
def _clean_accumulator():
"""Each test starts with an empty accumulator and restores it afterwards."""
api_server._access_updates = {}
yield
api_server._access_updates = {}
# --------------------------------------------------------------------------- #
# 5.1 hot path: cache hit records into the accumulator, no synchronous SQLite.
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_cache_hit_records_accumulator_no_sqlite(tmp_path, monkeypatch):
monkeypatch.chdir(tmp_path)
channel, post_id, fid = "testchan", 5, "fidHIT"
cache_dir = tmp_path / "data" / "cache" / channel / str(post_id)
cache_dir.mkdir(parents=True)
(cache_dir / fid).write_bytes(b"cached-bytes")
async def fake_get(_cid, _pid):
return _fake_plain_message()
monkeypatch.setattr(api_server.client, "safe_get_messages", fake_get)
# Spy: the single-row synchronous updater must NOT be called on the hot path.
called = []
monkeypatch.setattr(api_server, "update_media_file_access_sync",
lambda *a, **k: called.append(a))
path, delete_after = await api_server.download_media_file(channel, post_id, fid)
assert path == str(cache_dir / fid)
assert delete_after is False
assert called == [] # no synchronous SQLite access-write happened
assert (channel, post_id, fid) in api_server._access_updates
# --------------------------------------------------------------------------- #
# gotcha: str(channel) key discipline on the hot path (int-ish channel).
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_cache_hit_keys_channel_as_str(tmp_path, monkeypatch):
monkeypatch.chdir(tmp_path)
channel_int, post_id, fid = 12345, 7, "fidINT"
cache_dir = tmp_path / "data" / "cache" / str(channel_int) / str(post_id)
cache_dir.mkdir(parents=True)
(cache_dir / fid).write_bytes(b"x")
async def fake_get(_cid, _pid):
return _fake_plain_message()
monkeypatch.setattr(api_server.client, "safe_get_messages", fake_get)
await api_server.download_media_file(channel_int, post_id, fid)
assert ("12345", post_id, fid) in api_server._access_updates # string form
assert (channel_int, post_id, fid) not in api_server._access_updates # never the raw int
# --------------------------------------------------------------------------- #
# 5.1 flush: hit -> flush -> the accumulated timestamp is written to the DB.
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_flush_writes_accumulated_timestamps(tmp_path, monkeypatch):
db = str(tmp_path / "flush.db")
init_db_sync(db)
# Seed two existing rows with an old timestamp.
conn = sqlite3.connect(db)
conn.executemany(
"INSERT INTO media_file_ids (channel, post_id, file_unique_id, added) VALUES (?,?,?,?)",
[("chA", 1, "fA", 1.0), ("chB", 2, "fB", 2.0)],
)
conn.commit()
conn.close()
monkeypatch.setattr(api_server, "DB_PATH", db)
api_server._access_updates = {
("chA", 1, "fA"): 111.0,
("chB", 2, "fB"): 222.0,
}
await api_server._flush_access_updates()
# Dict cleared after flush.
assert api_server._access_updates == {}
conn = sqlite3.connect(db)
rows = dict(((c, p, f), a) for c, p, f, a in
conn.execute("SELECT channel, post_id, file_unique_id, added FROM media_file_ids"))
conn.close()
assert rows[("chA", 1, "fA")] == 111.0
assert rows[("chB", 2, "fB")] == 222.0
@pytest.mark.asyncio
async def test_flush_empty_is_noop(monkeypatch):
api_server._access_updates = {}
# Must not raise and must not touch the threadpool/DB.
monkeypatch.setattr(api_server, "update_media_file_access_bulk_sync",
lambda *a, **k: pytest.fail("bulk sync should not run for an empty batch"))
await api_server._flush_access_updates()
# --------------------------------------------------------------------------- #
# 5.1 snapshot-then-clear: an update arriving DURING the flush is not lost.
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_snapshot_then_clear_does_not_lose_late_update(tmp_path, monkeypatch):
db = str(tmp_path / "race.db")
init_db_sync(db)
monkeypatch.setattr(api_server, "DB_PATH", db)
late_key = ("chLate", 9, "fLate")
def fake_bulk(_db, entries):
# Simulates a cache hit landing WHILE the flush's to_thread runs: because the flush
# already replaced the module dict with a fresh one, this write goes into the NEW dict.
api_server._access_updates[late_key] = 999.0
monkeypatch.setattr(api_server, "update_media_file_access_bulk_sync", fake_bulk)
api_server._access_updates = {("chA", 1, "fA"): 111.0}
await api_server._flush_access_updates()
# The snapshot (chA) was flushed and cleared; the late update survives in the new dict.
assert api_server._access_updates == {late_key: 999.0}
# --------------------------------------------------------------------------- #
# 5.1 re-queue on failure: a failed bulk write is not lost; setdefault keeps a
# fresher write that arrived during the flush.
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_failed_flush_requeues_batch_without_clobbering_fresh(tmp_path, monkeypatch):
db = str(tmp_path / "fail.db")
init_db_sync(db)
monkeypatch.setattr(api_server, "DB_PATH", db)
stale_key = ("chStale", 1, "fS") # only in the failing snapshot
fresh_key = ("chFresh", 2, "fF") # re-written FRESHER during the flush
def fake_bulk(_db, _entries):
# A newer cache hit for fresh_key lands in the fresh dict WHILE the bulk write runs,
# then the write fails. The re-queue must restore stale_key but must NOT overwrite
# the newer fresh_key value (setdefault, not assignment).
api_server._access_updates[fresh_key] = 999.0
raise sqlite3.OperationalError("database is locked")
monkeypatch.setattr(api_server, "update_media_file_access_bulk_sync", fake_bulk)
api_server._access_updates = {stale_key: 1.0, fresh_key: 2.0}
with pytest.raises(sqlite3.OperationalError):
await api_server._flush_access_updates()
# stale_key restored with its snapshot value; fresh_key keeps the NEWER value (not 2.0).
assert api_server._access_updates == {stale_key: 1.0, fresh_key: 999.0}
# --------------------------------------------------------------------------- #
# 5.1 bulk fn: empty no-op, multi-row, and existing-row update via str channel key.
# --------------------------------------------------------------------------- #
def test_bulk_access_update_real_sql(tmp_path):
db = str(tmp_path / "bulk.db")
init_db_sync(db)
# Empty batch is a no-op (no crash).
update_media_file_access_bulk_sync(db, [])
# Seed existing rows.
conn = sqlite3.connect(db)
conn.executemany(
"INSERT INTO media_file_ids (channel, post_id, file_unique_id, added) VALUES (?,?,?,?)",
[("chA", 1, "fA", 1.0), ("chB", 2, "fB", 2.0)],
)
conn.commit()
conn.close()
# Multi-row update of EXISTING rows: WHERE matches on the str channel key.
update_media_file_access_bulk_sync(db, [
("chA", 1, "fA", 500.0),
("chB", 2, "fB", 600.0),
])
conn = sqlite3.connect(db)
rows = dict(((c, p, f), a) for c, p, f, a in
conn.execute("SELECT channel, post_id, file_unique_id, added FROM media_file_ids"))
# A row keyed by an int channel does NOT match the string "chA" WHERE (documents the gotcha).
n = conn.execute("SELECT COUNT(*) FROM media_file_ids WHERE channel = 1").fetchone()[0]
conn.close()
assert rows[("chA", 1, "fA")] == 500.0
assert rows[("chB", 2, "fB")] == 600.0
assert n == 0