tuwunel_database/engine/
context.rs

1use std::{
2	collections::BTreeMap,
3	fs::remove_dir_all,
4	path::Path,
5	sync::{Arc, Mutex},
6};
7
8use rocksdb::{Cache, LruCacheOptions};
9use tuwunel_core::{
10	Result, Server, debug,
11	utils::{math::usize_from_f64, result::LogErr},
12};
13
14use super::env::Env;
15use crate::pool::Pool;
16
17/// One block-cache pool, plus the column families participating in it.
18///
19/// Pools may be shared by multiple CFs (`SHARED_POOL`, symmetric
20/// `CacheDisp::SharedWith` pairs); the participant list lets the admin
21/// surface name them.
22pub(crate) struct ColCache {
23	pub(crate) cache: Cache,
24
25	pub(crate) participants: Vec<&'static str>,
26}
27
28/// Holds shared resources that must outlive an opened database.
29///
30/// The worker pool, caches, server handle, and RocksDB environment are prepared
31/// before the database opens. Keeping them in one shared context gives every
32/// engine component a common owner for those resources.
33pub(crate) struct Context {
34	pub(crate) pool: Arc<Pool>,
35
36	/// Retained because rust-rocksdb's `Cache` binding lacks `get_capacity`.
37	pub(crate) row_cache_capacity: usize,
38
39	pub(crate) row_cache: Mutex<Cache>,
40
41	pub(crate) col_cache: Mutex<ColCaches>,
42
43	pub(crate) server: Arc<Server>,
44
45	pub(super) env: Arc<Env>,
46}
47
48/// Map of block-cache pools keyed by pool name. The pool name is either
49/// `SHARED_POOL` or the first-arrival CF that created it.
50pub(crate) type ColCaches = BTreeMap<&'static str, ColCache>;
51
52/// Name under which the shared block cache (every CF with
53/// `CacheDisp::Shared`) is registered in [`Context::col_cache`].
54pub(crate) const SHARED_POOL: &str = "Shared";
55
56impl Context {
57	pub(crate) fn new(server: &Arc<Server>) -> Result<Arc<Self>> {
58		let config = &server.config;
59		let cache_capacity_bytes = config.db_cache_capacity_mb * 1024.0 * 1024.0;
60
61		let col_cache_shards: i32 = 128;
62		let col_shard_bits = col_cache_shards.ilog2().try_into()?;
63		let col_cache_capacity_bytes = usize_from_f64(cache_capacity_bytes * 0.50)?;
64
65		let row_cache_shards: i32 = 128;
66		let row_shard_bits = row_cache_shards.ilog2().try_into()?;
67		let row_cache_capacity_bytes = usize_from_f64(cache_capacity_bytes * 0.50)?;
68
69		let mut row_cache_opts = LruCacheOptions::default();
70		row_cache_opts.set_num_shard_bits(row_shard_bits);
71		row_cache_opts.set_capacity(row_cache_capacity_bytes);
72		let row_cache = Cache::new_lru_cache_opts(&row_cache_opts);
73
74		let mut col_cache_opts = LruCacheOptions::default();
75		col_cache_opts.set_num_shard_bits(col_shard_bits);
76		col_cache_opts.set_capacity(col_cache_capacity_bytes);
77		let col_cache = Cache::new_lru_cache_opts(&col_cache_opts);
78		let shared = ColCache {
79			cache: col_cache,
80			participants: Vec::new(),
81		};
82		let col_cache: ColCaches = [(SHARED_POOL, shared)].into();
83
84		Ok(Arc::new(Self {
85			pool: Pool::new(server)?,
86			row_cache_capacity: row_cache_capacity_bytes,
87			row_cache: row_cache.into(),
88			col_cache: col_cache.into(),
89			server: server.clone(),
90			env: Env::acquire(server)?,
91		}))
92	}
93}
94
95impl Drop for Context {
96	#[cold]
97	fn drop(&mut self) {
98		debug!("Closing frontend pool");
99		self.pool.close();
100
101		after_close(self, &self.server.config.database_path)
102			.expect("Failed to execute after_close handler");
103	}
104}
105
106/// For unit and integration tests the 'fresh' directive deletes found db.
107pub(super) fn before_open(ctx: &Arc<Context>, path: &Path) -> Result {
108	if ctx.server.config.test.contains("fresh") {
109		match delete_database_for_testing(ctx, path) {
110			| Err(e) if !e.is_not_found() => return Err(e),
111			| _ => (),
112		}
113	}
114
115	Ok(())
116}
117
118/// For unit and integration tests the 'cleanup' directive deletes after close
119/// to cleanup.
120fn after_close(ctx: &Context, path: &Path) -> Result {
121	if ctx.server.config.test.contains("cleanup") {
122		delete_database_for_testing(ctx, path)
123			.log_err()
124			.ok();
125	}
126
127	Ok(())
128}
129
130/// For unit and integration tests; removes the database directory when called.
131/// To prevent misuse, cfg!(test) must be true for a unit test or the
132/// integration test server is named localhost.
133#[tracing::instrument(level = "debug", skip_all)]
134fn delete_database_for_testing(ctx: &Context, path: &Path) -> Result {
135	let config = &ctx.server.config;
136	let localhost = config
137		.server_name
138		.as_str()
139		.starts_with("localhost");
140
141	if !cfg!(test) && !localhost {
142		return Ok(());
143	}
144
145	debug_assert!(
146		config.test.contains("cleanup") | config.test.contains("fresh"),
147		"missing any test directive legitimating this call.",
148	);
149
150	remove_dir_all(path).map_err(Into::into)
151}
tuwunel_database/engine/context.rs

tuwunel_database/engine/
context.rs
