Materialization claim needs an identity, not just CLEAN (#92)

glidefs/src/block/handler.rs

@@ -450,6 +450,17 @@ impl BlockHandler {
             // NOT_PRESENT (even if the data pwrite never completed, the
             // guest's WRITE_FIXED is the source of truth for the bytes the
             // caller is writing).
+            //
+            // A LIVE (merely slow) winner can't clobber us afterwards
+            // either. Its commit (`write_materialized`) takes the data_file
+            // WRITE lock, and our data lands under the rotation READ gate
+            // held from before WRITE_FIXED through the DIRTY commit — so the
+            // two are ordered, never interleaved: the winner either lands
+            // its merged block before ours (our bytes go on top) or finds
+            // the block DIRTY and aborts. And if the block is meanwhile
+            // evicted and RE-claimed, `pre_write`'s `set_present`
+            // invalidates the winner's claim token, so its stale pre-image
+            // aborts there too. See `MaterializationClaim::is_valid`.
             let clean_deadline = std::time::Instant::now()
                 + std::time::Duration::from_secs(5);
             loop {
@@ -549,9 +560,12 @@ impl BlockHandler {
                 Ok(data) => data,
                 Err(e) => {
                     // Revert the state claim — a dangling CLEAN block parks
-                    // every sibling writer on the CLEAN-wait. (The promote
-                    // claim releases when `claim` drops.)
-                    self.cache.unclaim_block(idx);
+                    // every sibling writer on the CLEAN-wait. Claim-checked:
+                    // if the block was re-claimed underneath us, the CLEAN is
+                    // a later writer's and must not be knocked back to
+                    // NOT_PRESENT. (The promote claim releases when `claim`
+                    // drops.)
+                    claim.unclaim();
                     return Err(e.into());
                 }
             };
@@ -573,16 +587,16 @@ impl BlockHandler {
             // claim (committed newer data) while we fetched. On a steal the
             // block is DIRTY with the guest's bytes; our stale pre-image
             // must NOT land. The caller's own WRITE_FIXED still follows.
-            match self.cache.write_materialized(block_start, &block_data, idx) {
+            match self.cache.write_materialized(block_start, &block_data, &claim) {
                 Ok(true) => {}
                 Ok(false) => {
                     tracing::debug!(
                         block = idx,
-                        "backfill: claim stolen by concurrent guest write; skipping stale pre-image"
+                        "backfill: claim stolen or re-claimed by a concurrent guest write; skipping stale pre-image"
                     );
                 }
                 Err(e) => {
-                    self.cache.unclaim_block(idx);
+                    claim.unclaim();
                     return Err(e.into());
                 }
             }
@@ -844,8 +858,10 @@ impl BlockHandler {
                     Err(e) => {
                         // Revert the state claim — a dangling CLEAN block
                         // parks every sibling writer on the CLEAN-wait.
-                        // (The promote claim releases when `claim` drops.)
-                        self.cache.unclaim_block(idx);
+                        // Claim-checked: a CLEAN that is no longer ours
+                        // belongs to a later writer. (The promote claim
+                        // releases when `claim` drops.)
+                        claim.unclaim();
                         tracing::warn!(
                             block = idx,
                             error = %e,
@@ -861,7 +877,16 @@ impl BlockHandler {
                 // block is CLEAN (claimed); cache.write transitions it DIRTY
                 // and the rest of the block stays sparse zeros — consistent
                 // with the all-zero prior.
+                //
+                // Only while the claim is still ours: if the block was
+                // stolen, flushed out and re-claimed while we fetched, "the
+                // remainder is legitimately zeros" no longer holds (the S3
+                // image moved on), so re-enter the state machine instead.
                 if prior.is_empty() || prior.iter().all(|&b| b == 0) {
+                    if !claim.is_valid() {
+                        drop(claim);
+                        continue 'block_retry;
+                    }
                     self.cache.write(write_start, &data[data_offset..data_offset + write_len])?;
                     drop(claim);
                     break 'block_retry;
@@ -870,7 +895,7 @@ impl BlockHandler {
                 // We hold the claim. Merge guest data onto prior block.
                 let mut block_buf = prior.to_vec();
                 if block_buf.len() != block_size {
-                    self.cache.unclaim_block(idx);
+                    claim.unclaim();
                     tracing::error!(
                         block = idx,
                         expected = block_size,
@@ -885,27 +910,28 @@ impl BlockHandler {
 
                 _backfill_gate!(BackfillStep::BeforeWrite, idx, self.cache.block_state(idx).raw());
 
-                // Materialized write (write lock + CLEAN re-check). A steal
-                // means a concurrent guest write committed newer data while
-                // we fetched — our merged block embeds the STALE pre-image,
-                // so discard it and re-enter the state machine: the block is
-                // now DIRTY, and the has_local_data arm writes only our own
-                // guest sub-range on top.
-                match self.cache.write_materialized(block_start, &block_buf, idx) {
+                // Materialized write (write lock + CLEAN + claim-identity
+                // re-check). A steal means a concurrent guest write committed
+                // newer data while we fetched — our merged block embeds the
+                // STALE pre-image, so discard it and re-enter the state
+                // machine: the block is now DIRTY (or re-claimed by a later
+                // writer), and the retry writes only our own guest sub-range
+                // on top of whatever materialized it.
+                match self.cache.write_materialized(block_start, &block_buf, &claim) {
                     Ok(true) => {
                         drop(claim);
                         break 'block_retry;
                     }
                     Ok(false) => {
                         tracing::debug!(
                             block = idx,
-                            "backfill_and_write: claim stolen by concurrent guest write; rewriting sub-range only"
+                            "backfill_and_write: claim stolen or re-claimed by a concurrent guest write; rewriting sub-range only"
                         );
                         drop(claim);
                         continue 'block_retry;
                     }
                     Err(e) => {
-                        self.cache.unclaim_block(idx);
+                        claim.unclaim();
                         return Err(e.into());
                     }
                 }

glidefs/src/block/write_cache/flush.rs

@@ -602,27 +602,81 @@ impl WriteCache<Active> {
     ///
     /// Drop order: finish the data write (`cache.write` → DIRTY) BEFORE
     /// dropping the guard, so waiters re-check and observe DIRTY.
+    ///
+    /// The claim carries a [`ClaimToken`](super::inner::ClaimToken) identity:
+    /// `CLEAN` alone does not prove the block is still the one that was
+    /// claimed (see [`MaterializationClaim::is_valid`]), so the commit
+    /// ([`WriteCache::write_materialized`]) takes the guard, not a bare index.
     pub fn claim_block_for_materialization(
         &self,
         block_idx: usize,
     ) -> Option<MaterializationClaim<'_>> {
-        if !self.inner.promote_claim.try_claim(block_idx) {
-            return None;
-        }
+        let token = self.inner.promote_claim.try_claim_token(block_idx)?;
         if !self.inner.try_set_present(block_idx) {
             self.inner.promote_claim.release(block_idx);
             return None;
         }
-        Some(MaterializationClaim { cache_inner: &self.inner, block_idx })
+        Some(MaterializationClaim { cache_inner: &self.inner, block_idx, token })
     }
 }
 
 /// RAII guard for a block claimed for S3 materialization — holds the block's
 /// promote claim from [`WriteCache::claim_block_for_materialization`] until
 /// drop. See that method for the protocol.
 pub struct MaterializationClaim<'a> {
-    cache_inner: &'a super::inner::CacheInner,
-    block_idx: usize,
+    pub(super) cache_inner: &'a super::inner::CacheInner,
+    pub(super) block_idx: usize,
+    pub(super) token: super::inner::ClaimToken,
+}
+
+impl MaterializationClaim<'_> {
+    /// The block this claim owns.
+    #[inline]
+    pub fn block_idx(&self) -> usize {
+        self.block_idx
+    }
+
+    /// Is this still the live claim on the block?
+    ///
+    /// A materialization claim is held across an S3 fetch, so `CLEAN` on its
+    /// own is not proof of ownership: the block can be stolen (CLEAN→DIRTY),
+    /// uploaded + evicted by a flush (→SYNCING→NOT_PRESENT) and re-claimed by
+    /// a later writer — landing back on CLEAN with someone else's data
+    /// pending. Committing the fetched pre-image then rolls that later,
+    /// already-durable write back to the S3 image. The token distinguishes
+    /// the two: any bypassing `NOT_PRESENT → CLEAN` (the ublk ZC `pre_write`)
+    /// invalidates the live claim before flipping the state.
+    ///
+    /// Must be evaluated in the same critical section as the state check —
+    /// `write_materialized` does both under the `data_file` write lock.
+    #[inline]
+    pub fn is_valid(&self) -> bool {
+        self.cache_inner
+            .promote_claim
+            .is_valid(self.block_idx, self.token)
+    }
+
+    /// Revert this claim (CAS CLEAN → NOT_PRESENT) for the backfill error
+    /// path: a claimed block whose S3 fetch failed must not stay CLEAN —
+    /// sibling writers park on the CLEAN-wait expecting a CLEAN→DIRTY
+    /// transition that would never come.
+    ///
+    /// No-op if the claim is no longer ours: the CLEAN we would revert is
+    /// then a *later* writer's claim, and knocking it back to NOT_PRESENT
+    /// would strand that writer's in-flight data write on an unclaimed block.
+    /// Returns true if this call performed the transition.
+    pub fn unclaim(&self) -> bool {
+        use crate::block::block_map::SparseBlockState;
+        self.cache_inner
+            .promote_claim
+            .if_valid(self.block_idx, self.token, || {
+                self.cache_inner
+                    .state_map
+                    .cas(self.block_idx, SparseBlockState::CLEAN, SparseBlockState::NOT_PRESENT)
+                    .is_ok()
+            })
+            .unwrap_or(false)
+    }
 }
 
 impl Drop for MaterializationClaim<'_> {

glidefs/src/block/write_cache/inner.rs

@@ -312,27 +312,105 @@ impl HandoffPhase {
 /// state for an occasionally-held flag. Sparse storage costs O(in-flight),
 /// independent of device size.
 ///
-/// Empty cost: `Mutex<HashSet>` ≈ 64 B per export.
+/// Empty cost: `Mutex<HashMap>` ≈ 64 B per export.
 /// Per-claim cost: one short mutex critical section (insert/remove). The
 /// mutex contends only during the promote handshake itself, never on the
 /// data plane.
+///
+/// ## Claim identity (`ClaimToken`)
+///
+/// A materialization claim ([`WriteCache::claim_block_for_materialization`])
+/// is held across an S3 fetch — seconds, not microseconds. `CLEAN` alone is
+/// therefore NOT proof that the block is still the one the holder claimed:
+/// the block can be stolen (CLEAN→DIRTY), flushed (→SYNCING→NOT_PRESENT) and
+/// re-claimed by an unrelated writer, ending back at CLEAN. A commit gated
+/// only on "state == CLEAN" would land the holder's now-stale pre-image on
+/// top of that later writer's data — resurrecting a pre-image over an
+/// already-durable page.
+///
+/// Every claim therefore carries a monotonically-increasing [`ClaimToken`],
+/// and any `NOT_PRESENT → CLEAN` transition that does NOT go through the
+/// claim ([`CacheInner::set_present`], i.e. the ublk ZC `pre_write` path)
+/// invalidates whatever claim is live for that block *before* it flips the
+/// state. `is_valid(idx, token)` is then an exact "this CLEAN is still
+/// mine" test, and the holder aborts instead of clobbering.
 pub(crate) struct PromoteClaimBitmap {
-    in_flight: parking_lot::Mutex<std::collections::HashSet<usize>>,
+    in_flight: parking_lot::Mutex<std::collections::HashMap<usize, ClaimSlot>>,
     released: parking_lot::Condvar,
+    next_token: AtomicU64,
+}
+
+/// Identity of one claim on a block. Compared by value — a claim that was
+/// invalidated or replaced never matches its holder's token again.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) struct ClaimToken(u64);
+
+struct ClaimSlot {
+    token: ClaimToken,
+    /// Cleared when someone else materializes/claims this block underneath
+    /// the holder. A stale holder must not commit.
+    valid: bool,
 }
 
 impl PromoteClaimBitmap {
     pub(super) fn new() -> Self {
         Self {
-            in_flight: parking_lot::Mutex::new(std::collections::HashSet::new()),
+            in_flight: parking_lot::Mutex::new(std::collections::HashMap::new()),
             released: parking_lot::Condvar::new(),
+            next_token: AtomicU64::new(1),
         }
     }
 
     /// Insert `idx` into the in-flight set. `true` iff this caller's
     /// insertion was new (i.e., they own the claim).
     pub(super) fn try_claim(&self, idx: usize) -> bool {
-        self.in_flight.lock().insert(idx)
+        self.try_claim_token(idx).is_some()
+    }
+
+    /// [`Self::try_claim`], returning the claim's identity. Callers that
+    /// hold the claim across an await point (S3 materialization) must use
+    /// this and gate their commit on [`Self::is_valid`].
+    pub(super) fn try_claim_token(&self, idx: usize) -> Option<ClaimToken> {
+        let token = ClaimToken(self.next_token.fetch_add(1, Ordering::Relaxed));
+        let mut g = self.in_flight.lock();
+        if g.contains_key(&idx) {
+            return None;
+        }
+        g.insert(idx, ClaimSlot { token, valid: true });
+        Some(token)
+    }
+
+    /// Is `token` still the live, un-invalidated claim on `idx`?
+    pub(super) fn is_valid(&self, idx: usize, token: ClaimToken) -> bool {
+        self.if_valid(idx, token, || ()).is_some()
+    }
+
+    /// Run `f` iff `token` is still the live claim on `idx`, holding the
+    /// claim map across it — so a bypassing `NOT_PRESENT → CLEAN` cannot
+    /// slip in between the validity test and the commit it guards.
+    ///
+    /// `f` must not re-enter the claim map (the mutex is not reentrant).
+    pub(super) fn if_valid<R>(
+        &self,
+        idx: usize,
+        token: ClaimToken,
+        f: impl FnOnce() -> R,
+    ) -> Option<R> {
+        let g = self.in_flight.lock();
+        match g.get(&idx) {
+            Some(slot) if slot.valid && slot.token == token => Some(f()),
+            _ => None,
+        }
+    }
+
+    /// Invalidate the live claim on `idx`, if any. Called by every
+    /// `NOT_PRESENT → CLEAN` transition that bypasses the claim, BEFORE the
+    /// state flips — so a holder that reads `CLEAN` and then finds its token
+    /// still valid is guaranteed to be looking at its own claim.
+    pub(super) fn invalidate(&self, idx: usize) {
+        if let Some(slot) = self.in_flight.lock().get_mut(&idx) {
+            slot.valid = false;
+        }
     }
 
     /// Park until another task `release`s `idx` (or `deadline` passes).
@@ -342,7 +420,7 @@ impl PromoteClaimBitmap {
     /// thread, but only for the actual wait, with no CPU burn.
     pub(super) fn wait_for_release(&self, idx: usize, deadline: std::time::Instant) -> bool {
         let mut g = self.in_flight.lock();
-        while g.contains(&idx) {
+        while g.contains_key(&idx) {
             let now = std::time::Instant::now();
             if now >= deadline {
                 return false;
@@ -362,7 +440,7 @@ impl PromoteClaimBitmap {
 
     /// Whether any task currently holds the claim for `idx`.
     pub(super) fn is_claimed(&self, idx: usize) -> bool {
-        self.in_flight.lock().contains(&idx)
+        self.in_flight.lock().contains_key(&idx)
     }
 }
 
@@ -854,11 +932,29 @@ impl CacheInner {
     }
 
     /// Mark block as present (lock-free CAS NOT_PRESENT -> CLEAN).
+    ///
+    /// Unlike [`WriteCache::try_claim_block`] this does NOT linearize through
+    /// the promote claim — the ublk ZC `pre_write` marks its blocks present
+    /// and only lands the guest bytes later (kernel `WRITE_FIXED`), so it
+    /// cannot hold the claim across the transition. That makes it the one
+    /// path that can install a CLEAN on a block some materializer still
+    /// believes it owns (its claim's CLEAN having been stolen + flushed +
+    /// evicted meanwhile). Invalidate that claim BEFORE flipping the state:
+    /// the holder's `write_materialized` then aborts instead of landing its
+    /// stale S3 pre-image over the newer, already-durable data.
+    ///
+    /// The claim map is only touched when the block is actually NOT_PRESENT
+    /// — the steady-state overwrite path (block already present) stays
+    /// lock-free.
     #[inline]
     pub(super) fn set_present(&self, block_num: usize) {
         if block_num >= self.num_blocks {
             return;
         }
+        if self.state_map.is_present(block_num) {
+            return;
+        }
+        self.promote_claim.invalidate(block_num);
         self.state_map.set_present(block_num);
     }