risingwave_storage/hummock/sstable/
block.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::cmp::Ordering;
16use std::fmt::Debug;
17use std::io::{Read, Write};
18use std::mem::size_of;
19use std::ops::Range;
20
21use bytes::{Buf, BufMut, Bytes, BytesMut};
22use risingwave_common::catalog::TableId;
23use risingwave_hummock_sdk::KeyComparator;
24use risingwave_hummock_sdk::key::{FullKey, TableKey};
25use serde::{Deserialize, Serialize};
26
27use super::utils::{CompressionAlgorithm, bytes_diff_below_max_key_length, xxhash64_verify};
28use crate::hummock::sstable::utils;
29use crate::hummock::sstable::utils::xxhash64_checksum;
30use crate::hummock::{HummockError, HummockResult};
31use crate::monitor::Hitmap;
32
33pub const DEFAULT_BLOCK_SIZE: usize = 4 * 1024;
34pub const DEFAULT_RESTART_INTERVAL: usize = 16;
35pub const DEFAULT_ENTRY_SIZE: usize = 24; // table_id(u64) + primary_key(u64) + epoch(u64)
36
37#[allow(non_camel_case_types)]
38#[derive(Clone, Copy, PartialEq, Eq, Debug)]
39pub enum LenType {
40    u8 = 1,
41    u16 = 2,
42    u32 = 3,
43}
44
45macro_rules! put_fn {
46    ($name:ident, $($value:ident: $type:ty),*) => {
47        fn $name<T: BufMut>(&self, buf: &mut T, $($value: $type),*) {
48            match *self {
49                LenType::u8 => {
50                    $(buf.put_u8($value as u8);)*
51                },
52
53                LenType::u16 => {
54                    $(buf.put_u16($value as u16);)*
55                },
56
57                LenType::u32 => {
58                    $(buf.put_u32($value as u32);)*
59                },
60            }
61        }
62    };
63}
64
65macro_rules! get_fn {
66    ($name:ident, $($type:ty),*) => {
67        #[allow(unused_parens)]
68        fn $name<T: Buf>(&self, buf: &mut T) -> ($($type), *) {
69            match *self {
70                LenType::u8 => {
71                    ($(buf.get_u8() as $type),*)
72                }
73                LenType::u16 => {
74                    ($(buf.get_u16() as $type),*)
75                }
76                LenType::u32 => {
77                    ($(buf.get_u32() as $type),*)
78                }
79            }
80        }
81    };
82}
83
84impl From<u8> for LenType {
85    fn from(value: u8) -> Self {
86        match value {
87            1 => LenType::u8,
88            2 => LenType::u16,
89            3 => LenType::u32,
90            _ => {
91                panic!("unexpected type {}", value)
92            }
93        }
94    }
95}
96
97impl LenType {
98    put_fn!(put, v1: usize);
99
100    put_fn!(put2, v1: usize, v2: usize);
101
102    get_fn!(get, usize);
103
104    get_fn!(get2, usize, usize);
105
106    fn new(len: usize) -> Self {
107        const U8_MAX: usize = u8::MAX as usize + 1;
108        const U16_MAX: usize = u16::MAX as usize + 1;
109        const U32_MAX: usize = u32::MAX as usize + 1;
110
111        match len {
112            0..U8_MAX => LenType::u8,
113            U8_MAX..U16_MAX => LenType::u16,
114            U16_MAX..U32_MAX => LenType::u32,
115            _ => unreachable!("unexpected LenType {}", len),
116        }
117    }
118
119    fn len(&self) -> usize {
120        match *self {
121            Self::u8 => size_of::<u8>(),
122            Self::u16 => size_of::<u16>(),
123            Self::u32 => size_of::<u32>(),
124        }
125    }
126}
127
128#[derive(Clone, Copy, Debug, PartialEq, Eq)]
129pub struct RestartPoint {
130    pub offset: u32,
131    pub key_len_type: LenType,
132    pub value_len_type: LenType,
133}
134
135impl RestartPoint {
136    fn size_of() -> usize {
137        // store key_len_type and value_len_type in u8 related to `BlockBuilder::build`
138        // encoding_value = (key_len_type << 4) | value_len_type
139        std::mem::size_of::<u32>() + std::mem::size_of::<LenType>()
140    }
141}
142
143pub struct Block {
144    /// Uncompressed entries data, with restart encoded restart points info.
145    data: Bytes,
146    /// Uncompressed entries data length.
147    data_len: usize,
148
149    /// Table id of this block.
150    table_id: TableId,
151
152    /// Restart points.
153    restart_points: Vec<RestartPoint>,
154
155    hitmap: Hitmap<{ Self::HITMAP_ELEMS }>,
156}
157
158impl Clone for Block {
159    fn clone(&self) -> Self {
160        Self {
161            data: self.data.clone(),
162            data_len: self.data_len,
163            table_id: self.table_id,
164            restart_points: self.restart_points.clone(),
165            hitmap: self.hitmap.clone(),
166        }
167    }
168}
169
170impl Debug for Block {
171    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
172        f.debug_struct("Block")
173            .field("data_len", &self.data_len)
174            .field("table_id", &self.table_id)
175            .field("restart_points", &self.restart_points)
176            .finish()
177    }
178}
179
180impl Serialize for Block {
181    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
182    where
183        S: serde::Serializer,
184    {
185        serde_bytes::serialize(&self.data[..], serializer)
186    }
187}
188
189impl<'de> Deserialize<'de> for Block {
190    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
191    where
192        D: serde::Deserializer<'de>,
193    {
194        let data: Vec<u8> = serde_bytes::deserialize(deserializer)?;
195        let data = Bytes::from(data);
196        Ok(Block::decode_from_raw(data))
197    }
198}
199
200impl Block {
201    pub const HITMAP_ELEMS: usize = 4;
202
203    pub fn get_algorithm(buf: &Bytes) -> HummockResult<CompressionAlgorithm> {
204        let compression = CompressionAlgorithm::decode(&mut &buf[buf.len() - 9..buf.len() - 8])?;
205        Ok(compression)
206    }
207
208    pub fn decode(buf: Bytes, uncompressed_capacity: usize) -> HummockResult<Self> {
209        Self::decode_with_copy(buf, uncompressed_capacity, false)
210    }
211
212    pub fn decode_with_copy(
213        buf: Bytes,
214        uncompressed_capacity: usize,
215        copy: bool,
216    ) -> HummockResult<Self> {
217        // Verify checksum.
218
219        let xxhash64_checksum = (&buf[buf.len() - 8..]).get_u64_le();
220        xxhash64_verify(&buf[..buf.len() - 8], xxhash64_checksum)?;
221
222        // Decompress.
223        let compression = CompressionAlgorithm::decode(&mut &buf[buf.len() - 9..buf.len() - 8])?;
224        let compressed_data = &buf[..buf.len() - 9];
225
226        let buf = match compression {
227            CompressionAlgorithm::None => {
228                if copy {
229                    Bytes::copy_from_slice(&buf[0..(buf.len() - 9)])
230                } else {
231                    buf.slice(0..(buf.len() - 9))
232                }
233            }
234            CompressionAlgorithm::Lz4 => {
235                let mut decoder = lz4::Decoder::new(compressed_data.reader())
236                    .map_err(HummockError::decode_error)?;
237                let mut decoded = Vec::with_capacity(uncompressed_capacity);
238                let read_size = decoder
239                    .read_to_end(&mut decoded)
240                    .map_err(HummockError::decode_error)?;
241                assert_eq!(read_size, uncompressed_capacity);
242                Bytes::from(decoded)
243            }
244            CompressionAlgorithm::Zstd => {
245                let mut decoder = zstd::Decoder::new(compressed_data.reader())
246                    .map_err(HummockError::decode_error)?;
247                let mut decoded = Vec::with_capacity(uncompressed_capacity);
248                let read_size = decoder
249                    .read_to_end(&mut decoded)
250                    .map_err(HummockError::decode_error)?;
251                assert_eq!(read_size, uncompressed_capacity);
252                Bytes::from(decoded)
253            }
254        };
255
256        Ok(Self::decode_from_raw(buf))
257    }
258
259    pub fn decode_from_raw(buf: Bytes) -> Self {
260        let table_id = (&buf[buf.len() - 4..]).get_u32_le();
261        let n_index = (&buf[buf.len() - 8..buf.len() - 4]).get_u32_le() as usize;
262        let index_start = buf.len() - 8 - n_index * RestartPoint::size_of();
263        let (index_entries, remainder) = buf[index_start..buf.len() - 8].as_chunks::<5>();
264        assert!(remainder.is_empty());
265        let mut index_entries = index_entries.iter();
266        let decode_index = |entry: &[u8; 5]| RestartPoint {
267            offset: u32::from_le_bytes(entry[..4].try_into().unwrap()),
268            key_len_type: LenType::from(entry[4] >> 4),
269            value_len_type: LenType::from(entry[4] & 0x0f),
270        };
271        let mut current = decode_index(index_entries.next().unwrap());
272        let mut next = index_entries.next().map(decode_index);
273        let n_restarts = (&buf[index_start - 4..index_start]).get_u32_le() as usize;
274        let data_len = index_start - 4 - n_restarts * size_of::<u32>();
275        let (offsets, remainder) = buf[data_len..index_start - 4].as_chunks::<4>();
276        assert!(remainder.is_empty());
277        let mut restart_points = Vec::with_capacity(n_restarts);
278        for bytes in offsets {
279            let offset = u32::from_le_bytes(*bytes);
280            if let Some(entry) = next.as_ref()
281                && offset >= entry.offset
282            {
283                current = *entry;
284                next = index_entries.next().map(decode_index);
285            }
286            restart_points.push(RestartPoint {
287                offset,
288                key_len_type: current.key_len_type,
289                value_len_type: current.value_len_type,
290            });
291        }
292        Block {
293            data: buf,
294            data_len,
295            restart_points,
296            table_id: TableId::new(table_id),
297            hitmap: Hitmap::default(),
298        }
299    }
300
301    /// Original decoder retained only as a Criterion comparison control.
302    #[cfg(feature = "test")]
303    #[doc(hidden)]
304    pub fn decode_from_raw_legacy_for_bench(buf: Bytes) -> Self {
305        let table_id = (&buf[buf.len() - 4..]).get_u32_le();
306        // decode restart_points_type_index
307        let n_index = ((&buf[buf.len() - 8..buf.len() - 4]).get_u32_le()) as usize;
308        let index_data_len = size_of::<u32>() + n_index * RestartPoint::size_of();
309        let data_len = buf.len() - 4 - index_data_len;
310        let mut restart_points_type_index_buf = &buf[data_len..buf.len() - 8];
311
312        let mut index_key_vec = Vec::with_capacity(n_index);
313        for _ in 0..n_index {
314            let offset = restart_points_type_index_buf.get_u32_le();
315            let value = restart_points_type_index_buf.get_u8();
316            let key_len_type = LenType::from(value >> 4);
317            let value_len_type = LenType::from(value & 0x0F);
318
319            index_key_vec.push(RestartPoint {
320                offset,
321                key_len_type,
322                value_len_type,
323            });
324        }
325
326        // Decode restart points.
327        let n_restarts = ((&buf[data_len - 4..]).get_u32_le()) as usize;
328        let restart_points_len = size_of::<u32>() + n_restarts * (size_of::<u32>());
329        let restarts_end = data_len - 4;
330        let data_len = data_len - restart_points_len;
331        let mut restart_points = Vec::with_capacity(n_restarts);
332        let mut restart_points_buf = &buf[data_len..restarts_end];
333
334        let mut type_index: usize = 0;
335
336        for _ in 0..n_restarts {
337            let offset = restart_points_buf.get_u32_le();
338            if type_index < index_key_vec.len() - 1
339                && offset >= index_key_vec[type_index + 1].offset
340            {
341                type_index += 1;
342            }
343
344            restart_points.push(RestartPoint {
345                offset,
346                key_len_type: index_key_vec[type_index].key_len_type,
347                value_len_type: index_key_vec[type_index].value_len_type,
348            });
349        }
350
351        Block {
352            data: buf,
353            data_len,
354            restart_points,
355            table_id: TableId::new(table_id),
356            hitmap: Hitmap::default(),
357        }
358    }
359
360    /// Entries data len.
361    #[expect(clippy::len_without_is_empty)]
362    pub fn len(&self) -> usize {
363        assert!(!self.data.is_empty());
364        self.data_len
365    }
366
367    pub fn capacity(&self) -> usize {
368        self.data.len()
369            + self.restart_points.capacity() * std::mem::size_of::<u32>()
370            + std::mem::size_of::<u32>()
371    }
372
373    pub fn estimated_memory_weight(&self) -> usize {
374        size_of::<Self>()
375            + self.data.len()
376            + self.restart_points.capacity() * size_of::<RestartPoint>()
377            + Hitmap::<{ Self::HITMAP_ELEMS }>::bytes()
378    }
379
380    pub fn table_id(&self) -> TableId {
381        self.table_id
382    }
383
384    /// Gets restart point by index.
385    pub fn restart_point(&self, index: usize) -> RestartPoint {
386        self.restart_points[index]
387    }
388
389    /// Gets restart point len.
390    pub fn restart_point_len(&self) -> usize {
391        self.restart_points.len()
392    }
393
394    /// Searches the index of the restart point by partition point.
395    pub fn search_restart_partition_point<P>(&self, pred: P) -> usize
396    where
397        P: FnMut(&RestartPoint) -> bool,
398    {
399        self.restart_points.partition_point(pred)
400    }
401
402    pub fn data(&self) -> &[u8] {
403        &self.data[..self.data_len]
404    }
405
406    pub fn raw(&self) -> &[u8] {
407        &self.data[..]
408    }
409
410    pub fn hitmap(&self) -> &Hitmap<{ Self::HITMAP_ELEMS }> {
411        &self.hitmap
412    }
413
414    pub fn efficiency(&self) -> f64 {
415        self.hitmap.ratio()
416    }
417}
418
419/// [`KeyPrefix`] contains info for prefix compression.
420#[derive(Debug)]
421pub struct KeyPrefix {
422    overlap: usize,
423    diff: usize,
424    value: usize,
425    /// Used for calculating range, won't be encoded.
426    offset: usize,
427
428    len: usize,
429}
430
431impl KeyPrefix {
432    // This function is used in BlockBuilder::add to provide a wrapper for encode since the
433    // KeyPrefix len field is only useful in the decode phase
434    pub fn new_without_len(overlap: usize, diff: usize, value: usize, offset: usize) -> Self {
435        KeyPrefix {
436            overlap,
437            diff,
438            value,
439            offset,
440            len: 0, // not used when encode
441        }
442    }
443}
444
445impl KeyPrefix {
446    pub fn encode(&self, buf: &mut impl BufMut, key_len_type: LenType, value_len_type: LenType) {
447        key_len_type.put2(buf, self.overlap, self.diff);
448        value_len_type.put(buf, self.value);
449    }
450
451    pub fn decode(
452        buf: &mut impl Buf,
453        offset: usize,
454        key_len_type: LenType,
455        value_len_type: LenType,
456    ) -> Self {
457        let (overlap, diff) = key_len_type.get2(buf);
458        let value = value_len_type.get(buf);
459
460        let len = key_len_type.len() * 2 + value_len_type.len();
461
462        Self {
463            overlap,
464            diff,
465            value,
466            offset,
467            len,
468        }
469    }
470
471    /// Encoded length.
472    fn len(&self) -> usize {
473        self.len
474    }
475
476    /// Gets overlap len.
477    pub fn overlap_len(&self) -> usize {
478        self.overlap
479    }
480
481    /// Gets diff key range.
482    pub fn diff_key_range(&self) -> Range<usize> {
483        self.offset + self.len()..self.offset + self.len() + self.diff
484    }
485
486    /// Gets value range.
487    pub fn value_range(&self) -> Range<usize> {
488        self.offset + self.len() + self.diff..self.offset + self.len() + self.diff + self.value
489    }
490
491    /// Gets entry len.
492    pub fn entry_len(&self) -> usize {
493        self.len() + self.diff + self.value
494    }
495}
496
497pub struct BlockBuilderOptions {
498    /// Reserved bytes size when creating buffer to avoid frequent allocating.
499    pub capacity: usize,
500    /// Compression algorithm.
501    pub compression_algorithm: CompressionAlgorithm,
502    /// Restart point interval.
503    pub restart_interval: usize,
504}
505
506impl Default for BlockBuilderOptions {
507    fn default() -> Self {
508        Self {
509            capacity: DEFAULT_BLOCK_SIZE,
510            compression_algorithm: CompressionAlgorithm::None,
511            restart_interval: DEFAULT_RESTART_INTERVAL,
512        }
513    }
514}
515
516/// [`BlockBuilder`] encodes and appends block to a buffer.
517pub struct BlockBuilder {
518    /// Write buffer.
519    buf: BytesMut,
520    /// Compress buffer
521    compress_buf: BytesMut,
522    /// Entry interval between restart points.
523    restart_count: usize,
524    /// Restart points.
525    restart_points: Vec<u32>,
526    /// Last key.
527    last_key: Vec<u8>,
528    /// Count of entries in current block.
529    entry_count: usize,
530    /// Compression algorithm.
531    compression_algorithm: CompressionAlgorithm,
532
533    table_id: Option<TableId>,
534    // restart_points_type_index stores only the restart_point corresponding to each type change,
535    // as an index, in order to reduce space usage
536    restart_points_type_index: Vec<RestartPoint>,
537}
538
539impl BlockBuilder {
540    pub fn new(options: BlockBuilderOptions) -> Self {
541        Self {
542            // add more space to avoid re-allocate space. (for restart_points and restart_points_type_index)
543            buf: BytesMut::with_capacity(Self::buf_reserve_size(&options)),
544            compress_buf: BytesMut::default(),
545            restart_count: options.restart_interval,
546            restart_points: Vec::with_capacity(
547                options.capacity / DEFAULT_ENTRY_SIZE / options.restart_interval + 1,
548            ),
549            last_key: vec![],
550            entry_count: 0,
551            compression_algorithm: options.compression_algorithm,
552            table_id: None,
553            restart_points_type_index: Vec::default(),
554        }
555    }
556
557    /// Appends a kv pair to the block using pre-encoded key.
558    ///
559    /// This method accepts a pre-encoded key (`table_key` + `epoch`, without `table_id`)
560    /// to avoid redundant encoding operations. This is useful when the caller
561    /// has already encoded the full key and wants to reuse it.
562    ///
563    /// NOTE: Key must be added in ASCEND order.
564    ///
565    /// # Format
566    ///
567    /// ```plain
568    /// entry (kv pair): | overlap len (len_type) | diff len (len_type) | value len(len_type) | diff key | value |
569    /// ```
570    ///
571    /// # Panics
572    ///
573    /// Panic if key is not added in ASCEND order or `table_id` mismatch.
574    pub fn add(&mut self, table_id: TableId, encoded_key_without_table_id: &[u8], value: &[u8]) {
575        match self.table_id {
576            Some(current_table_id) => assert_eq!(current_table_id, table_id),
577            None => self.table_id = Some(table_id),
578        }
579        #[cfg(debug_assertions)]
580        self.debug_valid();
581
582        // Call the core implementation with the pre-encoded key
583        self.add_encoded_impl(encoded_key_without_table_id, value);
584    }
585
586    #[cfg(any(test, feature = "test"))]
587    pub fn add_for_test(&mut self, full_key: FullKey<&[u8]>, value: &[u8]) {
588        use risingwave_hummock_sdk::key::TABLE_PREFIX_LEN;
589        let input_table_id = full_key.user_key.table_id;
590        match self.table_id {
591            Some(current_table_id) => assert_eq!(current_table_id, input_table_id),
592            None => self.table_id = Some(input_table_id),
593        }
594        #[cfg(debug_assertions)]
595        self.debug_valid();
596
597        let mut key: BytesMut = Default::default();
598        full_key.encode_into(&mut key);
599
600        // Slice off the table_id prefix to get the key without table_id
601        let encoded_key_without_table_id = &key[TABLE_PREFIX_LEN..];
602
603        // Call the core implementation with encoded key
604        self.add_encoded_impl(encoded_key_without_table_id, value);
605    }
606
607    /// Core implementation that accepts pre-encoded key.
608    /// This is used by both `add` and `add_for_test` to avoid code duplication.
609    fn add_encoded_impl(&mut self, key: &[u8], value: &[u8]) {
610        if self.entry_count > 0 {
611            debug_assert!(!key.is_empty());
612            debug_assert_eq!(
613                KeyComparator::compare_encoded_full_key(&self.last_key[..], key),
614                Ordering::Less,
615            );
616        }
617        // Update restart point if needed and calculate diff key.
618        let k_type = LenType::new(key.len());
619        let v_type = LenType::new(value.len());
620
621        let type_mismatch = if let Some(RestartPoint {
622            offset: _,
623            key_len_type: last_key_len_type,
624            value_len_type: last_value_len_type,
625        }) = self.restart_points_type_index.last()
626        {
627            k_type != *last_key_len_type || v_type != *last_value_len_type
628        } else {
629            true
630        };
631
632        let diff_key = if self.entry_count.is_multiple_of(self.restart_count) || type_mismatch {
633            let offset = utils::checked_into_u32(self.buf.len()).unwrap_or_else(|_| {
634                panic!(
635                    "WARN overflow can't convert buf_len {} into u32 table {:?}",
636                    self.buf.len(),
637                    self.table_id,
638                )
639            });
640
641            self.restart_points.push(offset);
642
643            if type_mismatch {
644                self.restart_points_type_index.push(RestartPoint {
645                    offset,
646                    key_len_type: k_type,
647                    value_len_type: v_type,
648                });
649            }
650
651            key
652        } else {
653            bytes_diff_below_max_key_length(&self.last_key, key)
654        };
655
656        let prefix = KeyPrefix::new_without_len(
657            key.len() - diff_key.len(),
658            diff_key.len(),
659            value.len(),
660            self.buf.len(),
661        );
662
663        prefix.encode(&mut self.buf, k_type, v_type);
664        self.buf.put_slice(diff_key);
665        self.buf.put_slice(value);
666
667        self.last_key.clear();
668        self.last_key.extend_from_slice(key);
669        self.entry_count += 1;
670    }
671
672    pub fn get_last_key(&self) -> &[u8] {
673        &self.last_key
674    }
675
676    pub fn is_empty(&self) -> bool {
677        self.buf.is_empty()
678    }
679
680    pub fn clear(&mut self) {
681        self.buf.clear();
682        self.restart_points.clear();
683        self.table_id = None;
684        self.restart_points_type_index.clear();
685        self.last_key.clear();
686        self.entry_count = 0;
687    }
688
689    /// Calculate block size without compression.
690    pub fn uncompressed_block_size(&mut self) -> usize {
691        self.buf.len()
692            + (self.restart_points.len() + 1) * std::mem::size_of::<u32>()
693            + (RestartPoint::size_of()) // (offset + len_type(u8)) * len
694                * self.restart_points_type_index.len()
695            + std::mem::size_of::<u32>() // restart_points_type_index len
696            + std::mem::size_of::<u32>() // table_id len
697    }
698
699    /// Finishes building block.
700    ///
701    /// # Format
702    ///
703    /// ```plain
704    /// compressed: | entries | restart point 0 (4B) | ... | restart point N-1 (4B) | N (4B) | restart point index 0 (5B)| ... | restart point index N-1 (5B) | N (4B) | table id (4B)
705    /// uncompressed: | compression method (1B) | xxhash64 checksum (8B) |
706    /// ```
707    ///
708    /// # Panics
709    ///
710    /// Panic if there is compression error.
711    pub fn build(&mut self) -> &[u8] {
712        assert!(
713            self.entry_count > 0,
714            "buf_len {} entry_count {} table {:?}",
715            self.buf.len(),
716            self.entry_count,
717            self.table_id
718        );
719
720        for restart_point in &self.restart_points {
721            self.buf.put_u32_le(*restart_point);
722        }
723
724        self.buf.put_u32_le(
725            utils::checked_into_u32(self.restart_points.len()).unwrap_or_else(|_| {
726                panic!(
727                    "WARN overflow can't convert restart_points_len {} into u32 table {:?}",
728                    self.restart_points.len(),
729                    self.table_id,
730                )
731            }),
732        );
733        for RestartPoint {
734            offset,
735            key_len_type,
736            value_len_type,
737        } in &self.restart_points_type_index
738        {
739            self.buf.put_u32_le(*offset);
740
741            let mut value: u8 = 0;
742            value |= *key_len_type as u8;
743            value <<= 4;
744            value |= *value_len_type as u8;
745
746            self.buf.put_u8(value);
747        }
748
749        self.buf.put_u32_le(
750            utils::checked_into_u32(self.restart_points_type_index.len()).unwrap_or_else(|_| {
751                panic!(
752                    "WARN overflow can't convert restart_points_type_index_len {} into u32 table {:?}",
753                    self.restart_points_type_index.len(),
754                    self.table_id,
755                )
756            }),
757        );
758
759        self.buf.put_u32_le(self.table_id.unwrap().as_raw_id());
760        let result_buf = if self.compression_algorithm != CompressionAlgorithm::None {
761            self.compress_buf.clear();
762            self.compress_buf = Self::compress(
763                &self.buf[..],
764                self.compression_algorithm,
765                std::mem::take(&mut self.compress_buf),
766            );
767
768            &mut self.compress_buf
769        } else {
770            &mut self.buf
771        };
772
773        self.compression_algorithm.encode(result_buf);
774        let checksum = xxhash64_checksum(result_buf);
775        result_buf.put_u64_le(checksum);
776        assert!(
777            result_buf.len() < (u32::MAX) as usize,
778            "buf_len {} entry_count {} table {:?}",
779            result_buf.len(),
780            self.entry_count,
781            self.table_id
782        );
783
784        if self.compression_algorithm != CompressionAlgorithm::None {
785            self.compress_buf.as_ref()
786        } else {
787            self.buf.as_ref()
788        }
789    }
790
791    pub fn compress_block(
792        buf: Bytes,
793        target_compression: CompressionAlgorithm,
794    ) -> HummockResult<Bytes> {
795        // Verify checksum.
796        let checksum = (&buf[buf.len() - 8..]).get_u64_le();
797        xxhash64_verify(&buf[..buf.len() - 8], checksum)?;
798        // Decompress.
799        let compression = CompressionAlgorithm::decode(&mut &buf[buf.len() - 9..buf.len() - 8])?;
800        let compressed_data = &buf[..buf.len() - 9];
801        assert_eq!(compression, CompressionAlgorithm::None);
802        let mut compress_writer = Self::compress(
803            compressed_data,
804            target_compression,
805            BytesMut::with_capacity(buf.len()),
806        );
807
808        target_compression.encode(&mut compress_writer);
809        let checksum = xxhash64_checksum(&compress_writer);
810        compress_writer.put_u64_le(checksum);
811        Ok(compress_writer.freeze())
812    }
813
814    pub fn compress(
815        buf: &[u8],
816        compression_algorithm: CompressionAlgorithm,
817        compress_writer: BytesMut,
818    ) -> BytesMut {
819        match compression_algorithm {
820            CompressionAlgorithm::None => unreachable!(),
821            CompressionAlgorithm::Lz4 => {
822                let mut encoder = lz4::EncoderBuilder::new()
823                    .level(4)
824                    .build(compress_writer.writer())
825                    .map_err(HummockError::encode_error)
826                    .unwrap();
827                encoder
828                    .write_all(buf)
829                    .map_err(HummockError::encode_error)
830                    .unwrap();
831                let (writer, result) = encoder.finish();
832                result.map_err(HummockError::encode_error).unwrap();
833                writer.into_inner()
834            }
835            CompressionAlgorithm::Zstd => {
836                let mut encoder = zstd::Encoder::new(compress_writer.writer(), 4)
837                    .map_err(HummockError::encode_error)
838                    .unwrap();
839                encoder
840                    .write_all(buf)
841                    .map_err(HummockError::encode_error)
842                    .unwrap();
843                let writer = encoder
844                    .finish()
845                    .map_err(HummockError::encode_error)
846                    .unwrap();
847                writer.into_inner()
848            }
849        }
850    }
851
852    /// Approximate block len (uncompressed).
853    pub fn approximate_len(&self) -> usize {
854        // block + restart_points + restart_points.len + restart_points_type_indices +
855        // restart_points_type_indics.len compression_algorithm + checksum
856        self.buf.len()
857            + std::mem::size_of::<u32>() * self.restart_points.len() // restart_points
858            + std::mem::size_of::<u32>() // restart_points.len
859            + RestartPoint::size_of() * self.restart_points_type_index.len() // restart_points_type_indics
860            + std::mem::size_of::<u32>() // restart_points_type_indics.len
861            + std::mem::size_of::<CompressionAlgorithm>() // compression_algorithm
862            + std::mem::size_of::<u64>() // checksum
863            + std::mem::size_of::<u32>() // table_id
864    }
865
866    pub fn debug_valid(&self) {
867        if self.entry_count == 0 {
868            debug_assert!(self.buf.is_empty());
869            debug_assert!(self.restart_points.is_empty());
870            debug_assert!(self.restart_points_type_index.is_empty());
871            debug_assert!(self.last_key.is_empty());
872        }
873    }
874
875    pub fn table_id(&self) -> Option<TableId> {
876        self.table_id
877    }
878
879    fn buf_reserve_size(option: &BlockBuilderOptions) -> usize {
880        option.capacity + 1024 + 256
881    }
882}
883
884/// Attempts to shorten `block_smallest` key while preserving block boundary correctness.
885///
886/// Returns a shortened key if `prev_block_last < shortened <= block_smallest` can be satisfied
887/// with fewer bytes. This reduces block metadata size by exploiting common prefixes between
888/// adjacent blocks.
889///
890/// Returns `None` if the key cannot be shortened (e.g., keys differ only in the last byte).
891pub fn try_shorten_block_smallest_key(
892    prev_block_last: &FullKey<&[u8]>,
893    block_smallest: &FullKey<&[u8]>,
894) -> Option<FullKey<Vec<u8>>> {
895    /// Returns the length of the longest common prefix between two byte slices.
896    fn lcp_len(a: &[u8], b: &[u8]) -> usize {
897        let min_len = a.len().min(b.len());
898        for i in 0..min_len {
899            if a[i] != b[i] {
900                return i;
901            }
902        }
903        min_len
904    }
905
906    let prev_table_key = prev_block_last.user_key.table_key.as_ref();
907    let next_table_key = block_smallest.user_key.table_key.as_ref();
908    assert_eq!(
909        prev_block_last.user_key.table_id, block_smallest.user_key.table_id,
910        "table ids must match for shortening block smallest key"
911    );
912
913    let lcp = lcp_len(prev_table_key, next_table_key);
914
915    // Shortened key = LCP prefix + first differing byte.
916    // Only shorten if the result is strictly shorter than next_table_key.
917    let shortened_len = lcp + 1;
918    if shortened_len >= next_table_key.len() {
919        return None;
920    }
921
922    // Build candidate key: take LCP + 1 bytes from block_smallest
923    let cand = FullKey::new_with_gap_epoch(
924        block_smallest.user_key.table_id,
925        TableKey(&next_table_key[..shortened_len]),
926        block_smallest.epoch_with_gap,
927    );
928
929    // Invariant: prev_block_last < cand <= block_smallest
930    // This must hold by construction. If violated, it indicates a bug in the shortening logic.
931    // Use assert! (not debug_assert!) to fail fast and prevent writing incorrect keys to storage.
932    assert!(
933        prev_block_last.cmp(&cand).is_lt(),
934        "Invariant violated: prev_block_last >= cand. prev={:?}, cand={:?}",
935        prev_block_last,
936        cand
937    );
938    assert!(
939        cand.cmp(block_smallest).is_le(),
940        "Invariant violated: cand > block_smallest. cand={:?}, block_smallest={:?}",
941        cand,
942        block_smallest
943    );
944
945    Some(cand.copy_into())
946}
947
948#[cfg(test)]
949mod tests {
950
951    use risingwave_common::util::epoch::test_epoch;
952    use risingwave_hummock_sdk::key::MAX_KEY_LEN;
953
954    use super::*;
955    use crate::hummock::{BlockHolder, BlockIterator};
956
957    #[test]
958    fn test_block_enc_dec() {
959        let options = BlockBuilderOptions::default();
960        let mut builder = BlockBuilder::new(options);
961        builder.add_for_test(construct_full_key_struct_for_test(0, b"k1", 1), b"v01");
962        builder.add_for_test(construct_full_key_struct_for_test(0, b"k2", 2), b"v02");
963        builder.add_for_test(construct_full_key_struct_for_test(0, b"k3", 3), b"v03");
964        builder.add_for_test(construct_full_key_struct_for_test(0, b"k4", 4), b"v04");
965        let capacity = builder.uncompressed_block_size();
966        assert_eq!(capacity, builder.approximate_len() - 9);
967        let buf = builder.build().to_vec();
968        let block = Box::new(Block::decode(buf.into(), capacity).unwrap());
969        let mut bi = BlockIterator::new(BlockHolder::from_owned_block(block));
970
971        bi.seek_to_first();
972        assert!(bi.is_valid());
973        assert_eq!(construct_full_key_struct_for_test(0, b"k1", 1), bi.key());
974        assert_eq!(b"v01", bi.value());
975
976        bi.next();
977        assert!(bi.is_valid());
978        assert_eq!(construct_full_key_struct_for_test(0, b"k2", 2), bi.key());
979        assert_eq!(b"v02", bi.value());
980
981        bi.next();
982        assert!(bi.is_valid());
983        assert_eq!(construct_full_key_struct_for_test(0, b"k3", 3), bi.key());
984        assert_eq!(b"v03", bi.value());
985
986        bi.next();
987        assert!(bi.is_valid());
988        assert_eq!(construct_full_key_struct_for_test(0, b"k4", 4), bi.key());
989        assert_eq!(b"v04", bi.value());
990
991        bi.next();
992        assert!(!bi.is_valid());
993    }
994
995    #[test]
996    fn test_compressed_block_enc_dec() {
997        inner_test_compressed(CompressionAlgorithm::Lz4);
998        inner_test_compressed(CompressionAlgorithm::Zstd);
999    }
1000
1001    fn inner_test_compressed(algo: CompressionAlgorithm) {
1002        let options = BlockBuilderOptions {
1003            compression_algorithm: algo,
1004            ..Default::default()
1005        };
1006        let mut builder = BlockBuilder::new(options);
1007        builder.add_for_test(construct_full_key_struct_for_test(0, b"k1", 1), b"v01");
1008        builder.add_for_test(construct_full_key_struct_for_test(0, b"k2", 2), b"v02");
1009        builder.add_for_test(construct_full_key_struct_for_test(0, b"k3", 3), b"v03");
1010        builder.add_for_test(construct_full_key_struct_for_test(0, b"k4", 4), b"v04");
1011        let capacity = builder.uncompressed_block_size();
1012        assert_eq!(capacity, builder.approximate_len() - 9);
1013        let buf = builder.build().to_vec();
1014        let block = Box::new(Block::decode(buf.into(), capacity).unwrap());
1015        let mut bi = BlockIterator::new(BlockHolder::from_owned_block(block));
1016
1017        bi.seek_to_first();
1018        assert!(bi.is_valid());
1019        assert_eq!(construct_full_key_struct_for_test(0, b"k1", 1), bi.key());
1020        assert_eq!(b"v01", bi.value());
1021
1022        bi.next();
1023        assert!(bi.is_valid());
1024        assert_eq!(construct_full_key_struct_for_test(0, b"k2", 2), bi.key());
1025        assert_eq!(b"v02", bi.value());
1026
1027        bi.next();
1028        assert!(bi.is_valid());
1029        assert_eq!(construct_full_key_struct_for_test(0, b"k3", 3), bi.key());
1030        assert_eq!(b"v03", bi.value());
1031
1032        bi.next();
1033        assert!(bi.is_valid());
1034        assert_eq!(construct_full_key_struct_for_test(0, b"k4", 4), bi.key());
1035        assert_eq!(b"v04", bi.value());
1036
1037        bi.next();
1038        assert!(!bi.is_valid());
1039    }
1040
1041    pub fn construct_full_key_struct_for_test(
1042        table_id: u32,
1043        table_key: &[u8],
1044        epoch: u64,
1045    ) -> FullKey<&[u8]> {
1046        FullKey::for_test(TableId::new(table_id), table_key, test_epoch(epoch))
1047    }
1048
1049    #[test]
1050    fn test_block_enc_large_key() {
1051        let options = BlockBuilderOptions::default();
1052        let mut builder = BlockBuilder::new(options);
1053        let medium_key = vec![b'a'; MAX_KEY_LEN - 500];
1054        let large_key = vec![b'b'; MAX_KEY_LEN];
1055        let xlarge_key = vec![b'c'; MAX_KEY_LEN + 500];
1056
1057        builder.add_for_test(construct_full_key_struct_for_test(0, &medium_key, 1), b"v1");
1058        builder.add_for_test(construct_full_key_struct_for_test(0, &large_key, 2), b"v2");
1059        builder.add_for_test(construct_full_key_struct_for_test(0, &xlarge_key, 3), b"v3");
1060        let capacity = builder.uncompressed_block_size();
1061        assert_eq!(capacity, builder.approximate_len() - 9);
1062        let buf = builder.build().to_vec();
1063        let block = Box::new(Block::decode(buf.into(), capacity).unwrap());
1064        let mut bi = BlockIterator::new(BlockHolder::from_owned_block(block));
1065
1066        bi.seek_to_first();
1067        assert!(bi.is_valid());
1068        assert_eq!(
1069            construct_full_key_struct_for_test(0, &medium_key, 1),
1070            bi.key()
1071        );
1072        assert_eq!(b"v1", bi.value());
1073
1074        bi.next();
1075        assert!(bi.is_valid());
1076        assert_eq!(
1077            construct_full_key_struct_for_test(0, &large_key, 2),
1078            bi.key()
1079        );
1080        assert_eq!(b"v2", bi.value());
1081
1082        bi.next();
1083        assert!(bi.is_valid());
1084        assert_eq!(
1085            construct_full_key_struct_for_test(0, &xlarge_key, 3),
1086            bi.key()
1087        );
1088        assert_eq!(b"v3", bi.value());
1089
1090        bi.next();
1091        assert!(!bi.is_valid());
1092    }
1093
1094    #[test]
1095    fn test_block_restart_point() {
1096        let options = BlockBuilderOptions::default();
1097        let mut builder = BlockBuilder::new(options);
1098
1099        const KEY_COUNT: u8 = 100;
1100        const BUILDER_COUNT: u8 = 5;
1101
1102        for _ in 0..BUILDER_COUNT {
1103            for index in 0..KEY_COUNT {
1104                if index < 50 {
1105                    let mut medium_key = vec![b'A'; MAX_KEY_LEN - 500];
1106                    medium_key.push(index);
1107                    builder
1108                        .add_for_test(construct_full_key_struct_for_test(0, &medium_key, 1), b"v1");
1109                } else if index < 80 {
1110                    let mut large_key = vec![b'B'; MAX_KEY_LEN];
1111                    large_key.push(index);
1112                    builder
1113                        .add_for_test(construct_full_key_struct_for_test(0, &large_key, 2), b"v2");
1114                } else {
1115                    let mut xlarge_key = vec![b'C'; MAX_KEY_LEN + 500];
1116                    xlarge_key.push(index);
1117                    builder
1118                        .add_for_test(construct_full_key_struct_for_test(0, &xlarge_key, 3), b"v3");
1119                }
1120            }
1121
1122            let capacity = builder.uncompressed_block_size();
1123            assert_eq!(capacity, builder.approximate_len() - 9);
1124            let buf = builder.build().to_vec();
1125            let block = Box::new(Block::decode(buf.into(), capacity).unwrap());
1126            let mut bi = BlockIterator::new(BlockHolder::from_owned_block(block));
1127            bi.seek_to_first();
1128            assert!(bi.is_valid());
1129
1130            for index in 0..KEY_COUNT {
1131                if index < 50 {
1132                    let mut medium_key = vec![b'A'; MAX_KEY_LEN - 500];
1133                    medium_key.push(index);
1134                    assert_eq!(
1135                        construct_full_key_struct_for_test(0, &medium_key, 1),
1136                        bi.key()
1137                    );
1138                } else if index < 80 {
1139                    let mut large_key = vec![b'B'; MAX_KEY_LEN];
1140                    large_key.push(index);
1141                    assert_eq!(
1142                        construct_full_key_struct_for_test(0, &large_key, 2),
1143                        bi.key()
1144                    );
1145                } else {
1146                    let mut xlarge_key = vec![b'C'; MAX_KEY_LEN + 500];
1147                    xlarge_key.push(index);
1148                    assert_eq!(
1149                        construct_full_key_struct_for_test(0, &xlarge_key, 3),
1150                        bi.key()
1151                    );
1152                }
1153                bi.next();
1154            }
1155
1156            assert!(!bi.is_valid());
1157            builder.clear();
1158        }
1159    }
1160
1161    #[test]
1162    fn test_block_serde() {
1163        fn assmut_serde<'de, T: Serialize + Deserialize<'de>>() {}
1164
1165        assmut_serde::<Block>();
1166        assmut_serde::<Box<Block>>();
1167
1168        let options = BlockBuilderOptions::default();
1169        let mut builder = BlockBuilder::new(options);
1170        for i in 0..100 {
1171            builder.add_for_test(
1172                construct_full_key_struct_for_test(0, format!("k{i:8}").as_bytes(), i),
1173                format!("v{i:8}").as_bytes(),
1174            );
1175        }
1176
1177        let capacity = builder.uncompressed_block_size();
1178        assert_eq!(capacity, builder.approximate_len() - 9);
1179        let buf = builder.build().to_vec();
1180
1181        let block = Box::new(Block::decode(buf.into(), capacity).unwrap());
1182
1183        let buffer = bincode::serialize(&block).unwrap();
1184        let blk: Block = bincode::deserialize(&buffer).unwrap();
1185
1186        assert_eq!(block.data, blk.data);
1187        assert_eq!(block.data_len, blk.data_len);
1188        assert_eq!(block.table_id, blk.table_id,);
1189        assert_eq!(block.restart_points, blk.restart_points);
1190    }
1191
1192    #[test]
1193    fn test_try_shorten_block_smallest_key() {
1194        use risingwave_hummock_sdk::EpochWithGap;
1195
1196        fn make_full_key(table_id: u32, table_key: &[u8], epoch: u64) -> FullKey<&[u8]> {
1197            FullKey::new_with_gap_epoch(
1198                TableId::new(table_id),
1199                TableKey(table_key),
1200                EpochWithGap::new_from_epoch(test_epoch(epoch)),
1201            )
1202        }
1203
1204        /// Verifies the invariant: prev < shortened <= next using `FullKey::cmp`
1205        fn assert_invariant(
1206            prev: &FullKey<&[u8]>,
1207            shortened: &FullKey<Vec<u8>>,
1208            next: &FullKey<&[u8]>,
1209        ) {
1210            assert!(
1211                prev.cmp(&shortened.to_ref()).is_lt(),
1212                "Invariant violated: prev >= shortened. prev={:?}, shortened={:?}",
1213                prev,
1214                shortened
1215            );
1216            assert!(
1217                shortened.to_ref().cmp(next).is_le(),
1218                "Invariant violated: shortened > next. shortened={:?}, next={:?}",
1219                shortened,
1220                next
1221            );
1222        }
1223
1224        // Case 1: Basic shortening - prev="abc", next="abdef" -> cand="abd"
1225        {
1226            let prev = make_full_key(1, b"abc", 100);
1227            let next = make_full_key(1, b"abdef", 100);
1228            let result = try_shorten_block_smallest_key(&prev, &next);
1229            assert!(result.is_some());
1230            let shortened = result.unwrap();
1231            assert_eq!(shortened.user_key.table_key.as_ref(), b"abd");
1232            assert_invariant(&prev, &shortened, &next);
1233        }
1234
1235        // Case 2: Prefix case - prev="ab", next="abcd" -> cand="abc"
1236        {
1237            let prev = make_full_key(1, b"ab", 100);
1238            let next = make_full_key(1, b"abcd", 100);
1239            let result = try_shorten_block_smallest_key(&prev, &next);
1240            assert!(result.is_some());
1241            let shortened = result.unwrap();
1242            assert_eq!(shortened.user_key.table_key.as_ref(), b"abc");
1243            assert_invariant(&prev, &shortened, &next);
1244        }
1245
1246        // Case 3: Cannot shorten - only 1 char difference at end
1247        {
1248            let prev = make_full_key(1, b"abc", 100);
1249            let next = make_full_key(1, b"abd", 100);
1250            assert!(try_shorten_block_smallest_key(&prev, &next).is_none());
1251        }
1252
1253        // Case 4: Cannot shorten - next is only 1 char longer than LCP
1254        {
1255            let prev = make_full_key(1, b"abc", 100);
1256            let next = make_full_key(1, b"abcd", 100);
1257            assert!(try_shorten_block_smallest_key(&prev, &next).is_none());
1258        }
1259
1260        // Case 5: Same table_key, different epoch - cannot shorten
1261        {
1262            let prev = make_full_key(1, b"abc", 200);
1263            let next = make_full_key(1, b"abc", 100);
1264            assert!(try_shorten_block_smallest_key(&prev, &next).is_none());
1265        }
1266
1267        // Case 7: Single char keys - cannot shorten
1268        {
1269            let prev = make_full_key(1, b"a", 100);
1270            let next = make_full_key(1, b"b", 100);
1271            assert!(try_shorten_block_smallest_key(&prev, &next).is_none());
1272        }
1273
1274        // Case 8: Long common prefix with divergence
1275        {
1276            let prev = make_full_key(1, b"hello_world_abc", 100);
1277            let next = make_full_key(1, b"hello_world_xyz123", 100);
1278            let result = try_shorten_block_smallest_key(&prev, &next);
1279            assert!(result.is_some());
1280            let shortened = result.unwrap();
1281            assert_eq!(shortened.user_key.table_key.as_ref(), b"hello_world_x");
1282            assert_invariant(&prev, &shortened, &next);
1283        }
1284
1285        // Case 9: Empty common prefix (completely different keys)
1286        {
1287            let prev = make_full_key(1, b"aaa", 100);
1288            let next = make_full_key(1, b"bbbccc", 100);
1289            let result = try_shorten_block_smallest_key(&prev, &next);
1290            assert!(result.is_some());
1291            let shortened = result.unwrap();
1292            assert_eq!(shortened.user_key.table_key.as_ref(), b"b");
1293            assert_invariant(&prev, &shortened, &next);
1294        }
1295
1296        // Case 10: Minimal shortening - next is exactly 2 chars longer than LCP
1297        {
1298            let prev = make_full_key(1, b"abc", 100);
1299            let next = make_full_key(1, b"abcde", 100);
1300            let result = try_shorten_block_smallest_key(&prev, &next);
1301            assert!(result.is_some());
1302            let shortened = result.unwrap();
1303            assert_eq!(shortened.user_key.table_key.as_ref(), b"abcd");
1304            assert_invariant(&prev, &shortened, &next);
1305        }
1306    }
1307}
risingwave_storage/hummock/sstable/block.rs

risingwave_storage/hummock/sstable/
block.rs
