feat(tvix/store): implement PathInfoService with sled r/5649

This uses [sled](https://github.com/spacejam/sled) to store PathInfo
objects.

Change-Id: I12e8032e5562af8f884efa23a78049fd1108fdbc
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7726
Reviewed-by: tazjin <tazjin@tvl.su>
Tested-by: BuildkiteCI

7 files changed, 169 insertions, 37 deletions

diff --git a/tvix/store/Cargo.toml b/tvix/store/Cargo.toml
index abbd9b0dcbcd..0f1dd50e2b98 100644
--- a/tvix/store/Cargo.toml
+++ b/tvix/store/Cargo.toml
@@ -13,6 +13,7 @@ data-encoding = "2.3.3"
 lazy_static = "1.4.0"
 clap = { version = "4.0", features = ["derive", "env"] }
 prost = "0.11.2"
+sled = { version = "0.34.7", features = ["compression"] }
 thiserror = "1.0.38"
 tokio = { version = "1.23.0", features = ["rt-multi-thread"] }
 tokio-stream = "0.1.11"
@@ -30,6 +31,7 @@ tonic-build = "0.8.2"
 
 [dev-dependencies]
 test-case = "2.2.2"
+tempfile = "3.3.0"
 
 [features]
 default = [
diff --git a/tvix/store/src/dummy_path_info_service.rs b/tvix/store/src/dummy_path_info_service.rs
deleted file mode 100644
index 93359377f3dc..000000000000
--- a/tvix/store/src/dummy_path_info_service.rs
+++ /dev/null
@@ -1,35 +0,0 @@
-use crate::proto::path_info_service_server::PathInfoService;
-use crate::proto::CalculateNarResponse;
-use crate::proto::GetPathInfoRequest;
-use crate::proto::Node;
-use crate::proto::PathInfo;
-use tonic::{Request, Response, Result, Status};
-use tracing::{instrument, warn};
-
-pub struct DummyPathInfoService {}
-
-const NOT_IMPLEMENTED_MSG: &str = "not implemented";
-
-#[tonic::async_trait]
-impl PathInfoService for DummyPathInfoService {
-    #[instrument(skip(self))]
-    async fn get(&self, _request: Request<GetPathInfoRequest>) -> Result<Response<PathInfo>> {
-        warn!(NOT_IMPLEMENTED_MSG);
-        Err(Status::unimplemented(NOT_IMPLEMENTED_MSG))
-    }
-
-    #[instrument(skip(self))]
-    async fn put(&self, _request: Request<PathInfo>) -> Result<Response<PathInfo>> {
-        warn!(NOT_IMPLEMENTED_MSG);
-        Err(Status::unimplemented(NOT_IMPLEMENTED_MSG))
-    }
-
-    #[instrument(skip(self))]
-    async fn calculate_nar(
-        &self,
-        _request: Request<Node>,
-    ) -> Result<Response<CalculateNarResponse>> {
-        warn!(NOT_IMPLEMENTED_MSG);
-        Err(Status::unimplemented(NOT_IMPLEMENTED_MSG))
-    }
-}
diff --git a/tvix/store/src/lib.rs b/tvix/store/src/lib.rs
index 642d285990dd..8d7d45d028b9 100644
--- a/tvix/store/src/lib.rs
+++ b/tvix/store/src/lib.rs
@@ -4,7 +4,7 @@ pub mod store_path;
 
 pub mod dummy_blob_service;
 pub mod dummy_directory_service;
-pub mod dummy_path_info_service;
+pub mod sled_path_info_service;
 
 #[cfg(test)]
 mod tests;
diff --git a/tvix/store/src/main.rs b/tvix/store/src/main.rs
index 0ad1857db6bc..4fa09f28a285 100644
--- a/tvix/store/src/main.rs
+++ b/tvix/store/src/main.rs
@@ -36,7 +36,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let blob_service = tvix_store::dummy_blob_service::DummyBlobService {};
     let directory_service = tvix_store::dummy_directory_service::DummyDirectoryService {};
-    let path_info_service = tvix_store::dummy_path_info_service::DummyPathInfoService {};
+    let path_info_service =
+        tvix_store::sled_path_info_service::SledPathInfoService::new("pathinfo.sled".into())?;
 
     let mut router = server
         .add_service(BlobServiceServer::new(blob_service))
diff --git a/tvix/store/src/sled_path_info_service.rs b/tvix/store/src/sled_path_info_service.rs
new file mode 100644
index 000000000000..3215fd0e4730
--- /dev/null
+++ b/tvix/store/src/sled_path_info_service.rs
@@ -0,0 +1,89 @@
+use prost::Message;
+use std::path::PathBuf;
+
+use crate::proto::get_path_info_request::ByWhat;
+use crate::proto::path_info_service_server::PathInfoService;
+use crate::proto::CalculateNarResponse;
+use crate::proto::GetPathInfoRequest;
+use crate::proto::Node;
+use crate::proto::PathInfo;
+use crate::store_path::DIGEST_SIZE;
+use tonic::{Request, Response, Result, Status};
+use tracing::{instrument, warn};
+
+const NOT_IMPLEMENTED_MSG: &str = "not implemented";
+
+/// SledPathInfoService stores PathInfo in a [sled](https://github.com/spacejam/sled).
+///
+/// The PathInfo messages are stored as encoded protos, and keyed by their output hash,
+/// as that's currently the only request type available.
+pub struct SledPathInfoService {
+    db: sled::Db,
+}
+
+impl SledPathInfoService {
+    pub fn new(p: PathBuf) -> Result<Self, anyhow::Error> {
+        let config = sled::Config::default().use_compression(true).path(p);
+        let db = config.open()?;
+
+        Ok(Self { db })
+    }
+}
+
+#[tonic::async_trait]
+impl PathInfoService for SledPathInfoService {
+    #[instrument(skip(self))]
+    async fn get(&self, request: Request<GetPathInfoRequest>) -> Result<Response<PathInfo>> {
+        match request.into_inner().by_what {
+            None => Err(Status::unimplemented("by_what needs to be specified")),
+            Some(ByWhat::ByOutputHash(digest)) => {
+                if digest.len() != DIGEST_SIZE {
+                    return Err(Status::invalid_argument("invalid digest length"));
+                }
+
+                match self.db.get(digest) {
+                    Ok(None) => Err(Status::not_found("PathInfo not found")),
+                    Ok(Some(data)) => match PathInfo::decode(&*data) {
+                        Ok(path_info) => Ok(Response::new(path_info)),
+                        Err(e) => {
+                            warn!("failed to decode stored PathInfo: {}", e);
+                            Err(Status::internal("failed to decode stored PathInfo"))
+                        }
+                    },
+                    Err(e) => {
+                        warn!("failed to retrieve PathInfo: {}", e);
+                        Err(Status::internal("error during PathInfo lookup"))
+                    }
+                }
+            }
+        }
+    }
+
+    #[instrument(skip(self))]
+    async fn put(&self, request: Request<PathInfo>) -> Result<Response<PathInfo>> {
+        let path_info = request.into_inner();
+
+        // Call validate on the received PathInfo message.
+        match path_info.validate() {
+            Err(e) => Err(Status::invalid_argument(e.to_string())),
+            // In case the PathInfo is valid, and we were able to extract a NixPath, store it in the database.
+            // This overwrites existing PathInfo objects.
+            Ok(nix_path) => match self.db.insert(nix_path.digest, path_info.encode_to_vec()) {
+                Ok(_) => Ok(Response::new(path_info)),
+                Err(e) => {
+                    warn!("failed to insert PathInfo: {}", e);
+                    Err(Status::internal("failed to insert PathInfo"))
+                }
+            },
+        }
+    }
+
+    #[instrument(skip(self))]
+    async fn calculate_nar(
+        &self,
+        _request: Request<Node>,
+    ) -> Result<Response<CalculateNarResponse>> {
+        warn!(NOT_IMPLEMENTED_MSG);
+        Err(Status::unimplemented(NOT_IMPLEMENTED_MSG))
+    }
+}
diff --git a/tvix/store/src/tests/mod.rs b/tvix/store/src/tests/mod.rs
index b5fb0648e632..3f5d044d658a 100644
--- a/tvix/store/src/tests/mod.rs
+++ b/tvix/store/src/tests/mod.rs
@@ -1,6 +1,7 @@
 use crate::proto::{Directory, DirectoryNode, FileNode, SymlinkNode, ValidateDirectoryError};
 use lazy_static::lazy_static;
 
+mod path_info_service;
 mod pathinfo;
 
 lazy_static! {
diff --git a/tvix/store/src/tests/path_info_service.rs b/tvix/store/src/tests/path_info_service.rs
new file mode 100644
index 000000000000..dca61fe12a56
--- /dev/null
+++ b/tvix/store/src/tests/path_info_service.rs
@@ -0,0 +1,74 @@
+use tempfile::TempDir;
+
+use crate::proto::path_info_service_server::PathInfoService;
+use crate::proto::GetPathInfoRequest;
+use crate::proto::{get_path_info_request, PathInfo};
+use crate::sled_path_info_service::SledPathInfoService;
+
+use lazy_static::lazy_static;
+
+lazy_static! {
+    static ref DUMMY_OUTPUT_HASH: Vec<u8> = vec![
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00, 0x00
+    ];
+}
+
+/// Trying to get a non-existent PathInfo should return a not found error.
+#[tokio::test]
+async fn not_found() -> anyhow::Result<()> {
+    let service = SledPathInfoService::new(TempDir::new()?.path().to_path_buf())?;
+
+    let resp = service
+        .get(tonic::Request::new(GetPathInfoRequest {
+            by_what: Some(get_path_info_request::ByWhat::ByOutputHash(
+                DUMMY_OUTPUT_HASH.to_vec(),
+            )),
+        }))
+        .await;
+
+    match resp {
+        Err(status) => {
+            assert_eq!(status.code(), tonic::Code::NotFound);
+        }
+        Ok(_) => panic!("must fail"),
+    };
+
+    Ok(())
+}
+
+/// Put a PathInfo into the store, get it back.
+#[tokio::test]
+async fn put_get() -> anyhow::Result<()> {
+    let service = SledPathInfoService::new(TempDir::new()?.path().to_path_buf())?;
+
+    let path_info = PathInfo {
+        node: Some(crate::proto::Node {
+            node: Some(crate::proto::node::Node::Symlink(
+                crate::proto::SymlinkNode {
+                    name: "00000000000000000000000000000000-foo".to_string(),
+                    target: "doesntmatter".to_string(),
+                },
+            )),
+        }),
+        ..Default::default()
+    };
+
+    let resp = service.put(tonic::Request::new(path_info.clone())).await;
+
+    assert!(resp.is_ok());
+    assert_eq!(resp.expect("must succeed").into_inner(), path_info);
+
+    let resp = service
+        .get(tonic::Request::new(GetPathInfoRequest {
+            by_what: Some(get_path_info_request::ByWhat::ByOutputHash(
+                DUMMY_OUTPUT_HASH.to_vec(),
+            )),
+        }))
+        .await;
+
+    assert!(resp.is_ok());
+    assert_eq!(resp.expect("must succeed").into_inner(), path_info);
+
+    Ok(())
+}