sui_rpc_store/schema/
transaction_bitmap.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! `(dimension_key, bucket)` → `BitmapBlob`.
5//!
6//! Inverted bitmap index over `tx_seq` space. The dimension key is
7//! a variable-length opaque token (e.g. `[tag][sender]`); each
8//! bucket holds the roaring bitmap of tx_seqs whose containing
9//! transaction matches that dimension.
10//!
11//! Indexer pipelines stage merge operands carrying a small bitmap
12//! (often a single bit per write); the merge operator unions every
13//! operand against the existing on-disk bitmap and emits a single
14//! consolidated value optimized for the on-disk encoding.
15//!
16//! A per-bucket compaction filter reads the database schema's
17//! `tx_seq` pruning floor and drops buckets whose entire `tx_seq`
18//! range sits below it.
19
20use std::sync::Arc;
21use std::sync::atomic::AtomicU64;
22use std::sync::atomic::Ordering;
23
24use bytes::Buf;
25use bytes::BufMut;
26use prost::Message;
27use roaring::RoaringBitmap;
28use sui_consistent_store::Decode;
29use sui_consistent_store::Encode;
30use sui_consistent_store::Iter;
31use sui_consistent_store::Protobuf;
32use sui_consistent_store::error::DecodeError;
33use sui_consistent_store::error::EncodeError;
34use sui_consistent_store::error::Error;
35use sui_consistent_store::reader::Reader;
36
37use crate::proto::BitmapBlob;
38
39pub const NAME: &str = "transaction_bitmap";
40
41/// Number of consecutive `tx_seq` values represented by one
42/// bucket. Sized to keep individual bitmaps small (~8 KiB at
43/// worst-case density) and the per-row read cost predictable.
44pub const TX_BUCKET_SIZE: u64 = 65_536;
45
46const _: () = assert!(TX_BUCKET_SIZE <= u32::MAX as u64);
47
48#[derive(Debug, Clone, PartialEq, Eq, Hash)]
49pub struct Key {
50    pub dimension_key: Vec<u8>,
51    pub bucket: u64,
52}
53
54pub type Value = Protobuf<BitmapBlob>;
55
56impl Encode for Key {
57    fn encode_into<B: BufMut>(&self, buf: &mut B) -> Result<(), EncodeError> {
58        buf.put_slice(&self.dimension_key);
59        buf.put_slice(&self.bucket.to_be_bytes());
60        Ok(())
61    }
62}
63
64impl Decode for Key {
65    fn decode<B: Buf>(buf: &mut B) -> Result<Self, DecodeError> {
66        if buf.remaining() < 8 {
67            return Err(DecodeError::msg(format!(
68                "{NAME} Key too short: {} bytes",
69                buf.remaining(),
70            )));
71        }
72        let dim_len = buf.remaining() - 8;
73        let dim_bytes = buf.copy_to_bytes(dim_len);
74        let bucket = buf.get_u64();
75        Ok(Key {
76            dimension_key: dim_bytes.to_vec(),
77            bucket,
78        })
79    }
80}
81
82/// CF options: install the bitmap-union merge operator and a
83/// per-bucket compaction filter that drops buckets whose entire
84/// `tx_seq` range sits below the pruning floor.
85pub fn options(
86    resolver: &sui_consistent_store::CfOptionsResolver,
87    tx_seq_pruning_floor: Arc<AtomicU64>,
88) -> rocksdb::Options {
89    let mut opts = resolver.options(NAME);
90    opts.set_merge_operator_associative("transaction_bitmap_merge", merge);
91    opts.set_compaction_filter("transaction_bitmap_pruning", move |_level, key, _value| {
92        let pruned_exclusive = tx_seq_pruning_floor.load(Ordering::Relaxed);
93        if should_remove_bucket(key, pruned_exclusive) {
94            rocksdb::CompactionDecision::Remove
95        } else {
96            rocksdb::CompactionDecision::Keep
97        }
98    });
99    opts
100}
101
102/// Pure logic of the compaction filter: decide whether the bucket
103/// identified by `key`'s trailing 8-byte big-endian `bucket_id`
104/// can be removed given the exclusive `tx_seq` pruning floor.
105///
106/// A bucket is removable iff every `tx_seq` it covers is strictly
107/// below the floor — i.e. `(bucket_id + 1) * TX_BUCKET_SIZE <=
108/// pruned_exclusive`. Arithmetic uses `checked_*` so a corrupted
109/// `bucket_id` can't overflow and cause spurious removal.
110///
111/// Kept rather than removed on any malformed input — silent data
112/// loss is worse than a stuck row.
113pub(crate) fn should_remove_bucket(key: &[u8], pruned_exclusive: u64) -> bool {
114    if key.len() < 8 {
115        return false;
116    }
117    let bucket_id_bytes: [u8; 8] = key[key.len() - 8..].try_into().expect("slice length");
118    let bucket_id = u64::from_be_bytes(bucket_id_bytes);
119    bucket_id
120        .checked_add(1)
121        .and_then(|b| b.checked_mul(TX_BUCKET_SIZE))
122        .is_some_and(|highest_plus_one| highest_plus_one <= pruned_exclusive)
123}
124
125/// The bucket that owns a given `tx_seq`.
126pub fn bucket_of(tx_seq: u64) -> u64 {
127    tx_seq / TX_BUCKET_SIZE
128}
129
130/// The bit position within a bucket for a given `tx_seq`. The
131/// cast is safe because `TX_BUCKET_SIZE <= u32::MAX` (enforced
132/// at compile time above).
133pub fn bit_of(tx_seq: u64) -> u32 {
134    (tx_seq % TX_BUCKET_SIZE) as u32
135}
136
137/// Build a `(Key, Value)` pair that adds `tx_seq` to the bitmap
138/// for `(dimension_key, bucket_of(tx_seq))`. The merge operator
139/// unions this single-bit operand with whatever's already on
140/// disk.
141pub fn store_match(dimension_key: Vec<u8>, tx_seq: u64) -> (Key, Value) {
142    let mut bitmap = RoaringBitmap::new();
143    bitmap.insert(bit_of(tx_seq));
144    store_bitmap(dimension_key, bucket_of(tx_seq), bitmap)
145}
146
147/// Build a `(Key, Value)` pair that stages the given bitmap as a
148/// merge operand against the existing on-disk bitmap. Useful for
149/// pipelines that batch many tx_seqs into one bucket per
150/// dimension before writing.
151pub fn store_bitmap(dimension_key: Vec<u8>, bucket: u64, bitmap: RoaringBitmap) -> (Key, Value) {
152    (
153        Key {
154            dimension_key,
155            bucket,
156        },
157        Protobuf(BitmapBlob {
158            data: serialize_bitmap(&bitmap).into(),
159        }),
160    )
161}
162
163impl<R: Reader> super::RpcStoreSchema<R> {
164    /// Look up the bitmap for `(dimension_key, bucket)` and
165    /// return it deserialized.
166    pub fn get_transaction_bitmap(
167        &self,
168        dimension_key: Vec<u8>,
169        bucket: u64,
170    ) -> Result<Option<RoaringBitmap>, Error> {
171        let Some(stored) = self.transaction_bitmap.get(&Key {
172            dimension_key,
173            bucket,
174        })?
175        else {
176            return Ok(None);
177        };
178        let bytes = stored.into_inner().data;
179        let bitmap = RoaringBitmap::deserialize_from(bytes.as_ref())
180            .map_err(|e| DecodeError::with_source("deserialize RoaringBitmap", e))?;
181        Ok(Some(bitmap))
182    }
183
184    /// Iterate every bucket recorded against `dimension_key`, in
185    /// ascending bucket order.
186    pub fn iter_transaction_bitmap_buckets(
187        &self,
188        dimension_key: Vec<u8>,
189    ) -> Result<Iter<'_, Key, Value>, Error> {
190        self.transaction_bitmap
191            .iter_prefix(&DimensionPrefix(dimension_key))
192    }
193}
194
195/// Prefix encoder for "all buckets recorded against
196/// `dimension_key`". Encodes as the raw dimension bytes — the
197/// leading bytes of every `Key` whose `dimension_key` matches.
198pub struct DimensionPrefix(pub Vec<u8>);
199
200impl Encode for DimensionPrefix {
201    fn encode_into<B: BufMut>(&self, buf: &mut B) -> Result<(), EncodeError> {
202        buf.put_slice(&self.0);
203        Ok(())
204    }
205}
206
207/// Serialize a roaring bitmap for on-disk storage. Run-encodes
208/// dense containers first so a bucket that matches many
209/// consecutive `tx_seq` values compresses well.
210fn serialize_bitmap(bitmap: &RoaringBitmap) -> Vec<u8> {
211    let mut buf = Vec::with_capacity(bitmap.serialized_size());
212    bitmap
213        .serialize_into(&mut buf)
214        .expect("RoaringBitmap::serialize_into on Vec cannot fail");
215    buf
216}
217
218/// Associative merge: union every operand bitmap with the
219/// existing on-disk bitmap, then optimize the accumulator before
220/// re-serializing.
221///
222/// Encode / decode failures panic — this CF is written only by
223/// the crate's `store_*` helpers, so a parse failure indicates
224/// corruption rather than a recoverable situation.
225fn merge(
226    _key: &[u8],
227    existing_val: Option<&[u8]>,
228    operands: &rocksdb::MergeOperands,
229) -> Option<Vec<u8>> {
230    let mut acc = match existing_val {
231        Some(bytes) => decode_bitmap(bytes),
232        None => RoaringBitmap::new(),
233    };
234
235    for operand in operands {
236        let bitmap = decode_bitmap(operand);
237        acc |= bitmap;
238    }
239
240    // Convert dense containers to runs before persisting. The
241    // operands are typically tiny (one bit per call) so there's
242    // nothing for run-encoding to collapse on them; the
243    // accumulator is what RocksDB writes back to disk.
244    acc.optimize();
245    Some(encode_bitmap_blob(&acc))
246}
247
248fn decode_bitmap(bytes: &[u8]) -> RoaringBitmap {
249    let blob = BitmapBlob::decode(bytes).expect("decode BitmapBlob");
250    RoaringBitmap::deserialize_from(blob.data.as_ref()).expect("deserialize RoaringBitmap")
251}
252
253fn encode_bitmap_blob(bitmap: &RoaringBitmap) -> Vec<u8> {
254    let blob = BitmapBlob {
255        data: serialize_bitmap(bitmap).into(),
256    };
257    blob.encode_to_vec()
258}
259
260#[cfg(test)]
261mod tests {
262    use std::collections::BTreeSet;
263
264    use sui_consistent_store::Db;
265    use sui_consistent_store::DbOptions;
266
267    use super::*;
268    use crate::RpcStoreSchema;
269
270    fn fresh_db() -> (tempfile::TempDir, sui_consistent_store::Db, RpcStoreSchema) {
271        let dir = tempfile::tempdir().unwrap();
272        let (db, schema) = Db::open::<RpcStoreSchema>(dir.path(), DbOptions::default()).unwrap();
273        (dir, db, schema)
274    }
275
276    #[test]
277    fn bucket_and_bit_math() {
278        assert_eq!(bucket_of(0), 0);
279        assert_eq!(bit_of(0), 0);
280        assert_eq!(bucket_of(TX_BUCKET_SIZE - 1), 0);
281        assert_eq!(bit_of(TX_BUCKET_SIZE - 1), (TX_BUCKET_SIZE - 1) as u32);
282        assert_eq!(bucket_of(TX_BUCKET_SIZE), 1);
283        assert_eq!(bit_of(TX_BUCKET_SIZE), 0);
284        assert_eq!(bucket_of(3 * TX_BUCKET_SIZE + 7), 3);
285        assert_eq!(bit_of(3 * TX_BUCKET_SIZE + 7), 7);
286    }
287
288    #[test]
289    fn get_returns_none_for_unknown_bucket() {
290        let (_dir, _db, schema) = fresh_db();
291        assert!(
292            schema
293                .get_transaction_bitmap(b"sender:alice".to_vec(), 0)
294                .unwrap()
295                .is_none()
296        );
297    }
298
299    #[test]
300    fn single_match_round_trips_through_merge() {
301        let (_dir, db, schema) = fresh_db();
302        let (k, v) = store_match(b"sender:alice".to_vec(), 42);
303
304        let mut batch = db.batch();
305        batch.merge(&schema.transaction_bitmap, &k, &v).unwrap();
306        batch.commit().unwrap();
307
308        let bitmap = schema
309            .get_transaction_bitmap(b"sender:alice".to_vec(), bucket_of(42))
310            .unwrap()
311            .expect("bitmap present");
312        let bits: Vec<u32> = bitmap.iter().collect();
313        assert_eq!(bits, vec![42]);
314    }
315
316    #[test]
317    fn many_matches_in_one_bucket_union() {
318        let (_dir, db, schema) = fresh_db();
319        let dim = b"sender:alice".to_vec();
320
321        let mut batch = db.batch();
322        for tx_seq in [1u64, 17, 256, 9_999] {
323            let (k, v) = store_match(dim.clone(), tx_seq);
324            batch.merge(&schema.transaction_bitmap, &k, &v).unwrap();
325        }
326        batch.commit().unwrap();
327
328        let bitmap = schema
329            .get_transaction_bitmap(dim, 0)
330            .unwrap()
331            .expect("bitmap present");
332        let bits: BTreeSet<u32> = bitmap.iter().collect();
333        assert_eq!(bits, BTreeSet::from([1, 17, 256, 9_999]));
334    }
335
336    #[test]
337    fn distinct_dimensions_do_not_alias() {
338        let (_dir, db, schema) = fresh_db();
339        let (k_a, v_a) = store_match(b"sender:alice".to_vec(), 42);
340        let (k_b, v_b) = store_match(b"sender:bob".to_vec(), 100);
341        let mut batch = db.batch();
342        batch.merge(&schema.transaction_bitmap, &k_a, &v_a).unwrap();
343        batch.merge(&schema.transaction_bitmap, &k_b, &v_b).unwrap();
344        batch.commit().unwrap();
345
346        let alice = schema
347            .get_transaction_bitmap(b"sender:alice".to_vec(), 0)
348            .unwrap()
349            .unwrap();
350        let bob = schema
351            .get_transaction_bitmap(b"sender:bob".to_vec(), 0)
352            .unwrap()
353            .unwrap();
354        assert_eq!(alice.iter().collect::<Vec<u32>>(), vec![42]);
355        assert_eq!(bob.iter().collect::<Vec<u32>>(), vec![100]);
356    }
357
358    #[test]
359    fn should_remove_bucket_drops_only_fully_pruned_ranges() {
360        let dim = b"sender:alice";
361
362        // A bucket whose highest tx_seq is exactly at the floor:
363        // the floor is *exclusive*, so this bucket is still
364        // partially live and must not be removed.
365        let just_at_floor_key = Key {
366            dimension_key: dim.to_vec(),
367            bucket: 0,
368        }
369        .encode()
370        .unwrap();
371        assert!(!should_remove_bucket(
372            &just_at_floor_key,
373            TX_BUCKET_SIZE - 1
374        ));
375
376        // Move the floor one past the bucket's highest tx_seq:
377        // every entry it could hold is pruned, removable.
378        assert!(should_remove_bucket(&just_at_floor_key, TX_BUCKET_SIZE));
379
380        // Bucket 3 covers `tx_seq` in `[3 * BUCKET, 4 * BUCKET)`.
381        // Floor sitting in the middle of the bucket keeps it.
382        let middle_key = Key {
383            dimension_key: dim.to_vec(),
384            bucket: 3,
385        }
386        .encode()
387        .unwrap();
388        assert!(!should_remove_bucket(
389            &middle_key,
390            3 * TX_BUCKET_SIZE + (TX_BUCKET_SIZE / 2),
391        ));
392
393        // Floor past the bucket's high end → removable.
394        assert!(should_remove_bucket(&middle_key, 4 * TX_BUCKET_SIZE));
395
396        // Key shorter than 8 bytes → kept.
397        assert!(!should_remove_bucket(&[0u8; 4], u64::MAX));
398
399        // Floor of 0 → nothing removable.
400        assert!(!should_remove_bucket(&middle_key, 0));
401    }
402
403    #[test]
404    fn iter_walks_buckets_for_one_dimension_in_order() {
405        let (_dir, db, schema) = fresh_db();
406        let dim = b"sender:alice".to_vec();
407        let other = b"sender:bob".to_vec();
408
409        let mut batch = db.batch();
410        for tx_seq in [1u64, TX_BUCKET_SIZE + 5, 3 * TX_BUCKET_SIZE + 9] {
411            let (k, v) = store_match(dim.clone(), tx_seq);
412            batch.merge(&schema.transaction_bitmap, &k, &v).unwrap();
413        }
414        // Unrelated dimension — must not appear in our iter.
415        let (k_other, v_other) = store_match(other, 7);
416        batch
417            .merge(&schema.transaction_bitmap, &k_other, &v_other)
418            .unwrap();
419        batch.commit().unwrap();
420
421        let buckets: Vec<u64> = schema
422            .iter_transaction_bitmap_buckets(dim)
423            .unwrap()
424            .map(|res| res.unwrap().0.bucket)
425            .collect();
426        assert_eq!(buckets, vec![0, 1, 3]);
427    }
428}
sui_rpc_store/schema/transaction_bitmap.rs

sui_rpc_store/schema/
transaction_bitmap.rs
