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 the position `lower` gives within the prefix rather than at the
852    /// first key under the prefix.
853    ///
854    /// The bound value is **not** a full key: it is the entire remainder of
855    /// the key that follows `prefix` (all the trailing fields of a tuple key,
856    /// not a partial suffix). Both `prefix` and the bound value are
857    /// serialized with `be_fix_int_ser` and concatenated, mirroring how a
858    /// composite key is encoded. `Included` starts the scan at
859    /// `prefix ++ lower`, `Excluded` strictly after that position (also when
860    /// no such key exists), and `Unbounded` at the first key under the
861    /// prefix. The scan bounds never leave the prefix (excluding the
862    /// maximum remainder yields an empty scan, not the next prefix's keys),
863    /// so a paginated scan can resume from a cursor without an artificial
864    /// maximum key for the upper bound.
865    pub fn safe_iter_with_prefix_from<P, C>(
866        &self,
867        prefix: &P,
868        lower: Bound<&C>,
869    ) -> DbIterator<'_, (K, V)>
870    where
871        P: ?Sized + Serialize,
872        C: ?Sized + Serialize,
873        K: DeserializeOwned,
874        V: DeserializeOwned,
875    {
876        let (lower_bound, upper_bound) =
877            prefix_iterator_bounds_with_range::<P, C>(prefix, (lower, Bound::Unbounded));
878        self.iter_forward_raw(lower_bound, upper_bound)
879    }
880
881    /// Forward iterator over entries whose key begins with `prefix`, resuming
882    /// strictly after `cursor`, or from the first key under the prefix when
883    /// there is no cursor. `cursor` is the remainder of the key after
884    /// `prefix`, as for [`Self::safe_iter_with_prefix_from`].
885    pub fn safe_iter_with_prefix_after<P, C>(
886        &self,
887        prefix: &P,
888        cursor: Option<&C>,
889    ) -> DbIterator<'_, (K, V)>
890    where
891        P: ?Sized + Serialize,
892        C: ?Sized + Serialize,
893        K: DeserializeOwned,
894        V: DeserializeOwned,
895    {
896        self.safe_iter_with_prefix_from(prefix, cursor.map_or(Bound::Unbounded, Bound::Excluded))
897    }
898
899    /// Reverse counterpart of [`Self::safe_iter_with_prefix`]: the matching
900    /// entries in descending key order (e.g. for "latest entry under a
901    /// prefix").
902    pub fn safe_iter_with_prefix_reversed<P>(&self, prefix: &P) -> DbIterator<'_, (K, V)>
903    where
904        P: ?Sized + Serialize,
905        K: DeserializeOwned,
906        V: DeserializeOwned,
907    {
908        let (lower_bound, upper_bound) = prefix_iterator_bounds(prefix);
909        self.iter_reversed_raw(lower_bound, upper_bound)
910    }
911}
912
913/// Provides a mutable struct to form a collection of database write operations,
914/// and execute them.
915///
916/// Batching write and delete operations is faster than performing them one by
917/// one and ensures their atomicity,  ie. they are all written or none is.
918/// This is also true of operations across column families in the same database.
919///
920/// Serializations / Deserialization, and naming of column families is performed
921/// by passing a DBMap<K,V> with each operation.
922///
923/// ```
924/// use core::fmt::Error;
925/// use std::sync::Arc;
926///
927/// use prometheus_filtered::Registry;
928/// use tempfile::tempdir;
929/// use typed_store::{Map, metrics::DBMetrics, rocks::*};
930///
931/// #[tokio::main(flavor = "current_thread")]
932/// async fn main() -> Result<(), Error> {
933///     let rocks = open_cf_opts(
934///         tempfile::tempdir().unwrap(),
935///         None,
936///         MetricConf::default(),
937///         &[
938///             ("First_CF", rocksdb::Options::default()),
939///             ("Second_CF", rocksdb::Options::default()),
940///         ],
941///     )
942///     .unwrap();
943///
944///     let db_cf_1 = DBMap::reopen(
945///         &rocks,
946///         Some("First_CF"),
947///         &ReadWriteOptions::default(),
948///         false,
949///     )
950///     .expect("Failed to open storage");
951///     let keys_vals_1 = (1..100).map(|i| (i, i.to_string()));
952///
953///     let db_cf_2 = DBMap::reopen(
954///         &rocks,
955///         Some("Second_CF"),
956///         &ReadWriteOptions::default(),
957///         false,
958///     )
959///     .expect("Failed to open storage");
960///     let keys_vals_2 = (1000..1100).map(|i| (i, i.to_string()));
961///
962///     let mut batch = db_cf_1.batch();
963///     batch
964///         .insert_batch(&db_cf_1, keys_vals_1.clone())
965///         .expect("Failed to batch insert")
966///         .insert_batch(&db_cf_2, keys_vals_2.clone())
967///         .expect("Failed to batch insert");
968///
969///     let _ = batch.write().expect("Failed to execute batch");
970///     for (k, v) in keys_vals_1 {
971///         let val = db_cf_1.get(&k).expect("Failed to get inserted key");
972///         assert_eq!(Some(v), val);
973///     }
974///
975///     for (k, v) in keys_vals_2 {
976///         let val = db_cf_2.get(&k).expect("Failed to get inserted key");
977///         assert_eq!(Some(v), val);
978///     }
979///     Ok(())
980/// }
981/// ```
982pub struct DBBatch {
983    database: Arc<Database>,
984    batch: StorageWriteBatch,
985    db_metrics: Arc<DBMetrics>,
986    write_sample_interval: SamplingInterval,
987}
988
989impl DBBatch {
990    /// Create a new batch associated with a DB reference.
991    ///
992    /// Use `open_cf` to get the DB reference or an existing open database.
993    pub fn new(
994        dbref: &Arc<Database>,
995        batch: StorageWriteBatch,
996        db_metrics: &Arc<DBMetrics>,
997        write_sample_interval: &SamplingInterval,
998    ) -> Self {
999        DBBatch {
1000            database: dbref.clone(),
1001            batch,
1002            db_metrics: db_metrics.clone(),
1003            write_sample_interval: write_sample_interval.clone(),
1004        }
1005    }
1006
1007    /// Consume the batch and write its operations to the database
1008    #[instrument(level = "trace", skip_all, err)]
1009    pub fn write(self) -> Result<(), TypedStoreError> {
1010        self.write_opt(&rocksdb::WriteOptions::default())
1011    }
1012
1013    /// Consume the batch and write its operations to the database with custom
1014    /// write options
1015    #[instrument(level = "trace", skip_all, err)]
1016    pub fn write_opt(self, write_options: &rocksdb::WriteOptions) -> Result<(), TypedStoreError> {
1017        let db_name = self.database.db_name();
1018        let timer = self
1019            .db_metrics
1020            .op_metrics
1021            .rocksdb_batch_commit_latency_seconds
1022            .with_label_values(&[&db_name])
1023            .start_timer();
1024        let batch_size_bytes = self.size_in_bytes();
1025
1026        let perf_ctx = if self.write_sample_interval.sample() {
1027            Some(RocksDBPerfContext)
1028        } else {
1029            None
1030        };
1031        self.database.write_opt(self.batch, write_options)?;
1032        self.db_metrics
1033            .op_metrics
1034            .rocksdb_batch_commit_bytes
1035            .with_label_values(&[&db_name])
1036            .observe(batch_size_bytes as f64);
1037
1038        if perf_ctx.is_some() {
1039            self.db_metrics
1040                .write_perf_ctx_metrics
1041                .report_metrics(&db_name);
1042        }
1043        let elapsed_secs = timer.stop_and_record();
1044        let threshold_secs = very_slow_batch_write_threshold_secs(batch_size_bytes);
1045        if elapsed_secs > threshold_secs {
1046            warn!(
1047                elapsed_secs,
1048                threshold_secs,
1049                batch_size_bytes,
1050                ?db_name,
1051                "very slow batch write"
1052            );
1053            self.db_metrics
1054                .op_metrics
1055                .rocksdb_very_slow_batch_writes_count
1056                .with_label_values(&[&db_name])
1057                .inc();
1058            self.db_metrics
1059                .op_metrics
1060                .rocksdb_very_slow_batch_writes_duration_ms
1061                .with_label_values(&[&db_name])
1062                .inc_by((elapsed_secs * 1000.0) as u64);
1063        }
1064        Ok(())
1065    }
1066
1067    pub fn size_in_bytes(&self) -> usize {
1068        match self.batch {
1069            StorageWriteBatch::Rocks(ref b) => b.size_in_bytes(),
1070            StorageWriteBatch::InMemory(_) => 0,
1071        }
1072    }
1073
1074    pub fn delete_batch<J: Borrow<K>, K: Serialize, V>(
1075        &mut self,
1076        db: &DBMap<K, V>,
1077        purged_vals: impl IntoIterator<Item = J>,
1078    ) -> Result<(), TypedStoreError> {
1079        self.delete_batch_raw_keys(
1080            db,
1081            purged_vals
1082                .into_iter()
1083                .map(|key| be_fix_int_ser(key.borrow())),
1084        )
1085    }
1086
1087    /// [`Self::delete_batch`] for keys already in the column family's
1088    /// `be_fix_int_ser` encoding.
1089    fn delete_batch_raw_keys<K, V>(
1090        &mut self,
1091        db: &DBMap<K, V>,
1092        purged_vals: impl IntoIterator<Item = Vec<u8>>,
1093    ) -> Result<(), TypedStoreError> {
1094        if !Arc::ptr_eq(&db.db, &self.database) {
1095            return Err(TypedStoreError::CrossDBBatch);
1096        }
1097
1098        purged_vals
1099            .into_iter()
1100            .try_for_each::<_, Result<_, TypedStoreError>>(|k_buf| {
1101                match (&mut self.batch, &db.column_family) {
1102                    (StorageWriteBatch::Rocks(b), ColumnFamily::Rocks(name)) => {
1103                        b.delete_cf(&rocks_cf_from_db(&self.database, name)?, k_buf)
1104                    }
1105                    (StorageWriteBatch::InMemory(b), ColumnFamily::InMemory(name)) => {
1106                        b.delete_cf(name, k_buf)
1107                    }
1108                    _ => Err(TypedStoreError::RocksDB(
1109                        "typed store invariant violation".to_string(),
1110                    ))?,
1111                }
1112                Ok(())
1113            })?;
1114        Ok(())
1115    }
1116
1117    /// Deletes a range of keys between `from` (inclusive) and `to`
1118    /// (non-inclusive) by writing a range delete tombstone in the db map.
1119    /// The effect of this write is visible immediately, i.e. you won't see
1120    /// old values when you do a lookup or scan.
1121    pub fn schedule_delete_range<K: Serialize, V>(
1122        &mut self,
1123        db: &DBMap<K, V>,
1124        from: &K,
1125        to: &K,
1126    ) -> Result<(), TypedStoreError> {
1127        self.schedule_delete_range_raw(db, be_fix_int_ser(from), be_fix_int_ser(to))
1128    }
1129
1130    /// [`Self::schedule_delete_range`] for a range whose bounds are not keys,
1131    /// such as a whole key prefix. The bounds are raw bytes ordered like the
1132    /// map's keys, i.e. in the `be_fix_int_ser` encoding; `from` is inclusive
1133    /// and `to` exclusive.
1134    fn schedule_delete_range_raw<K, V>(
1135        &mut self,
1136        db: &DBMap<K, V>,
1137        from: Vec<u8>,
1138        to: Vec<u8>,
1139    ) -> Result<(), TypedStoreError> {
1140        if !Arc::ptr_eq(&db.db, &self.database) {
1141            return Err(TypedStoreError::CrossDBBatch);
1142        }
1143
1144        if let StorageWriteBatch::Rocks(b) = &mut self.batch {
1145            b.delete_range_cf(&rocks_cf_from_db(&self.database, db.cf_name())?, from, to);
1146        }
1147        Ok(())
1148    }
1149
1150    /// inserts a range of (key, value) pairs given as an iterator
1151    pub fn insert_batch<J: Borrow<K>, K: Serialize, U: Borrow<V>, V: Serialize>(
1152        &mut self,
1153        db: &DBMap<K, V>,
1154        new_vals: impl IntoIterator<Item = (J, U)>,
1155    ) -> Result<&mut Self, TypedStoreError> {
1156        self.insert_batch_raw_keys(
1157            db,
1158            new_vals
1159                .into_iter()
1160                .map(|(key, value)| (be_fix_int_ser(key.borrow()), value)),
1161        )
1162    }
1163
1164    /// [`Self::insert_batch`] for keys already in the column family's
1165    /// `be_fix_int_ser` encoding.
1166    fn insert_batch_raw_keys<K, U: Borrow<V>, V: Serialize>(
1167        &mut self,
1168        db: &DBMap<K, V>,
1169        new_vals: impl IntoIterator<Item = (Vec<u8>, U)>,
1170    ) -> Result<&mut Self, TypedStoreError> {
1171        if !Arc::ptr_eq(&db.db, &self.database) {
1172            return Err(TypedStoreError::CrossDBBatch);
1173        }
1174        let mut total = 0usize;
1175        new_vals
1176            .into_iter()
1177            .try_for_each::<_, Result<_, TypedStoreError>>(|(k_buf, v)| {
1178                let v_buf = bcs::to_bytes(v.borrow()).map_err(typed_store_err_from_bcs_err)?;
1179                total += k_buf.len() + v_buf.len();
1180                if db.opts.log_value_hash {
1181                    let key_hash = default_hash(&k_buf);
1182                    let value_hash = default_hash(&v_buf);
1183                    debug!(
1184                        "Insert to DB table: {:?}, key_hash: {:?}, value_hash: {:?}",
1185                        db.cf_name(),
1186                        key_hash,
1187                        value_hash
1188                    );
1189                }
1190                match (&mut self.batch, &db.column_family) {
1191                    (StorageWriteBatch::Rocks(b), ColumnFamily::Rocks(name)) => {
1192                        b.put_cf(&rocks_cf_from_db(&self.database, name)?, k_buf, v_buf)
1193                    }
1194                    (StorageWriteBatch::InMemory(b), ColumnFamily::InMemory(name)) => {
1195                        b.put_cf(name, k_buf, v_buf)
1196                    }
1197                    _ => Err(TypedStoreError::RocksDB(
1198                        "typed store invariant violation".to_string(),
1199                    ))?,
1200                }
1201                Ok(())
1202            })?;
1203        self.db_metrics
1204            .op_metrics
1205            .rocksdb_batch_put_bytes
1206            .with_label_values(&[db.cf_name()])
1207            .observe(total as f64);
1208        Ok(self)
1209    }
1210
1211    /// [`Self::insert_batch`] for a map sharing its column family by tag.
1212    pub fn insert_batch_tagged<J: Borrow<K>, K: Serialize, U: Borrow<V>, V: Serialize>(
1213        &mut self,
1214        map: &TaggedDBMap<K, V>,
1215        new_vals: impl IntoIterator<Item = (J, U)>,
1216    ) -> Result<&mut Self, TypedStoreError> {
1217        self.insert_batch_raw_keys(
1218            &map.map,
1219            new_vals
1220                .into_iter()
1221                .map(|(key, value)| (be_fix_int_ser(&(map.tag, key.borrow())), value)),
1222        )
1223    }
1224
1225    /// [`Self::delete_batch`] for a map sharing its column family by tag.
1226    pub fn delete_batch_tagged<J: Borrow<K>, K: Serialize, V>(
1227        &mut self,
1228        map: &TaggedDBMap<K, V>,
1229        purged_vals: impl IntoIterator<Item = J>,
1230    ) -> Result<(), TypedStoreError> {
1231        self.delete_batch_raw_keys(
1232            &map.map,
1233            purged_vals
1234                .into_iter()
1235                .map(|key| be_fix_int_ser(&(map.tag, key.borrow()))),
1236        )
1237    }
1238}
1239
1240impl<K, V> DBMap<K, V> {
1241    #[instrument(level = "trace", skip_all, err)]
1242    pub(crate) fn contains_raw_key(&self, key_buf: Vec<u8>) -> Result<bool, TypedStoreError> {
1243        let readopts = self.opts.readopts();
1244        Ok(self
1245            .db
1246            .key_may_exist_cf(&self.column_family, &key_buf, &readopts)
1247            && self
1248                .db
1249                .get(&self.column_family, &key_buf, &readopts)?
1250                .is_some())
1251    }
1252
1253    #[instrument(level = "trace", skip_all, err)]
1254    pub(crate) fn multi_contains_raw_keys(
1255        &self,
1256        keys: impl IntoIterator<Item = Vec<u8>>,
1257    ) -> Result<Vec<bool>, TypedStoreError> {
1258        let values = self.multi_get_pinned(keys)?;
1259        Ok(values.into_iter().map(|v| v.is_some()).collect())
1260    }
1261
1262    #[instrument(level = "trace", skip_all, err)]
1263    pub(crate) fn get_raw_key(&self, key_buf: Vec<u8>) -> Result<Option<V>, TypedStoreError>
1264    where
1265        V: DeserializeOwned,
1266    {
1267        let _timer = self
1268            .db_metrics
1269            .op_metrics
1270            .rocksdb_get_latency_seconds
1271            .with_label_values(&[self.cf_name()])
1272            .start_timer();
1273        let perf_ctx = if self.get_sample_interval.sample() {
1274            Some(RocksDBPerfContext)
1275        } else {
1276            None
1277        };
1278        let res = self
1279            .db
1280            .get(&self.column_family, &key_buf, &self.opts.readopts())?;
1281        self.db_metrics
1282            .op_metrics
1283            .rocksdb_get_bytes
1284            .with_label_values(&[self.cf_name()])
1285            .observe(res.as_ref().map_or(0.0, |v| v.len() as f64));
1286        if perf_ctx.is_some() {
1287            self.db_metrics
1288                .read_perf_ctx_metrics
1289                .report_metrics(self.cf_name());
1290        }
1291        match res {
1292            Some(data) => {
1293                let value = bcs::from_bytes(&data).map_err(typed_store_err_from_bcs_err);
1294                if value.is_err() {
1295                    let key_hash = default_hash(&key_buf);
1296                    let value_hash = default_hash(&data);
1297                    debug_fatal!(
1298                        "Failed to deserialize value from DB table {:?}, key_hash: {:?}, value_hash: {:?}, error: {:?}",
1299                        self.cf_name(),
1300                        key_hash,
1301                        value_hash,
1302                        value.as_ref().err().unwrap()
1303                    );
1304                }
1305                Ok(Some(value?))
1306            }
1307            None => Ok(None),
1308        }
1309    }
1310
1311    #[instrument(level = "trace", skip_all, err)]
1312    pub(crate) fn insert_raw_key(&self, key_buf: Vec<u8>, value: &V) -> Result<(), TypedStoreError>
1313    where
1314        V: Serialize,
1315    {
1316        let timer = self
1317            .db_metrics
1318            .op_metrics
1319            .rocksdb_put_latency_seconds
1320            .with_label_values(&[self.cf_name()])
1321            .start_timer();
1322        let perf_ctx = if self.write_sample_interval.sample() {
1323            Some(RocksDBPerfContext)
1324        } else {
1325            None
1326        };
1327        let value_buf = bcs::to_bytes(value).map_err(typed_store_err_from_bcs_err)?;
1328        self.db_metrics
1329            .op_metrics
1330            .rocksdb_put_bytes
1331            .with_label_values(&[self.cf_name()])
1332            .observe((key_buf.len() + value_buf.len()) as f64);
1333        if perf_ctx.is_some() {
1334            self.db_metrics
1335                .write_perf_ctx_metrics
1336                .report_metrics(self.cf_name());
1337        }
1338        self.db.put_cf(&self.column_family, key_buf, value_buf)?;
1339
1340        let elapsed_secs = timer.stop_and_record();
1341        if elapsed_secs > 1.0 {
1342            warn!(elapsed_secs, cf = ?self.cf_name(), "very slow insert");
1343            self.db_metrics
1344                .op_metrics
1345                .rocksdb_very_slow_puts_count
1346                .with_label_values(&[self.cf_name()])
1347                .inc();
1348            self.db_metrics
1349                .op_metrics
1350                .rocksdb_very_slow_puts_duration_ms
1351                .with_label_values(&[self.cf_name()])
1352                .inc_by((elapsed_secs * 1000.0) as u64);
1353        }
1354
1355        Ok(())
1356    }
1357
1358    #[instrument(level = "trace", skip_all, err)]
1359    pub(crate) fn remove_raw_key(&self, key_buf: Vec<u8>) -> Result<(), TypedStoreError> {
1360        let _timer = self
1361            .db_metrics
1362            .op_metrics
1363            .rocksdb_delete_latency_seconds
1364            .with_label_values(&[self.cf_name()])
1365            .start_timer();
1366        let perf_ctx = if self.write_sample_interval.sample() {
1367            Some(RocksDBPerfContext)
1368        } else {
1369            None
1370        };
1371        self.db.delete_cf(&self.column_family, key_buf)?;
1372        self.db_metrics
1373            .op_metrics
1374            .rocksdb_deletes
1375            .with_label_values(&[self.cf_name()])
1376            .inc();
1377        if perf_ctx.is_some() {
1378            self.db_metrics
1379                .write_perf_ctx_metrics
1380                .report_metrics(self.cf_name());
1381        }
1382        Ok(())
1383    }
1384
1385    #[instrument(level = "trace", skip_all, err)]
1386    pub(crate) fn multi_get_raw_keys(
1387        &self,
1388        keys: impl IntoIterator<Item = Vec<u8>>,
1389    ) -> Result<Vec<Option<V>>, TypedStoreError>
1390    where
1391        V: DeserializeOwned,
1392    {
1393        let results = self.multi_get_pinned(keys)?;
1394        let values_parsed: Result<Vec<_>, TypedStoreError> = results
1395            .into_iter()
1396            .map(|value_byte| match value_byte {
1397                Some(data) => Ok(Some(
1398                    bcs::from_bytes(&data).map_err(typed_store_err_from_bcs_err)?,
1399                )),
1400                None => Ok(None),
1401            })
1402            .collect();
1403
1404        values_parsed
1405    }
1406}
1407
1408impl<'a, K, V> Map<'a, K, V> for DBMap<K, V>
1409where
1410    K: Serialize + DeserializeOwned,
1411    V: Serialize + DeserializeOwned,
1412{
1413    type Error = TypedStoreError;
1414
1415    fn contains_key(&self, key: &K) -> Result<bool, TypedStoreError> {
1416        self.contains_raw_key(be_fix_int_ser(key))
1417    }
1418
1419    fn multi_contains_keys<J>(
1420        &self,
1421        keys: impl IntoIterator<Item = J>,
1422    ) -> Result<Vec<bool>, Self::Error>
1423    where
1424        J: Borrow<K>,
1425    {
1426        self.multi_contains_raw_keys(keys.into_iter().map(|k| be_fix_int_ser(k.borrow())))
1427    }
1428
1429    fn get(&self, key: &K) -> Result<Option<V>, TypedStoreError> {
1430        self.get_raw_key(be_fix_int_ser(key))
1431    }
1432
1433    fn insert(&self, key: &K, value: &V) -> Result<(), TypedStoreError> {
1434        self.insert_raw_key(be_fix_int_ser(key), value)
1435    }
1436
1437    fn remove(&self, key: &K) -> Result<(), TypedStoreError> {
1438        self.remove_raw_key(be_fix_int_ser(key))
1439    }
1440
1441    /// Writes a range delete tombstone to delete all entries in the db map.
1442    /// The effect of this write is visible immediately, i.e. you won't see
1443    /// old values when you do a lookup or scan.
1444    #[instrument(level = "trace", skip_all, err)]
1445    fn schedule_delete_all(&self) -> Result<(), TypedStoreError> {
1446        let Some(last_key) = self
1447            .safe_range_iter_reversed(..)
1448            .next()
1449            .transpose()?
1450            .map(|(k, _v)| k)
1451        else {
1452            return Ok(());
1453        };
1454        // The tombstone excludes its upper bound, so it has to end past the last key:
1455        // appending a zero byte gives the first key sorting above it.
1456        let mut to = be_fix_int_ser(&last_key);
1457        to.push(0);
1458        // The empty key sorts below every key, so the range starts before the first
1459        // entry.
1460        let mut batch = self.batch();
1461        batch.schedule_delete_range_raw(self, Vec::new(), to)?;
1462        batch.write()
1463    }
1464
1465    fn is_empty(&self) -> bool {
1466        self.safe_iter().next().is_none()
1467    }
1468
1469    fn safe_iter(&'a self) -> DbIterator<'a, (K, V)> {
1470        match &self.db.storage {
1471            Storage::Rocks(db) => {
1472                let db_iter = db.underlying.raw_iterator_cf_opt(
1473                    &rocks_cf(db, self.column_family.name()),
1474                    self.opts.readopts(),
1475                );
1476                let (_timer, bytes_scanned, keys_scanned, _perf_ctx) = self.create_iter_context();
1477                Box::new(SafeIter::new(
1478                    self.cf_name().to_string(),
1479                    db_iter,
1480                    _timer,
1481                    _perf_ctx,
1482                    bytes_scanned,
1483                    keys_scanned,
1484                    Some(self.db_metrics.clone()),
1485                ))
1486            }
1487            Storage::InMemory(db) => db.iterator(self.column_family.name(), None, None, false),
1488        }
1489    }
1490
1491    fn safe_iter_with_bounds(
1492        &'a self,
1493        lower_bound: Option<K>,
1494        upper_bound: Option<K>,
1495    ) -> DbIterator<'a, (K, V)> {
1496        let range = (
1497            lower_bound.map(Bound::Included).unwrap_or(Bound::Unbounded),
1498            upper_bound.map(Bound::Excluded).unwrap_or(Bound::Unbounded),
1499        );
1500        self.safe_range_iter(range)
1501    }
1502
1503    fn safe_range_iter(&'a self, range: impl RangeBounds<K>) -> DbIterator<'a, (K, V)> {
1504        let (lower_bound, upper_bound) = iterator_bounds_with_range(range);
1505        self.iter_forward_raw(lower_bound, upper_bound)
1506    }
1507
1508    /// Both directions derive their bounds from the same
1509    /// `iterator_bounds_with_range`, so they are guaranteed to cover the
1510    /// identical set of keys regardless of the bound inclusivity.
1511    fn safe_range_iter_reversed(&'a self, range: impl RangeBounds<K>) -> DbIterator<'a, (K, V)> {
1512        let (lower_bound, upper_bound) = iterator_bounds_with_range(range);
1513        self.iter_reversed_raw(lower_bound, upper_bound)
1514    }
1515
1516    /// Returns a vector of values corresponding to the keys provided.
1517    fn multi_get<J>(
1518        &self,
1519        keys: impl IntoIterator<Item = J>,
1520    ) -> Result<Vec<Option<V>>, TypedStoreError>
1521    where
1522        J: Borrow<K>,
1523    {
1524        self.multi_get_raw_keys(keys.into_iter().map(|k| be_fix_int_ser(k.borrow())))
1525    }
1526
1527    /// Convenience method for batch insertion
1528    #[instrument(level = "trace", skip_all, err)]
1529    fn multi_insert<J, U>(
1530        &self,
1531        key_val_pairs: impl IntoIterator<Item = (J, U)>,
1532    ) -> Result<(), Self::Error>
1533    where
1534        J: Borrow<K>,
1535        U: Borrow<V>,
1536    {
1537        let mut batch = self.batch();
1538        batch.insert_batch(self, key_val_pairs)?;
1539        batch.write()
1540    }
1541
1542    /// Convenience method for batch removal
1543    #[instrument(level = "trace", skip_all, err)]
1544    fn multi_remove<J>(&self, keys: impl IntoIterator<Item = J>) -> Result<(), Self::Error>
1545    where
1546        J: Borrow<K>,
1547    {
1548        let mut batch = self.batch();
1549        batch.delete_batch(self, keys)?;
1550        batch.write()
1551    }
1552
1553    /// Try to catch up with primary when running as secondary
1554    #[instrument(level = "trace", skip_all, err)]
1555    fn try_catch_up_with_primary(&self) -> Result<(), Self::Error> {
1556        self.db.try_catch_up_with_primary()
1557    }
1558}
1559
1560/// A typed map sharing one column family with other typed maps,
1561/// distinguished by a tag byte prefixed to every key: the big-endian fixint
1562/// serialization of `(tag, key)` keeps each map a contiguous key range of
1563/// the column family. Useful for sets of tables that are always created and
1564/// dropped together — one column family instead of one per table keeps the
1565/// column-family and SST-file counts low.
1566///
1567/// Every read is scoped to the map's tag: full scans iterate the tag's key
1568/// prefix, and ranges are two-sided within it, so rows of other tags never
1569/// surface and their (differently typed) values are never deserialized.
1570///
1571/// Tags of existing maps must never change or be reused, or old rows would
1572/// be read under the wrong types.
1573///
1574/// Metrics are labelled by column family, so all the maps of one column
1575/// family report together, with no way to tell the tags apart.
1576///
1577/// # Examples
1578///
1579/// ```
1580/// use tempfile::tempdir;
1581/// use typed_store::{Map, rocks::*};
1582///
1583/// #[tokio::main(flavor = "current_thread")]
1584/// async fn main() {
1585///     let db = open_cf_opts(
1586///         tempdir().unwrap(),
1587///         None,
1588///         MetricConf::default(),
1589///         &[("shared", rocksdb::Options::default())],
1590///     )
1591///     .unwrap();
1592///
1593///     let numbers: TaggedDBMap<u32, String> =
1594///         TaggedDBMap::reopen(&db, "shared", 0, &ReadWriteOptions::default(), false).unwrap();
1595///     let words: TaggedDBMap<String, u64> =
1596///         TaggedDBMap::reopen(&db, "shared", 1, &ReadWriteOptions::default(), false).unwrap();
1597///
1598///     let mut batch = numbers.batch();
1599///     batch
1600///         .insert_batch_tagged(&numbers, [(1, "one".to_string())])
1601///         .unwrap()
1602///         .insert_batch_tagged(&words, [("one".to_string(), 1)])
1603///         .unwrap();
1604///     batch.write().unwrap();
1605///
1606///     assert_eq!(numbers.get(&1).unwrap(), Some("one".to_string()));
1607///     assert_eq!(words.get(&"one".to_string()).unwrap(), Some(1));
1608///     // Each map sees only its own row.
1609///     assert_eq!(numbers.safe_iter().count(), 1);
1610/// }
1611/// ```
1612pub struct TaggedDBMap<K, V> {
1613    tag: u8,
1614    map: DBMap<(u8, K), V>,
1615}
1616
1617impl<K, V> TaggedDBMap<K, V> {
1618    /// Drops the tag prefixed to a key of the underlying map, turning its row
1619    /// into a row of this map.
1620    fn strip_tag(row: Result<((u8, K), V), TypedStoreError>) -> Result<(K, V), TypedStoreError> {
1621        row.map(|((_, key), value)| (key, value))
1622    }
1623
1624    /// Reopens an open database as a tagged map operating under `cf_name`.
1625    /// See [`DBMap::reopen`] for the remaining parameters.
1626    pub fn reopen(
1627        db: &Arc<Database>,
1628        cf_name: &str,
1629        tag: u8,
1630        rw_options: &ReadWriteOptions,
1631        skip_metrics_reporting: bool,
1632    ) -> Result<Self, TypedStoreError> {
1633        Ok(Self {
1634            tag,
1635            map: DBMap::reopen(db, Some(cf_name), rw_options, skip_metrics_reporting)?,
1636        })
1637    }
1638
1639    /// Create a batch associated with the map's database.
1640    pub fn batch(&self) -> DBBatch {
1641        self.map.batch()
1642    }
1643}
1644
1645impl<'a, K, V> Map<'a, K, V> for TaggedDBMap<K, V>
1646where
1647    K: Serialize + DeserializeOwned,
1648    V: Serialize + DeserializeOwned,
1649{
1650    type Error = TypedStoreError;
1651
1652    fn contains_key(&self, key: &K) -> Result<bool, TypedStoreError> {
1653        self.map.contains_raw_key(be_fix_int_ser(&(self.tag, key)))
1654    }
1655
1656    fn multi_contains_keys<J>(
1657        &self,
1658        keys: impl IntoIterator<Item = J>,
1659    ) -> Result<Vec<bool>, TypedStoreError>
1660    where
1661        J: Borrow<K>,
1662    {
1663        self.map.multi_contains_raw_keys(
1664            keys.into_iter()
1665                .map(|key| be_fix_int_ser(&(self.tag, key.borrow()))),
1666        )
1667    }
1668
1669    fn get(&self, key: &K) -> Result<Option<V>, TypedStoreError> {
1670        self.map.get_raw_key(be_fix_int_ser(&(self.tag, key)))
1671    }
1672
1673    fn multi_get<J>(
1674        &self,
1675        keys: impl IntoIterator<Item = J>,
1676    ) -> Result<Vec<Option<V>>, TypedStoreError>
1677    where
1678        J: Borrow<K>,
1679    {
1680        self.map.multi_get_raw_keys(
1681            keys.into_iter()
1682                .map(|key| be_fix_int_ser(&(self.tag, key.borrow()))),
1683        )
1684    }
1685
1686    fn insert(&self, key: &K, value: &V) -> Result<(), TypedStoreError> {
1687        self.map
1688            .insert_raw_key(be_fix_int_ser(&(self.tag, key)), value)
1689    }
1690
1691    fn remove(&self, key: &K) -> Result<(), TypedStoreError> {
1692        self.map.remove_raw_key(be_fix_int_ser(&(self.tag, key)))
1693    }
1694
1695    #[instrument(level = "trace", skip_all, err)]
1696    fn multi_insert<J, U>(
1697        &self,
1698        key_val_pairs: impl IntoIterator<Item = (J, U)>,
1699    ) -> Result<(), TypedStoreError>
1700    where
1701        J: Borrow<K>,
1702        U: Borrow<V>,
1703    {
1704        let mut batch = self.batch();
1705        batch.insert_batch_tagged(self, key_val_pairs)?;
1706        batch.write()
1707    }
1708
1709    #[instrument(level = "trace", skip_all, err)]
1710    fn multi_remove<J>(&self, keys: impl IntoIterator<Item = J>) -> Result<(), TypedStoreError>
1711    where
1712        J: Borrow<K>,
1713    {
1714        let mut batch = self.batch();
1715        batch.delete_batch_tagged(self, keys)?;
1716        batch.write()
1717    }
1718
1719    /// Deletes every entry of the map with a single range tombstone over the
1720    /// tag's key range, leaving the other tags of the column family untouched.
1721    ///
1722    /// The entries stay on disk until compaction, and until then iteration
1723    /// is slowed because of them.
1724    #[instrument(level = "trace", skip_all, err)]
1725    fn schedule_delete_all(&self) -> Result<(), TypedStoreError> {
1726        let from = be_fix_int_ser(&self.tag);
1727        let to = match prefix_iterator_bounds(&self.tag).1 {
1728            Some(to) => {
1729                if self.is_empty() {
1730                    return Ok(());
1731                }
1732                to
1733            }
1734            // The maximum tag has no following tag to bound the range
1735            // against, so end it just past the last entry. That read is why a
1736            // row written in between survives the clear, and why a value that
1737            // does not deserialize fails it.
1738            None => match self.safe_range_iter_reversed(..).next().transpose()? {
1739                Some((last_key, _)) => {
1740                    let mut to = be_fix_int_ser(&(self.tag, last_key));
1741                    to.push(0);
1742                    to
1743                }
1744                None => return Ok(()),
1745            },
1746        };
1747
1748        let mut batch = self.batch();
1749        batch.schedule_delete_range_raw(&self.map, from, to)?;
1750        batch.write()
1751    }
1752
1753    fn is_empty(&self) -> bool {
1754        self.safe_iter().next().is_none()
1755    }
1756
1757    fn safe_iter(&'a self) -> DbIterator<'a, (K, V)> {
1758        Box::new(
1759            self.map
1760                .safe_iter_with_prefix(&self.tag)
1761                .map(Self::strip_tag),
1762        )
1763    }
1764
1765    fn safe_iter_with_bounds(
1766        &'a self,
1767        lower_bound: Option<K>,
1768        upper_bound: Option<K>,
1769    ) -> DbIterator<'a, (K, V)> {
1770        let range = (
1771            lower_bound.map(Bound::Included).unwrap_or(Bound::Unbounded),
1772            upper_bound.map(Bound::Excluded).unwrap_or(Bound::Unbounded),
1773        );
1774        self.safe_range_iter(range)
1775    }
1776
1777    fn safe_range_iter(&'a self, range: impl RangeBounds<K>) -> DbIterator<'a, (K, V)> {
1778        let (lower_bound, upper_bound) = prefix_iterator_bounds_with_range(&self.tag, range);
1779        Box::new(
1780            self.map
1781                .iter_forward_raw(lower_bound, upper_bound)
1782                .map(Self::strip_tag),
1783        )
1784    }
1785
1786    fn safe_range_iter_reversed(&'a self, range: impl RangeBounds<K>) -> DbIterator<'a, (K, V)> {
1787        let (lower_bound, upper_bound) = prefix_iterator_bounds_with_range(&self.tag, range);
1788        Box::new(
1789            self.map
1790                .iter_reversed_raw(lower_bound, upper_bound)
1791                .map(Self::strip_tag),
1792        )
1793    }
1794
1795    fn try_catch_up_with_primary(&self) -> Result<(), TypedStoreError> {
1796        self.map.try_catch_up_with_primary()
1797    }
1798}
1799
1800fn default_hash(value: &[u8]) -> Digest<32> {
1801    let mut hasher = fastcrypto::hash::Blake2b256::default();
1802    hasher.update(value);
1803    hasher.finalize()
1804}
1805
1806// The throughput floor is used to compute the threshold for when a batch write
1807// is considered "very slow".
1808const THROUGHPUT_FLOOR_BYTES_PER_SEC: f64 = 32.0 * 1024.0 * 1024.0; // 32 MiB/s
1809
1810// A batch write that took less than this is never reported as very slow.
1811const MIN_VERY_SLOW_BATCH_WRITE_SECS: f64 = 1.0;
1812
1813/// Returns the elapsed time above which a batch write of `batch_size_bytes` is
1814/// reported as very slow: what it would take at
1815/// [`THROUGHPUT_FLOOR_BYTES_PER_SEC`], and never less than
1816/// [`MIN_VERY_SLOW_BATCH_WRITE_SECS`].
1817fn very_slow_batch_write_threshold_secs(batch_size_bytes: usize) -> f64 {
1818    (batch_size_bytes as f64 / THROUGHPUT_FLOOR_BYTES_PER_SEC).max(MIN_VERY_SLOW_BATCH_WRITE_SECS)
1819}
1820
1821#[cfg(test)]
1822mod tests {
1823    use super::very_slow_batch_write_threshold_secs;
1824
1825    const MIB: usize = 1024 * 1024;
1826
1827    #[test]
1828    fn very_slow_batch_write_threshold_never_below_one_second() {
1829        assert_eq!(very_slow_batch_write_threshold_secs(0), 1.0);
1830        assert_eq!(very_slow_batch_write_threshold_secs(1), 1.0);
1831        assert_eq!(very_slow_batch_write_threshold_secs(MIB), 1.0);
1832        assert_eq!(very_slow_batch_write_threshold_secs(32 * MIB), 1.0);
1833    }
1834
1835    #[test]
1836    fn very_slow_batch_write_threshold_grows_with_batch_size() {
1837        assert_eq!(very_slow_batch_write_threshold_secs(48 * MIB), 1.5);
1838        assert_eq!(very_slow_batch_write_threshold_secs(64 * MIB), 2.0);
1839        assert_eq!(very_slow_batch_write_threshold_secs(320 * MIB), 10.0);
1840    }
1841}
typed_store/database.rs

typed_store/
database.rs
