Repository: qdrant/qdrant
File: lib/gridstore/src/gridstore.rs (Rust)

Code before:
        Box::new(move || {
            let (Some(is_alive_flush_guard), Some(pages), Some(tracker), Some(bitmask)) = (
                is_alive_flush_lock.lock_if_alive(),
                pages.upgrade(),
                tracker.upgrade(),
                bitmask.upgrade(),
            ) else {
                log::trace!("Gridstore was cleared, cancelling flush");
                return Err(GridstoreError::FlushCancelled);
            };

            let bitmask_flusher = bitmask.read().flusher();
            bitmask_flusher()?;

            let page_flushers: Vec<_> = pages.read().iter().map(|p| p.flusher()).collect();
            for flusher in page_flushers {
                flusher()?;
            }

            let (old_pointers, tracker_flusher) = {
                let mut guard = tracker.write();
                let old_pointers = guard.write_pending(pending_updates);
                let flusher = guard.flusher();
                (old_pointers, flusher)
            };
            tracker_flusher()?;

            if old_pointers.is_empty() {
                // Nothing to do flush here
                return Ok(());
            }

            // Update all free blocks in the bitmask
            let bitmask_flusher = {
                let mut guard = bitmask.write();
                for (page_id, pointer_group) in
                    &old_pointers.into_iter().chunk_by(|pointer| pointer.page_id)
                {
                    let local_ranges = pointer_group.map(|pointer| {
                        let start = pointer.block_offset;
                        let end = pointer.block_offset
                            + Self::blocks_for_value(pointer.length as usize, block_size_bytes);
                        start as usize..end as usize
                    });
                    guard.mark_blocks_batch(page_id, local_ranges, false);
                }
                guard.flusher()
            };
            bitmask_flusher()?;

            // Keep the guard till the end of the flush to prevent concurrent drop/flushes

Review comment (refactor):
In my opinion it'd be nicer to abstract all of this away into `write_pending` and maybe rename it to `write_pending_flusher(pending)` which then returns the old pointers and the flush closure to us.

No strong opinion though. Feel free to keep it as is.

Code after:

        Box::new(move || {
            let (Some(is_alive_flush_guard), Some(pages), Some(tracker), Some(bitmask)) = (
                is_alive_flush_lock.lock_if_alive(),
                pages.upgrade(),
                tracker.upgrade(),
                bitmask.upgrade(),
            ) else {
                log::trace!("Gridstore was cleared, cancelling flush");
                return Err(GridstoreError::FlushCancelled);
            };

            let bitmask_flusher = bitmask.read().flusher();
            bitmask_flusher()?;

            let page_flushers: Vec<_> = pages.read().iter().map(|p| p.flusher()).collect();
            for flusher in page_flushers {
                flusher()?;
            }

            let old_pointers = Self::flush_tracker(&tracker, pending_updates)?;

            if old_pointers.is_empty() {
                // Nothing to do flush here
                return Ok(());
            }

            Self::flush_free_blocks(&bitmask, old_pointers, block_size_bytes)?;

            // Keep the guard till the end of the flush to prevent concurrent drop/flushes
            drop(is_alive_flush_guard);

            Ok(())
        })
    }

    /// Write pending updates to the tracker and flush it.
    /// Returns the old pointers that were overwritten and need to be freed.
    fn flush_tracker(
        tracker: &Arc<RwLock<Tracker>>,
        pending_updates: AHashMap<PointOffset, PointerUpdates>,
    ) -> crate::Result<Vec<ValuePointer>> {
        let (old_pointers, tracker_flusher) = {
            let mut guard = tracker.write();
            let old_pointers = guard.write_pending(pending_updates);
            let flusher = guard.flusher();
            (old_pointers, flusher)
        };
        tracker_flusher()?;
        Ok(old_pointers)
    }