Skip to main content

tempest_kv/base/
mod.rs

1//! This module contains base types for the key-value storage layer.
2//!
3//! - [`SeqNum`]: A sequence number new-type, since they are 56 bits, due to bit-packing.
4//! - [`KeyKind`]: A one-byte key-kind that identifies different operations in our storage layer.
5//! - [`KeyTrailer`]: This type packs the [`SeqNum`] in the higher 7 bytes, with [`KeyKind`] at the
6//!   lowest byte. Ordering by it therefore prioritizes the sequence number over the kind.
7//! - [`InternalKey`]: This is the key representation that our storage layer uses. It implements
8//!   custom ordering through the [`Comparer`] trait, allowing for key suffixes that encode custom
9//!   data. Tempest uses the suffix for `HlcTimestamps`, to allow ordering of data across storages.
10
11use std::{cmp, marker::PhantomData};
12
13use bincode::Options as BincodeOptions;
14use bytes::Bytes;
15use derive_more::Display;
16use nonmax::NonMaxU64;
17use num_enum::{IntoPrimitive, TryFromPrimitive};
18use serde::{Deserialize, Serialize};
19use zerocopy::{Immutable, IntoBytes, KnownLayout, LittleEndian, U64};
20
21pub mod comparer;
22
23pub use comparer::*;
24
25use crate::StorageError;
26
27/// The crate wide used [`bincode`] encoding options.
28#[doc(hidden)]
29pub fn bincode_options() -> impl BincodeOptions {
30    bincode::options()
31        .with_fixint_encoding() // Important: no variable length ints
32        .with_little_endian() // Ensure consistency across platforms
33        .allow_trailing_bytes()
34}
35
36#[derive(
37    derive_more::Debug,
38    Display,
39    Clone,
40    Copy,
41    PartialEq,
42    Eq,
43    PartialOrd,
44    Ord,
45    Hash,
46    Serialize,
47    Deserialize,
48)]
49#[debug("SeqNum({_0:?})")]
50pub(crate) struct SeqNum(NonMaxU64);
51
52impl SeqNum {
53    /// The zero seqnum represent the absence of any resonable seqnum, used as the lowest bound.
54    pub const ZERO: Self = unsafe { Self::new_unchecked(0) };
55
56    /// The sequence number one below the start.
57    /// Sometimes exists when searching in a silo, that had no inserts yet.
58    pub const MIN: Self = unsafe { Self::new_unchecked(15) };
59    /// The first sequence number used for keys.
60    /// 1-14 are reserved for later use
61    pub const START: Self = unsafe { Self::new_unchecked(16) };
62    /// The maximum sequence number available. Sequence numbers may only use 56 bits,
63    /// since the [`KeyTrailer`] encodes the [`SeqNum`] in the upper 7 bytes, with the
64    /// [`KeyKind`] in the lowest byte, to achieve the best bit packing.
65    pub const MAX: Self = unsafe { Self::new_unchecked((1 << 56) - 1) };
66
67    #[cfg(test)]
68    pub(crate) const TEST: Self = Self::ZERO;
69
70    pub(crate) const fn new(val: u64) -> Option<Self> {
71        if val > Self::MAX.get() {
72            return None;
73        }
74        // SAFETY: Just checked `val` is in valid range
75        Some(unsafe { SeqNum::new_unchecked(val) })
76    }
77
78    /// Creates a new `SeqNum` without verifying that it is within bounds.
79    ///
80    /// # Safety
81    ///
82    /// Caller has to ensure that `val` is at most [`SeqNum::MAX`].
83    #[inline]
84    pub(crate) const unsafe fn new_unchecked(val: u64) -> Self {
85        // SAFETY: User has to ensure that `val <= Self::MAX`
86        Self(unsafe { NonMaxU64::new_unchecked(val) })
87    }
88
89    // Returns the value as a u64 primitive type.
90    #[inline]
91    pub(crate) const fn get(&self) -> u64 {
92        self.0.get()
93    }
94}
95
96impl Default for SeqNum {
97    fn default() -> Self {
98        Self::START
99    }
100}
101
102// These values are part of the file format and shall never be changed.
103#[repr(u8)]
104#[derive(
105    Debug,
106    Display,
107    Clone,
108    Copy,
109    PartialEq,
110    Eq,
111    PartialOrd,
112    Ord,
113    Hash,
114    IntoPrimitive,
115    TryFromPrimitive,
116)]
117pub enum KeyKind {
118    Delete = 0,
119    Put = 1,
120}
121
122impl KeyKind {
123    /// The lowest value (byte-wise).
124    pub const MIN: Self = Self::Delete;
125    /// The highest value (byte-wise).
126    pub const MAX: Self = Self::Put;
127
128    #[cfg(test)]
129    pub(crate) const TEST: Self = Self::MIN;
130}
131
132#[derive(
133    derive_more::Debug,
134    Clone,
135    Copy,
136    PartialEq,
137    Eq,
138    PartialOrd,
139    Ord,
140    Hash,
141    IntoBytes,
142    KnownLayout,
143    Immutable,
144)]
145#[repr(transparent)]
146#[debug("KeyTrailer({}: seqnum={:?}, kind={})", _0, self.seqnum(), self.kind())]
147pub(crate) struct KeyTrailer(U64<LittleEndian>);
148
149impl KeyTrailer {
150    pub(crate) const MAX: Self = Self::new(SeqNum::MAX, KeyKind::MAX);
151
152    #[cfg(test)]
153    pub(crate) const TEST: Self = Self::new(SeqNum::TEST, KeyKind::TEST);
154
155    pub(crate) const fn new(seqnum: SeqNum, kind: KeyKind) -> Self {
156        Self(U64::new((seqnum.get() << 8) | (kind as u64)))
157    }
158
159    pub(crate) const fn seqnum(&self) -> SeqNum {
160        // SAFETY: we right shift by 8, so it's less than or equal to SeqNum::MAX
161        unsafe { SeqNum::new_unchecked(self.0.get() >> 8) }
162    }
163
164    pub(crate) const fn kind(&self) -> KeyKind {
165        // SAFETY: we mask by 0xFF to get the key kind bits.
166        // The result is always a valid KeyKind because KeyTrailer has no FromBytes impl;
167        // all instances go through new() or TryFrom<[u8; 8]>, both of which validate
168        // the KeyKind byte.
169        unsafe { std::mem::transmute((self.0.get() & 0xFF) as u8) }
170    }
171}
172
173impl TryFrom<[u8; 8]> for KeyTrailer {
174    type Error = StorageError;
175
176    fn try_from(bytes: [u8; 8]) -> Result<Self, Self::Error> {
177        let val = u64::from_le_bytes(bytes);
178        let kind_byte = (val & 0xFF) as u8;
179        let _ = KeyKind::try_from(kind_byte)?;
180        Ok(Self(U64::new(val)))
181    }
182}
183
184#[derive(derive_more::Debug, Clone)]
185#[debug("InternalKey(key={}, seqnum={:?}, kind={:?})", C::format(key.as_ref()), trailer.seqnum(), trailer.kind())]
186pub struct InternalKey<C: Comparer = DefaultComparer, K: AsRef<[u8]> = Bytes> {
187    key: K,
188    trailer: KeyTrailer,
189    _marker: PhantomData<C>,
190}
191
192impl<C: Comparer, K: AsRef<[u8]>> InternalKey<C, K> {
193    pub(crate) fn new(key: K, trailer: KeyTrailer) -> Self {
194        Self {
195            key,
196            trailer,
197            _marker: PhantomData,
198        }
199    }
200
201    pub(crate) const fn trailer(&self) -> KeyTrailer {
202        self.trailer
203    }
204
205    pub const fn key(&self) -> &K {
206        &self.key
207    }
208
209    pub(crate) fn compare_logical(
210        &self,
211        other: &InternalKey<C, impl AsRef<[u8]>>,
212    ) -> cmp::Ordering {
213        C::compare_logical(self.key().as_ref(), other.key().as_ref())
214    }
215
216    /// Converts this key to an `InternalKey<C, &[u8]>`, by slicing the key (borrow).
217    pub(crate) fn slice_key(&self) -> InternalKey<C, &[u8]> {
218        InternalKey::new(self.key.as_ref(), self.trailer)
219    }
220
221    /// Converts this key to an `InternalKey<C, Bytes>`, by copying the key (clone).
222    pub(crate) fn byte_key(&self) -> InternalKey<C, Bytes> {
223        InternalKey::new(Bytes::copy_from_slice(self.key.as_ref()), self.trailer)
224    }
225
226    /// Split up this key into `(prefix, suffix)` based on the comparer.
227    pub fn split_up(&self) -> (&[u8], &[u8]) {
228        let key = self.key.as_ref();
229        C::split_up(key)
230    }
231}
232
233impl<C: Comparer> InternalKey<C, Bytes> {
234    /// Generate an internal key for testing other parts of Tempest.
235    /// Test keys with a higher ID will be greater than ones with a lower ID.
236    /// This is useful to test correctness of internal ordering.
237    #[cfg(test)]
238    pub(crate) fn test(id: u64) -> Self {
239        let test_key = Bytes::from(id.to_be_bytes().to_vec());
240        Self::new(test_key, KeyTrailer::TEST)
241    }
242
243    #[cfg(test)]
244    pub(crate) fn test_with_seqnum(id: u64, seqnum: u64) -> Self {
245        let test_key = Bytes::from(id.to_be_bytes().to_vec());
246        Self::new(
247            test_key,
248            KeyTrailer::new(SeqNum::new(seqnum).unwrap(), KeyKind::TEST),
249        )
250    }
251
252    // Test helper that gets the key as an easily comparable `u64` value.
253    #[cfg(test)]
254    pub(crate) fn test_key_as_u64(&self) -> u64 {
255        let mut buf = [0u8; 8];
256        buf.copy_from_slice(self.key.as_ref());
257        u64::from_be_bytes(buf)
258    }
259}
260
261impl<C: Comparer, K: AsRef<[u8]>> cmp::PartialEq for InternalKey<C, K> {
262    fn eq(&self, other: &Self) -> bool {
263        self.cmp(other).is_eq()
264    }
265}
266
267impl<C: Comparer, K: AsRef<[u8]>> cmp::Eq for InternalKey<C, K> {}
268
269impl<C: Comparer, K: AsRef<[u8]>> cmp::PartialOrd for InternalKey<C, K> {
270    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
271        Some(self.cmp(other))
272    }
273}
274
275impl<C: Comparer, K: AsRef<[u8]>> cmp::Ord for InternalKey<C, K> {
276    fn cmp(&self, other: &Self) -> cmp::Ordering {
277        match C::compare_physical(self.key.as_ref(), other.key.as_ref()) {
278            cmp::Ordering::Equal => other.trailer.cmp(&self.trailer),
279            ord => ord,
280        }
281    }
282}