siftlode/backend/tests/test_download_job_locks.py

"""Overlapping requests on ONE job — DB-backed, two real sessions.

The routes that move a job's asset ref count (resume, cancel, delete, redownload) used to decide on an
unlocked read: two overlapping requests for one job (a double-click, a second tab, a replay after a
gateway timeout) both saw it holding its asset and both released the hold. The count then reaches 0
while another job still holds the asset, and GC or the ready-file delete removes it under that job.
A pause racing a cancel did the same one step later: its unlocked write turned the released job back
into a holder, which a later delete releases again. `own_job_for_update` makes the second request
wait and decide on the first one's committed result.
"""

import threading

import pytest
from fastapi import HTTPException
from sqlalchemy import select

from app.apitokens import FilesCaller
from app.db import SessionLocal
from app.downloads import service
from app.models import User
from app.downloads.models import DownloadJob, MediaAsset
from app.routes.downloads import cancel_download, pause_download
from app.routes.files import delete_file, own_job_for_update

_REF = "https://example.com/video"


def _asset_held_by_two_jobs(db, status="error"):
    mine, theirs = User(email="lock-mine@example.com"), User(email="lock-theirs@example.com")
    db.add_all([mine, theirs])
    db.flush()
    asset = MediaAsset(
        source_kind="url", source_ref=_REF, format_sig="shared", status="error", ref_count=2
    )
    db.add(asset)
    db.flush()
    jobs = [
        DownloadJob(
            user_id=owner.id,
            asset_id=asset.id,
            source_kind="url",
            source_ref=_REF,
            status=status if owner is mine else "error",
            queue_pos=0,
            profile_snapshot={},
        )
        for owner in (mine, theirs)
    ]
    db.add_all(jobs)
    db.commit()
    return mine, asset, jobs[0]


def delete_file_as(job_id, *, user, db):
    """`delete_file` as the user's own request (it takes the Files caller, not the user)."""
    return delete_file(job_id, caller=FilesCaller(user=user), db=db)


@pytest.mark.parametrize(
    ("status", "transition", "route", "left"),
    [
        ("error", service.cancel, cancel_download, "canceled"),
        ("error", service.delete, delete_file_as, None),
        ("queued", service.cancel, pause_download, "canceled"),
    ],
)
def test_overlapping_requests_for_one_job_release_its_hold_once(
    db, status, transition, route, left
):
    user, asset, job = _asset_held_by_two_jobs(db, status)
    job_id = job.id  # read before the delete case removes the row under the expired instance
    first, second = SessionLocal(), SessionLocal()
    finished = threading.Event()

    def the_other_request():
        try:
            route(job_id, user=user, db=second)
        except HTTPException:
            pass  # a request after the delete: the job is gone, a 404
        finally:
            second.close()
            finished.set()

    try:
        locked = own_job_for_update(first, user, job_id)
        other = threading.Thread(target=the_other_request)
        other.start()
        # The second request waits for the first one's locks instead of deciding on the same row.
        assert not finished.wait(0.5), "the second request did not wait for the first"
        transition(first, locked)
        other.join(10)
        assert finished.is_set()
    finally:
        first.close()

    db.expire_all()
    mine = select(DownloadJob.status).where(DownloadJob.id == job_id)
    assert db.execute(mine).scalar_one_or_none() == left
    assert db.get(MediaAsset, asset.id).ref_count == 1  # the other user's job still holds it


def test_losing_the_asset_insert_race_keeps_the_callers_own_writes(db, monkeypatch):
    # Retry clears the job's error, THEN may create the cookie asset. When another request inserted that
    # asset first, the loser used to roll the whole session back — and requeue the job with its stale
    # error, which a stall or crash later leaves showing under the new failure.
    user = User(email="race@example.com")
    db.add(user)
    db.flush()
    job = DownloadJob(
        user_id=user.id,
        source_kind="url",
        source_ref=_REF,
        status="error",
        error="the old failure",
        queue_pos=0,
        profile_snapshot={},
    )
    db.add(job)
    db.commit()
    job.error = None  # the caller's pending write

    real_flush, raced = db.flush, []

    def flush_after_the_other_request_inserted(*args, **kwargs):
        if not raced:  # the other request commits the same asset between our lookup and our insert
            raced.append(True)
            with SessionLocal() as other:
                other.add(
                    MediaAsset(source_kind="url", source_ref=_REF, format_sig="raced", ref_count=0)
                )
                other.commit()
        return real_flush(*args, **kwargs)

    monkeypatch.setattr(db, "flush", flush_after_the_other_request_inserted)

    asset = service.get_or_create_asset(db, "url", _REF, "raced")
    db.commit()

    assert raced and asset.format_sig == "raced"
    db.expire_all()
    assert db.get(DownloadJob, job.id).error is None