sui_rpc_store/schema/
event_bitmap.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! `(dimension_key, bucket)` → `BitmapBlob`.
5//!
6//! Same wire shape as [`super::transaction_bitmap`]
7//! but indexes packed-event-seq space — each set bit identifies a
8//! single event by `(tx_seq << EVENT_BITS) | event_idx`.
9//!
10//! The merge operator is identical in structure to the
11//! transaction-bitmap one (union + optimize). The per-bucket
12//! compaction filter translates the database schema's `tx_seq`
13//! pruning floor into packed-event-seq space
14//! (`tx_seq << EVENT_BITS`) and drops buckets that fit entirely
15//! below it.
16
17use std::sync::Arc;
18use std::sync::atomic::AtomicU64;
19use std::sync::atomic::Ordering;
20
21use bytes::Buf;
22use bytes::BufMut;
23use prost::Message;
24use roaring::RoaringBitmap;
25use sui_consistent_store::Decode;
26use sui_consistent_store::Encode;
27use sui_consistent_store::Iter;
28use sui_consistent_store::Protobuf;
29use sui_consistent_store::error::DecodeError;
30use sui_consistent_store::error::EncodeError;
31use sui_consistent_store::error::Error;
32use sui_consistent_store::reader::Reader;
33use sui_inverted_index::event_seq::{EVENT_BITS, encode_event_seq};
34
35use crate::proto::BitmapBlob;
36
37pub const NAME: &str = "event_bitmap";
38
39/// Number of consecutive `packed_event_seq` values represented by
40/// one bucket. Sized so each bucket covers
41/// `EVENT_BUCKET_SIZE >> EVENT_BITS = 4096` consecutive
42/// transactions worth of events.
43pub const EVENT_BUCKET_SIZE: u64 = 1 << 28;
44
45const _: () = assert!(EVENT_BUCKET_SIZE <= u32::MAX as u64);
46
47#[derive(Debug, Clone, PartialEq, Eq, Hash)]
48pub struct Key {
49    pub dimension_key: Vec<u8>,
50    pub bucket: u64,
51}
52
53pub type Value = Protobuf<BitmapBlob>;
54
55impl Encode for Key {
56    fn encode_into<B: BufMut>(&self, buf: &mut B) -> Result<(), EncodeError> {
57        buf.put_slice(&self.dimension_key);
58        buf.put_slice(&self.bucket.to_be_bytes());
59        Ok(())
60    }
61}
62
63impl Decode for Key {
64    fn decode<B: Buf>(buf: &mut B) -> Result<Self, DecodeError> {
65        if buf.remaining() < 8 {
66            return Err(DecodeError::msg(format!(
67                "{NAME} Key too short: {} bytes",
68                buf.remaining(),
69            )));
70        }
71        let dim_len = buf.remaining() - 8;
72        let dim_bytes = buf.copy_to_bytes(dim_len);
73        let bucket = buf.get_u64();
74        Ok(Key {
75            dimension_key: dim_bytes.to_vec(),
76            bucket,
77        })
78    }
79}
80
81/// CF options: install the bitmap-union merge operator and a
82/// per-bucket compaction filter that drops buckets whose entire
83/// packed-event-seq range sits below the pruning floor.
84pub fn options(
85    resolver: &sui_consistent_store::CfOptionsResolver,
86    tx_seq_pruning_floor: Arc<AtomicU64>,
87) -> rocksdb::Options {
88    let mut opts = resolver.options(NAME);
89    opts.set_merge_operator_associative("event_bitmap_merge", merge);
90    opts.set_compaction_filter("event_bitmap_pruning", move |_level, key, _value| {
91        let tx_seq_pruned = tx_seq_pruning_floor.load(Ordering::Relaxed);
92        if should_remove_bucket(key, tx_seq_pruned) {
93            rocksdb::CompactionDecision::Remove
94        } else {
95            rocksdb::CompactionDecision::Keep
96        }
97    });
98    opts
99}
100
101/// Pure logic of the compaction filter.
102///
103/// Translates the `tx_seq` floor into packed-event-seq space and
104/// asks whether every packed event the bucket could hold is
105/// strictly below the translated floor. Kept on any malformed
106/// input — silent data loss is worse than a stuck row.
107pub(crate) fn should_remove_bucket(key: &[u8], tx_seq_pruned_exclusive: u64) -> bool {
108    if key.len() < 8 {
109        return false;
110    }
111    let bucket_id_bytes: [u8; 8] = key[key.len() - 8..].try_into().expect("slice length");
112    let bucket_id = u64::from_be_bytes(bucket_id_bytes);
113    let packed_floor = packed_pruning_floor(tx_seq_pruned_exclusive);
114    bucket_id
115        .checked_add(1)
116        .and_then(|b| b.checked_mul(EVENT_BUCKET_SIZE))
117        .is_some_and(|highest_plus_one| highest_plus_one <= packed_floor)
118}
119
120/// Translate the `tx_seq` floor into packed-event-seq space.
121///
122/// `packed_event_seq = tx_seq << EVENT_BITS`. For
123/// `tx_seq >= 2^(64 - EVENT_BITS)` the shift would overflow a
124/// `u64`; we saturate to `u64::MAX`, which represents "every
125/// possible event has been pruned" — the conservative direction
126/// for a removal decision.
127fn packed_pruning_floor(tx_seq_pruned_exclusive: u64) -> u64 {
128    const OVERFLOW_THRESHOLD: u64 = 1u64 << (64 - EVENT_BITS);
129    if tx_seq_pruned_exclusive < OVERFLOW_THRESHOLD {
130        tx_seq_pruned_exclusive << EVENT_BITS
131    } else {
132        u64::MAX
133    }
134}
135
136/// The bucket that owns a given packed event sequence.
137pub fn bucket_of(packed: u64) -> u64 {
138    packed / EVENT_BUCKET_SIZE
139}
140
141/// The bit position within a bucket for a given packed event
142/// sequence. The cast is safe because `EVENT_BUCKET_SIZE`
143/// fits in a `u32` (enforced at compile time above).
144pub fn bit_of(packed: u64) -> u32 {
145    (packed % EVENT_BUCKET_SIZE) as u32
146}
147
148/// Build a `(Key, Value)` pair that adds the event identified by
149/// `(tx_seq, event_idx)` to the bitmap for its dimension and
150/// bucket. The merge operator unions this single-bit operand
151/// with whatever's already on disk.
152pub fn store_match(dimension_key: Vec<u8>, tx_seq: u64, event_idx: u32) -> (Key, Value) {
153    let packed = encode_event_seq(tx_seq, event_idx);
154    let mut bitmap = RoaringBitmap::new();
155    bitmap.insert(bit_of(packed));
156    store_bitmap(dimension_key, bucket_of(packed), bitmap)
157}
158
159/// Build a `(Key, Value)` pair that stages the given bitmap as a
160/// merge operand against the existing on-disk bitmap. Useful for
161/// pipelines that batch many events into one bucket per dimension
162/// before writing.
163pub fn store_bitmap(dimension_key: Vec<u8>, bucket: u64, bitmap: RoaringBitmap) -> (Key, Value) {
164    (
165        Key {
166            dimension_key,
167            bucket,
168        },
169        Protobuf(BitmapBlob {
170            data: serialize_bitmap(&bitmap).into(),
171        }),
172    )
173}
174
175impl<R: Reader> super::RpcStoreSchema<R> {
176    /// Look up the event bitmap for `(dimension_key, bucket)` and
177    /// return it deserialized.
178    pub fn get_event_bitmap(
179        &self,
180        dimension_key: Vec<u8>,
181        bucket: u64,
182    ) -> Result<Option<RoaringBitmap>, Error> {
183        let Some(stored) = self.event_bitmap.get(&Key {
184            dimension_key,
185            bucket,
186        })?
187        else {
188            return Ok(None);
189        };
190        let bytes = stored.into_inner().data;
191        let bitmap = RoaringBitmap::deserialize_from(bytes.as_ref())
192            .map_err(|e| DecodeError::with_source("deserialize RoaringBitmap", e))?;
193        Ok(Some(bitmap))
194    }
195
196    /// Iterate every bucket recorded against `dimension_key`, in
197    /// ascending bucket order.
198    pub fn iter_event_bitmap_buckets(
199        &self,
200        dimension_key: Vec<u8>,
201    ) -> Result<Iter<'_, Key, Value>, Error> {
202        self.event_bitmap
203            .iter_prefix(&DimensionPrefix(dimension_key))
204    }
205}
206
207/// Prefix encoder for "all buckets recorded against
208/// `dimension_key`". Encodes as the raw dimension bytes — the
209/// leading bytes of every `Key` whose `dimension_key` matches.
210pub struct DimensionPrefix(pub Vec<u8>);
211
212impl Encode for DimensionPrefix {
213    fn encode_into<B: BufMut>(&self, buf: &mut B) -> Result<(), EncodeError> {
214        buf.put_slice(&self.0);
215        Ok(())
216    }
217}
218
219/// Serialize a roaring bitmap for on-disk storage. Run-encodes
220/// dense containers first so a bucket that matches many
221/// consecutive packed event sequences compresses well.
222fn serialize_bitmap(bitmap: &RoaringBitmap) -> Vec<u8> {
223    let mut buf = Vec::with_capacity(bitmap.serialized_size());
224    bitmap
225        .serialize_into(&mut buf)
226        .expect("RoaringBitmap::serialize_into on Vec cannot fail");
227    buf
228}
229
230/// Associative merge: union every operand bitmap with the
231/// existing on-disk bitmap, then optimize the accumulator before
232/// re-serializing.
233///
234/// Encode / decode failures panic — this CF is written only by
235/// the crate's `store_*` helpers, so a parse failure indicates
236/// corruption rather than a recoverable situation.
237fn merge(
238    _key: &[u8],
239    existing_val: Option<&[u8]>,
240    operands: &rocksdb::MergeOperands,
241) -> Option<Vec<u8>> {
242    let mut acc = match existing_val {
243        Some(bytes) => decode_bitmap(bytes),
244        None => RoaringBitmap::new(),
245    };
246
247    for operand in operands {
248        let bitmap = decode_bitmap(operand);
249        acc |= bitmap;
250    }
251
252    acc.optimize();
253    Some(encode_bitmap_blob(&acc))
254}
255
256fn decode_bitmap(bytes: &[u8]) -> RoaringBitmap {
257    let blob = BitmapBlob::decode(bytes).expect("decode BitmapBlob");
258    RoaringBitmap::deserialize_from(blob.data.as_ref()).expect("deserialize RoaringBitmap")
259}
260
261fn encode_bitmap_blob(bitmap: &RoaringBitmap) -> Vec<u8> {
262    let blob = BitmapBlob {
263        data: serialize_bitmap(bitmap).into(),
264    };
265    blob.encode_to_vec()
266}
267
268#[cfg(test)]
269mod tests {
270    use std::collections::BTreeSet;
271
272    use sui_consistent_store::Db;
273    use sui_consistent_store::DbOptions;
274
275    use super::*;
276    use crate::RpcStoreSchema;
277
278    fn fresh_db() -> (tempfile::TempDir, sui_consistent_store::Db, RpcStoreSchema) {
279        let dir = tempfile::tempdir().unwrap();
280        let (db, schema) = Db::open::<RpcStoreSchema>(dir.path(), DbOptions::default()).unwrap();
281        (dir, db, schema)
282    }
283
284    #[test]
285    fn pack_bucket_and_bit_math() {
286        // tx_seq=0, event_idx=0 → packed 0 → bucket 0 / bit 0.
287        let p = encode_event_seq(0, 0);
288        assert_eq!(p, 0);
289        assert_eq!(bucket_of(p), 0);
290        assert_eq!(bit_of(p), 0);
291
292        // tx_seq=1, event_idx=0 → packed `1 << 16` = 65_536.
293        let p = encode_event_seq(1, 0);
294        assert_eq!(p, 1 << EVENT_BITS);
295        assert_eq!(bucket_of(p), 0);
296        assert_eq!(bit_of(p), 1 << EVENT_BITS);
297
298        // The first packed value of the next bucket sits at the
299        // boundary `EVENT_BUCKET_SIZE` — that's
300        // `EVENT_BUCKET_SIZE >> EVENT_BITS = 4096` transactions in.
301        let first_in_next_bucket = encode_event_seq(EVENT_BUCKET_SIZE >> EVENT_BITS, 0);
302        assert_eq!(first_in_next_bucket, EVENT_BUCKET_SIZE);
303        assert_eq!(bucket_of(first_in_next_bucket), 1);
304        assert_eq!(bit_of(first_in_next_bucket), 0);
305    }
306
307    #[test]
308    fn get_returns_none_for_unknown_bucket() {
309        let (_dir, _db, schema) = fresh_db();
310        assert!(
311            schema
312                .get_event_bitmap(b"emitting_module:coin".to_vec(), 0)
313                .unwrap()
314                .is_none()
315        );
316    }
317
318    #[test]
319    fn single_match_round_trips_through_merge() {
320        let (_dir, db, schema) = fresh_db();
321        let (k, v) = store_match(b"emitting_module:coin".to_vec(), 42, 3);
322
323        let mut batch = db.batch();
324        batch.merge(&schema.event_bitmap, &k, &v).unwrap();
325        batch.commit().unwrap();
326
327        let packed = encode_event_seq(42, 3);
328        let bitmap = schema
329            .get_event_bitmap(b"emitting_module:coin".to_vec(), bucket_of(packed))
330            .unwrap()
331            .expect("bitmap present");
332        let bits: Vec<u32> = bitmap.iter().collect();
333        assert_eq!(bits, vec![bit_of(packed)]);
334    }
335
336    #[test]
337    fn many_matches_in_one_bucket_union() {
338        let (_dir, db, schema) = fresh_db();
339        let dim = b"emitting_module:coin".to_vec();
340        let entries: Vec<(u64, u32)> = vec![(1, 0), (1, 7), (2, 0), (5, 12)];
341
342        let mut batch = db.batch();
343        for (tx, idx) in &entries {
344            let (k, v) = store_match(dim.clone(), *tx, *idx);
345            batch.merge(&schema.event_bitmap, &k, &v).unwrap();
346        }
347        batch.commit().unwrap();
348
349        let bitmap = schema
350            .get_event_bitmap(dim, 0)
351            .unwrap()
352            .expect("bitmap present");
353        let bits: BTreeSet<u32> = bitmap.iter().collect();
354        let expected: BTreeSet<u32> = entries
355            .iter()
356            .map(|(tx, idx)| bit_of(encode_event_seq(*tx, *idx)))
357            .collect();
358        assert_eq!(bits, expected);
359    }
360
361    #[test]
362    fn distinct_dimensions_do_not_alias() {
363        let (_dir, db, schema) = fresh_db();
364        let (k_a, v_a) = store_match(b"emitting_module:coin".to_vec(), 42, 1);
365        let (k_b, v_b) = store_match(b"emitting_module:nft".to_vec(), 100, 2);
366        let mut batch = db.batch();
367        batch.merge(&schema.event_bitmap, &k_a, &v_a).unwrap();
368        batch.merge(&schema.event_bitmap, &k_b, &v_b).unwrap();
369        batch.commit().unwrap();
370
371        let coin = schema
372            .get_event_bitmap(b"emitting_module:coin".to_vec(), 0)
373            .unwrap()
374            .unwrap();
375        let nft = schema
376            .get_event_bitmap(b"emitting_module:nft".to_vec(), 0)
377            .unwrap()
378            .unwrap();
379        assert_eq!(
380            coin.iter().collect::<Vec<u32>>(),
381            vec![bit_of(encode_event_seq(42, 1))]
382        );
383        assert_eq!(
384            nft.iter().collect::<Vec<u32>>(),
385            vec![bit_of(encode_event_seq(100, 2))]
386        );
387    }
388
389    #[test]
390    fn should_remove_bucket_drops_only_fully_pruned_ranges() {
391        let dim = b"emitting_module:coin";
392        let bucket_0_key = Key {
393            dimension_key: dim.to_vec(),
394            bucket: 0,
395        }
396        .encode()
397        .unwrap();
398
399        // Floor 0 → nothing pruned.
400        assert!(!should_remove_bucket(&bucket_0_key, 0));
401
402        // EVENT_BUCKET_SIZE in packed-event-seq space corresponds
403        // to `EVENT_BUCKET_SIZE >> EVENT_BITS` transactions —
404        // anything below that tx_seq floor keeps bucket 0 alive.
405        let txs_per_bucket = EVENT_BUCKET_SIZE >> EVENT_BITS;
406        assert!(!should_remove_bucket(&bucket_0_key, txs_per_bucket - 1));
407        // At the tx_seq floor that translates to exactly
408        // EVENT_BUCKET_SIZE in packed space, bucket 0 becomes
409        // fully pruned.
410        assert!(should_remove_bucket(&bucket_0_key, txs_per_bucket));
411
412        // Bucket 5 needs floor past 6 * EVENT_BUCKET_SIZE in
413        // packed space, i.e. tx_seq past 6 * txs_per_bucket.
414        let bucket_5_key = Key {
415            dimension_key: dim.to_vec(),
416            bucket: 5,
417        }
418        .encode()
419        .unwrap();
420        assert!(!should_remove_bucket(&bucket_5_key, 6 * txs_per_bucket - 1));
421        assert!(should_remove_bucket(&bucket_5_key, 6 * txs_per_bucket));
422
423        // Key too short → kept.
424        assert!(!should_remove_bucket(&[0u8; 4], u64::MAX));
425    }
426
427    #[test]
428    fn packed_pruning_floor_saturates_on_overflow() {
429        assert_eq!(packed_pruning_floor(0), 0);
430        assert_eq!(packed_pruning_floor(1), 1u64 << EVENT_BITS);
431        // Just below the overflow threshold.
432        let just_below = (1u64 << (64 - EVENT_BITS)) - 1;
433        assert_eq!(packed_pruning_floor(just_below), just_below << EVENT_BITS,);
434        // At the threshold — `tx_seq << EVENT_BITS` would
435        // overflow, so we saturate.
436        assert_eq!(packed_pruning_floor(1u64 << (64 - EVENT_BITS)), u64::MAX);
437        assert_eq!(packed_pruning_floor(u64::MAX), u64::MAX);
438    }
439
440    #[test]
441    fn iter_walks_buckets_for_one_dimension_in_order() {
442        let (_dir, db, schema) = fresh_db();
443        let dim = b"emitting_module:coin".to_vec();
444        let other = b"emitting_module:nft".to_vec();
445        // Three events whose packed seqs land in distinct
446        // buckets: bucket 0, bucket 1 (just past 4096 txs), and
447        // bucket 3.
448        let txs_per_bucket = EVENT_BUCKET_SIZE >> EVENT_BITS;
449        let tx_seqs = [0u64, txs_per_bucket + 5, 3 * txs_per_bucket + 9];
450
451        let mut batch = db.batch();
452        for tx in tx_seqs {
453            let (k, v) = store_match(dim.clone(), tx, 0);
454            batch.merge(&schema.event_bitmap, &k, &v).unwrap();
455        }
456        // Unrelated dimension — must not appear in our iter.
457        let (k_other, v_other) = store_match(other, 0, 0);
458        batch
459            .merge(&schema.event_bitmap, &k_other, &v_other)
460            .unwrap();
461        batch.commit().unwrap();
462
463        let buckets: Vec<u64> = schema
464            .iter_event_bitmap_buckets(dim)
465            .unwrap()
466            .map(|res| res.unwrap().0.bucket)
467            .collect();
468        assert_eq!(buckets, vec![0, 1, 3]);
469    }
470}
sui_rpc_store/schema/event_bitmap.rs

sui_rpc_store/schema/
event_bitmap.rs
