Recover snapshot cache after canceled restores - #199
Conversation
|
🔍 OpenCodeReview found 1 issue(s) in this PR.
|
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| if wait_for_inflight_result(&mut existing.receiver).await? | ||
| == InflightWait::Abandoned | ||
| { | ||
| self.clear_abandoned_inflight(key, &existing.receiver).await; | ||
| } | ||
| continue; |
There was a problem hiding this comment.
When the materialization future is cancelled, dropping the watch sender only marks the inflight entry abandoned; it does not remove or roll back the file that the user-supplied fetch may already have created. On retry, ensure_cached_at checks fs::try_exists(&local_path) before electing a new producer and immediately calls pin_local_file, so a partially written artifact can be registered and returned as a successful cache hit instead of being rematerialized. Please write to a temporary path and atomically rename on success (and clean it up on cancellation), or otherwise invalidate/remove the target before retrying.
Suggestion:
| if wait_for_inflight_result(&mut existing.receiver).await? | |
| == InflightWait::Abandoned | |
| { | |
| self.clear_abandoned_inflight(key, &existing.receiver).await; | |
| } | |
| continue; | |
| if wait_for_inflight_result(&mut existing.receiver).await? | |
| == InflightWait::Abandoned | |
| { | |
| self.clear_abandoned_inflight(key, &existing.receiver).await; | |
| let _ = fs::remove_file(&local_path).await; | |
| } | |
| continue; |
|
Production validation on term.so:
|
There was a problem hiding this comment.
This PR fixes a real issue: when a materialization leader is canceled, it leaves a stale entry in the inflight map that permanently poisons the snapshot key, causing every later restore to fail until a restart. The Abandoned handling and the same_channel ABA protection are the right approach. That said, I found two correctness issues that I'd like to see fixed before approving — I've left inline comments at the corresponding locations.
Note: the first one roots in pre-existing code, but it only becomes reachable because of this PR's new retry path — previously the stale entry failed everything loudly, so a partial file could never be registered as a valid cache entry. Since this PR turns that loud failure into a potential silent success, I think it needs to be addressed here.
| let _ = sender.send(Some(result)); | ||
| let mut inflight = self.inflight.lock().await; | ||
| inflight.remove(key); |
There was a problem hiding this comment.
Issue 2: Cancellation window between send and remove
send → lock().await → remove leaves a cancellation window: if the finisher is canceled at the .await, the map keeps a dead entry that already carries Some(Success/Failed). Stale Failed makes every later request replay the same old error forever; stale Success spins ensure_cached_at in a miss → Finished → continue loop once the index entry is evicted.
Moving send inside the lock fixes it — no .await between send and remove:
let mut inflight = self.inflight.lock().await;
let _ = sender.send(Some(result));
inflight.remove(key);If the task is canceled before acquiring the lock, the sender is dropped without a result, which your Abandoned path already handles.
| let result = fetch(local_path.clone()) | ||
| .await | ||
| .with_context(|| format!("materialize '{key}' into local cache")); |
There was a problem hiding this comment.
Issue 1: Partial files are treated as successful fetches
fetch writes directly to the final path, but whether the file is usable is decided only by try_exists. These two don't compose safely:
- The leader is canceled mid-write (or fails with
Errafter writing partial data) → the final path exists, but the content is incomplete. - With this PR, the next waiter clears the abandoned inflight entry and retries.
- On retry,
try_exists(local_path)returns true →pin_local_fileregisters the partial file in the index and returns a handle — without ever re-runningfetch.
This turns an explicit failure into a silent success: e.g. resuming a sandbox from a truncated vm_state.bin, which will fail much later and be far harder to diagnose than a restore error. This is reachable in production — the OSS resolver's P2P export writes chunks straight to the target path.
2 participants
What changed
When the task leading a snapshot artifact materialization is canceled, the next waiter now removes that abandoned in-flight entry and retries. The cleanup compares the watch channel before removal so a concurrent waiter cannot remove a newer fetch.
A regression test cancels the materialization leader, then verifies a second caller materializes and reads the artifact.
Why
A client disconnect can cancel the leader after the cache records the fetch but before it publishes a result. The sender is dropped, while the receiver remains in the in-flight map. Every later restore of that snapshot then fails with
ended without a resultuntil the AgentENV server restarts.This change makes canceled restores self-healing and prevents a transient disconnect from permanently blocking sandbox recovery on that node.
Validation
cargo test snapshot::artifact_cache::tests— 7 passedcargo clippy --locked --all-targets -- -D warnings— passedcargo test --locked— 741 passed, 4 ignored; 5 unrelated ublk tests could not connect to the local ublk daemon (channel closed)