fix(scanner): preserve dirty usage convergence (#3844)

Co-authored-by: Henry Guo <marshawcoco@users.noreply.github.com>

crates/ecstore/src/data_usage.rs

@@ -50,6 +50,10 @@ struct CachedBucketUsage {
     usage: BucketUsageInfo,
     refreshed_at: SystemTime,
     usage_updated_at: SystemTime,
+    // Set by request-path mutations until a scanner snapshot catches up to the same core counts.
+    dirty: bool,
+    // Set when a newer scanner snapshot was observed but did not include the dirty counts yet.
+    stale_snapshot_pending: bool,
 }
 
 type UsageMemoryCache = Arc<RwLock<HashMap<String, CachedBucketUsage>>>;
@@ -481,6 +485,8 @@ fn cached_bucket_usage_from_backend(usage: BucketUsageInfo, updated_at: SystemTi
         usage,
         refreshed_at: SystemTime::now(),
         usage_updated_at: updated_at,
+        dirty: false,
+        stale_snapshot_pending: false,
     }
 }
 
@@ -490,6 +496,8 @@ fn cached_bucket_usage_now(usage: BucketUsageInfo) -> CachedBucketUsage {
         usage,
         refreshed_at: now,
         usage_updated_at: now,
+        dirty: false,
+        stale_snapshot_pending: false,
     }
 }
 
@@ -497,6 +505,13 @@ fn data_usage_info_updated_at(data_usage_info: &DataUsageInfo) -> SystemTime {
     data_usage_info.last_update.unwrap_or(SystemTime::UNIX_EPOCH)
 }
 
+fn bucket_usage_counts_match(left: &BucketUsageInfo, right: &BucketUsageInfo) -> bool {
+    left.size == right.size
+        && left.objects_count == right.objects_count
+        && left.versions_count == right.versions_count
+        && left.delete_markers_count == right.delete_markers_count
+}
+
 /// Fast in-memory update for immediate quota and admin usage consistency.
 pub async fn record_bucket_object_write_memory(bucket: &str, previous_current_size: Option<u64>, new_size: u64) {
     ensure_bucket_usage_cached(bucket).await;
@@ -520,6 +535,8 @@ pub async fn record_bucket_object_write_memory(bucket: &str, previous_current_si
     let now = SystemTime::now();
     entry.refreshed_at = now;
     entry.usage_updated_at = now;
+    entry.dirty = true;
+    entry.stale_snapshot_pending = false;
 }
 
 /// Fast in-memory increment for immediate quota consistency.
@@ -545,6 +562,8 @@ pub async fn record_bucket_object_delete_memory(bucket: &str, deleted_size: u64,
     let now = SystemTime::now();
     entry.refreshed_at = now;
     entry.usage_updated_at = now;
+    entry.dirty = true;
+    entry.stale_snapshot_pending = false;
 }
 
 /// Fast in-memory decrement for immediate quota consistency
@@ -643,19 +662,36 @@ pub async fn replace_bucket_usage_memory_from_info(data_usage_info: &DataUsageIn
     let mut next_cache = HashMap::new();
 
     for (bucket, bucket_usage) in data_usage_info.buckets_usage.iter() {
-        if let Some(existing) = cache.get(bucket)
-            && existing.usage_updated_at > usage_updated_at
-        {
-            next_cache.insert(bucket.clone(), existing.clone());
-            continue;
+        if let Some(existing) = cache.get(bucket) {
+            if existing.usage_updated_at > usage_updated_at {
+                next_cache.insert(bucket.clone(), existing.clone());
+                continue;
+            }
+
+            if existing.dirty && !bucket_usage_counts_match(&existing.usage, bucket_usage) {
+                // A scanner snapshot can be saved after newer writes but still miss them if it listed the bucket earlier.
+                let mut preserved = existing.clone();
+                preserved.stale_snapshot_pending = true;
+                next_cache.insert(bucket.clone(), preserved);
+                continue;
+            }
         }
 
         next_cache.insert(bucket.clone(), cached_bucket_usage_from_backend(bucket_usage.clone(), usage_updated_at));
     }
 
     for (bucket, existing) in cache.iter() {
-        if !data_usage_info.buckets_usage.contains_key(bucket) && existing.usage_updated_at > usage_updated_at {
-            next_cache.insert(bucket.clone(), existing.clone());
+        if !data_usage_info.buckets_usage.contains_key(bucket) {
+            if existing.usage_updated_at > usage_updated_at {
+                next_cache.insert(bucket.clone(), existing.clone());
+                continue;
+            }
+
+            if existing.dirty {
+                let mut preserved = existing.clone();
+                preserved.stale_snapshot_pending = true;
+                next_cache.insert(bucket.clone(), preserved);
+            }
         }
     }
 
@@ -672,7 +708,7 @@ pub async fn apply_bucket_usage_memory_overlay(data_usage_info: &mut DataUsageIn
     let mut changed = false;
 
     for (bucket, cached) in cache.iter() {
-        if persisted_update.is_some_and(|persisted| cached.usage_updated_at <= persisted) {
+        if !cached.stale_snapshot_pending && persisted_update.is_some_and(|persisted| cached.usage_updated_at <= persisted) {
             continue;
         }
 
@@ -1171,4 +1207,34 @@ mod tests {
             Some((0, 0))
         );
     }
+
+    #[tokio::test]
+    #[serial]
+    async fn scanner_sync_preserves_dirty_memory_update_with_later_partial_snapshot() {
+        clear_usage_memory_cache_for_test().await;
+
+        let now = SystemTime::now();
+        let old_persisted = data_usage_info_for_test("bucket-a", 900, 9_000, now - Duration::from_secs(10));
+        replace_bucket_usage_memory_from_info(&old_persisted).await;
+
+        for _ in 0..100 {
+            record_bucket_object_write_memory("bucket-a", None, 10).await;
+        }
+
+        let scanner_partial = data_usage_info_for_test("bucket-a", 950, 9_500, now + Duration::from_secs(10));
+        replace_bucket_usage_memory_from_info(&scanner_partial).await;
+
+        let mut response = scanner_partial.clone();
+        apply_bucket_usage_memory_overlay(&mut response).await;
+
+        assert_eq!(response.objects_total_count, 1000);
+        assert_eq!(response.objects_total_size, 10_000);
+        assert_eq!(
+            response
+                .buckets_usage
+                .get("bucket-a")
+                .map(|usage| (usage.objects_count, usage.size)),
+            Some((1000, 10_000))
+        );
+    }
 }

crates/scanner/src/scanner.rs

@@ -124,6 +124,10 @@ enum ScannerCycleWakeReason {
 }
 
 async fn wait_for_next_scanner_cycle(ctx: &CancellationToken, delay: Duration) -> ScannerCycleWakeReason {
+    if dirty_usage_buckets_pending() {
+        return ScannerCycleWakeReason::DirtyUsage;
+    }
+
     let sleep = tokio::time::sleep(delay);
     tokio::pin!(sleep);
 
@@ -1624,6 +1628,19 @@ mod tests {
         crate::scanner_io::clear_dirty_usage_bucket("photos");
     }
 
+    #[tokio::test]
+    #[serial]
+    async fn test_wait_for_next_scanner_cycle_sees_existing_dirty_usage() {
+        crate::scanner_io::clear_dirty_usage_bucket("photos");
+        crate::scanner_io::record_dirty_usage_bucket("photos");
+
+        let ctx = CancellationToken::new();
+        let reason = wait_for_next_scanner_cycle(&ctx, Duration::from_secs(60)).await;
+
+        assert_eq!(reason, ScannerCycleWakeReason::DirtyUsage);
+        crate::scanner_io::clear_dirty_usage_bucket("photos");
+    }
+
     #[test]
     #[serial]
     fn test_get_cycle_scan_mode_runs_deep_until_selection_window_completes() {

crates/scanner/src/scanner_io.rs

@@ -165,6 +165,15 @@ fn clear_dirty_usage_buckets(snapshot: &DirtyUsageBuckets) {
         .record_scanner_dirty_usage_cycle_clear(usize_to_u64_saturated(cleared_buckets), usize_to_u64_saturated(pending_buckets));
 }
 
+fn dirty_usage_snapshot_covers_current(snapshot: &DirtyUsageBuckets) -> bool {
+    let dirty_buckets = dirty_usage_buckets();
+    dirty_buckets.iter().all(|(bucket, generation)| {
+        snapshot
+            .get(bucket)
+            .is_some_and(|snapshot_generation| snapshot_generation == generation)
+    })
+}
+
 #[cfg(test)]
 fn dirty_usage_bucket_count() -> usize {
     dirty_usage_buckets().len()
@@ -421,9 +430,10 @@ fn completed_data_usage_info(
     all_buckets: &[String],
     budget_elapsed: bool,
     cancelled: bool,
+    dirty_usage_current: bool,
 ) -> Option<(DataUsageInfo, SystemTime)> {
     let completed_set_count = results.iter().filter(|result| result.info.last_update.is_some()).count();
-    if !should_publish_completed_snapshot(completed_set_count, results.len(), budget_elapsed, cancelled) {
+    if !should_publish_completed_snapshot(completed_set_count, results.len(), budget_elapsed, cancelled) || !dirty_usage_current {
         return None;
     }
 
@@ -478,12 +488,17 @@ mod publish_gate_tests {
         let first_set = completed_root_cache("bucket-a", 1, 10);
         let second_set = completed_root_cache("bucket-b", 2, 20);
 
-        assert!(completed_data_usage_info(&[first_set.clone(), DataUsageCache::default()], &all_buckets, false, false).is_none());
-        assert!(completed_data_usage_info(&[first_set.clone(), second_set.clone()], &all_buckets, true, false).is_none());
-        assert!(completed_data_usage_info(&[first_set.clone(), second_set.clone()], &all_buckets, false, true).is_none());
+        assert!(
+            completed_data_usage_info(&[first_set.clone(), DataUsageCache::default()], &all_buckets, false, false, true)
+                .is_none()
+        );
+        assert!(completed_data_usage_info(&[first_set.clone(), second_set.clone()], &all_buckets, true, false, true).is_none());
+        assert!(completed_data_usage_info(&[first_set.clone(), second_set.clone()], &all_buckets, false, true, true).is_none());
+        assert!(completed_data_usage_info(&[first_set.clone(), second_set.clone()], &all_buckets, false, false, false).is_none());
 
-        let (data_usage_info, last_update) = completed_data_usage_info(&[first_set, second_set], &all_buckets, false, false)
-            .expect("all completed sets should produce a publishable data usage snapshot");
+        let (data_usage_info, last_update) =
+            completed_data_usage_info(&[first_set, second_set], &all_buckets, false, false, true)
+                .expect("all completed sets should produce a publishable data usage snapshot");
         assert_eq!(last_update, SystemTime::UNIX_EPOCH + Duration::from_secs(20));
         assert_eq!(data_usage_info.objects_total_count, 3);
         assert_eq!(data_usage_info.buckets_usage.len(), 2);
@@ -773,6 +788,7 @@ impl ScannerIO for ECStore {
         let results_mutex_for_updates = results_mutex.clone();
         let budget_for_updates = budget.clone();
         let child_token_for_updates = child_token.clone();
+        let dirty_usage_buckets_for_updates = dirty_usage_buckets.clone();
         tokio::spawn(async move {
             let mut last_update = SystemTime::UNIX_EPOCH;
             let mut has_sent_once = false;
@@ -795,6 +811,7 @@ impl ScannerIO for ECStore {
                                 &all_buckets_clone,
                                 budget_for_updates.budget_elapsed(),
                                 child_token_for_updates.is_cancelled(),
+                                dirty_usage_snapshot_covers_current(dirty_usage_buckets_for_updates.as_ref()),
                             )
                         };
 
@@ -813,6 +830,7 @@ impl ScannerIO for ECStore {
                                 &all_buckets_clone,
                                 budget_for_updates.budget_elapsed(),
                                 child_token_for_updates.is_cancelled(),
+                                dirty_usage_snapshot_covers_current(dirty_usage_buckets_for_updates.as_ref()),
                             )
                         };
 
@@ -1554,6 +1572,22 @@ mod tests {
         clear_dirty_usage_buckets_for_tests();
     }
 
+    #[test]
+    #[serial]
+    fn dirty_usage_snapshot_detects_uncovered_generation() {
+        clear_dirty_usage_buckets_for_tests();
+        record_dirty_usage_bucket("photos");
+        let buckets = vec![bucket_info("photos")];
+        let snapshot = snapshot_dirty_usage_buckets(&buckets);
+
+        assert!(dirty_usage_snapshot_covers_current(&snapshot));
+
+        record_dirty_usage_bucket("photos");
+
+        assert!(!dirty_usage_snapshot_covers_current(&snapshot));
+        clear_dirty_usage_buckets_for_tests();
+    }
+
     #[test]
     #[serial]
     fn clear_dirty_usage_bucket_removes_deleted_bucket_marker() {