coordinode-lsm-tree 7.0.0

//! Definitions for compaction filters
//!
//! Compaction filters allow users to run custom logic during compactions, e.g. custom cleanup rules such as TTL.
//! Because compactions run in background workers, using compactions filters instead of scans can massively increase the efficiency of the storage engine.
use crate::path::Path;
use crate::{
    BlobIndirection, InternalValue, KvSeparationOptions, SeqNo, UserKey, UserValue, ValueType,
    coding::{Decode, Encode},
    compaction::{
        stream::{StreamFilter, StreamFilterVerdict},
        worker::Options,
    },
    key::InternalKey,
    version::Version,
    vlog::{Accessor, BlobFileWriter, ValueHandle},
};
#[cfg(not(feature = "std"))]
use alloc::boxed::Box;
use core::panic::RefUnwindSafe;

/// Verdict returned by a [`CompactionFilter`].
#[non_exhaustive]
#[derive(Debug, Default)]
pub enum Verdict {
    /// Keeps the item.
    #[default]
    Keep,

    /// Removes the item.
    Remove,

    /// Removes the item and replace it with a weak tombstone.
    ///
    /// This may cause old versions of this item to be resurrected.
    /// The semantics of this operation are identical to [`remove_weak`](crate::AbstractTree::remove_weak).
    RemoveWeak,

    /// Replaces the value of the item.
    ReplaceValue(UserValue),

    /// Destroys a value - does not leave behind a tombstone.
    ///
    /// Only use in situations where you absolutely 100% know your
    /// item key is never written or updated multiple times.
    Destroy,
}

/// Trait to implement custom compaction filters
pub trait CompactionFilter: Send {
    /// Returns whether an item should be kept during compaction.
    ///
    /// # Panicking
    ///
    /// This function should NOT panic.
    ///
    /// # Errors
    ///
    /// Returning an error will abort the running compaction.
    /// This should only be done when **strictly** necessary, such as when fetching a value fails.
    fn filter_item(&mut self, item: ItemAccessor<'_>, ctx: &Context) -> crate::Result<Verdict>;

    // TODO: how would we go about adding custom properties to tables?
    // a function that gets called every time a table is cut...? e.g.:
    // fn on_table_cut(&mut self, ctx: ...) {}

    /// Called when compaction is finished.
    ///
    /// # Panicking
    ///
    /// This function should NOT panic.
    fn finish(self: Box<Self>) {}
}

/// Context passed to compaction filters and factories for each compaction run
#[non_exhaustive]
#[derive(Debug)]
pub struct Context {
    /// Whether we are compacting into the last level.
    pub is_last_level: bool,
}

/// Trait that creates compaction filter objects for each compaction
pub trait Factory: Send + Sync + RefUnwindSafe {
    /// Returns the compaction filter name.
    ///
    /// This is currently only used for logging purposes.
    fn name(&self) -> &str;

    /// Returns a new compaction filter.
    ///
    /// # Panicking
    ///
    /// This function should NOT panic.
    fn make_filter(&self, ctx: &Context) -> Box<dyn CompactionFilter>;
}

struct AccessorShared<'a> {
    opts: &'a Options,
    version: &'a Version,
    blobs_folder: &'a Path,
}

impl AccessorShared<'_> {
    /// Fetches a value from the blob store.
    fn get_indirect_value(
        &self,
        user_key: &[u8],
        vhandle: &ValueHandle,
    ) -> crate::Result<Option<UserValue>> {
        // No dictionary passed: a compaction reads blob files of every
        // generation the tree holds, and each carries its own. Not counted:
        // this read is part of the compaction, which is maintenance, and the
        // read counters leave compaction's input out.
        let accessor = Accessor::new(
            &self.version.blob_files,
            self.opts.runtime_config.load().blob_cache_fill,
            #[cfg(feature = "metrics")]
            None,
        );

        accessor.get(
            self.opts.tree_id,
            user_key,
            vhandle,
            &self.opts.config.cache,
        )
    }
}

/// Accessor for the key/value from a compaction filter
pub struct ItemAccessor<'a> {
    item: &'a InternalValue,
    shared: &'a AccessorShared<'a>,
}

impl<'a> ItemAccessor<'a> {
    /// Get the key of this item.
    #[must_use]
    pub fn key(&self) -> &'a UserKey {
        &self.item.key.user_key
    }

    /// Get the sequence number of this item.
    #[must_use]
    pub fn seqno(&self) -> SeqNo {
        self.item.key.seqno
    }

    /// Returns whether this item's value is stored separately.
    #[must_use]
    #[doc(hidden)]
    pub fn is_indirection(&self) -> bool {
        self.item.key.value_type.is_indirection()
    }

    /// Get the value of this item.
    ///
    /// # Errors
    ///
    /// This method will return an error if blob retrieval fails.
    pub fn value(&self) -> crate::Result<UserValue> {
        match self.item.key.value_type {
            // MergeOperand: return raw operand bytes so filters can inspect them.
            // The compaction pipeline preserves MergeOperand type on Replace.
            crate::ValueType::Value | crate::ValueType::MergeOperand => Ok(self.item.value.clone()),
            crate::ValueType::Indirection => {
                // Resolve and read the value from a blob file.
                let mut reader = &self.item.value[..];
                let indirection = BlobIndirection::decode_from(&mut reader)?;
                let vhandle = indirection.vhandle;

                let value = self
                    .shared
                    .get_indirect_value(&self.item.key.user_key, &vhandle)?;

                if let Some(value) = value {
                    Ok(value)
                } else {
                    log::error!(
                        "failed to read referenced blob file during execution of compaction filter. key: {:?}, vptr: {indirection:?}",
                        self.item.key,
                    );
                    Err(crate::Error::Unrecoverable)
                }
            }
            crate::ValueType::CellRow => {
                let value = crate::blob_tree::field_row::resolve_row(&self.item.value, |vptr| {
                    self.shared
                        .get_indirect_value(&self.item.key.user_key, &vptr.vhandle)?
                        .ok_or_else(|| {
                            log::error!(
                                "failed to read a cell's blob during execution of compaction filter. key: {:?}, vptr: {vptr:?}",
                                self.item.key,
                            );
                            crate::Error::Unrecoverable
                        })
                })?;
                Ok(UserValue::from(value))
            }
            crate::ValueType::WeakTombstone | crate::ValueType::Tombstone => {
                unreachable!("tombstones are filtered out before calling filter")
            }
        }
    }
}

/// The transform counters a [`StreamFilterAdapter`] ticks on a non-`Keep`
/// verdict.
pub(crate) struct TransformCounters {
    /// Shared with the merge stream and the table writer: an output whose
    /// window saw a TRANSFORMATION is not derivable from its inputs, so the
    /// writer strips its compaction lineage (a `Keep`-only run keeps it, and
    /// the manifest-repair dedup stays effective).
    pub transform: Option<alloc::sync::Arc<portable_atomic::AtomicU64>>,
    /// The FILTER's non-`Keep` verdicts alone. The shared counter above also
    /// ticks on the merge stream's own watermark-gated drops, which the GC
    /// watermark already accounts for; a user filter removes or rewrites rows
    /// regardless of any watermark, so the install needs to know that the
    /// filter itself changed something to raise the persisted retention floor
    /// to its own seqno.
    pub filter: alloc::sync::Arc<portable_atomic::AtomicU64>,
}

/// Where a filter's separated replacement values go: a blob file in `folder`,
/// opened on the first one into `writer`, under the compaction's `runtime`
/// snapshot like everything else the compaction writes. Without a `writer`
/// the replacement stays inline, for a compaction whose own blob output
/// separates it where the row is written: a relocating one, whose relocated
/// frames and new values then share target-sized files in key order.
pub(crate) struct FilterBlobOutput<'a> {
    pub folder: &'a Path,
    pub writer: Option<&'a mut Option<BlobFileWriter>>,
    pub runtime: &'a crate::runtime_config::RuntimeConfig,
}

/// Adapts a [`CompactionFilter`] to a [`StreamFilter`]
//
// NOTE: this slightly helps insulate CompactionStream from lifetime spam
pub(crate) struct StreamFilterAdapter<'a, 'b: 'a> {
    filter: Option<&'a mut (dyn CompactionFilter + 'b)>,
    shared: AccessorShared<'a>,
    blob_opts: Option<&'a KvSeparationOptions>,
    /// `None` when the compaction's output separates the values itself.
    blob_writer: Option<&'a mut Option<BlobFileWriter>>,
    runtime: &'a crate::runtime_config::RuntimeConfig,
    ctx: &'a Context,
    counters: TransformCounters,
}

impl<'a, 'b: 'a> StreamFilterAdapter<'a, 'b> {
    pub fn new(
        filter: Option<&'a mut (dyn CompactionFilter + 'b)>,
        opts: &'a Options,
        version: &'a Version,
        FilterBlobOutput {
            folder,
            writer,
            runtime,
        }: FilterBlobOutput<'a>,
        ctx: &'a Context,
        counters: TransformCounters,
    ) -> Self {
        Self {
            filter,
            shared: AccessorShared {
                opts,
                version,
                blobs_folder: folder,
            },
            blob_opts: opts.config.kv_separation_opts.as_ref(),
            blob_writer: writer,
            runtime,
            ctx,
            counters,
        }
    }

    /// Redirects a write to a blob file if KV separation is enabled and
    /// the value meets the separation threshold.
    fn handle_write(
        &mut self,
        prev_key: &InternalKey,
        new_value: UserValue,
    ) -> crate::Result<(ValueType, UserValue)> {
        let Some(blob_opts) = self.blob_opts else {
            return Ok((ValueType::Value, new_value));
        };
        // Separated where the row is written, and paced there by the merge
        // loop, which sees the whole value inline.
        let Some(blob_writer) = self.blob_writer.as_deref_mut() else {
            return Ok((ValueType::Value, new_value));
        };

        #[expect(clippy::cast_possible_truncation, reason = "values are u32 length max")]
        let value_size = new_value.len() as u32;

        if value_size < blob_opts.separation_threshold {
            return Ok((ValueType::Value, new_value));
        }

        // The payload goes to the value log below, and the stream emits only
        // the small pointer, which is all the merge loop's limiter sees; the
        // write is debited here, before it is issued, as the merge loop
        // debits an item before writing it. A stop that cuts the wait short
        // aborts the compaction the way the merge loop does.
        let opts = self.shared.opts;
        if opts
            .rate_limiter
            .request_interruptible(u64::from(value_size), || opts.stop_signal.is_stopped())
        {
            return Err(super::worker::cancelled_compaction());
        }

        let writer = if let Some(writer) = blob_writer {
            writer
        } else {
            // Instantiate writer as necessary, under the compaction's snapshot:
            // the SSTs it writes and the blob files its filter writes follow one
            // configuration, and the compaction's hold on that snapshot keeps
            // the dictionary these files name until they are installed.
            let rc = self.runtime;
            let writer = BlobFileWriter::new(
                self.shared.opts.blob_file_id_generator.clone(),
                self.shared.blobs_folder,
                self.shared.opts.tree_id,
                self.shared.opts.config.descriptor_table.clone(),
                self.shared.opts.config.fs.clone(),
            )?
            .use_output_ledger(self.shared.opts.outputs.clone())
            .use_target_size(blob_opts.file_target_size)
            .use_compression(rc.blob_compression)
            .use_sync_mode(self.shared.opts.config.sync_mode)
            .use_writeback_bytes(self.shared.opts.config.writeback_bytes)
            .use_lifetime_groups(blob_opts.lifetime_groups)
            .use_clock(self.shared.opts.config.clock.clone());

            // A filter that rewrites a separated value writes a NEW blob file
            // under the tree's blob policy, so it needs both halves: the
            // dictionary to compress with, and the set to pin on the file it
            // produces. Without the first, a `ZstdDict` policy fails the write.
            #[cfg(zstd_any)]
            let writer = {
                let dicts = self.shared.opts.config.current_zstd_dictionaries();
                writer
                    .use_zstd_dictionary(dicts.for_compression(rc.blob_compression)?)
                    .use_zstd_two_pass_seed(rc.zstd_two_pass_seed)
                    .use_zstd_dictionaries(dicts)
            };

            blob_writer.insert(writer)
        };

        // A value the filter writes is new: it has outlived nothing yet.
        writer.select_lifetime_class(
            blob_opts
                .lifetime_classifier()
                .class_of(&prev_key.user_key, false),
        )?;
        let vhandle = writer.write(&prev_key.user_key, prev_key.seqno, &new_value)?;

        let indirection = BlobIndirection {
            vhandle,
            size: value_size,
        };

        Ok((ValueType::Indirection, indirection.encode_into_vec().into()))
    }
}

impl<'a, 'b: 'a> StreamFilter for StreamFilterAdapter<'a, 'b> {
    /// The adapter is installed on every compaction, so without a user filter
    /// behind it there is nothing that can return a verdict.
    fn keeps_everything(&self) -> bool {
        self.filter.is_none()
    }

    fn read_separated_base(&mut self, base: &InternalValue) -> crate::Result<Option<UserValue>> {
        ItemAccessor {
            item: base,
            shared: &self.shared,
        }
        .value()
        .map(Some)
    }

    fn place_merged_value(
        &mut self,
        key: &InternalKey,
        value: UserValue,
    ) -> crate::Result<(ValueType, UserValue)> {
        self.handle_write(key, value)
    }

    fn filter_item(&mut self, item: &InternalValue) -> crate::Result<StreamFilterVerdict> {
        let Some(filter) = self.filter.as_mut() else {
            return Ok(StreamFilterVerdict::Keep);
        };

        let verdict = filter.filter_item(
            ItemAccessor {
                item,
                shared: &self.shared,
            },
            self.ctx,
        )?;
        if !matches!(verdict, Verdict::Keep) {
            self.counters
                .filter
                .fetch_add(1, core::sync::atomic::Ordering::Relaxed);
            if let Some(marker) = &self.counters.transform {
                marker.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
            }
        }
        match verdict {
            Verdict::Destroy => Ok(StreamFilterVerdict::Drop),
            Verdict::Keep => Ok(StreamFilterVerdict::Keep),
            Verdict::Remove => Ok(StreamFilterVerdict::Replace((
                ValueType::Tombstone,
                UserValue::empty(),
            ))),
            Verdict::RemoveWeak => Ok(StreamFilterVerdict::Replace((
                ValueType::WeakTombstone,
                UserValue::empty(),
            ))),
            // An operand stays inline: an indirection reads as a put, so an
            // operand stored as one would stop merging onto its base. The
            // stream keeps the operand type for a `Value` replacement.
            Verdict::ReplaceValue(new_value) if item.key.value_type.is_merge_operand() => {
                Ok(StreamFilterVerdict::Replace((ValueType::Value, new_value)))
            }
            Verdict::ReplaceValue(new_value) => self
                .handle_write(&item.key, new_value)
                .map(StreamFilterVerdict::Replace),
        }
    }
}