spacetimedb/
replica_context.rs

1use super::database_logger::DatabaseLogger;
2use crate::db::relational_db::RelationalDB;
3use crate::error::DBError;
4use crate::messages::control_db::Database;
5use crate::subscription::module_subscription_actor::ModuleSubscriptions;
6use std::io;
7use std::ops::Deref;
8use std::sync::Arc;
9
10pub type Result<T> = anyhow::Result<T>;
11
12/// A "live" database.
13#[derive(Clone)]
14pub struct ReplicaContext {
15    pub database: Database,
16    pub replica_id: u64,
17    pub logger: Arc<DatabaseLogger>,
18    pub subscriptions: ModuleSubscriptions,
19    pub relational_db: Arc<RelationalDB>,
20}
21
22impl ReplicaContext {
23    /// The number of bytes on disk occupied by the database's durability layer.
24    ///
25    /// An in-memory database will return `Ok(0)`.
26    pub fn durability_size_on_disk(&self) -> io::Result<u64> {
27        self.relational_db.size_on_disk()
28    }
29
30    /// The size of the log file.
31    pub fn log_file_size(&self) -> std::result::Result<u64, DBError> {
32        Ok(self.logger.size()?)
33    }
34
35    /// Obtain an array which can be summed to obtain the total disk usage.
36    ///
37    /// Some sources of size-on-disk may error, in which case the corresponding array element will be None.
38    pub fn total_disk_usage(&self) -> TotalDiskUsage {
39        TotalDiskUsage {
40            durability: self.durability_size_on_disk().ok(),
41            logs: self.log_file_size().ok(),
42        }
43    }
44
45    /// The size in bytes of all of the in-memory data of the database.
46    pub fn mem_usage(&self) -> usize {
47        self.relational_db.size_in_memory()
48    }
49
50    /// Update data size stats.
51    pub fn update_gauges(&self) {
52        self.relational_db.update_data_size_metrics();
53        self.subscriptions.update_gauges();
54    }
55}
56
57impl Deref for ReplicaContext {
58    type Target = Database;
59
60    fn deref(&self) -> &Self::Target {
61        &self.database
62    }
63}
64
65#[derive(Copy, Clone, Default)]
66pub struct TotalDiskUsage {
67    pub durability: Option<u64>,
68    pub logs: Option<u64>,
69}
70
71impl TotalDiskUsage {
72    /// Returns self, but if any of the sources are None then we take it from fallback
73    pub fn or(self, fallback: TotalDiskUsage) -> Self {
74        Self {
75            durability: self.durability.or(fallback.durability),
76            logs: self.logs.or(fallback.logs),
77        }
78    }
79
80    pub fn sum(&self) -> u64 {
81        self.durability.unwrap_or(0) + self.logs.unwrap_or(0)
82    }
83}