feat(worker): smooth download % from adapter byte-level progress
The Floor bar used a file-count poller that jumps 0→~99% for small/fast albums (streamrip reports only 0/1). But yt-dlp and slskd already emit real byte-level progress via the on_progress callback the pipeline was discarding (lambda _pct: None). - Thread a throttled on_progress → Job.downloadProgress (write on a >=1% move or every 0.5s; clamped 0..1; write failures swallowed). Called synchronously from adapter.download on the pipeline thread, so it reuses the main conn safely. - The file-count poller stays as the FALLBACK: it skips writing while a `reports` flag is set (an adapter is driving real progress), and the flag is cleared per attempt so a fall-through to streamrip still gets the estimate. - Smooth bars for YouTube + Soulseek now; Qobuz/streamrip keeps the file-count estimate until its per-track callback is hooked later. worker 216 tests / 7-skip (on_progress write/throttle/clamp covered; existing pipeline tests now exercise the fake's on_progress end-to-end). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
import os
|
||||
import shutil
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import replace
|
||||
from typing import Sequence
|
||||
|
||||
@@ -26,20 +27,24 @@ def _count_staged_audio(staging: str) -> int:
|
||||
return n
|
||||
|
||||
|
||||
def _poll_download_progress(job_id: str, staging: str, expected: int, stop: "threading.Event") -> None:
|
||||
def _poll_download_progress(job_id: str, staging: str, expected: int, stop: "threading.Event",
|
||||
reports: "threading.Event") -> None:
|
||||
"""Until `stop`, periodically write Job.downloadProgress = files-in-staging / expected
|
||||
(capped 0.99). Uses its own short-lived connection (psycopg conns aren't shareable across
|
||||
threads). A missing DSN or expected<=0 just no-ops."""
|
||||
(capped 0.99) — the FALLBACK for adapters (streamrip/Qobuz) that report no byte-level
|
||||
progress. Skips writing while `reports` is set, i.e. the in-flight adapter is driving
|
||||
progress itself via on_progress (yt-dlp/slskd). Uses its own short-lived connection
|
||||
(psycopg conns aren't shareable across threads). A missing DSN or expected<=0 just no-ops."""
|
||||
dsn = os.environ.get("DATABASE_URL")
|
||||
if not dsn or expected <= 0:
|
||||
return
|
||||
conn = psycopg.connect(dsn)
|
||||
try:
|
||||
while not stop.is_set():
|
||||
frac = min(_count_staged_audio(staging) / expected, 0.99)
|
||||
with conn.cursor() as cur:
|
||||
cur.execute('UPDATE "Job" SET "downloadProgress" = %s WHERE id = %s', (frac, job_id))
|
||||
conn.commit()
|
||||
if not reports.is_set(): # the adapter isn't reporting real progress → estimate
|
||||
frac = min(_count_staged_audio(staging) / expected, 0.99)
|
||||
with conn.cursor() as cur:
|
||||
cur.execute('UPDATE "Job" SET "downloadProgress" = %s WHERE id = %s', (frac, job_id))
|
||||
conn.commit()
|
||||
stop.wait(1.5)
|
||||
except Exception as e: # a progress poller must never affect the job
|
||||
print(f"pipeline: download progress poller error: {e}", flush=True)
|
||||
@@ -62,6 +67,29 @@ def _set_download_progress(conn: psycopg.Connection, job_id: str, frac: float) -
|
||||
conn.commit()
|
||||
|
||||
|
||||
def _make_on_progress(conn: psycopg.Connection, job_id: str, reports: "threading.Event"):
|
||||
"""Build the on_progress(pct) callback threaded into adapter.download. It marks the
|
||||
adapter as reporting real byte-level progress (so the file-count poller yields) and
|
||||
throttle-writes Job.downloadProgress. Called synchronously from adapter.download on the
|
||||
pipeline's own thread, so it safely reuses `conn`."""
|
||||
state = {"frac": -1.0, "t": 0.0}
|
||||
|
||||
def on_progress(pct: float) -> None:
|
||||
reports.set()
|
||||
frac = 0.0 if pct < 0 else 1.0 if pct > 1 else float(pct)
|
||||
now = time.monotonic()
|
||||
# throttle DB writes: on a >=1% move or every 0.5s, not once per received byte
|
||||
if frac - state["frac"] >= 0.01 or (now - state["t"]) >= 0.5:
|
||||
state["frac"] = frac
|
||||
state["t"] = now
|
||||
try:
|
||||
_set_download_progress(conn, job_id, frac)
|
||||
except Exception as e: # a progress write must never fail the download
|
||||
print(f"pipeline: on_progress write failed: {e}", flush=True)
|
||||
|
||||
return on_progress
|
||||
|
||||
|
||||
def _request_id(conn: psycopg.Connection, job_id: str) -> str:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute('SELECT "requestId" FROM "Job" WHERE id = %s', (job_id,))
|
||||
@@ -245,8 +273,9 @@ def run_pipeline(
|
||||
_set_download_progress(conn, job_id, 0.0) # reset for this run
|
||||
expected = target.track_count or (ranked[0].track_count if ranked else 0)
|
||||
_stop = threading.Event()
|
||||
_reports = threading.Event() # set while the in-flight adapter reports real byte progress
|
||||
_poller = threading.Thread(
|
||||
target=_poll_download_progress, args=(job_id, staging, expected, _stop), daemon=True
|
||||
target=_poll_download_progress, args=(job_id, staging, expected, _stop, _reports), daemon=True
|
||||
)
|
||||
_poller.start()
|
||||
try:
|
||||
@@ -257,7 +286,8 @@ def run_pipeline(
|
||||
continue
|
||||
_mark_chosen(conn, job_id, candidate.source_ref) # reflect the source now in flight
|
||||
shutil.rmtree(staging, ignore_errors=True) # clean slate per attempt
|
||||
result = adapter.download(candidate, staging, lambda _pct: None)
|
||||
_reports.clear() # this attempt hasn't reported yet → poller estimates until it does
|
||||
result = adapter.download(candidate, staging, _make_on_progress(conn, job_id, _reports))
|
||||
if result.ok:
|
||||
winner = candidate
|
||||
break
|
||||
|
||||
Reference in New Issue
Block a user