snix_castore/directoryservice/
redb.rs

1use futures::{StreamExt, TryStreamExt, stream::BoxStream};
2use prost::Message;
3use redb::{ReadableDatabase, TableDefinition};
4use std::{path::PathBuf, sync::Arc};
5use tonic::async_trait;
6use tracing::{instrument, warn};
7
8use super::{Directory, DirectoryPutter, DirectoryService, traversal};
9use crate::{
10    B3Digest,
11    composition::{CompositionContext, ServiceBuilder},
12    directoryservice::directory_graph::DirectoryGraphBuilder,
13    proto,
14};
15
16const DIRECTORY_TABLE: TableDefinition<[u8; B3Digest::LENGTH], Vec<u8>> =
17    TableDefinition::new("directory");
18
19enum Db {
20    ReadOnly(redb::ReadOnlyDatabase),
21    ReadWrite(redb::Database),
22}
23
24impl Db {
25    fn begin_read(&self) -> Result<redb::ReadTransaction, Error> {
26        match self {
27            Db::ReadOnly(db) => Ok(db.begin_read()?),
28            Db::ReadWrite(db) => Ok(db.begin_read()?),
29        }
30    }
31
32    fn begin_write(&self) -> Result<redb::WriteTransaction, Error> {
33        match self {
34            Db::ReadOnly(_) => Err(Error::OpenedReadonly),
35            Db::ReadWrite(db) => Ok(db.begin_write()?),
36        }
37    }
38}
39
40#[derive(Clone)]
41pub struct RedbDirectoryService {
42    instance_name: String,
43
44    /// An Arc'ed Database, read-only or writeable.
45    db: Arc<Db>,
46}
47
48impl RedbDirectoryService {
49    /// Constructs a new instance using the specified config.
50    pub async fn new(
51        instance_name: String,
52        config: RedbDirectoryServiceConfig,
53    ) -> Result<Self, Error> {
54        if let Some(path) = config.path.as_ref() {
55            if path == "" {
56                return Err(Error::WrongConfig("empty path is disallowed"));
57            }
58            if path == "/" {
59                return Err(Error::WrongConfig("cowardly refusing to open / with redb"));
60            }
61
62            if !config.read_only
63                && let Some(parent) = path.parent()
64            {
65                tokio::fs::create_dir_all(parent).await?;
66            }
67        }
68
69        let db = if config.path.is_some() {
70            tokio::task::spawn_blocking(move || {
71                let mut builder = redb::Database::builder();
72                configure_builder(&mut builder, &config);
73
74                let path = config.path.expect("Snix bug: path is Some");
75
76                if config.read_only {
77                    Ok::<_, Error>(Db::ReadOnly(builder.open_read_only(&path)?))
78                } else {
79                    let db = builder.create(&path)?;
80                    create_schema(&db)?;
81                    Ok(Db::ReadWrite(db))
82                }
83            })
84            .await??
85        } else {
86            if config.read_only {
87                return Err(Error::WrongConfig("in-memory database cannot be read-only"));
88            }
89            let mut builder = redb::Database::builder();
90            configure_builder(&mut builder, &config);
91
92            let db = builder
93                .create_with_backend(redb::backends::InMemoryBackend::new())
94                .expect("Snix bug: unable to create in-memory redb");
95
96            create_schema(&db)?;
97            Db::ReadWrite(db)
98        };
99
100        Ok(Self {
101            instance_name,
102            db: Arc::new(db),
103        })
104    }
105
106    /// Constructs a new instance using the in-memory backend.
107    /// Only used for testing purposes and mocks, use [Self::new] with a None
108    /// path in config for other usecases.
109    /// Sync, as there's no real IO happening.
110    #[cfg(any(test, feature = "mocks"))]
111    pub fn new_temporary(instance_name: String) -> Self {
112        let mut builder = redb::Database::builder();
113        configure_builder(
114            &mut builder,
115            &RedbDirectoryServiceConfig {
116                path: None,
117                cache_size: None,
118                read_only: false,
119            },
120        );
121
122        let db = builder
123            .create_with_backend(redb::backends::InMemoryBackend::new())
124            .expect("Snix bug: unable to create in-memory redb");
125
126        create_schema(&db).expect("Snix bug: unable to create schema for in-memory redb");
127
128        Self {
129            instance_name,
130            db: Arc::new(Db::ReadWrite(db)),
131        }
132    }
133}
134
135/// Applies options from [RedbDirectoryServiceConfig] to a [redb::Builder].
136fn configure_builder(builder: &mut redb::Builder, config: &RedbDirectoryServiceConfig) {
137    if let Some(cache_size) = config.cache_size {
138        builder.set_cache_size(cache_size);
139    }
140}
141
142/// Ensures all tables are present.
143/// Opens a write transaction and calls open_table on DIRECTORY_TABLE, which will
144/// create it if not present.
145#[allow(clippy::result_large_err)]
146fn create_schema(db: &redb::Database) -> Result<(), Error> {
147    let txn = db.begin_write()?;
148    txn.open_table(DIRECTORY_TABLE)?;
149    txn.commit()?;
150
151    Ok(())
152}
153
154#[async_trait]
155impl DirectoryService for RedbDirectoryService {
156    #[instrument(skip(self, digest), fields(directory.digest = %digest, instance_name = %self.instance_name))]
157    async fn get(&self, digest: &B3Digest) -> Result<Option<Directory>, super::Error> {
158        let db = self.db.clone();
159        let digest = *digest;
160        // Retrieves the protobuf-encoded Directory for the corresponding digest.
161        let directory_data = match tokio::task::spawn_blocking(move || -> Result<_, Error> {
162            let txn = db.begin_read()?;
163            let table = txn.open_table(DIRECTORY_TABLE)?;
164            Ok(table.get(*digest)?)
165        })
166        .await
167        .map_err(Error::TokioJoin)??
168        {
169            // The Directory was not found, return None.
170            None => return Ok(None),
171            Some(directory_data) => directory_data.value(),
172        };
173
174        // We check that the digest of the retrieved Directory matches the expected digest.
175        let actual = B3Digest::from(blake3::hash(&directory_data));
176        if actual != digest {
177            return Err(Error::WrongDigest {
178                expected: digest,
179                actual,
180            }
181            .into());
182        }
183
184        // Attempt to decode the retrieved protobuf-encoded Directory
185        let proto_directory =
186            proto::Directory::decode(directory_data.as_slice()).map_err(Error::ProtobufDecode)?;
187        let directory = Directory::try_from(proto_directory).map_err(Error::DirectoryValidation)?;
188
189        Ok(Some(directory))
190    }
191
192    #[instrument(skip(self, directory), fields(directory.digest = %directory.digest(), instance_name = %self.instance_name))]
193    async fn put(&self, directory: Directory) -> Result<B3Digest, super::Error> {
194        let db = self.db.clone();
195        let digest = tokio::task::spawn_blocking(move || -> Result<_, Error> {
196            let digest = directory.digest();
197
198            // Store the directory in the table.
199            let txn = db.begin_write()?;
200            {
201                let mut table = txn.open_table(DIRECTORY_TABLE)?;
202                table.insert(
203                    digest.as_ref(),
204                    proto::Directory::from(directory).encode_to_vec(),
205                )?;
206            }
207            txn.commit()?;
208
209            Ok(digest)
210        })
211        .await
212        .map_err(Error::TokioJoin)??;
213
214        Ok(digest)
215    }
216
217    #[instrument(skip_all, fields(directory.digest = %root_directory_digest, instance_name = %self.instance_name))]
218    fn get_recursive(
219        &self,
220        root_directory_digest: &B3Digest,
221    ) -> BoxStream<'static, Result<Directory, super::Error>> {
222        // FUTUREWORK: Ideally we should have all of the directory traversing happen in a single
223        // redb transaction to avoid constantly closing and opening new transactions for the
224        // database.
225        let svc = self.clone();
226        traversal::root_to_leaves(*root_directory_digest, move |digest| {
227            let svc = svc.clone();
228            async move { svc.get(&digest).await }
229        })
230        .map_err(Error::DirectoryTraversal)
231        .err_into()
232        .boxed()
233    }
234
235    #[instrument(skip_all)]
236    fn put_multiple_start(&self) -> Box<dyn DirectoryPutter> {
237        Box::new(RedbDirectoryPutter {
238            db: self.db.clone(),
239            builder: Some(DirectoryGraphBuilder::new_leaves_to_root()),
240        })
241    }
242}
243
244pub struct RedbDirectoryPutter {
245    db: Arc<Db>,
246
247    /// The directories (inside the directory validator) that we insert later,
248    /// or None, if they were already inserted.
249    builder: Option<DirectoryGraphBuilder>,
250}
251
252#[async_trait]
253impl DirectoryPutter for RedbDirectoryPutter {
254    #[instrument(level = "trace", skip_all, fields(directory.digest=%directory.digest()), err)]
255    async fn put(&mut self, directory: Directory) -> Result<(), super::Error> {
256        let builder = self
257            .builder
258            .as_mut()
259            .ok_or_else(|| Error::DirectoryPutterAlreadyClosed)?;
260
261        builder
262            .try_insert(directory)
263            .map_err(Error::DirectoryOrdering)?;
264
265        Ok(())
266    }
267
268    #[instrument(level = "trace", skip_all, ret, err)]
269    async fn close(&mut self) -> Result<B3Digest, super::Error> {
270        let builder = self
271            .builder
272            .take()
273            .ok_or_else(|| Error::DirectoryPutterAlreadyClosed)?;
274
275        // Insert all directories as a batch.
276        let db = self.db.clone();
277        let root_digest = tokio::task::spawn_blocking(move || {
278            // Retrieve the validated directories.
279            let directory_graph = builder.build().map_err(Error::DirectoryOrdering)?;
280            let root_digest = directory_graph.root().digest();
281
282            let txn = db.begin_write()?;
283            // Looping over all the verified directories, queuing them up for a
284            // batch insertion.
285            {
286                let mut table = txn.open_table(DIRECTORY_TABLE)?;
287                for directory in directory_graph.drain_leaves_to_root() {
288                    table.insert(
289                        directory.digest().as_ref(),
290                        proto::Directory::from(directory).encode_to_vec(),
291                    )?;
292                }
293            }
294            txn.commit()?;
295
296            Ok::<_, Error>(root_digest)
297        })
298        .await
299        .map_err(Error::TokioJoin)??;
300
301        Ok(root_digest)
302    }
303}
304
305#[derive(thiserror::Error, Debug)]
306pub enum Error {
307    #[error("wrong arguments: {0}")]
308    WrongConfig(&'static str),
309    #[error("serde-qs error: {0}")]
310    SerdeQS(#[from] serde_qs::Error),
311
312    #[error("Directory Graph ordering error")]
313    DirectoryOrdering(#[from] crate::directoryservice::OrderingError),
314
315    #[error("DirectoryPutter already closed")]
316    DirectoryPutterAlreadyClosed,
317
318    #[error("failure during directory traversal")]
319    DirectoryTraversal(#[source] traversal::Error),
320
321    #[error("requested directory has wrong digest, expected {expected}, actual {actual}")]
322    WrongDigest {
323        expected: B3Digest,
324        actual: B3Digest,
325    },
326    #[error("failed to decode protobuf: {0}")]
327    ProtobufDecode(#[from] prost::DecodeError),
328    #[error("failed to validate directory: {0}")]
329    DirectoryValidation(#[from] crate::DirectoryError),
330
331    #[error("unable to open write txn, database opened read-only")]
332    OpenedReadonly,
333    #[error("redb commit error: {0}")]
334    RedbCommit(#[from] redb::CommitError),
335    #[error("redb database error: {0}")]
336    RedbDatabase(#[from] redb::DatabaseError),
337    #[error("redb error: {0}")]
338    Redb(#[from] redb::Error),
339    #[error("redb storage error: {0}")]
340    RedbStorage(#[from] redb::StorageError),
341    #[error("redb table error: {0}")]
342    RedbTable(#[from] redb::TableError),
343    #[error("redb txn error: {0}")]
344    RedbTransaction(#[from] redb::TransactionError),
345
346    #[error("join error: {0}")]
347    TokioJoin(#[from] tokio::task::JoinError),
348    #[error("io error: {0}")]
349    IO(#[from] std::io::Error),
350}
351
352impl From<Error> for super::Error {
353    fn from(value: Error) -> Self {
354        Self(Box::new(value))
355    }
356}
357
358#[derive(Clone, Default, serde::Deserialize)]
359#[serde(deny_unknown_fields)]
360pub struct RedbDirectoryServiceConfig {
361    path: Option<PathBuf>,
362
363    /// The amount of memory (in bytes) used for caching data
364    cache_size: Option<usize>,
365
366    /// Whether to open read-only.
367    #[serde(default)]
368    read_only: bool,
369}
370
371impl TryFrom<url::Url> for RedbDirectoryServiceConfig {
372    type Error = Box<dyn std::error::Error + Send + Sync>;
373
374    fn try_from(url: url::Url) -> Result<Self, Self::Error> {
375        if url.has_host() {
376            return Err(Error::WrongConfig("no host allowed").into());
377        }
378
379        let path: Option<PathBuf> = match (url.scheme(), url.has_authority(), url.path()) {
380            ("redb+memory", false, "") => None,
381            ("redb+memory", false, _) => Err(Box::new(Error::WrongConfig(
382                "redb+memory with path is disallowed",
383            )))?,
384            ("redb+memory", true, _) => Err(Box::new(Error::WrongConfig(
385                "redb+memory may not have authority",
386            )))?,
387            ("redb", _, "") => Err(Box::new(Error::WrongConfig(
388                "redb without path is disallowed, use redb+memory if you want in-memory",
389            )))?,
390            ("redb", true, _path) => Err(Box::new(Error::WrongConfig("authority disallowed")))?,
391            ("redb", false, path) => Some(path.into()),
392            (_scheme, _, _) => Err(Box::new(Error::WrongConfig("unrecognized scheme")))?,
393        };
394
395        let mut config: RedbDirectoryServiceConfig =
396            serde_qs::from_str(url.query().unwrap_or_default())?;
397
398        config.path = path;
399
400        Ok(config)
401    }
402}
403
404#[async_trait]
405impl ServiceBuilder for RedbDirectoryServiceConfig {
406    type Output = dyn DirectoryService;
407    async fn build<'a>(
408        &'a self,
409        instance_name: &str,
410        _context: &CompositionContext,
411    ) -> Result<Arc<Self::Output>, Box<dyn std::error::Error + Send + Sync>> {
412        Ok(Arc::new(
413            RedbDirectoryService::new(instance_name.to_string(), self.to_owned()).await?,
414        ))
415    }
416}
417
418#[cfg(test)]
419mod tests {
420    use tempfile::TempDir;
421
422    use crate::{
423        directoryservice::{DirectoryService, RedbDirectoryService, RedbDirectoryServiceConfig},
424        fixtures::DIRECTORY_A,
425    };
426
427    #[tokio::test]
428    async fn reopen_as_read_only() {
429        let tempdir = TempDir::new().unwrap();
430        let path = tempdir.path().join("data.redb");
431
432        let config = RedbDirectoryServiceConfig {
433            path: Some(path),
434            cache_size: None,
435            read_only: false,
436        };
437
438        // Create a read-write directory service and insert some data.
439        {
440            let directory_service = RedbDirectoryService::new("rw".to_string(), config.clone())
441                .await
442                .expect("to construct");
443
444            directory_service
445                .put(DIRECTORY_A.clone())
446                .await
447                .expect("to insert");
448        } // we drop the rw database here.
449
450        // Re-open the same path in ro mode (twice)
451        let ro_config = RedbDirectoryServiceConfig {
452            read_only: true,
453            ..config
454        };
455
456        let directory_service_ro_1 =
457            RedbDirectoryService::new("ro1".to_string(), ro_config.clone())
458                .await
459                .expect("to construct");
460        let directory_service_ro_2 = RedbDirectoryService::new("ro2".to_string(), ro_config)
461            .await
462            .expect("to construct");
463
464        assert_eq!(
465            directory_service_ro_1
466                .get(&DIRECTORY_A.digest())
467                .await
468                .expect("get to succeed")
469                .expect("to be Some(_)")
470                .digest(),
471            DIRECTORY_A.digest()
472        );
473        assert_eq!(
474            directory_service_ro_2
475                .get(&DIRECTORY_A.digest())
476                .await
477                .expect("get to succeed")
478                .expect("to be Some(_)")
479                .digest(),
480            DIRECTORY_A.digest()
481        );
482    }
483
484    #[tokio::test]
485    async fn read_only_nonexistent() {
486        let tempdir = TempDir::new().unwrap();
487        let path = tempdir.path().join("data.redb");
488
489        let config = RedbDirectoryServiceConfig {
490            path: Some(path),
491            cache_size: None,
492            read_only: true,
493        };
494
495        // Opening a read-only redb should fail if the path doesn't exist.
496        assert!(
497            RedbDirectoryService::new("test".to_string(), config)
498                .await
499                .is_err(),
500            "opening new path r/o should fail"
501        );
502    }
503
504    #[tokio::test]
505    async fn open_rw_and_ro() {
506        let tempdir = TempDir::new().unwrap();
507        let path = tempdir.path().join("data.redb");
508
509        let config = RedbDirectoryServiceConfig {
510            path: Some(path),
511            cache_size: None,
512            read_only: false,
513        };
514
515        let _directory_service = RedbDirectoryService::new("rw".to_string(), config.clone())
516            .await
517            .expect("to construct");
518
519        // Opening a read-only redb should fail if it's already opened read-write.
520        assert!(
521            RedbDirectoryService::new(
522                "ro".to_string(),
523                RedbDirectoryServiceConfig {
524                    read_only: true,
525                    ..config
526                }
527            )
528            .await
529            .is_err(),
530            "opening r/o should fail if still open r/w"
531        );
532    }
533}
snix_castore/directoryservice/redb.rs

snix_castore/directoryservice/
redb.rs
