typed_store/
database.rs

1// Copyright (c) 2026 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{
5    borrow::Borrow,
6    marker::PhantomData,
7    ops::{Bound, Deref, RangeBounds},
8    path::Path,
9    sync::Arc,
10    time::Duration,
11};
12
13use fastcrypto::hash::{Digest, HashFunction};
14use iota_common::debug_fatal;
15use iota_macros::{fail_point, nondeterministic};
16use prometheus_filtered::{Histogram, HistogramTimer};
17use rocksdb::{DBPinnableSlice, Error, LiveFile, ReadOptions, WriteBatch, checkpoint::Checkpoint};
18use serde::{Serialize, de::DeserializeOwned};
19use tokio::sync::oneshot;
20use tracing::{debug, error, instrument, warn};
21use typed_store_error::TypedStoreError;
22
23use crate::{
24    DbIterator,
25    memstore::{InMemoryBatch, InMemoryDB},
26    metrics::{DBMetrics, RocksDBPerfContext, SamplingInterval},
27    rocks::{
28        RocksDB,
29        errors::{typed_store_err_from_bcs_err, typed_store_err_from_rocks_err},
30        options::ReadWriteOptions,
31        rocks_cf, rocks_util,
32        safe_iter::{SafeIter, SafeRevIter},
33    },
34    traits::{Map, TableSummary},
35    util::{
36        be_fix_int_ser, iterator_bounds_with_range, prefix_iterator_bounds,
37        prefix_iterator_bounds_with_range,
38    },
39};
40
41#[derive(Clone)]
42pub(crate) enum ColumnFamily {
43    Rocks(String),
44    #[allow(dead_code)]
45    InMemory(String),
46}
47
48impl std::fmt::Debug for ColumnFamily {
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        match self {
51            ColumnFamily::Rocks(name) => write!(f, "RocksDB cf: {name}"),
52            ColumnFamily::InMemory(name) => write!(f, "InMemory cf: {name}"),
53        }
54    }
55}
56
57impl ColumnFamily {
58    pub(crate) fn name(&self) -> &str {
59        match self {
60            ColumnFamily::Rocks(name) => name,
61            ColumnFamily::InMemory(name) => name,
62        }
63    }
64}
65
66pub(crate) enum Storage {
67    Rocks(RocksDB),
68    #[allow(dead_code)]
69    InMemory(InMemoryDB),
70}
71
72impl std::fmt::Debug for Storage {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        match self {
75            Storage::Rocks(db) => write!(f, "RocksDB Storage {db:?}"),
76            Storage::InMemory(db) => write!(f, "InMemoryDB Storage {db:?}"),
77        }
78    }
79}
80
81pub(crate) enum GetResult<'a> {
82    Rocks(DBPinnableSlice<'a>),
83    InMemory(Vec<u8>),
84}
85
86impl Deref for GetResult<'_> {
87    type Target = [u8];
88    fn deref(&self) -> &[u8] {
89        match self {
90            GetResult::Rocks(d) => d.deref(),
91            GetResult::InMemory(d) => d.deref(),
92        }
93    }
94}
95
96pub enum StorageWriteBatch {
97    Rocks(rocksdb::WriteBatch),
98    InMemory(InMemoryBatch),
99}
100
101#[derive(Debug, Default)]
102pub struct MetricConf {
103    pub db_name: String,
104    pub read_sample_interval: SamplingInterval,
105    pub write_sample_interval: SamplingInterval,
106    pub iter_sample_interval: SamplingInterval,
107}
108
109impl MetricConf {
110    pub fn new(db_name: &str) -> Self {
111        if db_name.is_empty() {
112            error!("A meaningful db name should be used for metrics reporting.")
113        }
114        Self {
115            db_name: db_name.to_string(),
116            read_sample_interval: SamplingInterval::default(),
117            write_sample_interval: SamplingInterval::default(),
118            iter_sample_interval: SamplingInterval::default(),
119        }
120    }
121
122    pub fn with_sampling(self, read_interval: SamplingInterval) -> Self {
123        Self {
124            db_name: self.db_name,
125            read_sample_interval: read_interval,
126            write_sample_interval: SamplingInterval::default(),
127            iter_sample_interval: SamplingInterval::default(),
128        }
129    }
130}
131
132const CF_METRICS_REPORT_PERIOD_SECS: u64 = 30;
133
134#[derive(Debug)]
135pub struct Database {
136    pub(crate) storage: Storage,
137    pub(crate) metric_conf: MetricConf,
138}
139
140impl Drop for Database {
141    fn drop(&mut self) {
142        DBMetrics::get().decrement_num_active_dbs(&self.metric_conf.db_name);
143    }
144}
145
146impl Database {
147    pub(crate) fn new(storage: Storage, metric_conf: MetricConf) -> Self {
148        DBMetrics::get().increment_num_active_dbs(&metric_conf.db_name);
149        Self {
150            storage,
151            metric_conf,
152        }
153    }
154
155    pub(crate) fn get<K: AsRef<[u8]>>(
156        &self,
157        cf: &ColumnFamily,
158        key: K,
159        readopts: &ReadOptions,
160    ) -> Result<Option<GetResult<'_>>, TypedStoreError> {
161        match (&self.storage, cf) {
162            (Storage::Rocks(db), ColumnFamily::Rocks(_)) => Ok(db
163                .underlying
164                .get_pinned_cf_opt(&rocks_cf(db, cf.name()), key, readopts)
165                .map_err(typed_store_err_from_rocks_err)?
166                .map(GetResult::Rocks)),
167            (Storage::InMemory(db), ColumnFamily::InMemory(cf_name)) => {
168                Ok(db.get(cf_name, key).map(GetResult::InMemory))
169            }
170
171            _ => Err(TypedStoreError::RocksDB(
172                "typed store invariant violation".to_string(),
173            )),
174        }
175    }
176
177    pub(crate) fn multi_get<I, K>(
178        &self,
179        cf: &ColumnFamily,
180        keys: I,
181        readopts: &ReadOptions,
182    ) -> Vec<Result<Option<GetResult<'_>>, TypedStoreError>>
183    where
184        I: IntoIterator<Item = K>,
185        K: AsRef<[u8]>,
186    {
187        match (&self.storage, cf) {
188            (Storage::Rocks(db), ColumnFamily::Rocks(_)) => {
189                let keys_vec: Vec<K> = keys.into_iter().collect();
190                let res = db.underlying.batched_multi_get_cf_opt(
191                    &rocks_cf(db, cf.name()),
192                    keys_vec.iter(),
193                    // sorted_input
194                    false,
195                    readopts,
196                );
197                res.into_iter()
198                    .map(|r| {
199                        r.map_err(typed_store_err_from_rocks_err)
200                            .map(|item| item.map(GetResult::Rocks))
201                    })
202                    .collect()
203            }
204            (Storage::InMemory(db), ColumnFamily::InMemory(cf_name)) => db
205                .multi_get(cf_name, keys)
206                .into_iter()
207                .map(|r| Ok(r.map(GetResult::InMemory)))
208                .collect(),
209            _ => unreachable!("typed store invariant violation"),
210        }
211    }
212
213    pub fn cf_handle(&self, name: &str) -> Option<()> {
214        match &self.storage {
215            Storage::Rocks(db) => db.underlying.cf_handle(name).map(|_| ()),
216            Storage::InMemory(db) => db.has_cf(name).then_some(()),
217        }
218    }
219
220    /// Creates a new column family at runtime. Fails if a column family with
221    /// this name already exists.
222    ///
223    /// `options` hold until the database is next opened, after which a
224    /// column family the caller does not declare is opened with the
225    /// crate's options instead. Declare it to keep them.
226    #[instrument(level = "debug", skip(self, options), err)]
227    pub fn create_cf(&self, name: &str, options: &rocksdb::Options) -> Result<(), TypedStoreError> {
228        match &self.storage {
229            Storage::Rocks(db) => nondeterministic!(db.underlying.create_cf(name, options))
230                .map_err(typed_store_err_from_rocks_err),
231            Storage::InMemory(db) => {
232                db.create_cf(name);
233                Ok(())
234            }
235        }
236    }
237
238    /// Maps on the column family have to be dropped first: a map resolves
239    /// it on every operation and panics once it is gone.
240    #[instrument(level = "debug", skip(self), err)]
241    pub fn drop_cf(&self, name: &str) -> Result<(), TypedStoreError> {
242        if name == rocksdb::DEFAULT_COLUMN_FAMILY_NAME {
243            // rocksdb refuses this, but not before its wrapper has taken the
244            // handle out of its own map, losing it for the rest of the process.
245            return Err(TypedStoreError::RocksDB(format!(
246                "the {name} column family cannot be dropped"
247            )));
248        }
249        match &self.storage {
250            Storage::Rocks(db) => nondeterministic!(db.underlying.drop_cf(name))
251                .map_err(typed_store_err_from_rocks_err),
252            Storage::InMemory(db) => {
253                db.drop_cf(name);
254                Ok(())
255            }
256        }
257    }
258
259    pub(crate) fn delete_cf<K: AsRef<[u8]>>(
260        &self,
261        cf: &ColumnFamily,
262        key: K,
263    ) -> Result<(), TypedStoreError> {
264        fail_point!("delete-cf-before");
265        let ret = match (&self.storage, cf) {
266            (Storage::Rocks(db), ColumnFamily::Rocks(_)) => db
267                .underlying
268                .delete_cf(&rocks_cf(db, cf.name()), key)
269                .map_err(typed_store_err_from_rocks_err),
270            (Storage::InMemory(db), ColumnFamily::InMemory(cf_name)) => {
271                db.delete(cf_name, key.as_ref());
272                Ok(())
273            }
274            _ => Err(TypedStoreError::RocksDB(
275                "typed store invariant violation".to_string(),
276            )),
277        };
278        fail_point!("delete-cf-after");
279        #[allow(clippy::let_and_return)]
280        ret
281    }
282
283    pub fn path_for_pruning(&self) -> &Path {
284        match &self.storage {
285            Storage::Rocks(rocks) => rocks.underlying.path(),
286            _ => unimplemented!("method is only supported for rocksdb backend"),
287        }
288    }
289
290    pub(crate) fn put_cf(
291        &self,
292        cf: &ColumnFamily,
293        key: Vec<u8>,
294        value: Vec<u8>,
295    ) -> Result<(), TypedStoreError> {
296        fail_point!("put-cf-before");
297        let ret = match (&self.storage, cf) {
298            (Storage::Rocks(db), ColumnFamily::Rocks(_)) => db
299                .underlying
300                .put_cf(&rocks_cf(db, cf.name()), key, value)
301                .map_err(typed_store_err_from_rocks_err),
302            (Storage::InMemory(db), ColumnFamily::InMemory(cf_name)) => {
303                db.put(cf_name, key, value);
304                Ok(())
305            }
306            _ => Err(TypedStoreError::RocksDB(
307                "typed store invariant violation".to_string(),
308            )),
309        };
310        fail_point!("put-cf-after");
311        #[allow(clippy::let_and_return)]
312        ret
313    }
314
315    pub(crate) fn key_may_exist_cf<K: AsRef<[u8]>>(
316        &self,
317        cf: &ColumnFamily,
318        key: K,
319        readopts: &ReadOptions,
320    ) -> bool {
321        match &self.storage {
322            // [`rocksdb::DBWithThreadMode::key_may_exist_cf`] can have false positives,
323            // but no false negatives. We use it to short-circuit the absent case
324            Storage::Rocks(rocks) => {
325                rocks
326                    .underlying
327                    .key_may_exist_cf_opt(&rocks_cf(rocks, cf.name()), key, readopts)
328            }
329            _ => true,
330        }
331    }
332
333    /// Flush the memtable of a single column family to SST files on disk.
334    pub(crate) fn flush_cf(&self, cf: &ColumnFamily) -> Result<(), TypedStoreError> {
335        match &self.storage {
336            // A flush blocks on RocksDB background threads; under the simulator it
337            // must run off the test thread, like the opens in `rocks/mod.rs`, or
338            // real flush durations leak into the simulated schedule. No-op outside
339            // msim.
340            Storage::Rocks(rocks) => nondeterministic!({
341                let cf_name = cf.name();
342                if let Some(handle) = rocks.underlying.cf_handle(cf_name) {
343                    rocks.underlying.flush_cf(&handle).map_err(|e| {
344                        TypedStoreError::RocksDB(format!(
345                            "Failed to flush column family {cf_name}: {e}"
346                        ))
347                    })?;
348                }
349                Ok(())
350            }),
351            // InMemory databases don't need flushing.
352            Storage::InMemory(_) => Ok(()),
353        }
354    }
355
356    /// Flush the memtables of every column family of the database to SST files
357    /// on disk.
358    pub fn flush_all(&self) -> Result<(), TypedStoreError> {
359        match &self.storage {
360            // See `flush_cf` for why the flushes run off the test thread under
361            // the simulator.
362            Storage::Rocks(rocks) => nondeterministic!({
363                for cf_name in rocks.cf_names().map_err(|e| {
364                    TypedStoreError::RocksDB(format!(
365                        "Failed to list column families of {}: {e}",
366                        rocks.underlying.path().display()
367                    ))
368                })? {
369                    if let Some(cf) = rocks.underlying.cf_handle(&cf_name) {
370                        rocks.underlying.flush_cf(&cf).map_err(|e| {
371                            TypedStoreError::RocksDB(format!(
372                                "Failed to flush column family {cf_name}: {e}"
373                            ))
374                        })?;
375                    }
376                }
377                Ok(())
378            }),
379            // InMemory databases don't need flushing.
380            Storage::InMemory(_) => Ok(()),
381        }
382    }
383
384    pub fn write(&self, batch: StorageWriteBatch) -> Result<(), TypedStoreError> {
385        self.write_opt(batch, &rocksdb::WriteOptions::default())
386    }
387
388    pub fn write_opt(
389        &self,
390        batch: StorageWriteBatch,
391        write_options: &rocksdb::WriteOptions,
392    ) -> Result<(), TypedStoreError> {
393        fail_point!("batch-write-before");
394        let ret = match (&self.storage, batch) {
395            (Storage::Rocks(rocks), StorageWriteBatch::Rocks(batch)) => rocks
396                .underlying
397                .write_opt(batch, write_options)
398                .map_err(typed_store_err_from_rocks_err),
399            (Storage::InMemory(db), StorageWriteBatch::InMemory(batch)) => {
400                // InMemory doesn't support write options.
401                db.write(batch);
402                Ok(())
403            }
404            _ => Err(TypedStoreError::RocksDB(
405                "using invalid batch type for the database".to_string(),
406            )),
407        };
408        fail_point!("batch-write-after");
409
410        #[allow(clippy::let_and_return)]
411        ret
412    }
413
414    pub(crate) fn compact_range_cf<K: AsRef<[u8]>>(
415        &self,
416        cf: &ColumnFamily,
417        start: Option<K>,
418        end: Option<K>,
419    ) {
420        if let Storage::Rocks(rocksdb) = &self.storage {
421            rocksdb
422                .underlying
423                .compact_range_cf(&rocks_cf(rocksdb, cf.name()), start, end);
424        }
425    }
426
427    pub fn checkpoint(&self, path: &Path) -> Result<(), TypedStoreError> {
428        // TODO: implement for other storage types
429        if let Storage::Rocks(rocks) = &self.storage {
430            let checkpoint =
431                Checkpoint::new(&rocks.underlying).map_err(typed_store_err_from_rocks_err)?;
432            checkpoint
433                .create_checkpoint(path)
434                .map_err(|e| TypedStoreError::RocksDB(e.to_string()))?;
435        }
436        Ok(())
437    }
438
439    pub fn get_sampling_interval(&self) -> SamplingInterval {
440        self.metric_conf.read_sample_interval.new_from_self()
441    }
442
443    pub fn multiget_sampling_interval(&self) -> SamplingInterval {
444        self.metric_conf.read_sample_interval.new_from_self()
445    }
446
447    pub fn write_sampling_interval(&self) -> SamplingInterval {
448        self.metric_conf.write_sample_interval.new_from_self()
449    }
450
451    pub fn iter_sampling_interval(&self) -> SamplingInterval {
452        self.metric_conf.iter_sample_interval.new_from_self()
453    }
454
455    pub(crate) fn db_name(&self) -> String {
456        let name = &self.metric_conf.db_name;
457        if name.is_empty() {
458            "default".to_string()
459        } else {
460            name.clone()
461        }
462    }
463
464    pub fn live_files(&self) -> Result<Vec<LiveFile>, Error> {
465        match &self.storage {
466            Storage::Rocks(rocks) => rocks.underlying.live_files(),
467            _ => Ok(vec![]),
468        }
469    }
470
471    pub(crate) fn try_catch_up_with_primary(&self) -> Result<(), TypedStoreError> {
472        if let Storage::Rocks(rocks) = &self.storage {
473            rocks
474                .underlying
475                .try_catch_up_with_primary()
476                .map_err(typed_store_err_from_rocks_err)?;
477        }
478        Ok(())
479    }
480}
481
482fn rocks_cf_from_db<'a>(
483    db: &'a Database,
484    cf_name: &str,
485) -> Result<Arc<rocksdb::BoundColumnFamily<'a>>, TypedStoreError> {
486    match &db.storage {
487        Storage::Rocks(rocksdb) => Ok(rocksdb
488            .underlying
489            .cf_handle(cf_name)
490            .expect("the column family was deleted unexpectedly")),
491        _ => Err(TypedStoreError::RocksDB(
492            "using invalid batch type for the database".to_string(),
493        )),
494    }
495}
496
497/// An interface to a rocksDB database, keyed by a columnfamily
498#[derive(Clone, Debug)]
499pub struct DBMap<K, V> {
500    pub db: Arc<Database>,
501    _phantom: PhantomData<fn(K) -> V>,
502    column_family: ColumnFamily,
503    pub opts: ReadWriteOptions,
504    db_metrics: Arc<DBMetrics>,
505    get_sample_interval: SamplingInterval,
506    multiget_sample_interval: SamplingInterval,
507    write_sample_interval: SamplingInterval,
508    iter_sample_interval: SamplingInterval,
509    _metrics_task_cancel_handle: Arc<oneshot::Sender<()>>,
510}
511
512impl<K, V> DBMap<K, V> {
513    pub(crate) fn new(
514        db: Arc<Database>,
515        opts: &ReadWriteOptions,
516        column_family: ColumnFamily,
517        skip_metrics_reporting: bool,
518    ) -> Self {
519        let db_cloned = Arc::downgrade(&db);
520        let db_metrics = DBMetrics::get();
521        let db_metrics_cloned = db_metrics.clone();
522        let cf = column_family.name().to_string();
523
524        let (sender, mut recv) = tokio::sync::oneshot::channel();
525        if !skip_metrics_reporting && matches!(db.storage, Storage::Rocks(_)) {
526            tokio::task::spawn(async move {
527                let mut interval =
528                    tokio::time::interval(Duration::from_secs(CF_METRICS_REPORT_PERIOD_SECS));
529                loop {
530                    tokio::select! {
531                        _ = interval.tick() => {
532                            if let Some(db) = db_cloned.upgrade() {
533                                let cf = cf.clone();
534                                let db_metrics = db_metrics.clone();
535                                if let Err(e) = tokio::task::spawn_blocking(move || {
536                                    Self::report_rocksdb_metrics(&db, &cf, &db_metrics);
537                                }).await {
538                                    error!("Failed to log metrics with error: {}", e);
539                                }
540                            } else {
541                                break;
542                            }
543                        }
544                        _ = &mut recv => break,
545                    }
546                }
547                debug!("Returning the cf metric logging task for DBMap: {}", &cf);
548            });
549        }
550        DBMap {
551            db: db.clone(),
552            opts: opts.clone(),
553            _phantom: PhantomData,
554            column_family,
555            db_metrics: db_metrics_cloned,
556            _metrics_task_cancel_handle: Arc::new(sender),
557            get_sample_interval: db.get_sampling_interval(),
558            multiget_sample_interval: db.multiget_sampling_interval(),
559            write_sample_interval: db.write_sampling_interval(),
560            iter_sample_interval: db.iter_sampling_interval(),
561        }
562    }
563
564    /// Reopens an open database as a typed map operating under a specific
565    /// column family. if no column family is passed, the default column
566    /// family is used.
567    ///
568    /// When `skip_metrics_reporting` is true, no periodic per-column-family
569    /// metrics task is spawned; use this for deprecated tables and for large
570    /// sets of rarely-touched column families.
571    #[instrument(level = "debug", skip(db), err)]
572    pub fn reopen(
573        db: &Arc<Database>,
574        opt_cf: Option<&str>,
575        rw_options: &ReadWriteOptions,
576        skip_metrics_reporting: bool,
577    ) -> Result<Self, TypedStoreError> {
578        let cf_key = opt_cf
579            .unwrap_or(rocksdb::DEFAULT_COLUMN_FAMILY_NAME)
580            .to_owned();
581        if db.cf_handle(&cf_key).is_none() {
582            return Err(TypedStoreError::UnregisteredColumn(cf_key));
583        }
584
585        let column_family = match &db.storage {
586            Storage::Rocks(_) => ColumnFamily::Rocks(cf_key),
587            Storage::InMemory(_) => ColumnFamily::InMemory(cf_key),
588        };
589        Ok(DBMap::new(
590            db.clone(),
591            rw_options,
592            column_family,
593            skip_metrics_reporting,
594        ))
595    }
596
597    pub fn cf_name(&self) -> &str {
598        self.column_family.name()
599    }
600
601    pub fn batch(&self) -> DBBatch {
602        let batch = match &self.db.storage {
603            Storage::Rocks(_) => StorageWriteBatch::Rocks(WriteBatch::default()),
604            Storage::InMemory(_) => StorageWriteBatch::InMemory(InMemoryBatch::default()),
605        };
606        DBBatch::new(
607            &self.db,
608            batch,
609            &self.db_metrics,
610            &self.write_sample_interval,
611        )
612    }
613
614    /// Flush the memtable of this table's column family to SST files on disk.
615    pub fn flush(&self) -> Result<(), TypedStoreError> {
616        self.db.flush_cf(&self.column_family)
617    }
618
619    /// Flush the memtables of every column family of the database to SST files
620    /// on disk.
621    pub fn flush_all(&self) -> Result<(), TypedStoreError> {
622        self.db.flush_all()
623    }
624
625    pub fn compact_range<J: Serialize>(&self, start: &J, end: &J) -> Result<(), TypedStoreError> {
626        let from_buf = be_fix_int_ser(start);
627        let to_buf = be_fix_int_ser(end);
628        self.db
629            .compact_range_cf(&self.column_family, Some(from_buf), Some(to_buf));
630        Ok(())
631    }
632
633    pub fn compact_range_raw(
634        &self,
635        cf_name: &str,
636        start: Vec<u8>,
637        end: Vec<u8>,
638    ) -> Result<(), TypedStoreError> {
639        let cf = match &self.db.storage {
640            Storage::Rocks(_) => ColumnFamily::Rocks(cf_name.to_string()),
641            Storage::InMemory(_) => ColumnFamily::InMemory(cf_name.to_string()),
642        };
643        self.db.compact_range_cf(&cf, Some(start), Some(end));
644        Ok(())
645    }
646
647    /// Returns a vector of raw values corresponding to the keys provided.
648    fn multi_get_pinned(
649        &self,
650        keys_bytes: impl IntoIterator<Item = Vec<u8>>,
651    ) -> Result<Vec<Option<GetResult<'_>>>, TypedStoreError> {
652        let _timer = self
653            .db_metrics
654            .op_metrics
655            .rocksdb_multiget_latency_seconds
656            .with_label_values(&[self.cf_name()])
657            .start_timer();
658        let perf_ctx = if self.multiget_sample_interval.sample() {
659            Some(RocksDBPerfContext)
660        } else {
661            None
662        };
663        let results: Result<Vec<_>, TypedStoreError> = self
664            .db
665            .multi_get(&self.column_family, keys_bytes, &self.opts.readopts())
666            .into_iter()
667            .collect();
668        let entries = results?;
669        let entry_size = entries
670            .iter()
671            .flatten()
672            .map(|entry| entry.len())
673            .sum::<usize>();
674        self.db_metrics
675            .op_metrics
676            .rocksdb_multiget_bytes
677            .with_label_values(&[self.cf_name()])
678            .observe(entry_size as f64);
679        if perf_ctx.is_some() {
680            self.db_metrics
681                .read_perf_ctx_metrics
682                .report_metrics(self.cf_name());
683        }
684        Ok(entries)
685    }
686
687    pub fn checkpoint_db(&self, path: &Path) -> Result<(), TypedStoreError> {
688        self.db.checkpoint(path)
689    }
690
691    pub fn table_summary(&self) -> eyre::Result<TableSummary>
692    where
693        K: Serialize + DeserializeOwned,
694        V: Serialize + DeserializeOwned,
695    {
696        let mut num_keys = 0;
697        let mut key_bytes_total = 0;
698        let mut value_bytes_total = 0;
699        let mut key_hist = hdrhistogram::Histogram::<u64>::new_with_max(100000, 2).unwrap();
700        let mut value_hist = hdrhistogram::Histogram::<u64>::new_with_max(100000, 2).unwrap();
701        for item in self.safe_iter() {
702            let (key, value) = item?;
703            num_keys += 1;
704            let key_len = be_fix_int_ser(key.borrow()).len();
705            let value_len = bcs::to_bytes(value.borrow())?.len();
706            key_bytes_total += key_len;
707            value_bytes_total += value_len;
708            key_hist.record(key_len as u64)?;
709            value_hist.record(value_len as u64)?;
710        }
711        Ok(TableSummary {
712            num_keys,
713            key_bytes_total,
714            value_bytes_total,
715            key_hist,
716            value_hist,
717        })
718    }
719
720    // Creates metrics and context for tracking an iterator usage and performance.
721    fn create_iter_context(
722        &self,
723    ) -> (
724        Option<HistogramTimer>,
725        Option<Histogram>,
726        Option<Histogram>,
727        Option<RocksDBPerfContext>,
728    ) {
729        let timer = self
730            .db_metrics
731            .op_metrics
732            .rocksdb_iter_latency_seconds
733            .with_label_values(&[self.cf_name()])
734            .start_timer();
735        let bytes_scanned = self
736            .db_metrics
737            .op_metrics
738            .rocksdb_iter_bytes
739            .with_label_values(&[self.cf_name()]);
740        let keys_scanned = self
741            .db_metrics
742            .op_metrics
743            .rocksdb_iter_keys
744            .with_label_values(&[self.cf_name()]);
745        let perf_ctx = if self.iter_sample_interval.sample() {
746            Some(RocksDBPerfContext)
747        } else {
748            None
749        };
750        (
751            Some(timer),
752            Some(bytes_scanned),
753            Some(keys_scanned),
754            perf_ctx,
755        )
756    }
757
758    /// Shared rocksdb `SafeIter` construction (raw iterator + scan metrics)
759    /// used by both the forward and reverse raw iterators.
760    fn rocks_safe_iter<'a>(&self, db: &'a RocksDB, readopts: ReadOptions) -> SafeIter<'a, K, V>
761    where
762        K: DeserializeOwned,
763        V: DeserializeOwned,
764    {
765        let db_iter = db
766            .underlying
767            .raw_iterator_cf_opt(&rocks_cf(db, self.column_family.name()), readopts);
768        let (_timer, bytes_scanned, keys_scanned, _perf_ctx) = self.create_iter_context();
769        SafeIter::new(
770            self.cf_name().to_string(),
771            db_iter,
772            _timer,
773            _perf_ctx,
774            bytes_scanned,
775            keys_scanned,
776            Some(self.db_metrics.clone()),
777        )
778    }
779
780    /// Forward iterator over the raw byte bounds `[lower_bound, upper_bound)`;
781    /// both bounds (already serialized, `upper_bound` exclusive) are applied to
782    /// the read options.
783    fn iter_forward_raw(
784        &self,
785        lower_bound: Option<Vec<u8>>,
786        upper_bound: Option<Vec<u8>>,
787    ) -> DbIterator<'_, (K, V)>
788    where
789        K: DeserializeOwned,
790        V: DeserializeOwned,
791    {
792        match &self.db.storage {
793            Storage::Rocks(db) => {
794                let readopts =
795                    rocks_util::apply_range_bounds(self.opts.readopts(), lower_bound, upper_bound);
796                Box::new(self.rocks_safe_iter(db, readopts))
797            }
798            Storage::InMemory(db) => {
799                db.iterator(self.column_family.name(), lower_bound, upper_bound, false)
800            }
801        }
802    }
803
804    /// Reverse counterpart of [`Self::iter_forward_raw`], yielding the same
805    /// keys in descending order. Only the lower bound is applied to the
806    /// read options; the (exclusive) `upper_bound` is enforced by
807    /// [`SafeRevIter`]'s initial seek, because setting an
808    /// iterate-upper-bound would stop `seek_for_prev` from landing on the
809    /// boundary key.
810    fn iter_reversed_raw(
811        &self,
812        lower_bound: Option<Vec<u8>>,
813        upper_bound: Option<Vec<u8>>,
814    ) -> DbIterator<'_, (K, V)>
815    where
816        K: DeserializeOwned,
817        V: DeserializeOwned,
818    {
819        match &self.db.storage {
820            Storage::Rocks(db) => {
821                let readopts =
822                    rocks_util::apply_range_bounds(self.opts.readopts(), lower_bound, None);
823                Box::new(SafeRevIter::new(
824                    self.rocks_safe_iter(db, readopts),
825                    upper_bound,
826                ))
827            }
828            Storage::InMemory(db) => {
829                db.iterator(self.column_family.name(), lower_bound, upper_bound, true)
830            }
831        }
832    }
833
834    /// Forward iterator over every entry whose key begins with `prefix`.
835    ///
836    /// `prefix` is serialized with `be_fix_int_ser` and must form a prefix of
837    /// the column family's key encoding — typically the leading field(s) of a
838    /// tuple key. This avoids constructing artificial maximum keys to bound a
839    /// composite-key scan.
840    pub fn safe_iter_with_prefix<P>(&self, prefix: &P) -> DbIterator<'_, (K, V)>
841    where
842        P: ?Sized + Serialize,
843        K: DeserializeOwned,
844        V: DeserializeOwned,
845    {
846        let (lower_bound, upper_bound) = prefix_iterator_bounds(prefix);
847        self.iter_forward_raw(lower_bound, upper_bound)
848    }
849
850    /// Forward iterator over entries whose key begins with `prefix`, starting
851    /// at `prefix ++ cursor` rather than at the first key under the prefix.
852    ///
853    /// `cursor` is **not** a full key — it is the remainder of the key that
854    /// follows `prefix` (typically the trailing field(s) of a tuple key). Both
855    /// `prefix` and `cursor` are serialized with `be_fix_int_ser` and
856    /// concatenated, mirroring how a composite key is encoded; the scan then
857    /// covers `[prefix ++ cursor, end-of-prefix)`. This lets a paginated scan
858    /// resume from a cursor without an artificial maximum key for the upper
859    /// bound.
860    pub fn safe_iter_with_prefix_from<P, C>(&self, prefix: &P, cursor: &C) -> DbIterator<'_, (K, V)>
861    where
862        P: ?Sized + Serialize,
863        C: ?Sized + Serialize,
864        K: DeserializeOwned,
865        V: DeserializeOwned,
866    {
867        let (lower_bound, upper_bound) = prefix_iterator_bounds(prefix);
868        let lower_bound = lower_bound.map(|mut lower| {
869            lower.extend_from_slice(&be_fix_int_ser(cursor));
870            lower
871        });
872        self.iter_forward_raw(lower_bound, upper_bound)
873    }
874
875    /// Reverse counterpart of [`Self::safe_iter_with_prefix`]: the matching
876    /// entries in descending key order (e.g. for "latest entry under a
877    /// prefix").
878    pub fn safe_iter_with_prefix_reversed<P>(&self, prefix: &P) -> DbIterator<'_, (K, V)>
879    where
880        P: ?Sized + Serialize,
881        K: DeserializeOwned,
882        V: DeserializeOwned,
883    {
884        let (lower_bound, upper_bound) = prefix_iterator_bounds(prefix);
885        self.iter_reversed_raw(lower_bound, upper_bound)
886    }
887}
888
889/// Provides a mutable struct to form a collection of database write operations,
890/// and execute them.
891///
892/// Batching write and delete operations is faster than performing them one by
893/// one and ensures their atomicity,  ie. they are all written or none is.
894/// This is also true of operations across column families in the same database.
895///
896/// Serializations / Deserialization, and naming of column families is performed
897/// by passing a DBMap<K,V> with each operation.
898///
899/// ```
900/// use core::fmt::Error;
901/// use std::sync::Arc;
902///
903/// use prometheus_filtered::Registry;
904/// use tempfile::tempdir;
905/// use typed_store::{Map, metrics::DBMetrics, rocks::*};
906///
907/// #[tokio::main(flavor = "current_thread")]
908/// async fn main() -> Result<(), Error> {
909///     let rocks = open_cf_opts(
910///         tempfile::tempdir().unwrap(),
911///         None,
912///         MetricConf::default(),
913///         &[
914///             ("First_CF", rocksdb::Options::default()),
915///             ("Second_CF", rocksdb::Options::default()),
916///         ],
917///     )
918///     .unwrap();
919///
920///     let db_cf_1 = DBMap::reopen(
921///         &rocks,
922///         Some("First_CF"),
923///         &ReadWriteOptions::default(),
924///         false,
925///     )
926///     .expect("Failed to open storage");
927///     let keys_vals_1 = (1..100).map(|i| (i, i.to_string()));
928///
929///     let db_cf_2 = DBMap::reopen(
930///         &rocks,
931///         Some("Second_CF"),
932///         &ReadWriteOptions::default(),
933///         false,
934///     )
935///     .expect("Failed to open storage");
936///     let keys_vals_2 = (1000..1100).map(|i| (i, i.to_string()));
937///
938///     let mut batch = db_cf_1.batch();
939///     batch
940///         .insert_batch(&db_cf_1, keys_vals_1.clone())
941///         .expect("Failed to batch insert")
942///         .insert_batch(&db_cf_2, keys_vals_2.clone())
943///         .expect("Failed to batch insert");
944///
945///     let _ = batch.write().expect("Failed to execute batch");
946///     for (k, v) in keys_vals_1 {
947///         let val = db_cf_1.get(&k).expect("Failed to get inserted key");
948///         assert_eq!(Some(v), val);
949///     }
950///
951///     for (k, v) in keys_vals_2 {
952///         let val = db_cf_2.get(&k).expect("Failed to get inserted key");
953///         assert_eq!(Some(v), val);
954///     }
955///     Ok(())
956/// }
957/// ```
958pub struct DBBatch {
959    database: Arc<Database>,
960    batch: StorageWriteBatch,
961    db_metrics: Arc<DBMetrics>,
962    write_sample_interval: SamplingInterval,
963}
964
965impl DBBatch {
966    /// Create a new batch associated with a DB reference.
967    ///
968    /// Use `open_cf` to get the DB reference or an existing open database.
969    pub fn new(
970        dbref: &Arc<Database>,
971        batch: StorageWriteBatch,
972        db_metrics: &Arc<DBMetrics>,
973        write_sample_interval: &SamplingInterval,
974    ) -> Self {
975        DBBatch {
976            database: dbref.clone(),
977            batch,
978            db_metrics: db_metrics.clone(),
979            write_sample_interval: write_sample_interval.clone(),
980        }
981    }
982
983    /// Consume the batch and write its operations to the database
984    #[instrument(level = "trace", skip_all, err)]
985    pub fn write(self) -> Result<(), TypedStoreError> {
986        self.write_opt(&rocksdb::WriteOptions::default())
987    }
988
989    /// Consume the batch and write its operations to the database with custom
990    /// write options
991    #[instrument(level = "trace", skip_all, err)]
992    pub fn write_opt(self, write_options: &rocksdb::WriteOptions) -> Result<(), TypedStoreError> {
993        let db_name = self.database.db_name();
994        let timer = self
995            .db_metrics
996            .op_metrics
997            .rocksdb_batch_commit_latency_seconds
998            .with_label_values(&[&db_name])
999            .start_timer();
1000        let batch_size_bytes = self.size_in_bytes();
1001
1002        let perf_ctx = if self.write_sample_interval.sample() {
1003            Some(RocksDBPerfContext)
1004        } else {
1005            None
1006        };
1007        self.database.write_opt(self.batch, write_options)?;
1008        self.db_metrics
1009            .op_metrics
1010            .rocksdb_batch_commit_bytes
1011            .with_label_values(&[&db_name])
1012            .observe(batch_size_bytes as f64);
1013
1014        if perf_ctx.is_some() {
1015            self.db_metrics
1016                .write_perf_ctx_metrics
1017                .report_metrics(&db_name);
1018        }
1019        let elapsed_secs = timer.stop_and_record();
1020        let threshold_secs = very_slow_batch_write_threshold_secs(batch_size_bytes);
1021        if elapsed_secs > threshold_secs {
1022            warn!(
1023                elapsed_secs,
1024                threshold_secs,
1025                batch_size_bytes,
1026                ?db_name,
1027                "very slow batch write"
1028            );
1029            self.db_metrics
1030                .op_metrics
1031                .rocksdb_very_slow_batch_writes_count
1032                .with_label_values(&[&db_name])
1033                .inc();
1034            self.db_metrics
1035                .op_metrics
1036                .rocksdb_very_slow_batch_writes_duration_ms
1037                .with_label_values(&[&db_name])
1038                .inc_by((elapsed_secs * 1000.0) as u64);
1039        }
1040        Ok(())
1041    }
1042
1043    pub fn size_in_bytes(&self) -> usize {
1044        match self.batch {
1045            StorageWriteBatch::Rocks(ref b) => b.size_in_bytes(),
1046            StorageWriteBatch::InMemory(_) => 0,
1047        }
1048    }
1049
1050    pub fn delete_batch<J: Borrow<K>, K: Serialize, V>(
1051        &mut self,
1052        db: &DBMap<K, V>,
1053        purged_vals: impl IntoIterator<Item = J>,
1054    ) -> Result<(), TypedStoreError> {
1055        self.delete_batch_raw_keys(
1056            db,
1057            purged_vals
1058                .into_iter()
1059                .map(|key| be_fix_int_ser(key.borrow())),
1060        )
1061    }
1062
1063    /// [`Self::delete_batch`] for keys already in the column family's
1064    /// `be_fix_int_ser` encoding.
1065    fn delete_batch_raw_keys<K, V>(
1066        &mut self,
1067        db: &DBMap<K, V>,
1068        purged_vals: impl IntoIterator<Item = Vec<u8>>,
1069    ) -> Result<(), TypedStoreError> {
1070        if !Arc::ptr_eq(&db.db, &self.database) {
1071            return Err(TypedStoreError::CrossDBBatch);
1072        }
1073
1074        purged_vals
1075            .into_iter()
1076            .try_for_each::<_, Result<_, TypedStoreError>>(|k_buf| {
1077                match (&mut self.batch, &db.column_family) {
1078                    (StorageWriteBatch::Rocks(b), ColumnFamily::Rocks(name)) => {
1079                        b.delete_cf(&rocks_cf_from_db(&self.database, name)?, k_buf)
1080                    }
1081                    (StorageWriteBatch::InMemory(b), ColumnFamily::InMemory(name)) => {
1082                        b.delete_cf(name, k_buf)
1083                    }
1084                    _ => Err(TypedStoreError::RocksDB(
1085                        "typed store invariant violation".to_string(),
1086                    ))?,
1087                }
1088                Ok(())
1089            })?;
1090        Ok(())
1091    }
1092
1093    /// Deletes a range of keys between `from` (inclusive) and `to`
1094    /// (non-inclusive) by writing a range delete tombstone in the db map.
1095    /// The effect of this write is visible immediately, i.e. you won't see
1096    /// old values when you do a lookup or scan.
1097    pub fn schedule_delete_range<K: Serialize, V>(
1098        &mut self,
1099        db: &DBMap<K, V>,
1100        from: &K,
1101        to: &K,
1102    ) -> Result<(), TypedStoreError> {
1103        self.schedule_delete_range_raw(db, be_fix_int_ser(from), be_fix_int_ser(to))
1104    }
1105
1106    /// [`Self::schedule_delete_range`] for a range whose bounds are not keys,
1107    /// such as a whole key prefix. The bounds are raw bytes ordered like the
1108    /// map's keys, i.e. in the `be_fix_int_ser` encoding; `from` is inclusive
1109    /// and `to` exclusive.
1110    fn schedule_delete_range_raw<K, V>(
1111        &mut self,
1112        db: &DBMap<K, V>,
1113        from: Vec<u8>,
1114        to: Vec<u8>,
1115    ) -> Result<(), TypedStoreError> {
1116        if !Arc::ptr_eq(&db.db, &self.database) {
1117            return Err(TypedStoreError::CrossDBBatch);
1118        }
1119
1120        if let StorageWriteBatch::Rocks(b) = &mut self.batch {
1121            b.delete_range_cf(&rocks_cf_from_db(&self.database, db.cf_name())?, from, to);
1122        }
1123        Ok(())
1124    }
1125
1126    /// inserts a range of (key, value) pairs given as an iterator
1127    pub fn insert_batch<J: Borrow<K>, K: Serialize, U: Borrow<V>, V: Serialize>(
1128        &mut self,
1129        db: &DBMap<K, V>,
1130        new_vals: impl IntoIterator<Item = (J, U)>,
1131    ) -> Result<&mut Self, TypedStoreError> {
1132        self.insert_batch_raw_keys(
1133            db,
1134            new_vals
1135                .into_iter()
1136                .map(|(key, value)| (be_fix_int_ser(key.borrow()), value)),
1137        )
1138    }
1139
1140    /// [`Self::insert_batch`] for keys already in the column family's
1141    /// `be_fix_int_ser` encoding.
1142    fn insert_batch_raw_keys<K, U: Borrow<V>, V: Serialize>(
1143        &mut self,
1144        db: &DBMap<K, V>,
1145        new_vals: impl IntoIterator<Item = (Vec<u8>, U)>,
1146    ) -> Result<&mut Self, TypedStoreError> {
1147        if !Arc::ptr_eq(&db.db, &self.database) {
1148            return Err(TypedStoreError::CrossDBBatch);
1149        }
1150        let mut total = 0usize;
1151        new_vals
1152            .into_iter()
1153            .try_for_each::<_, Result<_, TypedStoreError>>(|(k_buf, v)| {
1154                let v_buf = bcs::to_bytes(v.borrow()).map_err(typed_store_err_from_bcs_err)?;
1155                total += k_buf.len() + v_buf.len();
1156                if db.opts.log_value_hash {
1157                    let key_hash = default_hash(&k_buf);
1158                    let value_hash = default_hash(&v_buf);
1159                    debug!(
1160                        "Insert to DB table: {:?}, key_hash: {:?}, value_hash: {:?}",
1161                        db.cf_name(),
1162                        key_hash,
1163                        value_hash
1164                    );
1165                }
1166                match (&mut self.batch, &db.column_family) {
1167                    (StorageWriteBatch::Rocks(b), ColumnFamily::Rocks(name)) => {
1168                        b.put_cf(&rocks_cf_from_db(&self.database, name)?, k_buf, v_buf)
1169                    }
1170                    (StorageWriteBatch::InMemory(b), ColumnFamily::InMemory(name)) => {
1171                        b.put_cf(name, k_buf, v_buf)
1172                    }
1173                    _ => Err(TypedStoreError::RocksDB(
1174                        "typed store invariant violation".to_string(),
1175                    ))?,
1176                }
1177                Ok(())
1178            })?;
1179        self.db_metrics
1180            .op_metrics
1181            .rocksdb_batch_put_bytes
1182            .with_label_values(&[db.cf_name()])
1183            .observe(total as f64);
1184        Ok(self)
1185    }
1186
1187    /// [`Self::insert_batch`] for a map sharing its column family by tag.
1188    pub fn insert_batch_tagged<J: Borrow<K>, K: Serialize, U: Borrow<V>, V: Serialize>(
1189        &mut self,
1190        map: &TaggedDBMap<K, V>,
1191        new_vals: impl IntoIterator<Item = (J, U)>,
1192    ) -> Result<&mut Self, TypedStoreError> {
1193        self.insert_batch_raw_keys(
1194            &map.map,
1195            new_vals
1196                .into_iter()
1197                .map(|(key, value)| (be_fix_int_ser(&(map.tag, key.borrow())), value)),
1198        )
1199    }
1200
1201    /// [`Self::delete_batch`] for a map sharing its column family by tag.
1202    pub fn delete_batch_tagged<J: Borrow<K>, K: Serialize, V>(
1203        &mut self,
1204        map: &TaggedDBMap<K, V>,
1205        purged_vals: impl IntoIterator<Item = J>,
1206    ) -> Result<(), TypedStoreError> {
1207        self.delete_batch_raw_keys(
1208            &map.map,
1209            purged_vals
1210                .into_iter()
1211                .map(|key| be_fix_int_ser(&(map.tag, key.borrow()))),
1212        )
1213    }
1214}
1215
1216impl<K, V> DBMap<K, V> {
1217    #[instrument(level = "trace", skip_all, err)]
1218    pub(crate) fn contains_raw_key(&self, key_buf: Vec<u8>) -> Result<bool, TypedStoreError> {
1219        let readopts = self.opts.readopts();
1220        Ok(self
1221            .db
1222            .key_may_exist_cf(&self.column_family, &key_buf, &readopts)
1223            && self
1224                .db
1225                .get(&self.column_family, &key_buf, &readopts)?
1226                .is_some())
1227    }
1228
1229    #[instrument(level = "trace", skip_all, err)]
1230    pub(crate) fn multi_contains_raw_keys(
1231        &self,
1232        keys: impl IntoIterator<Item = Vec<u8>>,
1233    ) -> Result<Vec<bool>, TypedStoreError> {
1234        let values = self.multi_get_pinned(keys)?;
1235        Ok(values.into_iter().map(|v| v.is_some()).collect())
1236    }
1237
1238    #[instrument(level = "trace", skip_all, err)]
1239    pub(crate) fn get_raw_key(&self, key_buf: Vec<u8>) -> Result<Option<V>, TypedStoreError>
1240    where
1241        V: DeserializeOwned,
1242    {
1243        let _timer = self
1244            .db_metrics
1245            .op_metrics
1246            .rocksdb_get_latency_seconds
1247            .with_label_values(&[self.cf_name()])
1248            .start_timer();
1249        let perf_ctx = if self.get_sample_interval.sample() {
1250            Some(RocksDBPerfContext)
1251        } else {
1252            None
1253        };
1254        let res = self
1255            .db
1256            .get(&self.column_family, &key_buf, &self.opts.readopts())?;
1257        self.db_metrics
1258            .op_metrics
1259            .rocksdb_get_bytes
1260            .with_label_values(&[self.cf_name()])
1261            .observe(res.as_ref().map_or(0.0, |v| v.len() as f64));
1262        if perf_ctx.is_some() {
1263            self.db_metrics
1264                .read_perf_ctx_metrics
1265                .report_metrics(self.cf_name());
1266        }
1267        match res {
1268            Some(data) => {
1269                let value = bcs::from_bytes(&data).map_err(typed_store_err_from_bcs_err);
1270                if value.is_err() {
1271                    let key_hash = default_hash(&key_buf);
1272                    let value_hash = default_hash(&data);
1273                    debug_fatal!(
1274                        "Failed to deserialize value from DB table {:?}, key_hash: {:?}, value_hash: {:?}, error: {:?}",
1275                        self.cf_name(),
1276                        key_hash,
1277                        value_hash,
1278                        value.as_ref().err().unwrap()
1279                    );
1280                }
1281                Ok(Some(value?))
1282            }
1283            None => Ok(None),
1284        }
1285    }
1286
1287    #[instrument(level = "trace", skip_all, err)]
1288    pub(crate) fn insert_raw_key(&self, key_buf: Vec<u8>, value: &V) -> Result<(), TypedStoreError>
1289    where
1290        V: Serialize,
1291    {
1292        let timer = self
1293            .db_metrics
1294            .op_metrics
1295            .rocksdb_put_latency_seconds
1296            .with_label_values(&[self.cf_name()])
1297            .start_timer();
1298        let perf_ctx = if self.write_sample_interval.sample() {
1299            Some(RocksDBPerfContext)
1300        } else {
1301            None
1302        };
1303        let value_buf = bcs::to_bytes(value).map_err(typed_store_err_from_bcs_err)?;
1304        self.db_metrics
1305            .op_metrics
1306            .rocksdb_put_bytes
1307            .with_label_values(&[self.cf_name()])
1308            .observe((key_buf.len() + value_buf.len()) as f64);
1309        if perf_ctx.is_some() {
1310            self.db_metrics
1311                .write_perf_ctx_metrics
1312                .report_metrics(self.cf_name());
1313        }
1314        self.db.put_cf(&self.column_family, key_buf, value_buf)?;
1315
1316        let elapsed_secs = timer.stop_and_record();
1317        if elapsed_secs > 1.0 {
1318            warn!(elapsed_secs, cf = ?self.cf_name(), "very slow insert");
1319            self.db_metrics
1320                .op_metrics
1321                .rocksdb_very_slow_puts_count
1322                .with_label_values(&[self.cf_name()])
1323                .inc();
1324            self.db_metrics
1325                .op_metrics
1326                .rocksdb_very_slow_puts_duration_ms
1327                .with_label_values(&[self.cf_name()])
1328                .inc_by((elapsed_secs * 1000.0) as u64);
1329        }
1330
1331        Ok(())
1332    }
1333
1334    #[instrument(level = "trace", skip_all, err)]
1335    pub(crate) fn remove_raw_key(&self, key_buf: Vec<u8>) -> Result<(), TypedStoreError> {
1336        let _timer = self
1337            .db_metrics
1338            .op_metrics
1339            .rocksdb_delete_latency_seconds
1340            .with_label_values(&[self.cf_name()])
1341            .start_timer();
1342        let perf_ctx = if self.write_sample_interval.sample() {
1343            Some(RocksDBPerfContext)
1344        } else {
1345            None
1346        };
1347        self.db.delete_cf(&self.column_family, key_buf)?;
1348        self.db_metrics
1349            .op_metrics
1350            .rocksdb_deletes
1351            .with_label_values(&[self.cf_name()])
1352            .inc();
1353        if perf_ctx.is_some() {
1354            self.db_metrics
1355                .write_perf_ctx_metrics
1356                .report_metrics(self.cf_name());
1357        }
1358        Ok(())
1359    }
1360
1361    #[instrument(level = "trace", skip_all, err)]
1362    pub(crate) fn multi_get_raw_keys(
1363        &self,
1364        keys: impl IntoIterator<Item = Vec<u8>>,
1365    ) -> Result<Vec<Option<V>>, TypedStoreError>
1366    where
1367        V: DeserializeOwned,
1368    {
1369        let results = self.multi_get_pinned(keys)?;
1370        let values_parsed: Result<Vec<_>, TypedStoreError> = results
1371            .into_iter()
1372            .map(|value_byte| match value_byte {
1373                Some(data) => Ok(Some(
1374                    bcs::from_bytes(&data).map_err(typed_store_err_from_bcs_err)?,
1375                )),
1376                None => Ok(None),
1377            })
1378            .collect();
1379
1380        values_parsed
1381    }
1382}
1383
1384impl<'a, K, V> Map<'a, K, V> for DBMap<K, V>
1385where
1386    K: Serialize + DeserializeOwned,
1387    V: Serialize + DeserializeOwned,
1388{
1389    type Error = TypedStoreError;
1390
1391    fn contains_key(&self, key: &K) -> Result<bool, TypedStoreError> {
1392        self.contains_raw_key(be_fix_int_ser(key))
1393    }
1394
1395    fn multi_contains_keys<J>(
1396        &self,
1397        keys: impl IntoIterator<Item = J>,
1398    ) -> Result<Vec<bool>, Self::Error>
1399    where
1400        J: Borrow<K>,
1401    {
1402        self.multi_contains_raw_keys(keys.into_iter().map(|k| be_fix_int_ser(k.borrow())))
1403    }
1404
1405    fn get(&self, key: &K) -> Result<Option<V>, TypedStoreError> {
1406        self.get_raw_key(be_fix_int_ser(key))
1407    }
1408
1409    fn insert(&self, key: &K, value: &V) -> Result<(), TypedStoreError> {
1410        self.insert_raw_key(be_fix_int_ser(key), value)
1411    }
1412
1413    fn remove(&self, key: &K) -> Result<(), TypedStoreError> {
1414        self.remove_raw_key(be_fix_int_ser(key))
1415    }
1416
1417    /// Writes a range delete tombstone to delete all entries in the db map.
1418    /// The effect of this write is visible immediately, i.e. you won't see
1419    /// old values when you do a lookup or scan.
1420    #[instrument(level = "trace", skip_all, err)]
1421    fn schedule_delete_all(&self) -> Result<(), TypedStoreError> {
1422        let Some(last_key) = self
1423            .safe_range_iter_reversed(..)
1424            .next()
1425            .transpose()?
1426            .map(|(k, _v)| k)
1427        else {
1428            return Ok(());
1429        };
1430        // The tombstone excludes its upper bound, so it has to end past the last key:
1431        // appending a zero byte gives the first key sorting above it.
1432        let mut to = be_fix_int_ser(&last_key);
1433        to.push(0);
1434        // The empty key sorts below every key, so the range starts before the first
1435        // entry.
1436        let mut batch = self.batch();
1437        batch.schedule_delete_range_raw(self, Vec::new(), to)?;
1438        batch.write()
1439    }
1440
1441    fn is_empty(&self) -> bool {
1442        self.safe_iter().next().is_none()
1443    }
1444
1445    fn safe_iter(&'a self) -> DbIterator<'a, (K, V)> {
1446        match &self.db.storage {
1447            Storage::Rocks(db) => {
1448                let db_iter = db.underlying.raw_iterator_cf_opt(
1449                    &rocks_cf(db, self.column_family.name()),
1450                    self.opts.readopts(),
1451                );
1452                let (_timer, bytes_scanned, keys_scanned, _perf_ctx) = self.create_iter_context();
1453                Box::new(SafeIter::new(
1454                    self.cf_name().to_string(),
1455                    db_iter,
1456                    _timer,
1457                    _perf_ctx,
1458                    bytes_scanned,
1459                    keys_scanned,
1460                    Some(self.db_metrics.clone()),
1461                ))
1462            }
1463            Storage::InMemory(db) => db.iterator(self.column_family.name(), None, None, false),
1464        }
1465    }
1466
1467    fn safe_iter_with_bounds(
1468        &'a self,
1469        lower_bound: Option<K>,
1470        upper_bound: Option<K>,
1471    ) -> DbIterator<'a, (K, V)> {
1472        let range = (
1473            lower_bound.map(Bound::Included).unwrap_or(Bound::Unbounded),
1474            upper_bound.map(Bound::Excluded).unwrap_or(Bound::Unbounded),
1475        );
1476        self.safe_range_iter(range)
1477    }
1478
1479    fn safe_range_iter(&'a self, range: impl RangeBounds<K>) -> DbIterator<'a, (K, V)> {
1480        let (lower_bound, upper_bound) = iterator_bounds_with_range(range);
1481        self.iter_forward_raw(lower_bound, upper_bound)
1482    }
1483
1484    /// Both directions derive their bounds from the same
1485    /// `iterator_bounds_with_range`, so they are guaranteed to cover the
1486    /// identical set of keys regardless of the bound inclusivity.
1487    fn safe_range_iter_reversed(&'a self, range: impl RangeBounds<K>) -> DbIterator<'a, (K, V)> {
1488        let (lower_bound, upper_bound) = iterator_bounds_with_range(range);
1489        self.iter_reversed_raw(lower_bound, upper_bound)
1490    }
1491
1492    /// Returns a vector of values corresponding to the keys provided.
1493    fn multi_get<J>(
1494        &self,
1495        keys: impl IntoIterator<Item = J>,
1496    ) -> Result<Vec<Option<V>>, TypedStoreError>
1497    where
1498        J: Borrow<K>,
1499    {
1500        self.multi_get_raw_keys(keys.into_iter().map(|k| be_fix_int_ser(k.borrow())))
1501    }
1502
1503    /// Convenience method for batch insertion
1504    #[instrument(level = "trace", skip_all, err)]
1505    fn multi_insert<J, U>(
1506        &self,
1507        key_val_pairs: impl IntoIterator<Item = (J, U)>,
1508    ) -> Result<(), Self::Error>
1509    where
1510        J: Borrow<K>,
1511        U: Borrow<V>,
1512    {
1513        let mut batch = self.batch();
1514        batch.insert_batch(self, key_val_pairs)?;
1515        batch.write()
1516    }
1517
1518    /// Convenience method for batch removal
1519    #[instrument(level = "trace", skip_all, err)]
1520    fn multi_remove<J>(&self, keys: impl IntoIterator<Item = J>) -> Result<(), Self::Error>
1521    where
1522        J: Borrow<K>,
1523    {
1524        let mut batch = self.batch();
1525        batch.delete_batch(self, keys)?;
1526        batch.write()
1527    }
1528
1529    /// Try to catch up with primary when running as secondary
1530    #[instrument(level = "trace", skip_all, err)]
1531    fn try_catch_up_with_primary(&self) -> Result<(), Self::Error> {
1532        self.db.try_catch_up_with_primary()
1533    }
1534}
1535
1536/// A typed map sharing one column family with other typed maps,
1537/// distinguished by a tag byte prefixed to every key: the big-endian fixint
1538/// serialization of `(tag, key)` keeps each map a contiguous key range of
1539/// the column family. Useful for sets of tables that are always created and
1540/// dropped together — one column family instead of one per table keeps the
1541/// column-family and SST-file counts low.
1542///
1543/// Every read is scoped to the map's tag: full scans iterate the tag's key
1544/// prefix, and ranges are two-sided within it, so rows of other tags never
1545/// surface and their (differently typed) values are never deserialized.
1546///
1547/// Tags of existing maps must never change or be reused, or old rows would
1548/// be read under the wrong types.
1549///
1550/// Metrics are labelled by column family, so all the maps of one column
1551/// family report together, with no way to tell the tags apart.
1552///
1553/// # Examples
1554///
1555/// ```
1556/// use tempfile::tempdir;
1557/// use typed_store::{Map, rocks::*};
1558///
1559/// #[tokio::main(flavor = "current_thread")]
1560/// async fn main() {
1561///     let db = open_cf_opts(
1562///         tempdir().unwrap(),
1563///         None,
1564///         MetricConf::default(),
1565///         &[("shared", rocksdb::Options::default())],
1566///     )
1567///     .unwrap();
1568///
1569///     let numbers: TaggedDBMap<u32, String> =
1570///         TaggedDBMap::reopen(&db, "shared", 0, &ReadWriteOptions::default(), false).unwrap();
1571///     let words: TaggedDBMap<String, u64> =
1572///         TaggedDBMap::reopen(&db, "shared", 1, &ReadWriteOptions::default(), false).unwrap();
1573///
1574///     let mut batch = numbers.batch();
1575///     batch
1576///         .insert_batch_tagged(&numbers, [(1, "one".to_string())])
1577///         .unwrap()
1578///         .insert_batch_tagged(&words, [("one".to_string(), 1)])
1579///         .unwrap();
1580///     batch.write().unwrap();
1581///
1582///     assert_eq!(numbers.get(&1).unwrap(), Some("one".to_string()));
1583///     assert_eq!(words.get(&"one".to_string()).unwrap(), Some(1));
1584///     // Each map sees only its own row.
1585///     assert_eq!(numbers.safe_iter().count(), 1);
1586/// }
1587/// ```
1588pub struct TaggedDBMap<K, V> {
1589    tag: u8,
1590    map: DBMap<(u8, K), V>,
1591}
1592
1593impl<K, V> TaggedDBMap<K, V> {
1594    /// Drops the tag prefixed to a key of the underlying map, turning its row
1595    /// into a row of this map.
1596    fn strip_tag(row: Result<((u8, K), V), TypedStoreError>) -> Result<(K, V), TypedStoreError> {
1597        row.map(|((_, key), value)| (key, value))
1598    }
1599
1600    /// Reopens an open database as a tagged map operating under `cf_name`.
1601    /// See [`DBMap::reopen`] for the remaining parameters.
1602    pub fn reopen(
1603        db: &Arc<Database>,
1604        cf_name: &str,
1605        tag: u8,
1606        rw_options: &ReadWriteOptions,
1607        skip_metrics_reporting: bool,
1608    ) -> Result<Self, TypedStoreError> {
1609        Ok(Self {
1610            tag,
1611            map: DBMap::reopen(db, Some(cf_name), rw_options, skip_metrics_reporting)?,
1612        })
1613    }
1614
1615    /// Create a batch associated with the map's database.
1616    pub fn batch(&self) -> DBBatch {
1617        self.map.batch()
1618    }
1619}
1620
1621impl<'a, K, V> Map<'a, K, V> for TaggedDBMap<K, V>
1622where
1623    K: Serialize + DeserializeOwned,
1624    V: Serialize + DeserializeOwned,
1625{
1626    type Error = TypedStoreError;
1627
1628    fn contains_key(&self, key: &K) -> Result<bool, TypedStoreError> {
1629        self.map.contains_raw_key(be_fix_int_ser(&(self.tag, key)))
1630    }
1631
1632    fn multi_contains_keys<J>(
1633        &self,
1634        keys: impl IntoIterator<Item = J>,
1635    ) -> Result<Vec<bool>, TypedStoreError>
1636    where
1637        J: Borrow<K>,
1638    {
1639        self.map.multi_contains_raw_keys(
1640            keys.into_iter()
1641                .map(|key| be_fix_int_ser(&(self.tag, key.borrow()))),
1642        )
1643    }
1644
1645    fn get(&self, key: &K) -> Result<Option<V>, TypedStoreError> {
1646        self.map.get_raw_key(be_fix_int_ser(&(self.tag, key)))
1647    }
1648
1649    fn multi_get<J>(
1650        &self,
1651        keys: impl IntoIterator<Item = J>,
1652    ) -> Result<Vec<Option<V>>, TypedStoreError>
1653    where
1654        J: Borrow<K>,
1655    {
1656        self.map.multi_get_raw_keys(
1657            keys.into_iter()
1658                .map(|key| be_fix_int_ser(&(self.tag, key.borrow()))),
1659        )
1660    }
1661
1662    fn insert(&self, key: &K, value: &V) -> Result<(), TypedStoreError> {
1663        self.map
1664            .insert_raw_key(be_fix_int_ser(&(self.tag, key)), value)
1665    }
1666
1667    fn remove(&self, key: &K) -> Result<(), TypedStoreError> {
1668        self.map.remove_raw_key(be_fix_int_ser(&(self.tag, key)))
1669    }
1670
1671    #[instrument(level = "trace", skip_all, err)]
1672    fn multi_insert<J, U>(
1673        &self,
1674        key_val_pairs: impl IntoIterator<Item = (J, U)>,
1675    ) -> Result<(), TypedStoreError>
1676    where
1677        J: Borrow<K>,
1678        U: Borrow<V>,
1679    {
1680        let mut batch = self.batch();
1681        batch.insert_batch_tagged(self, key_val_pairs)?;
1682        batch.write()
1683    }
1684
1685    #[instrument(level = "trace", skip_all, err)]
1686    fn multi_remove<J>(&self, keys: impl IntoIterator<Item = J>) -> Result<(), TypedStoreError>
1687    where
1688        J: Borrow<K>,
1689    {
1690        let mut batch = self.batch();
1691        batch.delete_batch_tagged(self, keys)?;
1692        batch.write()
1693    }
1694
1695    /// Deletes every entry of the map with a single range tombstone over the
1696    /// tag's key range, leaving the other tags of the column family untouched.
1697    ///
1698    /// The entries stay on disk until compaction, and until then iteration
1699    /// is slowed because of them.
1700    #[instrument(level = "trace", skip_all, err)]
1701    fn schedule_delete_all(&self) -> Result<(), TypedStoreError> {
1702        let from = be_fix_int_ser(&self.tag);
1703        let to = match prefix_iterator_bounds(&self.tag).1 {
1704            Some(to) => {
1705                if self.is_empty() {
1706                    return Ok(());
1707                }
1708                to
1709            }
1710            // The maximum tag has no following tag to bound the range
1711            // against, so end it just past the last entry. That read is why a
1712            // row written in between survives the clear, and why a value that
1713            // does not deserialize fails it.
1714            None => match self.safe_range_iter_reversed(..).next().transpose()? {
1715                Some((last_key, _)) => {
1716                    let mut to = be_fix_int_ser(&(self.tag, last_key));
1717                    to.push(0);
1718                    to
1719                }
1720                None => return Ok(()),
1721            },
1722        };
1723
1724        let mut batch = self.batch();
1725        batch.schedule_delete_range_raw(&self.map, from, to)?;
1726        batch.write()
1727    }
1728
1729    fn is_empty(&self) -> bool {
1730        self.safe_iter().next().is_none()
1731    }
1732
1733    fn safe_iter(&'a self) -> DbIterator<'a, (K, V)> {
1734        Box::new(
1735            self.map
1736                .safe_iter_with_prefix(&self.tag)
1737                .map(Self::strip_tag),
1738        )
1739    }
1740
1741    fn safe_iter_with_bounds(
1742        &'a self,
1743        lower_bound: Option<K>,
1744        upper_bound: Option<K>,
1745    ) -> DbIterator<'a, (K, V)> {
1746        let range = (
1747            lower_bound.map(Bound::Included).unwrap_or(Bound::Unbounded),
1748            upper_bound.map(Bound::Excluded).unwrap_or(Bound::Unbounded),
1749        );
1750        self.safe_range_iter(range)
1751    }
1752
1753    fn safe_range_iter(&'a self, range: impl RangeBounds<K>) -> DbIterator<'a, (K, V)> {
1754        let (lower_bound, upper_bound) = prefix_iterator_bounds_with_range(&self.tag, range);
1755        Box::new(
1756            self.map
1757                .iter_forward_raw(lower_bound, upper_bound)
1758                .map(Self::strip_tag),
1759        )
1760    }
1761
1762    fn safe_range_iter_reversed(&'a self, range: impl RangeBounds<K>) -> DbIterator<'a, (K, V)> {
1763        let (lower_bound, upper_bound) = prefix_iterator_bounds_with_range(&self.tag, range);
1764        Box::new(
1765            self.map
1766                .iter_reversed_raw(lower_bound, upper_bound)
1767                .map(Self::strip_tag),
1768        )
1769    }
1770
1771    fn try_catch_up_with_primary(&self) -> Result<(), TypedStoreError> {
1772        self.map.try_catch_up_with_primary()
1773    }
1774}
1775
1776fn default_hash(value: &[u8]) -> Digest<32> {
1777    let mut hasher = fastcrypto::hash::Blake2b256::default();
1778    hasher.update(value);
1779    hasher.finalize()
1780}
1781
1782// The throughput floor is used to compute the threshold for when a batch write
1783// is considered "very slow".
1784const THROUGHPUT_FLOOR_BYTES_PER_SEC: f64 = 32.0 * 1024.0 * 1024.0; // 32 MiB/s
1785
1786// A batch write that took less than this is never reported as very slow.
1787const MIN_VERY_SLOW_BATCH_WRITE_SECS: f64 = 1.0;
1788
1789/// Returns the elapsed time above which a batch write of `batch_size_bytes` is
1790/// reported as very slow: what it would take at
1791/// [`THROUGHPUT_FLOOR_BYTES_PER_SEC`], and never less than
1792/// [`MIN_VERY_SLOW_BATCH_WRITE_SECS`].
1793fn very_slow_batch_write_threshold_secs(batch_size_bytes: usize) -> f64 {
1794    (batch_size_bytes as f64 / THROUGHPUT_FLOOR_BYTES_PER_SEC).max(MIN_VERY_SLOW_BATCH_WRITE_SECS)
1795}
1796
1797#[cfg(test)]
1798mod tests {
1799    use super::very_slow_batch_write_threshold_secs;
1800
1801    const MIB: usize = 1024 * 1024;
1802
1803    #[test]
1804    fn very_slow_batch_write_threshold_never_below_one_second() {
1805        assert_eq!(very_slow_batch_write_threshold_secs(0), 1.0);
1806        assert_eq!(very_slow_batch_write_threshold_secs(1), 1.0);
1807        assert_eq!(very_slow_batch_write_threshold_secs(MIB), 1.0);
1808        assert_eq!(very_slow_batch_write_threshold_secs(32 * MIB), 1.0);
1809    }
1810
1811    #[test]
1812    fn very_slow_batch_write_threshold_grows_with_batch_size() {
1813        assert_eq!(very_slow_batch_write_threshold_secs(48 * MIB), 1.5);
1814        assert_eq!(very_slow_batch_write_threshold_secs(64 * MIB), 2.0);
1815        assert_eq!(very_slow_batch_write_threshold_secs(320 * MIB), 10.0);
1816    }
1817}
typed_store/database.rs

typed_store/
database.rs
