chore(tvix/store): move NAR rendering logic into Renderer struct r/5915

This moves the logic rendering NARs to a struct using the
previously introduced, more granular BlobService, ChunkService and
DirectoryService.

Instead of passing them around to the helper functions, they're kept as
members of a struct.

Remove the async invocations in the nar_renderer tests, there's nothing
async in here.

Change-Id: Ic6d24aaad68a1fda46ce29f2cdb5f7b87f481d5c
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8095
Reviewed-by: raitobezarius <tvl@lahfa.xyz>
Tested-by: BuildkiteCI

5 files changed, 275 insertions, 181 deletions

diff --git a/tvix/store/src/nar.rs b/tvix/store/src/nar.rs
deleted file mode 100644
index efcaf652e138..000000000000
--- a/tvix/store/src/nar.rs
+++ /dev/null
@@ -1,70 +0,0 @@
-//! This provides some common "client-side" libraries to interact with a tvix-
-//! store, in this case to render NAR.
-use crate::{
-    client::StoreClient,
-    proto::{self, NamedNode},
-};
-use anyhow::Result;
-use nix_compat::nar;
-
-/// Consumes a [proto::node::Node] pointing to the root of a (store) path,
-/// and writes the contents in NAR serialization to the passed
-/// [std::io::Write].
-///
-/// It uses a [StoreClient] to do the necessary lookups as it traverses the
-/// structure.
-pub fn write_nar<W: std::io::Write, SC: StoreClient>(
-    w: &mut W,
-    proto_root_node: proto::node::Node,
-    store_client: &mut SC,
-) -> Result<()> {
-    // Initialize NAR writer
-    let nar_root_node = nar::writer::open(w)?;
-
-    walk_node(nar_root_node, proto_root_node, store_client)
-}
-
-/// Process an intermediate node in the structure.
-/// This consumes the node.
-fn walk_node<SC: StoreClient>(
-    nar_node: nar::writer::Node,
-    proto_node: proto::node::Node,
-    store_client: &mut SC,
-) -> Result<()> {
-    match proto_node {
-        proto::node::Node::Symlink(proto_symlink_node) => {
-            nar_node.symlink(&proto_symlink_node.target)?;
-        }
-        proto::node::Node::File(proto_file_node) => {
-            nar_node.file(
-                proto_file_node.executable,
-                proto_file_node.size.into(),
-                &mut store_client.open_blob(proto_file_node.digest)?,
-            )?;
-        }
-        proto::node::Node::Directory(proto_directory_node) => {
-            // look up that node from the store client
-            let proto_directory = store_client.get_directory(proto_directory_node.digest)?;
-
-            // if it's None, that's an error!
-            if proto_directory.is_none() {
-                // TODO: proper error handling
-                panic!("not found!")
-            }
-
-            // start a directory node
-            let mut nar_node_directory = nar_node.directory()?;
-
-            // for each node in the directory, create a new entry with its name,
-            // and then invoke walk_node on that entry.
-            for proto_node in proto_directory.unwrap().nodes() {
-                let child_node = nar_node_directory.entry(proto_node.get_name())?;
-                walk_node(child_node, proto_node, store_client)?;
-            }
-
-            // close the directory
-            nar_node_directory.close()?;
-        }
-    }
-    Ok(())
-}
diff --git a/tvix/store/src/nar/mod.rs b/tvix/store/src/nar/mod.rs
new file mode 100644
index 000000000000..d7d2cec4d803
--- /dev/null
+++ b/tvix/store/src/nar/mod.rs
@@ -0,0 +1,22 @@
+use data_encoding::BASE64;
+use thiserror::Error;
+
+mod renderer;
+
+pub use renderer::NARRenderer;
+
+/// Errors that can encounter while rendering NARs.
+#[derive(Debug, Error)]
+pub enum RenderError {
+    #[error("failure talking to a backing store client: {0}")]
+    StoreError(crate::Error),
+
+    #[error("unable to find directory {}, referred from {}", BASE64.encode(.0), .1)]
+    DirectoryNotFound(Vec<u8>, String),
+
+    #[error("unable to find blob {}, referred from {}", BASE64.encode(.0), .1)]
+    BlobNotFound(Vec<u8>, String),
+
+    #[error("failure using the NAR writer: {0}")]
+    NARWriterError(std::io::Error),
+}
diff --git a/tvix/store/src/nar/renderer.rs b/tvix/store/src/nar/renderer.rs
new file mode 100644
index 000000000000..d8d9886b315d
--- /dev/null
+++ b/tvix/store/src/nar/renderer.rs
@@ -0,0 +1,136 @@
+use crate::{
+    blobservice::BlobService,
+    chunkservice::ChunkService,
+    directoryservice::DirectoryService,
+    proto::{self, NamedNode},
+    BlobReader,
+};
+use nix_compat::nar;
+
+use super::RenderError;
+
+/// A NAR renderer, using a blob_service, chunk_service and directory_service
+/// to render a NAR to a writer.
+#[derive(Clone)]
+pub struct NARRenderer<BS: BlobService, CS: ChunkService + Clone, DS: DirectoryService> {
+    blob_service: BS,
+    chunk_service: CS,
+    directory_service: DS,
+}
+
+impl<BS: BlobService, CS: ChunkService + Clone, DS: DirectoryService> NARRenderer<BS, CS, DS> {
+    pub fn new(blob_service: BS, chunk_service: CS, directory_service: DS) -> Self {
+        Self {
+            blob_service,
+            chunk_service,
+            directory_service,
+        }
+    }
+
+    /// Consumes a [proto::node::Node] pointing to the root of a (store) path,
+    /// and writes the contents in NAR serialization to the passed
+    /// [std::io::Write].
+    ///
+    /// It uses the different clients in the struct to perform the necessary
+    /// lookups as it traverses the structure.
+    pub fn write_nar<W: std::io::Write>(
+        &self,
+        w: &mut W,
+        proto_root_node: proto::node::Node,
+    ) -> Result<(), RenderError> {
+        // Initialize NAR writer
+        let nar_root_node = nar::writer::open(w).map_err(RenderError::NARWriterError)?;
+
+        self.walk_node(nar_root_node, proto_root_node)
+    }
+
+    /// Process an intermediate node in the structure.
+    /// This consumes the node.
+    fn walk_node(
+        &self,
+        nar_node: nar::writer::Node,
+        proto_node: proto::node::Node,
+    ) -> Result<(), RenderError> {
+        match proto_node {
+            proto::node::Node::Symlink(proto_symlink_node) => {
+                nar_node
+                    .symlink(&proto_symlink_node.target)
+                    .map_err(RenderError::NARWriterError)?;
+            }
+            proto::node::Node::File(proto_file_node) => {
+                // get the digest we're referring to
+                let digest = proto_file_node.digest;
+                // query blob_service for blob_meta
+                let resp = self
+                    .blob_service
+                    .stat(&proto::StatBlobRequest {
+                        digest: digest.to_vec(),
+                        include_chunks: true,
+                        ..Default::default()
+                    })
+                    .map_err(RenderError::StoreError)?;
+
+                match resp {
+                    // if it's None, that's an error!
+                    None => {
+                        return Err(RenderError::BlobNotFound(digest, proto_file_node.name));
+                    }
+                    Some(blob_meta) => {
+                        let mut blob_reader = std::io::BufReader::new(BlobReader::open(
+                            &self.chunk_service,
+                            blob_meta,
+                        ));
+                        nar_node
+                            .file(
+                                proto_file_node.executable,
+                                proto_file_node.size.into(),
+                                &mut blob_reader,
+                            )
+                            .map_err(RenderError::NARWriterError)?;
+                    }
+                }
+            }
+            proto::node::Node::Directory(proto_directory_node) => {
+                // get the digest we're referring to
+                let digest = proto_directory_node.digest;
+                // look it up with the directory service
+                let resp = self
+                    .directory_service
+                    .get(&proto::get_directory_request::ByWhat::Digest(
+                        digest.to_vec(),
+                    ))
+                    .map_err(RenderError::StoreError)?;
+
+                match resp {
+                    // if it's None, that's an error!
+                    None => {
+                        return Err(RenderError::DirectoryNotFound(
+                            digest,
+                            proto_directory_node.name,
+                        ))
+                    }
+                    Some(proto_directory) => {
+                        // start a directory node
+                        let mut nar_node_directory =
+                            nar_node.directory().map_err(RenderError::NARWriterError)?;
+
+                        // for each node in the directory, create a new entry with its name,
+                        // and then invoke walk_node on that entry.
+                        for proto_node in proto_directory.nodes() {
+                            let child_node = nar_node_directory
+                                .entry(proto_node.get_name())
+                                .map_err(RenderError::NARWriterError)?;
+                            self.walk_node(child_node, proto_node)?;
+                        }
+
+                        // close the directory
+                        nar_node_directory
+                            .close()
+                            .map_err(RenderError::NARWriterError)?;
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+}
diff --git a/tvix/store/src/tests/mod.rs b/tvix/store/src/tests/mod.rs
index b945763f3686..6947b277d481 100644
--- a/tvix/store/src/tests/mod.rs
+++ b/tvix/store/src/tests/mod.rs
@@ -1,3 +1,3 @@
 mod directory_service;
-mod nar;
+mod nar_renderer;
 mod path_info_service;
diff --git a/tvix/store/src/tests/nar.rs b/tvix/store/src/tests/nar_renderer.rs
index 5a865f52b492..43afdefbf930 100644
--- a/tvix/store/src/tests/nar.rs
+++ b/tvix/store/src/tests/nar_renderer.rs
@@ -1,12 +1,17 @@
-use data_encoding::BASE64;
-
-use crate::client::StoreClient;
-use crate::nar::write_nar;
+use crate::blobservice::BlobService;
+use crate::blobservice::SledBlobService;
+use crate::chunkservice::ChunkService;
+use crate::chunkservice::SledChunkService;
+use crate::directoryservice::DirectoryService;
+use crate::directoryservice::SledDirectoryService;
+use crate::nar::NARRenderer;
 use crate::proto;
 use crate::proto::DirectoryNode;
 use crate::proto::FileNode;
 use crate::proto::SymlinkNode;
 use lazy_static::lazy_static;
+use std::path::Path;
+use tempfile::TempDir;
 
 const HELLOWORLD_BLOB_CONTENTS: &[u8] = b"Hello World!";
 const EMPTY_BLOB_CONTENTS: &[u8] = b"";
@@ -44,93 +49,39 @@ lazy_static! {
     };
 }
 
-/// A Store client that fails if you ask it for a blob or a directory
-#[derive(Default)]
-struct FailingStoreClient {}
-
-impl StoreClient for FailingStoreClient {
-    fn open_blob(&self, digest: Vec<u8>) -> std::io::Result<Box<dyn std::io::BufRead>> {
-        panic!(
-            "open_blob should never be called, but was called with {}",
-            BASE64.encode(&digest),
-        );
-    }
-
-    fn get_directory(&self, digest: Vec<u8>) -> std::io::Result<Option<proto::Directory>> {
-        panic!(
-            "get_directory should never be called, but was called with {}",
-            BASE64.encode(&digest),
-        );
-    }
+fn gen_blob_service(p: &Path) -> impl BlobService {
+    SledBlobService::new(p.join("blobs")).unwrap()
 }
 
-/// Only allow a request for a blob with [HELLOWORLD_BLOB_DIGEST]
-/// panic on everything else.
-#[derive(Default)]
-struct HelloWorldBlobStoreClient {}
-
-impl StoreClient for HelloWorldBlobStoreClient {
-    fn open_blob(&self, digest: Vec<u8>) -> std::io::Result<Box<dyn std::io::BufRead>> {
-        if digest != HELLOWORLD_BLOB_DIGEST.to_vec() {
-            panic!("open_blob called with {}", BASE64.encode(&digest));
-        }
-
-        let b: Box<&[u8]> = Box::new(&HELLOWORLD_BLOB_CONTENTS);
-
-        Ok(b)
-    }
-
-    fn get_directory(&self, digest: Vec<u8>) -> std::io::Result<Option<proto::Directory>> {
-        panic!(
-            "get_directory should never be called, but was called with {}",
-            BASE64.encode(&digest),
-        );
-    }
+fn gen_chunk_service(p: &Path) -> impl ChunkService + Clone {
+    SledChunkService::new(p.join("chunks")).unwrap()
 }
 
-/// Allow blob requests for [HELLOWORLD_BLOB_DIGEST] and EMPTY_BLOB_DIGEST, and
-/// allow DIRECTORY_WITH_KEEP and DIRECTORY_COMPLICATED.
-#[derive(Default)]
-struct SomeDirectoryStoreClient {}
-
-impl StoreClient for SomeDirectoryStoreClient {
-    fn open_blob(&self, digest: Vec<u8>) -> std::io::Result<Box<dyn std::io::BufRead>> {
-        if digest == HELLOWORLD_BLOB_DIGEST.to_vec() {
-            let b: Box<&[u8]> = Box::new(&HELLOWORLD_BLOB_CONTENTS);
-            return Ok(b);
-        }
-        if digest == EMPTY_BLOB_DIGEST.to_vec() {
-            let b: Box<&[u8]> = Box::new(&EMPTY_BLOB_CONTENTS);
-            return Ok(b);
-        }
-        panic!("open_blob called with {}", BASE64.encode(&digest));
-    }
-
-    fn get_directory(&self, digest: Vec<u8>) -> std::io::Result<Option<proto::Directory>> {
-        if digest == DIRECTORY_WITH_KEEP.digest() {
-            return Ok(Some(DIRECTORY_WITH_KEEP.clone()));
-        }
-        if digest == DIRECTORY_COMPLICATED.digest() {
-            return Ok(Some(DIRECTORY_COMPLICATED.clone()));
-        }
-        panic!("get_directory called with {}", BASE64.encode(&digest));
-    }
+fn gen_directory_service(p: &Path) -> impl DirectoryService {
+    SledDirectoryService::new(p.join("directories")).unwrap()
 }
 
-#[tokio::test]
-async fn single_symlink() -> anyhow::Result<()> {
+#[test]
+fn single_symlink() -> anyhow::Result<()> {
+    let tmpdir = TempDir::new()?;
+    let renderer = NARRenderer::new(
+        gen_blob_service(tmpdir.path()),
+        gen_chunk_service(tmpdir.path()),
+        gen_directory_service(tmpdir.path()),
+    );
+    // don't put anything in the stores, as we don't actually do any requests.
+
     let mut buf: Vec<u8> = vec![];
-    let mut store_client = FailingStoreClient::default();
 
-    write_nar(
-        &mut buf,
-        crate::proto::node::Node::Symlink(SymlinkNode {
-            name: "doesntmatter".to_string(),
-            target: "/nix/store/somewhereelse".to_string(),
-        }),
-        &mut store_client,
-    )
-    .expect("must succeed");
+    renderer
+        .write_nar(
+            &mut buf,
+            crate::proto::node::Node::Symlink(SymlinkNode {
+                name: "doesntmatter".to_string(),
+                target: "/nix/store/somewhereelse".to_string(),
+            }),
+        )
+        .expect("must succeed");
 
     assert_eq!(
         buf,
@@ -150,22 +101,48 @@ async fn single_symlink() -> anyhow::Result<()> {
     Ok(())
 }
 
-#[tokio::test]
-async fn single_file() -> anyhow::Result<()> {
+#[test]
+fn single_file() -> anyhow::Result<()> {
+    let tmpdir = TempDir::new()?;
+
+    let blob_service = gen_blob_service(tmpdir.path());
+    let chunk_service = gen_chunk_service(tmpdir.path());
+
+    chunk_service
+        .put(HELLOWORLD_BLOB_CONTENTS.to_vec())
+        .unwrap();
+
+    blob_service
+        .put(
+            &HELLOWORLD_BLOB_DIGEST,
+            proto::BlobMeta {
+                chunks: vec![proto::blob_meta::ChunkMeta {
+                    digest: HELLOWORLD_BLOB_DIGEST.to_vec(),
+                    size: HELLOWORLD_BLOB_CONTENTS.len() as u32,
+                }],
+                ..Default::default()
+            },
+        )
+        .unwrap();
+
+    let renderer = NARRenderer::new(
+        blob_service,
+        chunk_service,
+        gen_directory_service(tmpdir.path()),
+    );
     let mut buf: Vec<u8> = vec![];
-    let mut store_client = HelloWorldBlobStoreClient::default();
 
-    write_nar(
-        &mut buf,
-        crate::proto::node::Node::File(FileNode {
-            name: "doesntmatter".to_string(),
-            digest: HELLOWORLD_BLOB_DIGEST.to_vec(),
-            size: HELLOWORLD_BLOB_CONTENTS.len() as u32,
-            executable: false,
-        }),
-        &mut store_client,
-    )
-    .expect("must succeed");
+    renderer
+        .write_nar(
+            &mut buf,
+            crate::proto::node::Node::File(FileNode {
+                name: "doesntmatter".to_string(),
+                digest: HELLOWORLD_BLOB_DIGEST.to_vec(),
+                size: HELLOWORLD_BLOB_CONTENTS.len() as u32,
+                executable: false,
+            }),
+        )
+        .expect("must succeed");
 
     assert_eq!(
         buf,
@@ -185,21 +162,50 @@ async fn single_file() -> anyhow::Result<()> {
     Ok(())
 }
 
-#[tokio::test]
-async fn test_complicated() -> anyhow::Result<()> {
+#[test]
+fn test_complicated() -> anyhow::Result<()> {
+    let tmpdir = TempDir::new()?;
+
+    let blob_service = gen_blob_service(tmpdir.path());
+    let chunk_service = gen_chunk_service(tmpdir.path());
+    let directory_service = gen_directory_service(tmpdir.path());
+
+    // put all data into the stores.
+    for blob_contents in [HELLOWORLD_BLOB_CONTENTS, EMPTY_BLOB_CONTENTS] {
+        let digest = chunk_service.put(blob_contents.to_vec()).unwrap();
+
+        blob_service
+            .put(
+                &digest,
+                proto::BlobMeta {
+                    chunks: vec![proto::blob_meta::ChunkMeta {
+                        digest: digest.to_vec(),
+                        size: blob_contents.len() as u32,
+                    }],
+                    ..Default::default()
+                },
+            )
+            .unwrap();
+    }
+
+    directory_service.put(DIRECTORY_WITH_KEEP.clone()).unwrap();
+    directory_service
+        .put(DIRECTORY_COMPLICATED.clone())
+        .unwrap();
+
+    let renderer = NARRenderer::new(blob_service, chunk_service, directory_service);
     let mut buf: Vec<u8> = vec![];
-    let mut store_client = SomeDirectoryStoreClient::default();
 
-    write_nar(
-        &mut buf,
-        crate::proto::node::Node::Directory(DirectoryNode {
-            name: "doesntmatter".to_string(),
-            digest: DIRECTORY_COMPLICATED.digest(),
-            size: DIRECTORY_COMPLICATED.size() as u32,
-        }),
-        &mut store_client,
-    )
-    .expect("must succeed");
+    renderer
+        .write_nar(
+            &mut buf,
+            crate::proto::node::Node::Directory(DirectoryNode {
+                name: "doesntmatter".to_string(),
+                digest: DIRECTORY_COMPLICATED.digest(),
+                size: DIRECTORY_COMPLICATED.size(),
+            }),
+        )
+        .expect("must succeed");
 
     assert_eq!(
         buf,