diff --git a/comfy_cli/command/models/models.py b/comfy_cli/command/models/models.py index 789fc326..2483f974 100644 --- a/comfy_cli/command/models/models.py +++ b/comfy_cli/command/models/models.py @@ -865,21 +865,29 @@ def _submit_background_download( ) raise typer.Exit(code=1) from e - # Claim `dest` atomically. Creating the claim file with `O_CREAT | O_EXCL` - # *is* the decision: the kernel serializes the create, exactly one of any - # number of simultaneous submitters gets the file, and there is no - # write-then-re-scan window for the losers to slip through. That window is - # what the previous guard could only narrow — a re-scan cannot see a claim - # that has not been written yet, and nothing ordered a competitor's `write` - # before our scan, so when we scanned first and lost the `_claim_order` tie - # neither side withdrew and both streamed a full copy (reproduced at 12 - # simultaneous submits; `started_at` is second-resolution, so the tie that - # makes it likely is the common case rather than an exotic one). + # Claim `dest` with a *held* lock. Taking an exclusive OS lock on the + # destination's claim file is what decides ownership: the kernel serializes + # it, exactly one of any number of simultaneous submitters gets it, and + # there is no write-then-re-scan window for the losers to slip through. That + # window is what the record-is-the-claim guard could only narrow — a re-scan + # cannot see a record that has not landed yet, and nothing ordered a + # competitor's `write` before our scan, so when we scanned first and lost the + # `_claim_order` tie neither side withdrew and both streamed a full copy + # (reproduced at 12 simultaneous submits; `started_at` is second-resolution, + # so the tie that makes it likely is the common case rather than an exotic + # one). # - # The record is written *before* the claim, never after: the claim is only a - # pointer, and a competitor probing it for liveness must never find it - # pointing at a record of ours that does not exist yet — it would read as - # stale and be cleared out from under us. + # We hold the lock only long enough to arbitrate and stamp our id into the + # payload, then hand it to the worker, which re-takes it and holds it for the + # whole transfer. The gap between our release and the worker's lock is the + # *spawn gap*, and it is the one thing a bare lock cannot arbitrate — hence + # the payload check in `_acquire_dest_claim`, which refuses a competitor + # whose record is active and younger than the startup grace. + # + # The record is written *before* the claim, never after: the payload is only + # a pointer, and a competitor reading it during the spawn gap must never find + # it pointing at a record of ours that does not exist yet — it would read as + # stale and be stolen out from under us. # # Unchanged non-goals. A *foreground* transfer still writes no claim file — # it claims with its state record only — and the claim lives under @@ -915,11 +923,12 @@ def _submit_background_download( with contextlib.suppress(OSError): download_state.write(workspace, state) if claim_file is not None: - # No worker will ever reach a terminal transition for this record, so - # nothing else would ever release the claim. It would still read - # stale (the record is `failed`) and self-clear on the next submit, - # but leaving it is a needless round of cleanup for the next caller. - download_state.release_claim(claim_file, owner_id=state.id) + # No worker will ever come up to take the lock for this record, so + # the payload we stamped names a download that will never run. It + # would read stale to the next submitter anyway (the record is + # `failed`), but blanking it now means the spawn-gap arbitration + # never has to consider it at all. + _clear_dest_claim(claim_file, state.id) renderer.error( code="download_worker_spawn_failed", message=f"Could not start the background download worker: {e}", @@ -972,30 +981,65 @@ def _download_worker( cancel_marker = download_state.cancel_marker_for(path) - # The submitter took an `O_EXCL` claim on this destination and we are the - # only thing that reaches a terminal transition for it, so releasing it is - # ours: everything past the cancel-sentinel check runs under a - # `try/finally: release()`, so no exit — including an OSError out of the - # very first `write_path` — can strand the claim. Derived, never carried in - # the record — the claim is addressed by destination, not by download. - # `release_claim` no-ops unless the claim still names us, so a claim a later - # submitter already cleared and replaced is never unlinked from under it. + # The submitter stamped this destination's claim with our id and released + # the lock; taking it back is our first act, and we hold it for the whole + # transfer. Holding is what makes ownership true rather than inferred: while + # we have it, no submitter can conclude we are gone from a record that reads + # terminal (`download-cancel` writes `cancelled` even when `stop_worker` + # fails) or from a state file it momentarily could not read. And if we are + # SIGKILLed the kernel drops it, so the next submitter takes the destination + # immediately — no reconcile, no grace wait, no sweeper. + # + # A short blocking window rather than a single attempt: a competing submitter + # may hold the lock for the moment it takes to read and rewrite the payload, + # and losing the destination to that would be a spurious refusal. # - # Derived defensively: `state.dest` is read back off disk, so `_dest_key` - # (`realpath`) is being handed untrusted text and can raise on a shape the - # submitter's own `dest` never had. Releasing a claim must never be what - # stops the transfer from recording its result — a claim we cannot address - # is left behind and self-clears on the next submit, exactly like the SIGKILL - # case. + # Derived defensively, never carried in the record — the claim is addressed + # by destination, not by download. `state.dest` is read back off disk, so + # `_dest_key` (`realpath`) is being handed untrusted text and can raise on a + # shape the submitter's own `dest` never had. Bookkeeping must never be what + # stops a transfer from recording its result, so a claim we cannot address is + # a claim we proceed without. claim_file: pathlib.Path | None try: claim_file = download_state.claim_marker_for(path, _dest_key(state.dest)) except (OSError, ValueError): claim_file = None + claim_lock: download_state.ClaimLock | None = None + claim_lost = False + if claim_file is not None: + try: + claim_lock = download_state.lock_claim(claim_file, blocking_timeout=WORKER_CLAIM_LOCK_TIMEOUT_S) + except OSError: + # The filesystem cannot lock at all (the submitter will have warned + # about it already). Degrade to the pre-lock behavior — transfer + # anyway — rather than refuse a download over bookkeeping. + claim_lock = None + else: + # Held by somebody else for the whole window, or stamped with a + # different download's id while we were starting: either way the + # destination is not ours and we must not write a byte of it. + # + # A claim with *no* payload is not that case and must not be treated + # as one. Nobody can own a destination whose lock we are holding, so + # an unstamped claim means only that nothing stamped it — the + # submitter degraded past `write_payload` on a transient error, or a + # sweep removed an unowned file and our own `O_CREAT` re-made it. + # Refusing there would turn a recoverable submit-side hiccup into a + # download that never runs. Stamp it and carry on. + payload = claim_lock.read_payload() if claim_lock is not None else None + owner = payload.get("download_id") if payload is not None else None + claim_lost = claim_lock is None or (owner is not None and owner != state.id) + if claim_lock is not None and owner is None: + with contextlib.suppress(OSError): + claim_lock.write_payload(state.id, state.dest) + def release() -> None: - if claim_file is not None: - download_state.release_claim(claim_file, owner_id=state.id) + nonlocal claim_lock + if claim_lock is not None: + claim_lock.release() + claim_lock = None def cancelled() -> bool: return cancel_marker.exists() @@ -1016,6 +1060,19 @@ def cancelled() -> bool: release() raise typer.Exit(code=0) + if claim_lost: + # We lost the destination during the spawn gap. Exit refused without + # touching `dest`, and leave a terminal record behind exactly as the + # submit-side withdraw path does: a `starting` record we abandon would + # read as a live claim to `_active_download_for` and to `download-cancel` + # and refuse every later submission until it aged out of the grace. + release() + state.status = "failed" + state.error = "another download claimed this destination first" + with contextlib.suppress(OSError): + download_state.write_path(path, state) + raise typer.Exit(code=1) + try: state.pid = os.getpid() state.pid_create_time = download_state.process_create_time(os.getpid()) @@ -1171,9 +1228,9 @@ def _claim_order(state: download_state.DownloadState) -> tuple[str, str]: It no longer decides a background submit. Ordering records can only rank the ones a racer happens to *see*, and a racer that scans before the other's record lands sees no competitor at all — so two submits could rank each other - into a pair of winners. What decides a background submit now is the `O_EXCL` - claim file (:func:`_acquire_dest_claim`), where the kernel's create is the - order and nobody has to see anybody. + into a pair of winners. What decides a background submit now is the held + claim lock (:func:`_acquire_dest_claim`), where the kernel's arbitration is + the order and nobody has to see anybody. This order survives for the readers that are still advisory and still want a deterministic pick: :func:`_active_download_for`'s winner-pick (one live @@ -1220,7 +1277,7 @@ def _active_download_for( this scan *after* writing its own claim and must not find itself. The scan is scoped to ``get_workspace()``'s state directory, and that is the - residual neither this nor the `O_EXCL` claim file closes — the claim lives in + residual neither this nor the claim lock closes — the claim lives in the same per-workspace state directory. A destination can sit outside the workspace (``--relative-path`` is only ``expanduser``-ed, so it accepts ``..`` and absolute paths), so two invocations run against *different* @@ -1302,36 +1359,115 @@ def _dest_claim_path(dest: pathlib.Path) -> pathlib.Path | None: return None -def _claim_holder(claim_file: pathlib.Path) -> tuple[str | None, download_state.DownloadState | None]: - """Resolve a claim file to ``(recorded id, live record)``. - - A claim is live iff its ``download_id`` resolves to a state record whose - *reconciled* status is still active, or whose worker process is still - provably running. The first arm delegates the liveness question to the - record's existing `pid` + `pid_create_time` identity proof and - `STARTUP_GRACE_S` window: a SIGKILLed worker's record demotes to `failed` - on reconcile, so its claim reads stale here and the next submitter clears - it — which is why a SIGKILL needs no sweeper. The second arm covers the - gap the first cannot: a terminal status does not prove the process exited - (a cancelled worker is still mid-write until its own terminal transition - lands, and it releases the claim itself right after), so while - `worker_alive` still proves the recorded process is that worker, the claim - is its to release, not ours to clear. - - An unreadable or corrupt claim, and a claim whose record is gone, both come - back with no live record: stale. The id is still returned when it could be - read, because the retry path words its refusal from it. +def _record_or_none(download_id: str | None) -> download_state.DownloadState | None: + """``download_id``'s record, or None for every way the lookup can fail. + + :func:`download_state.read` resolves (and creates) the state directory, so it + can raise ``OSError`` as well as ``ValueError`` on a malformed id. Both are + tolerated here because every caller is on a *refusal* path, where the record + is only wanted to word a message: a lookup that fails must not turn a clean + refusal into a traceback. """ - download_id = download_state.read_claim(claim_file) - if download_id is None: - return None, None - record = download_state.read(get_workspace(), download_id) + if not download_id: + return None + try: + return download_state.read(get_workspace(), download_id) + except (OSError, ValueError): + return None + + +def _spawn_gap_holder(holder_id: str) -> download_state.DownloadState | None: + """The record ``holder_id`` names, when it is a submit still inside its spawn gap. + + The one case a held lock cannot arbitrate on its own. Between a submitter's + release and its worker taking the lock nobody holds it, so a competitor + arriving in that window acquires cleanly and would otherwise steal a + destination whose download is about to start. + + The payload is what closes it: a record that is still *active* after + reconcile and younger than :data:`download_state.STARTUP_GRACE_S` is a submit + whose worker has not come up yet, and the destination is its own. Anything + else — no record, a record reconcile has demoted, a terminal one, or one too + old for its worker to still be booting — is stale, and the caller takes the + destination under the lock it already holds. + + Bounded by the grace on purpose. Past it, a download that is genuinely live + holds the lock, so we would never have got here to ask; refusing on the + payload alone would resurrect exactly the derived-ownership guess this design + replaced. The residual is the same one the grace has always had: a submitter + that dies between writing its record and spawning its worker wedges its + destination for up to a minute. + """ + record = _record_or_none(holder_id) if record is None: - return download_id, None + return None fresh, _ = _reconciled(record) - if fresh.status not in download_state.ACTIVE_STATUSES and not download_state.worker_alive(fresh): - return download_id, None - return download_id, fresh + if fresh.status not in download_state.ACTIVE_STATUSES: + return None + if download_state.elapsed_seconds(fresh) >= download_state.STARTUP_GRACE_S: + return None + return fresh + + +def _refuse_held_claim( + state: download_state.DownloadState, + dest: pathlib.Path, + claim_file: pathlib.Path, +) -> typer.Exit: + """Withdraw our record and word the refusal for a claim somebody is holding. + + No liveness judgment is made or needed: a held lock is the liveness proof. + The payload is read *unlocked* here and purely to name the holder — it may be + a rewrite in progress, and on Windows the holder's byte-range lock makes the + read fail outright — so both the id and the record it resolves to are + best-effort. + + A resolvable record gets the usual `model_download_in_flight` refusal. One we + cannot resolve gets a distinct code rather than a `model_download_in_flight` + missing the status/kind fields that code documents: quoting either would mean + inventing it. + """ + holder_id = download_state.read_claim(claim_file) + holder = _record_or_none(holder_id) + _withdraw_record(state, dest, holder_id) + if holder is not None: + return _in_flight_failure(holder, dest) + named = f" ({holder_id})" if holder_id else "" + return _download_failure( + code="model_download_claim_contested", + message=f"Another download{named} is holding the claim on {dest}.", + hint="check `comfy model downloads`, then retry", + details={"path": str(dest), "claim_file": str(claim_file), "download_id": holder_id}, + ) + + +def _clear_dest_claim(claim_file: pathlib.Path, owner_id: str) -> None: + """Blank the claim payload we stamped, when it is still ours. Never raises. + + Not an unlink: the claim file is where the lock lives, and unlinking one out + from under a descriptor another process already opened is how two holders end + up locking two different inodes for one destination (see + :meth:`download_state.ClaimLock.release`). Blanking it under the lock is the + same outcome for every reader — an empty claim resolves to nobody — with none + of that risk. + + Silent on every failure. Clearing a claim is bookkeeping on an error path + that is already reporting a real problem to the user. + """ + try: + lock = download_state.lock_claim(claim_file) + except OSError: + return + if lock is None: + return + try: + payload = lock.read_payload() + if payload is not None and payload.get("download_id") == owner_id: + lock.clear_payload() + except OSError: + pass + finally: + lock.release() def _withdraw_record(state: download_state.DownloadState, dest: pathlib.Path, winner_id: str | None) -> None: @@ -1349,16 +1485,24 @@ def _withdraw_record(state: download_state.DownloadState, dest: pathlib.Path, wi _persist_record(state) -# `os.link` is what makes the claim atomic, and a filesystem without hard links -# (exFAT, FAT32, some network and container mounts) refuses it outright rather -# than transiently. Told apart from a transient failure only to word the warning -# accurately: both degrade identically, because a download that used to work must -# not become an error over bookkeeping. +# How long a freshly spawned worker retries the claim lock before concluding it +# lost the destination. Only a competitor *inspecting* the claim holds it while +# we start up, and that is a ~100 byte read and rewrite, so this covers process +# scheduling rather than any real work. A genuine owner holds it for its whole +# transfer, and waiting longer for that one would only delay a refusal. +WORKER_CLAIM_LOCK_TIMEOUT_S = 5.0 + +# The exclusive lock is what makes the claim decisive, and a filesystem that +# cannot take one (some network and container mounts, a kernel with no flock) +# refuses it outright rather than transiently. Told apart from a transient +# failure only to word the warning accurately: both degrade identically, because +# a download that used to work must not become an error over bookkeeping. +# Contention is *not* in here — `lock_claim` reports that by returning None. _CLAIMS_UNSUPPORTED_ERRNOS = frozenset( - getattr(errno, name) for name in ("ENOTSUP", "EOPNOTSUPP", "EPERM", "ENOSYS") if hasattr(errno, name) + getattr(errno, name) for name in ("ENOTSUP", "EOPNOTSUPP", "ENOLCK", "ENOSYS") if hasattr(errno, name) ) -# ERROR_INVALID_FUNCTION / ERROR_NOT_SUPPORTED — what Windows returns for -# `CreateHardLinkW` on a volume that has no hard links. +# ERROR_INVALID_FUNCTION / ERROR_NOT_SUPPORTED — what Windows returns when the +# volume behind the handle does not implement byte-range locking. _CLAIMS_UNSUPPORTED_WINERRORS = frozenset((1, 50)) _claims_degraded_reported = False @@ -1381,9 +1525,9 @@ def _report_claim_degraded(exc: OSError, dest: pathlib.Path) -> None: getattr(exc, "winerror", None) in _CLAIMS_UNSUPPORTED_WINERRORS ) why = ( - "this filesystem does not support the hard link that publishes them" + "this filesystem does not support the exclusive lock they are held with" if unsupported - else "the claim could not be written" + else "the claim could not be locked" ) logger.warning( "atomic destination claims are unavailable (%s: %s); falling back to the " @@ -1400,94 +1544,67 @@ def _acquire_dest_claim( dest: pathlib.Path, claim_file: pathlib.Path, ) -> None: - """Take the `O_EXCL` claim on ``dest``, or withdraw and refuse. - - Returns None when the claim is ours. Otherwise the destination belongs to - someone else: our own record comes back off disk (so we leave no phantom - claim) and the caller gets the usual `model_download_in_flight` refusal. - - A claim we lose to is only decisive while it is *live*. A stale one — its - record demoted by reconcile, deleted, or the file corrupt — is unlinked and - the create retried exactly ONCE. Once, not in a loop: a second collision - means another submitter won the retry race rather than that the claim is - wedged, and that submitter is a competitor to refuse to, not a lock to keep - fighting for. + """Take the claim lock on ``dest`` and stamp it, or withdraw and refuse. + + Returns None when the destination is ours. Otherwise it belongs to someone + else: our own record comes back off disk (so we leave no phantom claim) and + the caller gets the usual `model_download_in_flight` refusal. + + Three outcomes, and only the first two are collisions: + + * **The lock is held.** A live owner has it — that is the whole proof, and it + is why a worker whose record says `cancelled` while the worker is still + running no longer loses its destination. Refuse. + * **We got the lock and the payload names a live, just-submitted download.** + We won it during another submit's spawn gap, before its worker could take + it. Release and refuse (:func:`_spawn_gap_holder`). + * **We got the lock and the payload is anything else** — empty, unresolvable, + naming a terminal or reconciled-dead record. Stale: stamp our id in and + spawn. There is no clear-then-retry any more, because the steal happens + under the lock in one step; the retry loop existed only because unlink and + re-create were two. + + We release before spawning rather than handing the descriptor to the worker: + the worker re-takes the lock itself (with a short blocking window) and holds + it for the transfer, which keeps ownership with the process actually writing + the bytes even if this one exits first. """ - for attempt in (1, 2): + try: + lock = download_state.lock_claim(claim_file) + except OSError as e: + # Anything other than contention (a read-only state dir, a vanished + # claims directory, a filesystem that cannot lock). Degrade to the + # advisory guard, exactly as an unavailable claims directory does — but + # say so first: the advisory guard re-scans rather than arbitrates, so + # this silently gives up the atomicity this whole path exists to + # provide, and a destination on such a filesystem can be raced again. + _report_claim_degraded(e, dest) + _enforce_claim(state, dest) + return + + if lock is None: + raise _refuse_held_claim(state, dest, claim_file) + + try: + payload = lock.read_payload() + holder_id = payload.get("download_id") if payload is not None else None + if holder_id is not None and holder_id != state.id: + holder = _spawn_gap_holder(holder_id) + if holder is not None: + lock.release() + _withdraw_record(state, dest, holder.id) + raise _in_flight_failure(holder, dest) try: - if download_state.acquire_claim(claim_file, download_id=state.id, dest=str(dest)): - return + lock.write_payload(state.id, str(dest)) except OSError as e: - # Anything other than the collision (a read-only state dir, a - # vanished claims directory, a filesystem with no hard links). - # Degrade to the advisory guard, exactly as an unavailable claims - # directory does — but say so first: the advisory guard re-scans - # rather than arbitrates, so this silently gives up the atomicity - # this whole path exists to provide, and a destination on such a - # filesystem can be raced again. + # The lock is ours but the payload would not go down (ENOSPC, EIO). + # A lock we release without stamping decides nothing for the worker + # we are about to spawn, so this is the same degradation as a + # filesystem that cannot lock at all. _report_claim_degraded(e, dest) _enforce_claim(state, dest) - return - - holder_id, holder = _claim_holder(claim_file) - if holder is not None: - _withdraw_record(state, dest, holder.id) - raise _in_flight_failure(holder, dest) - - if attempt == 1: - # Stale: the claim outlived the download it points at. Clear it and - # try once more. Another submitter may clear it first and win the - # create — that is the second pass below, not a problem here. - # - # Cleared *conditionally*, by the id we just read: between reading a - # claim and deciding it is stale sits a state-file read and a - # `reconcile`, and in that gap its worker may finish, release it, and - # a fresh submitter take a live claim at the same path. An - # unconditional unlink would delete that live claim and leave two - # downloads owning one destination. `release_claim` re-reads and only - # unlinks while the id still matches, which narrows the window to the - # compare-and-unlink inside it — it does not close it (the filesystem - # offers no conditional unlink), but the surviving window no longer - # spans a reconcile. If the claim did change under us, the retry - # below collides with the new holder and refuses, which is right. - if not download_state.release_claim(claim_file, owner_id=holder_id): - # False for two very different reasons, told apart by a re-read. - # The claim changing hands under us is the takeover race above — - # fall through and collide with the new holder. The claim still - # naming the id we judged stale means the unlink itself failed - # (a claim file we cannot remove, or a directory sitting at the - # claim path): retrying would collide with the same corpse - # forever and report a phantom in-flight download, so name the - # real problem instead. - if download_state.read_claim(claim_file) == holder_id: - _withdraw_record(state, dest, holder_id) - raise _download_failure( - code="model_download_claim_unclearable", - message=f"A stale download claim on {dest} could not be cleared.", - hint=f"remove the claim file at {claim_file} and retry, or check its permissions", - details={ - "path": str(dest), - "claim_file": str(claim_file), - "download_id": holder_id, - }, - ) - continue - - # Second collision: somebody else took the claim we just cleared. Back - # off rather than clear theirs too — a retry loop over a contested claim - # is how two submitters livelock each other. Their id, when the claim was - # readable, is all we can honestly report: this claim did not resolve to - # a live record, so quoting a status or kind from one would be inventing - # it — hence a distinct code rather than a `model_download_in_flight` - # missing the fields that code documents. - _withdraw_record(state, dest, holder_id) - named = f" ({holder_id})" if holder_id else "" - raise _download_failure( - code="model_download_claim_contested", - message=f"Another download{named} claimed {dest} first.", - hint="check `comfy model downloads`, then retry", - details={"path": str(dest), "download_id": holder_id}, - ) + finally: + lock.release() def _enforce_claim(state: download_state.DownloadState, dest: pathlib.Path) -> None: @@ -1500,7 +1617,7 @@ def _enforce_claim(state: download_state.DownloadState, dest: pathlib.Path) -> N the ``--background`` split and (for a Hugging Face url) a whole ``check_unauthorized`` round trip between it and the write. - It is check-then-act and cannot be otherwise — hence the `O_EXCL` claim file + It is check-then-act and cannot be otherwise — hence the held claim lock that now decides the background path (:func:`_acquire_dest_claim`). This stays because it is the only guard the *foreground* path has (a foreground transfer writes no claim file), and because it is what lets a background diff --git a/comfy_cli/download_state.py b/comfy_cli/download_state.py index 968bf098..e149327d 100644 --- a/comfy_cli/download_state.py +++ b/comfy_cli/download_state.py @@ -68,6 +68,7 @@ from __future__ import annotations import contextlib +import errno import hashlib import json import os @@ -81,6 +82,11 @@ from pathlib import Path from typing import Any +if sys.platform == "win32": + import msvcrt +else: + import fcntl + STATE_SCHEMA = "download-state/1" STATE_DIRNAME = ".comfy-downloads" @@ -182,18 +188,87 @@ def request_cancel(path: Path) -> bool: # dir rather than next to the user's model files: a claim is our bookkeeping, # not something a user should find in `models/loras`, and `list_all` globs # `*.json` at the top level so the subdirectory stays invisible to every verb. +# +# **Ownership of a destination is HELD, not derived.** The process performing a +# transfer holds an exclusive OS lock (`flock` on POSIX, `msvcrt.locking` on +# Windows) on that destination's claim file for the whole transfer, and the lock +# — not the payload, and not the record the payload points at — is what proves +# the destination is taken. That inverts the previous design, in which ownership +# was re-derived on every access by resolving the payload's `download_id` to a +# state record and asking whether that record looked live. Deriving left three +# windows that holding closes: +# +# * a read-then-unlink race in the old conditional `release_claim`: the claim +# could change hands between the compare and the unlink; +# * a record that reads terminal while its worker is demonstrably still running +# (`download-cancel` writes `cancelled` even when `stop_worker` fails), and a +# transient read failure on the record, both of which made a live claim look +# stale; +# * a stale-clear that had to be retried, because clearing and re-creating were +# two separate steps another submitter could interleave with. +# +# A SIGKILLed holder needs no sweeper and no grace wait: the kernel drops its +# lock when the process dies, so the next submitter acquires immediately. +# +# The payload is still written — the `download_id` it names is how a submitter +# words its refusal, and how a freshly spawned worker checks that the +# destination it was handed is still its own — but it is only ever rewritten +# **in place on the locked descriptor**. A temp-file-plus-rename publish would +# swap the inode the lock lives on, which is exactly the thing the lock cannot +# survive. +# +# Claim files are never unlinked on release. One is ~100 bytes keyed by the +# hash of a destination, so the directory is bounded by the number of distinct +# destinations a workspace has ever downloaded to, and leaving them in place +# sidesteps the unlink-under-lock pitfall entirely (an unlink makes every fd +# already opened on that path a lock on an orphan inode). :func:`prune` still +# removes the ones it can both lock and prove dead, on POSIX only. +# +# **Network filesystems are best effort.** `flock` over NFS/SMB is advisory at +# best and may be a local-only lock, so two machines sharing a mount can still +# both win. That is no worse than the pid-based liveness this replaces, which +# was meaningless across machines to begin with; cross-machine correctness is +# out of scope here, exactly as it was before. CLAIMS_DIRNAME = "claims" -# Private staging suffix for `acquire_claim`'s write-then-link publish. Never -# matched by the `*.claim` readers; leftovers (a SIGKILL between write and link) -# are swept by `prune` once they are old enough to be provably dead. +# Legacy staging suffix. The previous implementation published a claim by +# writing a private `.claim...tmp` sibling and hard-linking it into +# place; nothing writes these any more, but a workspace last touched by that +# version can still hold one, so :func:`prune` keeps sweeping them. CLAIM_TMP_SUFFIX = ".tmp" -# How old a `.claim.*.tmp` staging file must be before `prune` treats it as a -# crashed acquire's leftover rather than an acquire in progress. An acquire -# holds its temp file for milliseconds; an hour is comfortably conservative. +# How old such a leftover must be before `prune` treats it as a crashed +# acquire's residue rather than an acquire in progress. CLAIM_TMP_MAX_AGE_S = 60 * 60 +# How many times :func:`lock_claim` re-opens after its re-stat guard rejects the +# file it just locked. A rejection means the claim was unlinked or replaced +# between our `open` and our lock, so the lock we hold is on an orphan inode and +# decides nothing. Bounded rather than unbounded: only a concurrent `prune` +# unlinks claims at all, so more than a couple of rounds means something else is +# churning the path and refusing is the safe direction. +CLAIM_LOCK_RESTAT_ATTEMPTS = 3 + +# Poll interval for `lock_claim`'s bounded blocking wait. The lock is only ever +# held across a ~100 byte read/rewrite by anyone who is not transferring, so the +# wait is short and a tight poll costs nothing. +CLAIM_LOCK_POLL_S = 0.02 + +# What a lock attempt raises when the lock is simply held elsewhere, as opposed +# to the filesystem being unable to lock at all. POSIX `flock` reports +# EWOULDBLOCK/EAGAIN; Windows `msvcrt.locking` reports EACCES for `LK_NBLCK` and +# EDEADLOCK when a blocking `LK_LOCK` gives up. Everything else (ENOLCK, +# ENOTSUP/EOPNOTSUPP on a filesystem with no locking) propagates, because it +# means no lock was taken and the caller has to degrade rather than conclude +# somebody else owns the destination. +_LOCK_CONTENDED_ERRNOS = frozenset( + getattr(errno, name) for name in ("EACCES", "EAGAIN", "EWOULDBLOCK", "EDEADLK", "EDEADLOCK") if hasattr(errno, name) +) + +# `os.open` needs O_BINARY on Windows or the descriptor does newline translation +# on a payload we byte-count. A no-op everywhere else. +_O_BINARY = getattr(os, "O_BINARY", 0) + def claims_dir(workspace: Path) -> Path: """Return ``/.comfy-downloads/claims`` and ensure it exists, owner-only.""" @@ -229,126 +304,239 @@ def claim_marker_for(state_file: Path, dest_key: str) -> Path: The worker is handed ``--state `` and never re-resolves a workspace (same reason as :func:`cancel_marker_for`), so it derives the claim from the - state file's own directory. Creates nothing - the worker only ever releases. + state file's own directory. Creates nothing - the worker locks what is there. """ return Path(state_file).parent / CLAIMS_DIRNAME / claim_filename(dest_key) -def acquire_claim(path: Path, *, download_id: str, dest: str) -> bool: - """Atomically create the claim at ``path``. False when it already exists. - - The payload is written to a private sibling first and *published* with - ``os.link``, which fails ``EEXIST`` exactly like ``O_CREAT | O_EXCL`` does — - so file creation is still the atomic decider (exactly one of any number of - simultaneous submitters gets True, with no check-then-act window), but the - claim is never visible at ``path`` until its payload is complete. Creating - the file at ``path`` directly and writing into it afterwards would open a - window in which a colliding submitter reads an empty claim, calls it stale, - and unlinks a live winner. It also means a failed payload write (``ENOSPC``, - ``EIO``) publishes nothing: the temp file is removed and the ``OSError`` - propagates with no orphan claim left at ``path``. - - The temp name carries the pid and the download id, both unique to this - acquire, so concurrent submitters never collide on it; a leftover from a - SIGKILL mid-acquire is swept by :func:`prune`. - - **Atomic publication depends on hard-link support.** ``os.link`` is the - thing that makes exactly one submitter win, and a filesystem without hard - links (exFAT, FAT32, some network and container mounts) fails it with - ``OSError`` — ``ENOTSUP``/``EOPNOTSUPP``/``EPERM``, or ``ERROR_NOT_SUPPORTED`` - on Windows — rather than ``EEXIST``. That error is *not* a collision and is - not reported as one: it propagates, and no atomic lock was taken. The caller - is expected to degrade to its advisory guard and to say so - (:func:`comfy_cli.command.models.models._acquire_dest_claim`), because the - advisory guard re-scans rather than arbitrates — concurrent submitters can - race again on such a filesystem. - - The payload records the ``download_id`` that owns the claim (the pointer a - later submitter follows to decide whether the claim is still live), the - destination for a human reading the directory, and when it was taken. No - url: a presigned url is a credential-shaped thing and the claim does not - need one. - - Raises ``OSError`` for anything other than the collision - the caller - decides whether that is fatal. - """ - payload = json.dumps( - {"download_id": download_id, "dest": str(dest), "created_at": _now_iso()}, - indent=2, - ).encode("utf-8") - path = Path(path) - tmp = path.with_name(f"{path.name}.{os.getpid()}.{download_id}{CLAIM_TMP_SUFFIX}") - fd = os.open(str(tmp), os.O_CREAT | os.O_TRUNC | os.O_WRONLY, STATE_FILE_MODE) - try: - try: - view = memoryview(payload) - while view: - view = view[os.write(fd, view) :] - finally: - os.close(fd) - if sys.platform != "win32": - # `os.open`'s mode is masked by the umask, exactly as `write_path`'s - # is. Fixed up before the link, so the published claim never appears - # with a looser mode. - with contextlib.suppress(OSError): - tmp.chmod(STATE_FILE_MODE) - try: - os.link(str(tmp), str(path)) - except FileExistsError: - return False - return True - finally: - with contextlib.suppress(OSError): - tmp.unlink() - - -def read_claim(path: Path) -> str | None: - """The ``download_id`` recorded in the claim at ``path``, or None. +def _parse_claim_payload(raw: bytes | str) -> dict[str, Any] | None: + """The claim payload in ``raw``, or None for every shape we cannot resolve. - None for every unreadable shape - absent, truncated mid-write, corrupt, or - carrying an id that could not name a state file. The caller treats all of - them as *stale*, which is the safe direction: a claim nobody can resolve - would otherwise wedge its destination forever, and the record it points at - (not the claim) is what actually proves a download is live. + None for absent, empty, truncated mid-rewrite, corrupt, or carrying an id + that could not name a state file. Every one of those is treated as *stale* + by the caller, which is the safe direction now for a different reason than + it used to be: the lock, not the payload, is what proves a destination is + taken, so a payload nobody can read costs nothing but the id in an error + message. A live holder is still refused on the strength of its lock. """ try: - payload = json.loads(Path(path).read_text(encoding="utf-8")) - except (OSError, ValueError): + text = raw.decode("utf-8") if isinstance(raw, bytes) else raw + payload = json.loads(text) + except (UnicodeDecodeError, ValueError): return None if not isinstance(payload, dict): return None download_id = payload.get("download_id") if not isinstance(download_id, str) or not _SAFE_ID.match(download_id): return None - return download_id + return payload -def release_claim(path: Path, *, owner_id: str | None) -> bool: - """Drop the claim at ``path``, but only if ``owner_id`` still holds it. +def read_claim(path: Path) -> str | None: + """The ``download_id`` recorded in the claim at ``path``, read *unlocked*. - The ownership check is what keeps a finishing worker from unlinking a claim - that is no longer its own: once our record goes terminal our claim reads - stale, so a competing submitter may clear it and create *its* own in the - window before we get here, and an unconditional unlink would delete a live - claim. Never raises - releasing is bookkeeping. + Advisory only, and there is exactly one legitimate caller: a submitter that + could not take the lock and wants to name the holder in its refusal. It + cannot be used to decide ownership — the payload it reads may be a rewrite + in progress (the holder rewrites in place, so there is no atomic publish any + more), and on Windows a locked byte range makes the read fail outright. Both + come back None, which only costs the id in a message. - ``owner_id`` may be None, and the None case is load-bearing: an unreadable - claim makes :func:`read_claim` return None, so passing that None back here - means "unlink the claim nobody can read" — which is how a corrupt claim - file gets cleared instead of wedging its destination. Since claims are - published atomically (see :func:`acquire_claim`), an unreadable claim is - corrupt, not mid-write. + Ownership questions go through :func:`lock_claim` and + :meth:`ClaimLock.read_payload`, which read the same bytes under the lock. """ - path = Path(path) - if read_claim(path) != owner_id: - return False try: - path.unlink(missing_ok=True) + raw = Path(path).read_bytes() except OSError: + return None + payload = _parse_claim_payload(raw) + return payload["download_id"] if payload is not None else None + + +class ClaimLock: + """An exclusive OS lock held on one destination's claim file. + + Held for as long as the holder owns the destination — for a worker, the + whole transfer. Dropped by :meth:`release`, and by the kernel if the holder + dies, which is what makes a SIGKILLed worker cost the next submitter nothing. + + The payload is read and rewritten **in place on the locked descriptor**. + Never temp-file-plus-rename: a rename swaps the inode, and the lock lives on + the inode, so the holder would be left guarding a file nobody can see. + """ + + def __init__(self, path: Path, fd: int) -> None: + self._path = Path(path) + self._fd: int | None = fd + + @property + def path(self) -> Path: + return self._path + + @property + def held(self) -> bool: + return self._fd is not None + + def _descriptor(self) -> int: + if self._fd is None: + raise ValueError("this claim lock has already been released") + return self._fd + + def read_payload(self) -> dict[str, Any] | None: + """The claim payload, read under the lock. None when unreadable. + + Tolerant in exactly the way :func:`_parse_claim_payload` is, plus the + read itself: a claim file we hold the lock on but cannot read tells us + nothing about who owns the destination, and the lock has already + answered that question. + """ + fd = self._descriptor() + try: + os.lseek(fd, 0, os.SEEK_SET) + chunks: list[bytes] = [] + while True: + block = os.read(fd, 65536) + if not block: + break + chunks.append(block) + except OSError: + return None + return _parse_claim_payload(b"".join(chunks)) + + def write_payload(self, download_id: str, dest: str) -> None: + """Rewrite the payload in place. Raises ``OSError`` if it could not be. + + Same shape the claim has always carried, and still no url: a resolved + download url can be presigned, and a claim needs an id and a path. + """ + fd = self._descriptor() + body = json.dumps( + {"download_id": download_id, "dest": str(dest), "created_at": _now_iso()}, + indent=2, + ).encode("utf-8") + os.lseek(fd, 0, os.SEEK_SET) + os.ftruncate(fd, 0) + view = memoryview(body) + while view: + view = view[os.write(fd, view) :] + + def clear_payload(self) -> None: + """Blank the payload, leaving the (lockable) file in place. + + An empty claim parses as no payload at all, which every reader treats as + stale. Used where the old code unlinked: a submit whose worker never + started has nothing to hand the destination to. + """ + os.ftruncate(self._descriptor(), 0) + + def release(self) -> None: + """Close the descriptor, which drops the lock. Idempotent. + + Deliberately **not** an unlink. Unlinking a claim while holding its lock + leaves every descriptor already opened on that path guarding an orphan + inode, so two submitters could each hold a "lock" on a different inode + for one destination. Bounded growth is the cheaper problem, and + :func:`prune` bounds it. + """ + fd, self._fd = self._fd, None + if fd is None: + return + with contextlib.suppress(OSError): + os.close(fd) + + def __enter__(self) -> ClaimLock: + return self + + def __exit__(self, *exc: Any) -> bool: + self.release() return False + + +def _lock_fd_nb(fd: int) -> bool: + """Take the exclusive lock on ``fd`` without blocking. False when held. + + Raises ``OSError`` for anything that is not contention — a filesystem with + no locking is a degradation the caller has to handle, not a competitor. + """ + try: + if sys.platform == "win32": + os.lseek(fd, 0, os.SEEK_SET) + msvcrt.locking(fd, msvcrt.LK_NBLCK, 1) + else: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError as e: + if e.errno in _LOCK_CONTENDED_ERRNOS: + return False + raise return True +def _is_same_file(fd: int, path: Path) -> bool: + """True while ``fd`` still names the file at ``path``.""" + try: + here = os.fstat(fd) + there = os.stat(path) + except OSError: + return False + return (here.st_dev, here.st_ino) == (there.st_dev, there.st_ino) + + +def lock_claim(path: Path, *, blocking_timeout: float = 0.0) -> ClaimLock | None: + """Take the exclusive lock on the claim at ``path``, or None if it is held. + + ``O_CREAT`` without ``O_EXCL``: the claim file is a lock target that outlives + every holder, so it usually already exists and creating it decides nothing. + What decides is the lock. + + ``blocking_timeout`` bounds a non-blocking retry loop rather than issuing a + blocking lock, so the wait can never outlive the caller's patience. Zero (the + default) is a single attempt — the submit path wants an immediate answer, + because a lock that is held *is* the answer. A freshly spawned worker passes + a few seconds, because a competitor may hold the lock for the moment it takes + to inspect the payload. + + **The re-stat guard.** Between our ``open`` and our lock, the file we opened + can be unlinked (only :func:`prune` does this, and only on POSIX) and a + different one created at the same path. We would then hold a real lock on an + orphan inode while another process holds a real lock on the live one. So + after locking, ``fstat`` the descriptor against ``stat`` of the path: a + different (st_dev, st_ino) — or a path that is gone — means we lost the race, + so close and retry, up to :data:`CLAIM_LOCK_RESTAT_ATTEMPTS` times. Windows + skips the guard (inode identity there is not reliable) and pays for it by + never unlinking a claim file at all, which removes the race instead. + + Raises ``OSError`` when the claim cannot be opened or the filesystem cannot + lock. That is not a collision and must not be reported as one: the caller + degrades to its advisory guard and says so. + """ + path = Path(path) + deadline = time.monotonic() + max(0.0, float(blocking_timeout)) + for _ in range(CLAIM_LOCK_RESTAT_ATTEMPTS): + fd = os.open(str(path), os.O_CREAT | os.O_RDWR | _O_BINARY, STATE_FILE_MODE) + keep = False + try: + if sys.platform != "win32": + # `os.open`'s mode is masked by the umask, exactly as + # `write_path`'s is, and the file may predate this code. + with contextlib.suppress(OSError): + os.fchmod(fd, STATE_FILE_MODE) + while not _lock_fd_nb(fd): + if time.monotonic() >= deadline: + return None + time.sleep(CLAIM_LOCK_POLL_S) + if sys.platform == "win32" or _is_same_file(fd, path): + keep = True + return ClaimLock(path, fd) + finally: + if not keep: + with contextlib.suppress(OSError): + os.close(fd) + # Every attempt raced a concurrent unlink. Reporting "held" rather than + # "acquired" is the safe direction: the caller refuses instead of letting a + # second transfer into a destination whose ownership we could not settle. + return None + + @dataclass class DownloadState: id: str @@ -628,20 +816,31 @@ def _remove_record(path: Path) -> bool: def _sweep_claims(workspace: Path) -> None: - """Drop claims that no longer point at a live download, and dead temp files. - - The liveness rule is the same one the submit path applies before clearing a - stale claim: a claim stays while its ``download_id`` resolves to a record - that is active after :func:`reconcile`, or whose worker process is still - provably alive (a terminal status does not prove the process is gone — a - cancelled worker may still be mid-write for a moment). The unlink goes - through :func:`release_claim` with the id we judged stale, so a claim that - changes hands between the read and the unlink is left alone. - - Temp files are :func:`acquire_claim`'s private staging names; one only - survives a SIGKILL inside the milliseconds between write and publish, so - anything older than :data:`CLAIM_TMP_MAX_AGE_S` is a crashed acquire's - leftover. Best effort throughout, like the rest of :func:`prune`. + """Drop claim files nothing owns any more, and legacy temp leftovers. + + A claim file is never unlinked on release (see :class:`ClaimLock`), so this + is the only thing that bounds ``claims/`` below "every destination this + workspace ever downloaded to". It is a nicety, not a correctness step: a + left-behind claim costs ~100 bytes and refuses nobody, because ownership is + the *lock*, which no longer exists once its holder is gone. + + Two conditions, both required, and the lock has to be held for the unlink: + + * we can take the lock — so no transfer owns the destination right now; and + * the payload does not name a download that is still live. That second check + is what keeps the sweep out of the **spawn gap**: between a submit's + release and its worker's lock nobody holds the claim, and unlinking it + there would make the worker come up, find an empty payload, and refuse the + destination it was just handed. + + POSIX only. On Windows a claim file is never unlinked at all — inode + identity is unreliable there, so :func:`lock_claim` cannot run its re-stat + guard, and without that guard an unlink is how two holders end up locking + two different inodes for one destination. + + Temp files are the previous implementation's private staging names, kept + only so a workspace upgraded from it does not carry them forever. Best + effort throughout, like the rest of :func:`prune`. """ base = Path(workspace) / STATE_DIRNAME / CLAIMS_DIRNAME try: @@ -657,14 +856,25 @@ def _sweep_claims(workspace: Path) -> None: if path.stat().st_mtime < tmp_cutoff: path.unlink() continue - if not path.name.endswith(".claim"): + if not path.name.endswith(".claim") or sys.platform == "win32": continue - download_id = read_claim(path) - if download_id is not None: - record = read(workspace, download_id) - if record is not None and (reconcile(record).status in ACTIVE_STATUSES or worker_alive(record)): - continue - release_claim(path, owner_id=download_id) + try: + lock = lock_claim(path) + except OSError: + continue + if lock is None: + continue + try: + payload = lock.read_payload() + download_id = payload["download_id"] if payload is not None else None + if download_id is not None: + record = read(workspace, download_id) + if record is not None and (reconcile(record).status in ACTIVE_STATUSES or worker_alive(record)): + continue + with contextlib.suppress(OSError): + path.unlink() + finally: + lock.release() def prune(workspace: Path) -> int: @@ -689,13 +899,12 @@ def prune(workspace: Path) -> int: An in-flight record (``starting``/``downloading``) is never touched at any age, and never counts toward — or is evicted by — the cap. - Also sweeps ``claims/`` (see :func:`_sweep_claims`): a claim is normally - released by its own worker, and a stranded one only self-clears on the next - submit *to the same destination*, so claims for destinations never - re-submitted would otherwise accumulate for the life of the workspace — - the same unbounded growth the record rules above exist to prevent. Swept - claims do not count toward the returned total, which stays "records - removed". + Also sweeps ``claims/`` (see :func:`_sweep_claims`): a claim file is never + unlinked on release — the lock lives on its inode — so without this the + directory would grow to one small file per destination the workspace has + ever downloaded to, the same unbounded growth the record rules above exist + to prevent. Swept claims do not count toward the returned total, which stays + "records removed". Every step is best effort, exactly like :func:`write_path`'s OSError handling: a read-only state directory, a permissions problem, or a file a diff --git a/comfy_cli/error_codes.py b/comfy_cli/error_codes.py index 4f71cf42..b1c28c28 100644 --- a/comfy_cli/error_codes.py +++ b/comfy_cli/error_codes.py @@ -797,25 +797,16 @@ class ErrorCode: ), ErrorCode( "model_download_claim_contested", - "`comfy model download --background` lost the race for a destination it had just judged " - "free: the stale claim it cleared was re-taken by another submitter before its own retry, " - "and that new claim does not (yet) resolve to a live download record. `details.path` is the " - "destination; `details.download_id` names the new claim's holder when its claim file was " - "readable, and is null otherwise. Unlike `model_download_in_flight` there is no `status`/" - "`kind` to report — the competitor's record was not visible at refusal time.", + "`comfy model download --background` found the destination's claim lock held by another " + "process, and that holder does not resolve to a readable download record. Holding the lock " + "is what proves the destination is taken, so the refusal stands regardless — but there is " + "no record to quote, which is why this is a distinct code rather than a " + "`model_download_in_flight` missing its `status`/`kind` fields. `details.path` is the " + "destination, `details.claim_file` the claim the lock is held on, and " + "`details.download_id` the holder's recorded id when the claim payload was readable " + "(null otherwise — it may be a rewrite in progress, or locked against reading on Windows).", "check `comfy model downloads`, then retry", ), - ErrorCode( - "model_download_claim_unclearable", - "`comfy model download --background` found a stale destination claim it could not remove " - "(`details.claim_file`): the file is not deletable by this user, or something else (e.g. a " - "directory) sits at the claim path. Every submission to `details.path` will be refused " - "until the claim file is cleared, so the command reports the real obstacle rather than a " - "phantom in-flight download. `details.download_id` is the stale claim's recorded holder, " - "null when the claim was unreadable.", - "remove the claim file by hand (check its ownership and the permissions on the `claims/` " - "directory), then retry", - ), ErrorCode( "model_download_foreground_cancel", "`comfy model download-cancel` refused to cancel a download that is running in the " diff --git a/tests/comfy_cli/command/test_model_download_background.py b/tests/comfy_cli/command/test_model_download_background.py index da2b5ca3..193a741d 100644 --- a/tests/comfy_cli/command/test_model_download_background.py +++ b/tests/comfy_cli/command/test_model_download_background.py @@ -9,10 +9,12 @@ from __future__ import annotations +import contextlib import errno import json import logging import os +import signal import subprocess import sys import threading @@ -1149,15 +1151,116 @@ def _claim_owner(workspace) -> str | None: return download_state.read_claim(files[0]) +def _plant_claim(path: Path, download_id: str, dest: str) -> None: + """Stamp a claim payload the way a submitter does, holding no lock afterwards.""" + lock = download_state.lock_claim(path) + assert lock is not None, f"the claim at {path} was already held" + try: + lock.write_payload(download_id, dest) + finally: + lock.release() + + +def _is_lockable(path: Path) -> bool: + """True when nobody is holding the claim at ``path`` — i.e. nobody owns the destination.""" + lock = download_state.lock_claim(path) + if lock is None: + return False + lock.release() + return True + + +# The repo root, so a helper process started with `-c` can import `comfy_cli` +# without depending on how pytest was invoked. +_REPO_ROOT = str(Path(download_state.__file__).resolve().parents[1]) + +# Holds a real OS lock on a claim file in a *separate process*, which is the only +# way to prove the thing the held-lock design rests on: the lock, not the record +# it points at, is what says a destination is taken — and the kernel drops it +# when the holder dies. +_CLAIM_HOLDER_SCRIPT = """ +import sys, time +sys.path.insert(0, sys.argv[1]) +from comfy_cli import download_state +lock = download_state.lock_claim(sys.argv[2], blocking_timeout=20) +if lock is None: + sys.stderr.write("could not take the lock\\n") + sys.exit(1) +lock.write_payload(sys.argv[3], sys.argv[4]) +sys.stdout.write("held\\n") +sys.stdout.flush() +# Self-bounded: the parent kills this, but it must never outlive a crashed test. +time.sleep(120) +""" + + +# One real process submitting one background download, used by the cross-process +# race probe. Threads share a process, and a process-local mistake (a module +# global, a re-entrant guard) can make a threaded probe pass while the real +# thing — several `comfy model download --background` invocations — still races. +# Only separate processes exercise the kernel-level arbitration the lock is. +_RACE_SUBMIT_SCRIPT = """ +import pathlib, sys, time +sys.path.insert(0, sys.argv[1]) +from comfy_cli.command.models import models +root, ws, dest, ready, gate = (pathlib.Path(a) for a in sys.argv[1:6]) +models.get_workspace = lambda: ws +models._spawn_download_worker = lambda state_file, log_file: 31337 +(ready / str(__import__("os").getpid())).touch() +# Self-bounded: never outlive a crashed parent that will not open the gate. +deadline = time.monotonic() + 60 +while not gate.exists(): + if time.monotonic() > deadline: + sys.exit(4) + time.sleep(0.001) +try: + models._submit_background_download( + url="https://example.com/m.safetensors", + dest=dest, + downloader="httpx", + needs_civitai_auth=False, + needs_hf_auth=False, + ) +except BaseException: + sys.exit(3) +sys.exit(0) +""" + + +@contextlib.contextmanager +def _claim_holder_process(claim: Path, download_id: str, dest: str): + """Yield a live process holding the lock on ``claim``. Always reaped by pid.""" + proc = subprocess.Popen( + [sys.executable, "-c", _CLAIM_HOLDER_SCRIPT, _REPO_ROOT, str(claim), download_id, str(dest)], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + if (proc.stdout.readline() or "").strip() != "held": + proc.kill() + pytest.fail(f"the claim holder never took the lock: {proc.stderr.read()!r}") + yield proc + finally: + if proc.poll() is None: + proc.kill() + with contextlib.suppress(subprocess.TimeoutExpired): + proc.wait(timeout=30) + proc.stdout.close() + proc.stderr.close() + + class TestAtomicDestinationClaim: - """The `O_EXCL` claim file: creating it is what decides who owns a destination. + """The held claim lock: holding it is what decides who owns a destination. The record-is-the-claim guard it replaces was check-then-act — a submitter wrote its record and then re-scanned for competitors — and a re-scan cannot see a record that has not landed yet. When A scanned before B wrote *and* `_claim_order` ranked B first, neither side withdrew and both streamed a full - copy. File creation with `O_CREAT | O_EXCL` has no such window: the kernel - hands the file to exactly one caller. + copy. An exclusive OS lock has no such window: the kernel hands it to exactly + one caller, and — unlike the `O_EXCL` create that came between — it stays + held for the transfer's lifetime, so ownership never has to be re-derived + from a record that may read terminal while its worker is still running. """ DEST = ("models/loras", "m.safetensors") @@ -1228,10 +1331,10 @@ def test_a_competitor_landing_inside_the_old_window_still_leaves_one_winner( The old guard's window is `[our re-scan -> the competitor's write]`: a competitor whose record lands *after* we look is invisible to us, so our re-scan finds nothing and we proceed. Here the competitor is planted in - exactly that window — after our re-scan, before our claim — with an + exactly that window — after our re-scan, before our lock — with an identical `started_at` and a lexically smaller id, so `_claim_order` ranks it first and the old code's re-scan (which never saw it) could not - have used that. Only the `O_EXCL` create can decide this, and it does. + have used that. The claim settles it instead. """ dest = self._dest(workspace) stamp = datetime.now(timezone.utc).isoformat(timespec="seconds") @@ -1241,86 +1344,68 @@ def test_a_competitor_landing_inside_the_old_window_still_leaves_one_winner( competitor = _state(id="aaaaaaaaaaaa", dest=str(dest), status="starting", pid=None) competitor.started_at = stamp - real_acquire = download_state.acquire_claim + real_lock = download_state.lock_claim planted: list = [] - def acquire(path, *, download_id, dest): + def lock_claim(path, **kwargs): if not planted: # The window the old guard could only narrow: our re-scan has - # already run and found nothing, and the competitor lands now. + # already run and found nothing, and the competitor lands now — + # record first, then its stamp on the claim. + planted.append(True) download_state.write(workspace, competitor) - planted.append(real_acquire(path, download_id=competitor.id, dest=dest)) - return real_acquire(path, download_id=download_id, dest=dest) + _plant_claim(path, competitor.id, str(dest)) + return real_lock(path, **kwargs) - monkeypatch.setattr(download_state, "acquire_claim", acquire) + monkeypatch.setattr(download_state, "lock_claim", lock_claim) - with patch("comfy_cli.utils.is_running", return_value=True): - with pytest.raises(typer.Exit) as exc: - self._submit(dest) + with pytest.raises(typer.Exit) as exc: + self._submit(dest) assert exc.value.exit_code == 1 - assert json_renderer()["error"]["code"] == "model_download_in_flight" - # Exactly one active record and exactly one claim, and the survivor is - # the competitor *regardless* of `_claim_order` — which ranks it first - # here, and would have ranked it second on the opposite id draw without - # changing the outcome. + env = json_renderer() + assert env["error"]["code"] == "model_download_in_flight" + assert env["error"]["details"]["download_id"] == competitor.id + # Exactly one active record and one claim, and the survivor is the + # competitor *regardless* of `_claim_order`, which ranks it first here. assert [s.id for s in download_state.list_all(workspace)] == [competitor.id] assert _claim_owner(workspace) == competitor.id ours = download_state.DownloadState(id="zzzzzzzzzzzz", url="u", dest=str(dest), started_at=stamp) assert models._claim_order(competitor) < models._claim_order(ours) - def test_a_stale_claim_replaced_by_a_live_one_is_not_cleared(self, workspace, monkeypatch, json_renderer): - """The stale-clear is conditional on the id we read, not unconditional. - - Between reading a claim and calling it stale sits a state-file read and a - `reconcile`, and in that gap the claim's worker can finish, release it, - and a fresh submitter take a *live* claim at the same path. Unlinking - that one would leave two downloads owning one destination. + @pytest.mark.skipif(sys.platform == "win32", reason="inode identity is not reliable on Windows") + def test_stealing_a_stale_claim_keeps_the_same_inode(self, workspace, monkeypatch, json_renderer): + """The held-lock replacement for the old conditional stale-clear. + + Clearing a stale claim used to be unlink-then-re-create, and between + those two steps another submitter could take a *live* claim at the same + path — which the old code then had to avoid unlinking by re-reading the + id it had just judged stale. There is no such window now: the steal is a + rewrite of the payload on the descriptor we already hold the lock on. The + inode is what the lock lives on, so the assertion that it never changes + is the assertion that no competitor could have been holding a lock on a + different file for this destination. """ dest = self._dest(workspace) dead = _state(dest=str(dest), status="downloading", pid=4242, total_bytes=4096) dead.started_at = _stamp(download_state.STARTUP_GRACE_S * 3) download_state.write(workspace, dead) claim = download_state.claim_path(workspace, models._dest_key(dest)) - assert download_state.acquire_claim(claim, download_id=dead.id, dest=str(dest)) - - successor = _state(dest=str(dest), status="downloading", pid=1234, total_bytes=4096) - - real_holder = models._claim_holder - swapped: list = [] - - def claim_holder(path): - result = real_holder(path) - if not swapped: - # We have just decided the claim is stale. The dead worker's - # record was released and a successor took it over in this exact - # instant — after our read, before our unlink. - swapped.append(True) - download_state.write(workspace, successor) - path.unlink() - assert download_state.acquire_claim(path, download_id=successor.id, dest=str(dest)) - return result - - monkeypatch.setattr(models, "_claim_holder", claim_holder) + _plant_claim(claim, dead.id, str(dest)) + before = claim.stat().st_ino self._no_spawn(monkeypatch) - - # `download` prunes on entry, and prune sweeps stale claims — which would - # clear the dead claim before `_acquire_dest_claim` ever collides with it, - # so the staged takeover above would never run. The sweep has its own - # tests; here it is disabled to keep the takeover window open. + # `download` prunes on entry and the sweep would unlink this claim before + # `_acquire_dest_claim` ever reaches it, so the steal under the lock — + # the thing under test — would never run. The sweep has its own tests. monkeypatch.setattr(download_state, "prune", lambda ws: 0) - with patch("comfy_cli.utils.is_running", side_effect=lambda pid: pid == successor.pid): - with pytest.raises(typer.Exit) as exc: - self._submit(dest) + with patch("comfy_cli.utils.is_running", return_value=False): + self._submit(dest) - assert exc.value.exit_code == 1 env = json_renderer() - assert env["error"]["code"] == "model_download_in_flight" - assert env["error"]["details"]["download_id"] == successor.id - # The successor's claim survived, and we left no record of our own. - assert _claim_owner(workspace) == successor.id - assert sorted(s.id for s in download_state.list_all(workspace)) == sorted([dead.id, successor.id]) + assert env["ok"] is True + assert _claim_owner(workspace) == env["data"]["download_id"] != dead.id + assert claim.stat().st_ino == before, "the claim was re-created rather than rewritten in place" @pytest.mark.parametrize("iteration", range(3)) def test_twelve_simultaneous_submits_accept_exactly_one(self, workspace, monkeypatch, json_renderer, iteration): @@ -1367,17 +1452,66 @@ def submit() -> None: assert len(live) == 1 and live[0].status in download_state.ACTIVE_STATUSES assert _claim_owner(workspace) == live[0].id - def test_a_claim_left_by_a_killed_worker_self_clears(self, workspace, monkeypatch, json_renderer): - """SIGKILL leaves the claim file on disk; nothing sweeps it and nothing - needs to. Liveness is the *record* the claim points at, and that record - reconciles to `failed` once its pid is gone — so the claim reads stale and - the next submitter clears it in passing.""" + def test_four_real_processes_racing_one_destination_accept_exactly_one(self, workspace, tmp_path): + """The same acceptance criterion as the threaded probe, across real + processes. Threads share a heap, so a threaded probe can pass on a guard + that is only process-local; separate processes are what the kernel's lock + actually arbitrates between, and what a user running several + `comfy model download --background` commands actually produces.""" + dest = workspace / self.DEST[0] / "cross-process.safetensors" + ready = tmp_path / "ready" + ready.mkdir() + gate = tmp_path / "gate" + + procs = [ + subprocess.Popen( + [ + sys.executable, + "-c", + _RACE_SUBMIT_SCRIPT, + _REPO_ROOT, + str(workspace), + str(dest), + str(ready), + str(gate), + ], + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + text=True, + ) + for _ in range(4) + ] + try: + deadline = time.monotonic() + 120 + while len(list(ready.iterdir())) < len(procs): + assert time.monotonic() < deadline, "the racers never came up" + assert all(p.poll() is None for p in procs), "a racer died before the gate opened" + time.sleep(0.02) + gate.touch() + codes = [p.wait(timeout=120) for p in procs] + finally: + for proc in procs: + if proc.poll() is None: + proc.kill() + with contextlib.suppress(subprocess.TimeoutExpired): + proc.wait(timeout=30) + proc.stderr.close() + + assert sorted(codes) == [0, 3, 3, 3], codes + live = [s for s in download_state.list_all(workspace) if s.dest == str(dest)] + assert len(live) == 1 and live[0].status in download_state.ACTIVE_STATUSES + assert download_state.read_claim(download_state.claim_path(workspace, models._dest_key(dest))) == live[0].id + + def test_a_claim_left_by_a_killed_worker_is_taken_over(self, workspace, monkeypatch, json_renderer): + """SIGKILL drops the lock in the kernel, so the claim file that is left on + disk owns nothing. The next submitter locks it, finds a payload naming a + record reconcile has already demoted, and stamps its own id in.""" dest = self._dest(workspace) dead = _state(dest=str(dest), status="downloading", pid=4242, total_bytes=4096) dead.started_at = _stamp(download_state.STARTUP_GRACE_S * 3) download_state.write(workspace, dead) claim = download_state.claim_path(workspace, models._dest_key(dest)) - assert download_state.acquire_claim(claim, download_id=dead.id, dest=str(dest)) + _plant_claim(claim, dead.id, str(dest)) self._no_spawn(monkeypatch) with patch("comfy_cli.utils.is_running", return_value=False): @@ -1396,7 +1530,7 @@ def test_a_claim_whose_record_is_gone_is_stale(self, workspace, monkeypatch, jso wedge the destination forever. There is no download to point a user at.""" dest = self._dest(workspace) claim = download_state.claim_path(workspace, models._dest_key(dest)) - assert download_state.acquire_claim(claim, download_id="deadbeefcafe", dest=str(dest)) + _plant_claim(claim, "deadbeefcafe", str(dest)) self._no_spawn(monkeypatch) self._submit(dest) @@ -1424,36 +1558,39 @@ def test_an_unreadable_claim_is_stale(self, workspace, monkeypatch, json_rendere assert json_renderer()["ok"] is True assert _claim_owner(workspace) is not None - def test_a_second_collision_refuses_instead_of_retrying_forever(self, workspace, monkeypatch, json_renderer): - """Stale claims get exactly one retry. A second collision means another - submitter won the create we just freed up, not that the claim is wedged - - and clearing theirs too is how two submitters livelock each other.""" + def test_a_held_claim_is_refused_on_the_first_attempt(self, workspace, monkeypatch, json_renderer): + """A lock somebody else is holding gets no retry and no liveness check. + + The old code cleared a claim it judged stale and re-created it, so it + needed exactly one retry and a distinct refusal for the second collision. + Holding removes both: the lock being unavailable *is* the answer, so one + attempt is all there is. + """ dest = self._dest(workspace) rival_id = "cccccccccccc" + claim = download_state.claim_path(workspace, models._dest_key(dest)) + _plant_claim(claim, rival_id, str(dest)) - real_acquire = download_state.acquire_claim calls: list[int] = [] - def acquire(path, *, download_id, dest): + def held(path, **kwargs): calls.append(1) - if len(calls) <= 2: - # First: a stale claim is already there. Second: the rival got in - # between our unlink and our retry. - path.write_text(json.dumps({"download_id": rival_id, "dest": dest}), encoding="utf-8") - return False - return real_acquire(path, download_id=download_id, dest=dest) + return None - monkeypatch.setattr(download_state, "acquire_claim", acquire) + monkeypatch.setattr(download_state, "lock_claim", held) self._no_spawn(monkeypatch) + # Prune's claim sweep locks too; disabled so the count below is the + # submit path's own attempts and nothing else. + monkeypatch.setattr(download_state, "prune", lambda ws: 0) with pytest.raises(typer.Exit) as exc: self._submit(dest) assert exc.value.exit_code == 1 - assert len(calls) == 2, "the retry must happen exactly once" + assert len(calls) == 1, "a held lock must not be retried" env = json_renderer() - # A distinct code, not `model_download_in_flight`: the rival's claim did - # not resolve to a live record, so the status/kind that code documents + # A distinct code, not `model_download_in_flight`: the holder's payload + # did not resolve to a record, so the status/kind that code documents # could only have been invented. assert env["error"]["code"] == "model_download_claim_contested" assert env["error"]["details"]["download_id"] == rival_id @@ -1483,7 +1620,7 @@ def test_an_unusable_claims_directory_degrades_to_the_advisory_guard(self, works def test_a_degraded_claim_is_reported_once(self, workspace, monkeypatch, json_renderer, caplog): """Degrading is right; degrading *silently* is not. The advisory guard - re-scans rather than arbitrates, so on a filesystem with no hard links + re-scans rather than arbitrates, so on a filesystem that cannot lock every submit quietly gives up the atomicity this path exists for. Once per process, because it is a property of the filesystem, not of the run. """ @@ -1491,7 +1628,7 @@ def test_a_degraded_claim_is_reported_once(self, workspace, monkeypatch, json_re monkeypatch.setattr(models, "_claims_degraded_reported", False) monkeypatch.setattr( download_state, - "acquire_claim", + "lock_claim", MagicMock(side_effect=OSError(errno.EOPNOTSUPP, "Operation not supported")), ) self._no_spawn(monkeypatch) @@ -1502,7 +1639,7 @@ def test_a_degraded_claim_is_reported_once(self, workspace, monkeypatch, json_re warnings = [r for r in caplog.records if "atomic destination claims are unavailable" in r.message] assert len(warnings) == 1 - assert "does not support the hard link" in warnings[0].getMessage() + assert "does not support the exclusive lock" in warnings[0].getMessage() # A second submit on the same process stays quiet. caplog.clear() @@ -1512,14 +1649,14 @@ def test_a_degraded_claim_is_reported_once(self, workspace, monkeypatch, json_re assert [r for r in caplog.records if "atomic destination claims" in r.message] == [] def test_a_transient_claim_failure_is_reported_as_such(self, workspace, monkeypatch, json_renderer, caplog): - """A read-only state dir is not a filesystem without hard links. Both + """A read-only state dir is not a filesystem that cannot lock. Both degrade identically; only the wording differs, so the message does not send someone hunting a filesystem capability that is not the problem.""" dest = self._dest(workspace) monkeypatch.setattr(models, "_claims_degraded_reported", False) monkeypatch.setattr( download_state, - "acquire_claim", + "lock_claim", MagicMock(side_effect=OSError(errno.EIO, "Input/output error")), ) self._no_spawn(monkeypatch) @@ -1530,11 +1667,13 @@ def test_a_transient_claim_failure_is_reported_as_such(self, workspace, monkeypa warnings = [r for r in caplog.records if "atomic destination claims are unavailable" in r.message] assert len(warnings) == 1 - assert "the claim could not be written" in warnings[0].getMessage() + assert "the claim could not be locked" in warnings[0].getMessage() - def test_a_failed_spawn_releases_the_claim(self, workspace, monkeypatch, json_renderer): - """Nothing else ever would: no worker starts, so no worker reaches the - terminal transition that releases it.""" + def test_a_failed_spawn_blanks_the_claim(self, workspace, monkeypatch, json_renderer): + """No worker ever comes up to take the lock for this record, so the id we + stamped names a download that will never run. Blanked rather than + unlinked: the claim file is where the lock lives, and an empty payload + resolves to nobody, which is the same thing every reader wanted.""" dest = self._dest(workspace) def boom(state_file, log_file): @@ -1545,7 +1684,11 @@ def boom(state_file, log_file): self._submit(dest) assert json_renderer()["error"]["code"] == "download_worker_spawn_failed" - assert _claim_files(workspace) == [] + claims = _claim_files(workspace) + assert len(claims) == 1 + assert claims[0].read_bytes() == b"" + assert download_state.read_claim(claims[0]) is None + assert _is_lockable(claims[0]) def test_a_live_foreground_competitor_is_refused_before_any_claim_is_taken( self, workspace, monkeypatch, json_renderer @@ -1600,20 +1743,38 @@ def test_the_claim_does_not_persist_the_url(self, workspace, monkeypatch, json_r assert "super-secret" not in body assert json.loads(body).keys() == {"download_id", "dest", "created_at"} - def test_a_cancelled_record_with_a_live_worker_still_blocks(self, workspace, monkeypatch, json_renderer): - """A terminal status does not prove the process exited: a cancelled - worker is still mid-write until its own terminal transition lands, and - it releases the claim itself right after. Clearing its claim on the - strength of the status alone would let a second transfer into a - destination a live worker may still be writing.""" + def test_a_cancelled_record_whose_worker_still_holds_the_lock_blocks(self, workspace, monkeypatch, json_renderer): + """Finding 2, first leg: the record reads terminal while its worker runs. + + `download-cancel` writes `cancelled` even when `stop_worker` fails, so a + record can say terminal while the worker is still running and still + writing the destination. The derived guard then asked `worker_alive`, + which needs the recorded pid *and* its start time to agree — and a start + time that no longer matches (a recycled pid, a record written by an + earlier run) answers "dead" for a process that is very much alive, so the + destination was handed to a second transfer. The lock cannot be wrong + about this: the worker is holding it. The holder is a real process here, + because nothing about that can be faked with a record on disk. + """ dest = self._dest(workspace) - holder = _state(dest=str(dest), status="cancelled", pid=4242) - download_state.write(workspace, holder) claim = download_state.claim_path(workspace, models._dest_key(dest)) - assert download_state.acquire_claim(claim, download_id=holder.id, dest=str(dest)) self._no_spawn(monkeypatch) - with patch("comfy_cli.utils.is_running", return_value=True): + with _claim_holder_process(claim, "eeeeeeeeeeee", str(dest)) as holder_proc: + holder = _state( + id="eeeeeeeeeeee", + dest=str(dest), + status="cancelled", + pid=holder_proc.pid, + pid_create_time=1.0, + ) + download_state.write(workspace, holder) + # Both halves of the derived liveness question answer "stale" here, + # which is the whole point: this is the state in which the previous + # design let a second transfer into the destination. + assert download_state.reconcile(holder).is_terminal + assert not download_state.worker_alive(holder) + with pytest.raises(typer.Exit) as exc: self._submit(dest) @@ -1621,21 +1782,83 @@ def test_a_cancelled_record_with_a_live_worker_still_blocks(self, workspace, mon env = json_renderer() assert env["error"]["code"] == "model_download_in_flight" assert env["error"]["details"]["download_id"] == holder.id + assert env["error"]["details"]["status"] == "cancelled" + # We withdrew our own record rather than stealing theirs. assert _claim_owner(workspace) == holder.id - # We withdrew our own record rather than clearing theirs. assert [s.id for s in download_state.list_all(workspace)] == [holder.id] - def test_an_unclearable_stale_claim_is_reported_not_a_phantom(self, workspace, monkeypatch, json_renderer): - """A stale claim whose unlink fails would otherwise collide again on the - retry and be reported as `model_download_in_flight` naming a download - that does not exist — forever, on every submission. Name the real - obstacle instead.""" + def test_a_live_holder_blocks_even_when_its_record_cannot_be_read(self, workspace, monkeypatch, json_renderer): + """Finding 2, second leg: the record the claim points at is unreadable. + + Deriving ownership meant resolving the payload's id to a record, so every + way that read could come back empty — the record pruned, deleted, or + transiently unreadable — was indistinguishable from "the download is + gone" and cleared a claim whose worker was still transferring. Holding + asks the kernel instead, and the kernel is not guessing. + """ + dest = self._dest(workspace) + claim = download_state.claim_path(workspace, models._dest_key(dest)) + self._no_spawn(monkeypatch) + + with _claim_holder_process(claim, "eeeeeeeeeeee", str(dest)): + assert download_state.read(workspace, "eeeeeeeeeeee") is None + + with pytest.raises(typer.Exit) as exc: + self._submit(dest) + + assert exc.value.exit_code == 1 + env = json_renderer() + # No record to quote a status or kind from, so the honest refusal is the + # contested code rather than an in-flight one missing its fields. + assert env["error"]["code"] == "model_download_claim_contested" + assert env["error"]["details"]["download_id"] == "eeeeeeeeeeee" + assert _claim_owner(workspace) == "eeeeeeeeeeee" + assert download_state.list_all(workspace) == [] + + @pytest.mark.skipif(sys.platform == "win32", reason="SIGKILL has no Windows equivalent") + def test_a_sigkilled_holder_frees_the_destination_with_no_grace_wait(self, workspace, monkeypatch, json_renderer): + """The kernel drops a dead process's lock. Nothing has to sweep it, and + nobody has to wait out `STARTUP_GRACE_S` — which the record deliberately + sits inside here, so a grace wait would be visible as a refusal.""" + dest = self._dest(workspace) + claim = download_state.claim_path(workspace, models._dest_key(dest)) + self._no_spawn(monkeypatch) + + with _claim_holder_process(claim, "ffffffffffff", str(dest)) as holder_proc: + holder = _state( + id="ffffffffffff", + dest=str(dest), + status="downloading", + pid=holder_proc.pid, + pid_create_time=download_state.process_create_time(holder_proc.pid), + total_bytes=4096, + ) + download_state.write(workspace, holder) + assert not _is_lockable(claim), "the holder is not actually holding the lock" + + os.kill(holder_proc.pid, signal.SIGKILL) + holder_proc.wait(timeout=30) + + started = time.monotonic() + self._submit(dest) + elapsed = time.monotonic() - started + + env = json_renderer() + assert env["ok"] is True + assert env["data"]["download_id"] != holder.id + assert _claim_owner(workspace) == env["data"]["download_id"] + assert elapsed < download_state.STARTUP_GRACE_S / 2, f"the submit waited {elapsed:.1f}s" + + def test_a_submit_inside_another_submits_spawn_gap_is_refused(self, workspace, monkeypatch, json_renderer): + """The one window a bare lock cannot arbitrate: between a submitter's + release and its worker's lock nobody holds the claim. The payload closes + it — an active record younger than the startup grace is a download whose + worker has not booted yet, and the destination is its own.""" dest = self._dest(workspace) + starting = _state(dest=str(dest), status="starting", pid=None) + download_state.write(workspace, starting) claim = download_state.claim_path(workspace, models._dest_key(dest)) - assert download_state.acquire_claim(claim, download_id="abcdefabcdef", dest=str(dest)) - # No record for that id, so the claim is stale; the unlink is made to - # fail while the claim keeps naming the same dead holder. - monkeypatch.setattr(download_state, "release_claim", lambda path, *, owner_id: False) + _plant_claim(claim, starting.id, str(dest)) self._no_spawn(monkeypatch) with pytest.raises(typer.Exit) as exc: @@ -1643,22 +1866,74 @@ def test_an_unclearable_stale_claim_is_reported_not_a_phantom(self, workspace, m assert exc.value.exit_code == 1 env = json_renderer() - assert env["error"]["code"] == "model_download_claim_unclearable" - assert env["error"]["details"]["claim_file"] == str(claim) - assert env["error"]["details"]["download_id"] == "abcdefabcdef" - # We withdrew our own record, so nothing phantom is left behind. - assert download_state.list_all(workspace) == [] + assert env["error"]["code"] == "model_download_in_flight" + assert env["error"]["details"]["download_id"] == starting.id + assert _claim_owner(workspace) == starting.id + assert [s.id for s in download_state.list_all(workspace)] == [starting.id] + + def test_a_spawn_gap_that_outlived_the_grace_is_stolen(self, workspace, monkeypatch, json_renderer): + """Past the grace the arbitration stops: a download that is genuinely + live holds the lock, so having got this far means it is not. Refusing on + the payload alone here would be the derived-ownership guess this design + replaced — and would wedge the destination of a submitter that died + before it could spawn anything.""" + dest = self._dest(workspace) + abandoned = _state(dest=str(dest), status="starting", pid=None) + abandoned.started_at = _stamp(download_state.STARTUP_GRACE_S * 2) + download_state.write(workspace, abandoned) + claim = download_state.claim_path(workspace, models._dest_key(dest)) + _plant_claim(claim, abandoned.id, str(dest)) + self._no_spawn(monkeypatch) + + self._submit(dest) + + env = json_renderer() + assert env["ok"] is True + assert _claim_owner(workspace) == env["data"]["download_id"] != abandoned.id + + def test_a_payload_we_cannot_write_degrades_to_the_advisory_guard( + self, workspace, monkeypatch, json_renderer, caplog + ): + """The lock is ours but the stamp will not go down (ENOSPC, EIO). A lock + we release without stamping tells the worker we are about to spawn + nothing, so this is the same degradation as a filesystem that cannot lock + at all — warn, fall back to the re-scan, and do not turn a download that + used to work into an error.""" + dest = self._dest(workspace) + monkeypatch.setattr(models, "_claims_degraded_reported", False) + monkeypatch.setattr( + download_state.ClaimLock, + "write_payload", + MagicMock(side_effect=OSError(errno.ENOSPC, "No space left on device")), + ) + self._no_spawn(monkeypatch) + + with caplog.at_level(logging.WARNING, logger=models.__name__): + self._submit(dest) + + assert json_renderer()["ok"] is True + assert any("atomic destination claims are unavailable" in r.message for r in caplog.records) + # The claim was left lockable rather than half-stamped. + assert _is_lockable(_claim_files(workspace)[0]) class TestWorkerReleasesTheClaim: - """Every terminal transition hands the destination back.""" + """The worker holds the claim lock for the transfer and hands it back on + every exit — terminal transition, cancel, or an OSError on the first write. + + "Released" no longer means "unlinked": the claim file stays, because the lock + lives on its inode and unlinking one out from under another process's open + descriptor is how two holders end up locking two different files for one + destination. What the assertions check is that the claim is *lockable* again, + which is the only thing ownership was ever about. + """ def _prepare(self, workspace, tmp_path, **overrides) -> tuple[download_state.DownloadState, Path, Path]: dest = tmp_path / "m.safetensors" state = _state(dest=str(dest), **overrides) path = download_state.write(workspace, state) claim = download_state.claim_path(workspace, models._dest_key(dest)) - assert download_state.acquire_claim(claim, download_id=state.id, dest=str(dest)) + _plant_claim(claim, state.id, str(dest)) return state, path, claim def test_a_completed_transfer_releases_it(self, workspace, monkeypatch, tmp_path): @@ -1672,7 +1947,7 @@ def test_a_completed_transfer_releases_it(self, workspace, monkeypatch, tmp_path models._download_worker(state_file=str(path)) assert download_state.read(workspace, state.id).status == "completed" - assert not claim.exists() + assert _is_lockable(claim) def test_a_failed_transfer_releases_it(self, workspace, monkeypatch, tmp_path): state, path, claim = self._prepare(workspace, tmp_path) @@ -1685,7 +1960,7 @@ def boom(*args, **kwargs): models._download_worker(state_file=str(path)) assert download_state.read(workspace, state.id).status == "failed" - assert not claim.exists() + assert _is_lockable(claim) def test_a_cancel_that_beat_the_worker_to_the_start_releases_it(self, workspace, monkeypatch, tmp_path): """`download-cancel` never touches the claim itself - the worker boots, @@ -1698,7 +1973,7 @@ def test_a_cancel_that_beat_the_worker_to_the_start_releases_it(self, workspace, models._download_worker(state_file=str(path)) assert download_state.read(workspace, state.id).status == "cancelled" - assert not claim.exists() + assert _is_lockable(claim) def test_a_mid_transfer_cancel_releases_it(self, workspace, monkeypatch, tmp_path): state, path, claim = self._prepare(workspace, tmp_path) @@ -1714,24 +1989,112 @@ def transfer(url, filepath, headers, downloader, progress_callback): models._download_worker(state_file=str(path)) assert download_state.read(workspace, state.id).status == "cancelled" - assert not claim.exists() + assert _is_lockable(claim) + + def test_the_lock_is_held_for_the_whole_transfer(self, workspace, monkeypatch, tmp_path): + """The held-lock replacement for "a claim that is no longer ours". - def test_a_claim_that_is_no_longer_ours_is_left_alone(self, workspace, monkeypatch, tmp_path): - """Our record goes terminal a moment before we release, so the claim reads - stale to a submitter racing us - it may clear ours and create its own in - that window. An unconditional unlink would delete a live claim.""" + The old worker re-derived ownership at the end and had to avoid unlinking + a claim that had changed hands while it was still running — a window that + only existed because it was not holding anything. Here the lock is taken + before the first byte and dropped after the last, so mid-transfer the + destination is provably unavailable to anybody else and no competitor can + have taken it in the first place. + """ state, path, claim = self._prepare(workspace, tmp_path) + seen: list[bool] = [] def transfer(url, filepath, headers, downloader, progress_callback): + seen.append(_is_lockable(claim)) filepath.write_bytes(b"ok") - claim.unlink() - assert download_state.acquire_claim(claim, download_id="ffffffffffff", dest=state.dest) monkeypatch.setattr(models, "download_file", transfer) models._download_worker(state_file=str(path)) - assert claim.exists() - assert download_state.read_claim(claim) == "ffffffffffff" + assert seen == [False], "the worker was not holding the claim lock mid-transfer" + assert _is_lockable(claim) + assert download_state.read_claim(claim) == state.id + + def test_a_worker_that_lost_its_destination_never_touches_it(self, workspace, monkeypatch, tmp_path): + """The spawn gap seen from the worker's side. A competitor may have won + the lock and stamped its own id in while our interpreter was booting; the + payload check under the lock is what stops us writing bytes into a file + somebody else owns. We exit refused and leave a terminal record, exactly + as the submit-side withdraw path does.""" + state, path, claim = self._prepare(workspace, tmp_path) + _plant_claim(claim, "999999999999", state.dest) + monkeypatch.setattr(models, "download_file", MagicMock(side_effect=AssertionError("a transfer started"))) + + with pytest.raises(typer.Exit) as exc: + models._download_worker(state_file=str(path)) + + assert exc.value.exit_code == 1 + assert download_state.read(workspace, state.id).status == "failed" + assert not Path(state.dest).exists() + # Their claim is untouched and lockable again — we held it only to look. + assert download_state.read_claim(claim) == "999999999999" + assert _is_lockable(claim) + + def test_a_worker_adopts_an_unstamped_claim_instead_of_refusing(self, workspace, monkeypatch, tmp_path): + """Nobody can own a destination whose lock we are holding, so a claim with + no payload means only that nothing stamped it — a submitter that degraded + past `write_payload` on a transient error, or a sweep that removed an + unowned file our own `O_CREAT` then re-made. Refusing there would turn a + recoverable submit-side hiccup into a download that never runs.""" + dest = tmp_path / "m.safetensors" + state = _state(dest=str(dest)) + path = download_state.write(workspace, state) + claim = download_state.claim_path(workspace, models._dest_key(dest)) + download_state.lock_claim(claim).release() # created, never stamped + assert download_state.read_claim(claim) is None + monkeypatch.setattr( + models, + "download_file", + lambda url, filepath, headers, downloader, progress_callback: filepath.write_bytes(b"ok"), + ) + + models._download_worker(state_file=str(path)) + + assert download_state.read(workspace, state.id).status == "completed" + assert dest.read_bytes() == b"ok" + # And it stamped itself in, so the claim names a real download again. + assert download_state.read_claim(claim) == state.id + + def test_a_worker_locked_out_for_the_whole_window_refuses(self, workspace, monkeypatch, tmp_path): + """A lock held by somebody else for the entire blocking window is the + same verdict as a payload naming somebody else: the destination is not + ours. Kept short here so the test does not pay the real timeout.""" + state, path, claim = self._prepare(workspace, tmp_path) + monkeypatch.setattr(models, "WORKER_CLAIM_LOCK_TIMEOUT_S", 0.05) + monkeypatch.setattr(download_state, "lock_claim", lambda path, **kwargs: None) + monkeypatch.setattr(models, "download_file", MagicMock(side_effect=AssertionError("a transfer started"))) + + with pytest.raises(typer.Exit) as exc: + models._download_worker(state_file=str(path)) + + assert exc.value.exit_code == 1 + assert download_state.read(workspace, state.id).status == "failed" + assert not Path(state.dest).exists() + + def test_a_filesystem_that_cannot_lock_still_downloads(self, workspace, monkeypatch, tmp_path): + """Bookkeeping must never be what stops a transfer. When the lock cannot + be taken at all — not held, *unsupported* — the worker degrades to the + pre-lock behavior and moves the bytes, exactly as the submitter degrades + to the advisory guard.""" + state, path, claim = self._prepare(workspace, tmp_path) + monkeypatch.setattr( + download_state, "lock_claim", MagicMock(side_effect=OSError(errno.ENOLCK, "No locks available")) + ) + monkeypatch.setattr( + models, + "download_file", + lambda url, filepath, headers, downloader, progress_callback: filepath.write_bytes(b"ok"), + ) + + models._download_worker(state_file=str(path)) + + assert download_state.read(workspace, state.id).status == "completed" + assert Path(state.dest).read_bytes() == b"ok" def test_a_finished_download_does_not_wedge_its_destination(self, workspace, monkeypatch, tmp_path): """The whole risk of adding a lock: one that is never handed back turns a @@ -1759,7 +2122,7 @@ def submit(): first = _claim_owner(workspace) models._download_worker(state_file=str(download_state.state_path(workspace, first))) assert download_state.read(workspace, first).status == "completed" - assert _claim_files(workspace) == [] + assert _is_lockable(_claim_files(workspace)[0]) _reset_envelope() submit() @@ -1777,30 +2140,32 @@ def test_a_failed_first_state_write_releases_it(self, workspace, monkeypatch, tm with pytest.raises(OSError): models._download_worker(state_file=str(path)) - assert not claim.exists() + assert _is_lockable(claim) -class TestClaimAtomicPublish: - """`acquire_claim` writes the payload to a private temp name and publishes - it with `os.link`: creation is still the atomic decider, but the claim is - never visible half-written — a colliding submitter that reads it can never - mistake a live winner mid-write for a stale claim and unlink it.""" +class TestClaimPayloadWrites: + """The payload is rewritten **in place on the locked descriptor**. + + Never temp-file-plus-rename, which is how the previous implementation + published a claim: a rename swaps the inode, and the lock lives on the inode, + so the holder would be left guarding a file nobody else can see while a + second process locks the new one. + """ @staticmethod def _patch_claim_writes(monkeypatch, claims_dir, transform): - """Apply ``transform`` to writes on the claim's staging fd, and only it. - - `acquire_claim` is exercised by breaking `os.write`, but `os.write` is - process-wide: an unscoped patch also truncates (or fails) pytest's own - fd-level capture flushes and anything else holding a raw descriptor. - So the fd is identified at `os.open` time by the directory it was opened - in, and every other descriptor writes normally. Tracked by fd number, - which the kernel reuses after a close, so `os.close` untracks it too. + """Apply ``transform`` to writes on a claim's descriptor, and only it. + + `os.write` is process-wide: an unscoped patch also truncates (or fails) + pytest's own fd-level capture flushes and anything else holding a raw + descriptor. So the fd is identified at `os.open` time by the directory it + was opened in, and every other descriptor writes normally. Tracked by fd + number, which the kernel reuses after a close, so `os.close` untracks it. """ real_open, real_write, real_close = os.open, os.write, os.close staged: set[int] = set() - def _is_staging(path): + def _is_claim(path): try: return os.path.dirname(os.fsdecode(path)) == str(claims_dir) except TypeError: @@ -1808,7 +2173,7 @@ def _is_staging(path): def tracking_open(path, flags, *args, **kwargs): fd = real_open(path, flags, *args, **kwargs) - if _is_staging(path): + if _is_claim(path): staged.add(fd) return fd @@ -1825,7 +2190,7 @@ def untracking_close(fd): monkeypatch.setattr(os, "write", scoped_write) monkeypatch.setattr(os, "close", untracking_close) - def test_short_writes_still_publish_a_complete_payload(self, workspace, monkeypatch): + def test_short_writes_still_land_a_complete_payload(self, workspace, monkeypatch): claim = download_state.claim_path(workspace, "/tmp/m.safetensors") shortened = [] @@ -1835,52 +2200,242 @@ def one_byte(write, fd, data): self._patch_claim_writes(monkeypatch, claim.parent, one_byte) - assert download_state.acquire_claim(claim, download_id="aaaaaaaaaaaa", dest="/tmp/m.safetensors") + lock = download_state.lock_claim(claim) + assert lock is not None + try: + lock.write_payload("aaaaaaaaaaaa", "/tmp/m.safetensors") + assert lock.read_payload()["download_id"] == "aaaaaaaaaaaa" + finally: + lock.release() assert download_state.read_claim(claim) == "aaaaaaaaaaaa" - # The scoping must not quietly stop matching the staging fd: that would - # leave this asserting nothing but that a normal write works. + # The scoping must not quietly stop matching the claim's descriptor: that + # would leave this asserting nothing but that a normal write works. assert len(shortened) > 1, "the short-write patch never reached the claim's descriptor" - def test_a_failed_payload_write_publishes_nothing(self, workspace, monkeypatch): - """ENOSPC mid-payload must not leave a claim — empty at the claim path - (every reader would call it stale and clear it under us) or orphaned as - a temp file. The OSError propagates and the caller degrades to the - advisory guard, exactly like an unusable claims directory.""" + def test_a_failed_payload_write_raises_and_leaves_no_temp_file(self, workspace, monkeypatch): + """ENOSPC mid-payload propagates so the caller can degrade to the + advisory guard. There is nothing to clean up after it: the write went + straight into the claim the lock is on, so no staging sibling exists to + be orphaned, and a short payload reads as stale like any other.""" claim = download_state.claim_path(workspace, "/tmp/m.safetensors") def _enospc(_write, _fd, _data): - raise OSError(28, "No space left on device") + raise OSError(errno.ENOSPC, "No space left on device") self._patch_claim_writes(monkeypatch, claim.parent, _enospc) + lock = download_state.lock_claim(claim) + assert lock is not None + try: + with pytest.raises(OSError): + lock.write_payload("aaaaaaaaaaaa", "/tmp/m.safetensors") + finally: + lock.release() + + assert [p.name for p in claim.parent.iterdir()] == [claim.name] + assert download_state.read_claim(claim) is None + + def test_the_payload_carries_no_url(self, workspace): + """A resolved download url can be presigned - a credential-shaped thing. + The claim needs an id and a path, so that is all it carries.""" + claim = download_state.claim_path(workspace, "/tmp/m.safetensors") + _plant_claim(claim, "aaaaaaaaaaaa", "/tmp/m.safetensors") + assert json.loads(claim.read_text(encoding="utf-8")).keys() == {"download_id", "dest", "created_at"} + + +class TestClaimLockPrimitive: + """`lock_claim` / `ClaimLock` on their own — the platform halves included.""" + + def test_the_lock_is_exclusive_and_released_by_close(self, workspace): + claim = download_state.claim_path(workspace, "/tmp/x.safetensors") + first = download_state.lock_claim(claim) + assert first is not None + assert download_state.lock_claim(claim) is None + first.release() + second = download_state.lock_claim(claim) + assert second is not None + second.release() + + def test_release_is_idempotent(self, workspace): + claim = download_state.claim_path(workspace, "/tmp/x.safetensors") + lock = download_state.lock_claim(claim) + lock.release() + lock.release() + assert not lock.held + with pytest.raises(ValueError): + lock.write_payload("aaaaaaaaaaaa", "/tmp/x.safetensors") + + def test_a_blocking_timeout_is_bounded(self, workspace): + """The worker waits out a competitor's inspection, but a bounded + non-blocking retry rather than a blocking lock: the wait can never + outlive the caller's patience.""" + claim = download_state.claim_path(workspace, "/tmp/x.safetensors") + held = download_state.lock_claim(claim) + try: + started = time.monotonic() + assert download_state.lock_claim(claim, blocking_timeout=0.15) is None + waited = time.monotonic() - started + finally: + held.release() + assert 0.1 <= waited < 5.0, waited + + def test_an_empty_or_corrupt_payload_reads_as_nobody(self, workspace): + claim = download_state.claim_path(workspace, "/tmp/x.safetensors") + claim.write_text("not json at all", encoding="utf-8") + lock = download_state.lock_claim(claim) + try: + assert lock.read_payload() is None + lock.write_payload("aaaaaaaaaaaa", "/tmp/x.safetensors") + lock.clear_payload() + assert lock.read_payload() is None + finally: + lock.release() + + def test_a_filesystem_that_cannot_lock_raises_rather_than_reporting_contention(self, workspace, monkeypatch): + """The distinction the caller's degradation hangs on: contention is None, + everything else is an OSError. Reading ENOLCK as "somebody owns this" + would refuse every download on such a filesystem, forever.""" + claim = download_state.claim_path(workspace, "/tmp/x.safetensors") + monkeypatch.setattr( + download_state, "_lock_fd_nb", MagicMock(side_effect=OSError(errno.ENOLCK, "No locks available")) + ) with pytest.raises(OSError): - download_state.acquire_claim(claim, download_id="aaaaaaaaaaaa", dest="/tmp/m.safetensors") + download_state.lock_claim(claim) - assert not claim.exists() - assert not list(claim.parent.iterdir()) + @pytest.mark.skipif(sys.platform == "win32", reason="the POSIX half") + def test_posix_takes_a_non_blocking_flock(self, workspace, monkeypatch): + import fcntl + + claim = download_state.claim_path(workspace, "/tmp/x.safetensors") + calls: list[int] = [] + real = fcntl.flock + monkeypatch.setattr(fcntl, "flock", lambda fd, op: (calls.append(op), real(fd, op))[1]) + + lock = download_state.lock_claim(claim) + lock.release() + assert calls == [fcntl.LOCK_EX | fcntl.LOCK_NB] + + @pytest.mark.skipif(sys.platform != "win32", reason="the Windows half") + def test_windows_takes_a_non_blocking_msvcrt_lock_at_offset_zero(self, workspace, monkeypatch): + import msvcrt + + claim = download_state.claim_path(workspace, "/tmp/x.safetensors") + calls: list[tuple[int, int, int]] = [] + real = msvcrt.locking + + def spy(fd, mode, nbytes): + calls.append((os.lseek(fd, 0, os.SEEK_CUR), mode, nbytes)) + return real(fd, mode, nbytes) + + monkeypatch.setattr(msvcrt, "locking", spy) + + lock = download_state.lock_claim(claim) + lock.release() + assert calls == [(0, msvcrt.LK_NBLCK, 1)] + + @pytest.mark.skipif(sys.platform == "win32", reason="inode identity is not reliable on Windows") + def test_the_payload_rewrite_keeps_the_inode(self, workspace): + """The single property the whole design rests on: the lock lives on the + inode, so a payload rewrite must never replace the file.""" + claim = download_state.claim_path(workspace, "/tmp/x.safetensors") + lock = download_state.lock_claim(claim) + try: + lock.write_payload("aaaaaaaaaaaa", "/tmp/x.safetensors") + first = claim.stat().st_ino + for i in range(5): + lock.write_payload(f"bbbbbbbbbbb{i}", "/tmp/x.safetensors") + assert claim.stat().st_ino == first + assert os.fstat(lock._fd).st_ino == first + finally: + lock.release() + + @pytest.mark.skipif(sys.platform == "win32", reason="only POSIX ever unlinks a claim file") + def test_a_claim_replaced_under_the_lock_is_locked_again(self, workspace, monkeypatch): + """The re-stat guard. Between our `open` and our lock the file can be + unlinked (only `prune` does this) and another created at the same path — + leaving us holding a real lock on an orphan inode while somebody else + holds one on the live file. `fstat` vs `stat` catches it; we close and + retry.""" + claim = download_state.claim_path(workspace, "/tmp/x.safetensors") + real = download_state._lock_fd_nb + calls: list[int] = [] + + def swap_the_file_once(fd): + took = real(fd) + calls.append(fd) + if len(calls) == 1: + claim.unlink() + download_state.lock_claim(claim).release() + return took + + monkeypatch.setattr(download_state, "_lock_fd_nb", swap_the_file_once) + + lock = download_state.lock_claim(claim) + assert lock is not None + try: + assert len(calls) > 1, "the guard did not retry" + assert os.fstat(lock._fd).st_ino == claim.stat().st_ino + finally: + lock.release() + + @pytest.mark.skipif(sys.platform == "win32", reason="only POSIX ever unlinks a claim file") + def test_a_claim_replaced_on_every_attempt_reports_held(self, workspace, monkeypatch): + """Refusing is the safe direction when ownership cannot be settled: the + caller backs off instead of letting a second transfer into a destination + whose claim something else is churning.""" + claim = download_state.claim_path(workspace, "/tmp/x.safetensors") + real = download_state._lock_fd_nb + + def always_swap(fd): + took = real(fd) + with contextlib.suppress(OSError): + claim.unlink() + return took + + monkeypatch.setattr(download_state, "_lock_fd_nb", always_swap) + assert download_state.lock_claim(claim) is None class TestClaimSweep: - """`prune` drops claims that no longer point at a live download. + """`prune` removes claim files nothing owns. + + A nicety, not a correctness step: a claim file is never unlinked on release + (the lock lives on its inode), so without this the directory would grow to + one ~100 byte file per destination the workspace ever downloaded to. Two + conditions, both under the lock — we can take it, and the payload does not + name a live download. The second is what keeps the sweep out of the spawn + gap, where unlinking would make a booting worker refuse the destination it + was just handed. + """ - A claim is normally released by its own worker, and a stranded one only - self-clears on the next submit *to the same destination* — so claims for - destinations never re-submitted would otherwise accumulate for the life of - the workspace.""" + posix_only = pytest.mark.skipif(sys.platform == "win32", reason="a claim file is never unlinked on Windows") + @posix_only def test_a_claim_with_no_record_is_swept(self, workspace): claim = download_state.claim_path(workspace, "/tmp/a.safetensors") - assert download_state.acquire_claim(claim, download_id="abcdefabcdef", dest="/tmp/a.safetensors") + _plant_claim(claim, "abcdefabcdef", "/tmp/a.safetensors") download_state.prune(workspace) assert not claim.exists() + @posix_only + def test_an_empty_claim_is_swept(self, workspace): + """The residue of a spawn that failed: blanked, owned by nobody, and with + no payload to consult.""" + claim = download_state.claim_path(workspace, "/tmp/a.safetensors") + download_state.lock_claim(claim).release() + + download_state.prune(workspace) + + assert not claim.exists() + + @posix_only def test_a_dead_workers_claim_is_swept(self, workspace): state = _state(status="downloading", pid=4242) download_state.write(workspace, state) claim = download_state.claim_path(workspace, "/tmp/b.safetensors") - assert download_state.acquire_claim(claim, download_id=state.id, dest="/tmp/b.safetensors") + _plant_claim(claim, state.id, "/tmp/b.safetensors") with patch("comfy_cli.utils.is_running", return_value=False): download_state.prune(workspace) @@ -1891,7 +2446,7 @@ def test_a_live_claim_is_kept(self, workspace): state = _state(status="downloading", pid=1234) download_state.write(workspace, state) claim = download_state.claim_path(workspace, "/tmp/c.safetensors") - assert download_state.acquire_claim(claim, download_id=state.id, dest="/tmp/c.safetensors") + _plant_claim(claim, state.id, "/tmp/c.safetensors") with patch("comfy_cli.utils.is_running", return_value=True): download_state.prune(workspace) @@ -1904,17 +2459,33 @@ def test_a_terminal_record_with_a_live_worker_keeps_its_claim(self, workspace): state = _state(status="cancelled", pid=1234) download_state.write(workspace, state) claim = download_state.claim_path(workspace, "/tmp/d.safetensors") - assert download_state.acquire_claim(claim, download_id=state.id, dest="/tmp/d.safetensors") + _plant_claim(claim, state.id, "/tmp/d.safetensors") with patch("comfy_cli.utils.is_running", return_value=True): download_state.prune(workspace) assert claim.exists() + def test_a_held_claim_is_never_unlinked_under_its_holder(self, workspace): + """The lock is checked first and the unlink happens under it, so a + transfer in progress cannot have its claim swept even when the record it + points at reads terminal.""" + state = _state(status="cancelled", pid=None) + download_state.write(workspace, state) + claim = download_state.claim_path(workspace, "/tmp/e.safetensors") + holder = download_state.lock_claim(claim) + try: + holder.write_payload(state.id, "/tmp/e.safetensors") + download_state.prune(workspace) + finally: + holder.release() + + assert claim.exists() + def test_temp_leftovers_age_out(self, workspace): - """A `.tmp` staging file only survives a SIGKILL inside the milliseconds - between write and publish; an aged one is a crashed acquire's leftover, - a fresh one may be an acquire in progress.""" + """A `.tmp` staging file is residue from the previous implementation's + write-then-link publish; nothing writes them any more, but a workspace + upgraded from that version can still be carrying one.""" claims = download_state.claims_dir(workspace) old_tmp = claims / f"x.claim.123.aaaaaaaaaaaa{download_state.CLAIM_TMP_SUFFIX}" old_tmp.write_text("{}", encoding="utf-8")