From 79317be214ce2f1e3347438319d3482bb773a649 Mon Sep 17 00:00:00 2001
From: Yureka <tvl@yuka.dev>
Date: Thu, 18 Jul 2024 20:59:30 +0200
Subject: [PATCH] feat(tvix/castore): add sled to composition registry

Change-Id: I03fa8dfabcee14c5f657380f86bb1a7aa00e08ef
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11977
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>
---
 tvix/castore/src/directoryservice/mod.rs  |  3 +-
 tvix/castore/src/directoryservice/sled.rs | 49 +++++++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/tvix/castore/src/directoryservice/mod.rs b/tvix/castore/src/directoryservice/mod.rs
index 815502c81..603e09d5d 100644
--- a/tvix/castore/src/directoryservice/mod.rs
+++ b/tvix/castore/src/directoryservice/mod.rs
@@ -25,7 +25,7 @@ pub use self::memory::{MemoryDirectoryService, MemoryDirectoryServiceConfig};
 pub use self::object_store::{ObjectStoreDirectoryService, ObjectStoreDirectoryServiceConfig};
 pub use self::order_validator::{LeavesToRootValidator, OrderValidator, RootToLeavesValidator};
 pub use self::simple_putter::SimplePutter;
-pub use self::sled::SledDirectoryService;
+pub use self::sled::{SledDirectoryService, SledDirectoryServiceConfig};
 pub use self::traverse::descend_to;
 pub use self::utils::traverse_directory;
 
@@ -134,6 +134,7 @@ pub(crate) fn register_directory_services(reg: &mut Registry) {
     reg.register::<Box<dyn ServiceBuilder<Output = dyn DirectoryService>>, super::directoryservice::MemoryDirectoryServiceConfig>("memory");
     reg.register::<Box<dyn ServiceBuilder<Output = dyn DirectoryService>>, super::directoryservice::CacheConfig>("cache");
     reg.register::<Box<dyn ServiceBuilder<Output = dyn DirectoryService>>, super::directoryservice::GRPCDirectoryServiceConfig>("grpc");
+    reg.register::<Box<dyn ServiceBuilder<Output = dyn DirectoryService>>, super::directoryservice::SledDirectoryServiceConfig>("sled");
     #[cfg(feature = "cloud")]
     {
         reg.register::<Box<dyn ServiceBuilder<Output = dyn DirectoryService>>, super::directoryservice::BigtableParameters>("bigtable");
diff --git a/tvix/castore/src/directoryservice/sled.rs b/tvix/castore/src/directoryservice/sled.rs
index bd98ed6b1..8e74227b3 100644
--- a/tvix/castore/src/directoryservice/sled.rs
+++ b/tvix/castore/src/directoryservice/sled.rs
@@ -4,11 +4,13 @@ use futures::stream::BoxStream;
 use prost::Message;
 use std::ops::Deref;
 use std::path::Path;
+use std::sync::Arc;
 use tonic::async_trait;
 use tracing::{instrument, warn};
 
 use super::utils::traverse_directory;
 use super::{DirectoryGraph, DirectoryPutter, DirectoryService, LeavesToRootValidator};
+use crate::composition::{CompositionContext, ServiceBuilder};
 
 #[derive(Clone)]
 pub struct SledDirectoryService {
@@ -17,6 +19,12 @@ pub struct SledDirectoryService {
 
 impl SledDirectoryService {
     pub fn new<P: AsRef<Path>>(p: P) -> Result<Self, sled::Error> {
+        if p.as_ref() == Path::new("/") {
+            return Err(sled::Error::Unsupported(
+                "cowardly refusing to open / with sled".to_string(),
+            ));
+        }
+
         let config = sled::Config::default()
             .use_compression(false) // is a required parameter
             .path(p);
@@ -128,6 +136,47 @@ impl DirectoryService for SledDirectoryService {
     }
 }
 
+#[derive(serde::Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct SledDirectoryServiceConfig {
+    is_temporary: bool,
+    #[serde(default)]
+    /// required when is_temporary = false
+    path: Option<String>,
+}
+
+#[async_trait]
+impl ServiceBuilder for SledDirectoryServiceConfig {
+    type Output = dyn DirectoryService;
+    async fn build<'a>(
+        &'a self,
+        _instance_name: &str,
+        _context: &CompositionContext<dyn DirectoryService>,
+    ) -> Result<Arc<dyn DirectoryService>, Box<dyn std::error::Error + Send + Sync + 'static>> {
+        match self {
+            SledDirectoryServiceConfig {
+                is_temporary: true,
+                path: None,
+            } => Ok(Arc::new(SledDirectoryService::new_temporary()?)),
+            SledDirectoryServiceConfig {
+                is_temporary: true,
+                path: Some(_),
+            } => Err(Error::StorageError(
+                "Temporary SledDirectoryService can not have path".into(),
+            )
+            .into()),
+            SledDirectoryServiceConfig {
+                is_temporary: false,
+                path: None,
+            } => Err(Error::StorageError("SledDirectoryService is missing path".into()).into()),
+            SledDirectoryServiceConfig {
+                is_temporary: false,
+                path: Some(path),
+            } => Ok(Arc::new(SledDirectoryService::new(path)?)),
+        }
+    }
+}
+
 /// Buffers Directory messages to be uploaded and inserts them in a batch
 /// transaction on close.
 pub struct SledDirectoryPutter {