Skip to main content

lsm_tree/compaction/
filter.rs

1//! Definitions for compaction filters
2//!
3//! Compaction filters allow users to run custom logic during compactions, e.g. custom cleanup rules such as TTL.
4//! Because compactions run in background workers, using compactions filters instead of scans can massively increase the efficiency of the storage engine.
5use crate::{
6    coding::{Decode, Encode},
7    compaction::{
8        stream::{StreamFilter, StreamFilterVerdict},
9        worker::Options,
10    },
11    key::InternalKey,
12    version::Version,
13    vlog::{Accessor, BlobFileWriter, ValueHandle},
14    BlobIndirection, InternalValue, KvSeparationOptions, UserKey, UserValue, ValueType,
15};
16use std::{panic::RefUnwindSafe, path::Path};
17
18/// Verdict returned by a [`CompactionFilter`].
19#[non_exhaustive]
20#[derive(Debug, Default)]
21pub enum Verdict {
22    /// Keeps the item.
23    #[default]
24    Keep,
25
26    /// Removes the item.
27    Remove,
28
29    /// Removes the item and replace it with a weak tombstone.
30    ///
31    /// This may cause old versions of this item to be resurrected.
32    /// The semantics of this operation are identical to [`remove_weak`](crate::AbstractTree::remove_weak).
33    RemoveWeak,
34
35    /// Replaces the value of the item.
36    ReplaceValue(UserValue),
37
38    /// Destroys a value - does not leave behind a tombstone.
39    ///
40    /// Only use in situations where you absolutely 100% know your
41    /// item key is never written or updated multiple times.
42    Destroy,
43}
44
45/// Trait to implement custom compaction filters
46pub trait CompactionFilter: Send {
47    /// Returns whether an item should be kept during compaction.
48    ///
49    /// # Panicking
50    ///
51    /// This function should NOT panic.
52    ///
53    /// # Errors
54    ///
55    /// Returning an error will abort the running compaction.
56    /// This should only be done when **strictly** necessary, such as when fetching a value fails.
57    fn filter_item(&mut self, item: ItemAccessor<'_>, ctx: &Context) -> crate::Result<Verdict>;
58
59    // TODO: how would we go about adding custom properties to tables?
60    // a function that gets called every time a table is cut...? e.g.:
61    // fn on_table_cut(&mut self, ctx: ...) {}
62
63    /// Called when compaction is finished.
64    ///
65    /// # Panicking
66    ///
67    /// This function should NOT panic.
68    fn finish(self: Box<Self>) {}
69}
70
71/// Context passed to compaction filters and factories for each compaction run
72#[non_exhaustive]
73#[derive(Debug)]
74pub struct Context {
75    /// Whether we are compacting into the last level.
76    pub is_last_level: bool,
77}
78
79/// Trait that creates compaction filter objects for each compaction
80pub trait Factory: Send + Sync + RefUnwindSafe {
81    /// Returns the compaction filter name.
82    ///
83    /// This is currently only used for logging purposes.
84    fn name(&self) -> &str;
85
86    /// Returns a new compaction filter.
87    ///
88    /// # Panicking
89    ///
90    /// This function should NOT panic.
91    fn make_filter(&self, ctx: &Context) -> Box<dyn CompactionFilter>;
92}
93
94struct AccessorShared<'a> {
95    opts: &'a Options,
96    version: &'a Version,
97    blobs_folder: &'a Path,
98}
99
100impl AccessorShared<'_> {
101    /// Fetches a value from the blob store.
102    fn get_indirect_value(
103        &self,
104        user_key: &[u8],
105        vhandle: &ValueHandle,
106    ) -> crate::Result<Option<UserValue>> {
107        Accessor::new(&self.version.blob_files).get(
108            self.opts.tree_id,
109            self.blobs_folder,
110            user_key,
111            vhandle,
112            &self.opts.config.cache,
113        )
114    }
115}
116
117/// Accessor for the key/value from a compaction filter
118pub struct ItemAccessor<'a> {
119    item: &'a InternalValue,
120    shared: &'a AccessorShared<'a>,
121}
122
123impl<'a> ItemAccessor<'a> {
124    /// Get the key of this item.
125    #[must_use]
126    pub fn key(&self) -> &'a UserKey {
127        &self.item.key.user_key
128    }
129
130    /// Returns whether this item's value is stored separately.
131    #[must_use]
132    #[doc(hidden)]
133    pub fn is_indirection(&self) -> bool {
134        self.item.key.value_type.is_indirection()
135    }
136
137    /// Get the value of this item.
138    ///
139    /// # Errors
140    ///
141    /// This method will return an error if blob retrieval fails.
142    pub fn value(&self) -> crate::Result<UserValue> {
143        match self.item.key.value_type {
144            crate::ValueType::Value => Ok(self.item.value.clone()),
145            crate::ValueType::Indirection => {
146                // Resolve and read the value from a blob file.
147                let mut reader = &self.item.value[..];
148                let indirection = BlobIndirection::decode_from(&mut reader)?;
149                let vhandle = indirection.vhandle;
150
151                let value = self
152                    .shared
153                    .get_indirect_value(&self.item.key.user_key, &vhandle)?;
154
155                if let Some(value) = value {
156                    Ok(value)
157                } else {
158                    log::error!(
159                        "failed to read referenced blob file during execution of compaction filter. key: {:?}, vptr: {indirection:?}",
160                        self.item.key,
161                    );
162                    Err(crate::Error::Unrecoverable)
163                }
164            }
165            crate::ValueType::WeakTombstone | crate::ValueType::Tombstone => {
166                unreachable!("tombstones are filtered out before calling filter")
167            }
168        }
169    }
170}
171
172/// Adapts a [`CompactionFilter`] to a [`StreamFilter`]
173//
174// NOTE: this slightly helps insulate CompactionStream from lifetime spam
175pub(crate) struct StreamFilterAdapter<'a, 'b: 'a> {
176    filter: Option<&'a mut (dyn CompactionFilter + 'b)>,
177    shared: AccessorShared<'a>,
178    blob_opts: Option<&'a KvSeparationOptions>,
179    blob_writer: &'a mut Option<BlobFileWriter>,
180    ctx: &'a Context,
181}
182
183impl<'a, 'b: 'a> StreamFilterAdapter<'a, 'b> {
184    pub fn new(
185        filter: Option<&'a mut (dyn CompactionFilter + 'b)>,
186        opts: &'a Options,
187        version: &'a Version,
188        blobs_folder: &'a Path,
189        blob_writer: &'a mut Option<BlobFileWriter>,
190        ctx: &'a Context,
191    ) -> Self {
192        Self {
193            filter,
194            shared: AccessorShared {
195                opts,
196                version,
197                blobs_folder,
198            },
199            blob_opts: opts.config.kv_separation_opts.as_ref(),
200            blob_writer,
201            ctx,
202        }
203    }
204
205    /// Redirects a write to a blob file if KV separation is enabled and
206    /// the value meets the separation threshold.
207    fn handle_write(
208        &mut self,
209        prev_key: &InternalKey,
210        new_value: UserValue,
211    ) -> crate::Result<(ValueType, UserValue)> {
212        let Some(blob_opts) = self.blob_opts else {
213            return Ok((ValueType::Value, new_value));
214        };
215
216        #[expect(clippy::cast_possible_truncation, reason = "values are u32 length max")]
217        let value_size = new_value.len() as u32;
218
219        if value_size < blob_opts.separation_threshold {
220            return Ok((ValueType::Value, new_value));
221        }
222
223        let writer = if let Some(writer) = self.blob_writer {
224            writer
225        } else {
226            // Instantiate writer as necessary
227            let writer = BlobFileWriter::new(
228                self.shared.opts.blob_file_id_generator.clone(),
229                self.shared.blobs_folder,
230                self.shared.opts.tree_id,
231                self.shared.opts.config.descriptor_table.clone(),
232            )?
233            .use_target_size(blob_opts.file_target_size)
234            .use_compression(blob_opts.compression);
235
236            self.blob_writer.insert(writer)
237        };
238
239        let vhandle = writer.write(&prev_key.user_key, prev_key.seqno, &new_value)?;
240
241        let indirection = BlobIndirection {
242            vhandle,
243            size: value_size,
244        };
245
246        Ok((ValueType::Indirection, indirection.encode_into_vec().into()))
247    }
248}
249
250impl<'a, 'b: 'a> StreamFilter for StreamFilterAdapter<'a, 'b> {
251    fn filter_item(&mut self, item: &InternalValue) -> crate::Result<StreamFilterVerdict> {
252        let Some(filter) = self.filter.as_mut() else {
253            return Ok(StreamFilterVerdict::Keep);
254        };
255
256        match filter.filter_item(
257            ItemAccessor {
258                item,
259                shared: &self.shared,
260            },
261            self.ctx,
262        )? {
263            Verdict::Destroy => Ok(StreamFilterVerdict::Drop),
264            Verdict::Keep => Ok(StreamFilterVerdict::Keep),
265            Verdict::Remove => Ok(StreamFilterVerdict::Replace((
266                ValueType::Tombstone,
267                UserValue::empty(),
268            ))),
269            Verdict::RemoveWeak => Ok(StreamFilterVerdict::Replace((
270                ValueType::WeakTombstone,
271                UserValue::empty(),
272            ))),
273            Verdict::ReplaceValue(new_value) => self
274                .handle_write(&item.key, new_value)
275                .map(StreamFilterVerdict::Replace),
276        }
277    }
278}