qdrant-edge 0.7.2

use std::sync::Arc;

use ahash::AHashMap;
use crate::common::stored_bitslice::StoredBitSlice;
use crate::common::universal_io::UniversalWrite;
use itertools::Itertools;
use parking_lot::RwLock;

use crate::segment::common::Flusher;
use crate::segment::common::operation_error::{OperationError, OperationResult};

/// A wrapper around [`StoredBitSlice`] that delays writing changes to the underlying storage
/// until they get flushed manually.
/// This expects the underlying storage not to grow in size.
#[derive(Debug)]
pub struct BufferedUpdateBitSlice<S> {
    bitslice: Arc<RwLock<StoredBitSlice<S>>>,
    len: usize,
    pending_updates: Arc<RwLock<AHashMap<usize, bool>>>,
    /// Lock to prevent concurrent flush and drop
    is_alive_flush_lock: crate::common::is_alive_lock::IsAliveLock,
}

impl<S: UniversalWrite + Send + Sync + 'static> BufferedUpdateBitSlice<S> {
    pub fn new(bitslice: StoredBitSlice<S>) -> Self {
        let len = bitslice.bit_len() as usize;
        Self {
            bitslice: Arc::new(RwLock::new(bitslice)),
            len,
            pending_updates: Arc::new(RwLock::new(AHashMap::new())),
            is_alive_flush_lock: crate::common::is_alive_lock::IsAliveLock::new(),
        }
    }

    /// Sets the bit at `index` to `value` buffered.
    ///
    /// ## Panics
    /// Panics if the index is out of bounds.
    pub fn set(&self, index: usize, value: bool) {
        assert!(index < self.len, "index {index} out of range: {}", self.len);
        self.pending_updates.write().insert(index, value);
    }

    pub fn get(&self, index: usize) -> Option<bool> {
        if index >= self.len {
            return None;
        }
        if let Some(value) = self.pending_updates.read().get(&index) {
            Some(*value)
        } else {
            match self.bitslice.read().get_bit(index as u64) {
                Ok(value) => value,
                Err(err) => {
                    log::error!("Error reading bit at index {index}: {err}");
                    debug_assert!(false, "Error reading bit at index {index}: {err}");
                    None
                }
            }
        }
    }

    pub fn len(&self) -> usize {
        self.len
    }

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// Removes from `pending_updates` all results that are flushed.
    /// If values in `pending_updates` are changed, do not remove them.
    fn reconcile_persisted_updates(
        pending_updates: &RwLock<AHashMap<usize, bool>>,
        persisted: AHashMap<usize, bool>,
    ) {
        pending_updates
            .write()
            .retain(|point_id, a| persisted.get(point_id).is_none_or(|b| a != b));
    }

    /// Hint to the OS that pages backing the underlying mmap can be reclaimed.
    pub fn clear_cache(&self) -> OperationResult<()> {
        let Self {
            bitslice,
            len: _,
            pending_updates: _,
            is_alive_flush_lock: _,
        } = self;
        bitslice.read().clear_ram_cache()?;
        Ok(())
    }

    pub fn flusher(&self) -> Flusher {
        let updates = {
            let updates_guard = self.pending_updates.read();
            if updates_guard.is_empty() {
                return Box::new(|| Ok(()));
            }
            updates_guard.clone()
        };

        let bitslice = Arc::downgrade(&self.bitslice);
        let pending_updates_weak = Arc::downgrade(&self.pending_updates);
        let is_alive_flush_lock = self.is_alive_flush_lock.handle();

        Box::new(move || {
            let (Some(is_alive_flush_guard), Some(bitslice), Some(pending_updates_arc)) = (
                is_alive_flush_lock.lock_if_alive(),
                bitslice.upgrade(),
                pending_updates_weak.upgrade(),
            ) else {
                // Already dropped, skip flush
                log::trace!("BufferedUpdateBitSlice was dropped, cancelling flush");
                return Err(OperationError::cancelled(
                    "Aborted flushing on a dropped BufferedUpdateBitSlice instance",
                ));
            };

            let mut storage_write = bitslice.write();

            storage_write.set_ascending_bits_batch(
                updates
                    .iter()
                    .map(|(idx, value)| (*idx as u64, *value))
                    .sorted_unstable_by_key(|(idx, _)| *idx),
            )?;
            storage_write.flusher()()?;

            // Keep the guard till here to prevent concurrent drop/flushes
            // We don't touch files from here on and can drop the alive guard
            drop(is_alive_flush_guard);

            Self::reconcile_persisted_updates(&pending_updates_arc, updates);

            Ok(())
        })
    }
}