Fix reconciliation in more flushers (#7805)

* Reconcile BufferedDynamicFlags flusher

* Reconcile DatabaseColumnScheduledDeleteWrapper flusher

* Rename function for consistency

lib/segment/src/common/flags/buffered_dynamic_flags.rs

@@ -51,26 +51,29 @@ impl BufferedDynamicFlags {
     }
 
     pub fn flusher(&self) -> Flusher {
-        // take pending changes
-        let (updates, required_len) = {
-            let mut buffer_guard = self.buffer.write();
-            let updates = std::mem::take(&mut *buffer_guard);
-            let Some(required_len) = updates.keys().max().map(|&max_id| max_id as usize + 1) else {
+        let updates = {
+            let buffer_guard = self.buffer.read();
+            if buffer_guard.is_empty() {
                 return Box::new(|| Ok(()));
-            };
-            (updates, required_len)
+            }
+            buffer_guard.clone()
+        };
+
+        let Some(required_len) = updates.keys().max().map(|&max_id| max_id as usize + 1) else {
+            return Box::new(|| Ok(()));
         };
 
         // Weak reference to detect when the storage has been deleted
         let flags_arc = Arc::downgrade(&self.storage);
+        let buffer = Arc::downgrade(&self.buffer);
         let is_alive_flush_lock = self.is_alive_flush_lock.handle();
 
         Box::new(move || {
-            let Some(is_alive_flush_guard) = is_alive_flush_lock.lock_if_alive() else {
-                return Ok(());
-            };
-
-            let Some(flags_arc) = flags_arc.upgrade() else {
+            let (Some(is_alive_flush_guard), Some(flags_arc), Some(buffer_arc)) = (
+                is_alive_flush_lock.lock_if_alive(),
+                flags_arc.upgrade(),
+                buffer.upgrade(),
+            ) else {
                 log::debug!("skipping flushing on deleted storage");
                 return Ok(());
             };
@@ -83,12 +86,14 @@ impl BufferedDynamicFlags {
                 flags_guard.set_len(required_len)?;
             }
 
-            for (index, value) in updates {
+            for (&index, &value) in &updates {
                 flags_guard.set(index as usize, value);
             }
 
             flags_guard.flusher()()?;
 
+            reconcile_persisted_buffer(&buffer_arc, updates);
+
             // Keep the guard till the end of the flush to prevent concurrent drop/flushes
             drop(is_alive_flush_guard);
 
@@ -97,6 +102,17 @@ impl BufferedDynamicFlags {
     }
 }
 
+/// Removes from `buffer` all results that are flushed.
+/// If values in `pending_updates` are changed, do not remove them.
+fn reconcile_persisted_buffer(
+    buffer: &RwLock<AHashMap<u32, bool>>,
+    persisted: AHashMap<u32, bool>,
+) {
+    buffer
+        .write()
+        .retain(|point_id, a| persisted.get(point_id).is_none_or(|b| a != b));
+}
+
 #[cfg(test)]
 mod tests {
 

lib/segment/src/common/rocksdb_buffered_delete_wrapper.rs

@@ -1,4 +1,3 @@
-use std::mem;
 use std::sync::Arc;
 
 use ahash::AHashSet;
@@ -70,16 +69,38 @@ impl DatabaseColumnScheduledDeleteWrapper {
     }
 
     pub fn flusher(&self) -> Flusher {
-        let ids_to_delete = mem::take(&mut *self.deleted_pending_persistence.lock());
+        let ids_to_delete = self.deleted_pending_persistence.lock().clone();
         let wrapper = self.db.clone();
+
+        let deleted_pending_persistence = Arc::downgrade(&self.deleted_pending_persistence);
+
         Box::new(move || {
-            for id in ids_to_delete {
+            let Some(deleted_pending_persistence_arc) = deleted_pending_persistence.upgrade()
+            else {
+                return Ok(());
+            };
+
+            for id in &ids_to_delete {
                 wrapper.remove(id)?;
             }
-            wrapper.flusher()()
+            wrapper.flusher()()?;
+
+            Self::reconcile_persisted_deletes(ids_to_delete, &deleted_pending_persistence_arc);
+
+            Ok(())
         })
     }
 
+    /// Removes from `deleted_pending_persistence` all results that are flushed.
+    fn reconcile_persisted_deletes(
+        persisted: AHashSet<Vec<u8>>,
+        pending_operations: &Mutex<AHashSet<Vec<u8>>>,
+    ) {
+        pending_operations
+            .lock()
+            .retain(|pending| !persisted.contains(pending));
+    }
+
     pub fn lock_db(&self) -> LockedDatabaseColumnScheduledDeleteWrapper<'_> {
         LockedDatabaseColumnScheduledDeleteWrapper {
             base: self.db.lock_db(),

lib/segment/src/common/rocksdb_buffered_update_wrapper.rs

@@ -59,7 +59,7 @@ impl DatabaseColumnScheduledUpdateWrapper {
 
     /// Removes from `pending_updates` all results that are flushed.
     /// If values in `pending_updates` are changed, do not remove them.
-    fn clear_flushed_updates(
+    fn reconcile_persisted_updates(
         flushed: PendingOperations,
         pending_operations: Arc<Mutex<PendingOperations>>,
     ) {
@@ -91,7 +91,7 @@ impl DatabaseColumnScheduledUpdateWrapper {
             }
             wrapper.flusher()()?;
 
-            Self::clear_flushed_updates(
+            Self::reconcile_persisted_updates(
                 PendingOperations { deleted, inserted },
                 pending_operations_arc,
             );