feat(tvix/store): seekable nar renderer r/8715

Co-authored-by: edef <edef@edef.eu>
Change-Id: I233206e8aae35504ca0519ac88178dfc5596bedb
Reviewed-on: https://cl.tvl.fyi/c/depot/+/12439
Reviewed-by: flokli <flokli@flokli.de>
Autosubmit: yuka <yuka@yuka.dev>
Tested-by: BuildkiteCI

2 files changed, 112 insertions, 0 deletions

diff --git a/tvix/store/src/tests/mod.rs b/tvix/store/src/tests/mod.rs
index 1e7fc3f6b451..c7e6fee0cc6c 100644
--- a/tvix/store/src/tests/mod.rs
+++ b/tvix/store/src/tests/mod.rs
@@ -1,2 +1,3 @@
 pub mod fixtures;
 mod nar_renderer;
+mod nar_renderer_seekable;
diff --git a/tvix/store/src/tests/nar_renderer_seekable.rs b/tvix/store/src/tests/nar_renderer_seekable.rs
new file mode 100644
index 000000000000..233b95d0b036
--- /dev/null
+++ b/tvix/store/src/tests/nar_renderer_seekable.rs
@@ -0,0 +1,111 @@
+use crate::nar::seekable::Reader;
+use crate::tests::fixtures::blob_service_with_contents as blob_service;
+use crate::tests::fixtures::directory_service_with_contents as directory_service;
+use crate::tests::fixtures::*;
+use rstest::*;
+use rstest_reuse::*;
+use std::io;
+use std::pin::Pin;
+use std::sync::Arc;
+use tokio::io::{AsyncReadExt, AsyncSeek, AsyncSeekExt};
+use tvix_castore::blobservice::BlobService;
+use tvix_castore::directoryservice::DirectoryService;
+use tvix_castore::Node;
+
+#[apply(castore_fixtures_template)]
+#[tokio::test]
+async fn read_to_end(
+    #[future] blob_service: Arc<dyn BlobService>,
+    #[future] directory_service: Arc<dyn DirectoryService>,
+    #[case] test_input: &Node,
+    #[case] test_output: Result<Result<&Vec<u8>, io::ErrorKind>, crate::nar::RenderError>,
+) {
+    let reader_result = Reader::new(
+        test_input.clone(),
+        // don't put anything in the stores, as we don't actually do any requests.
+        blob_service.await,
+        directory_service.await,
+    )
+    .await;
+
+    match (reader_result, test_output) {
+        (Ok(_), Err(_)) => panic!("creating reader should have failed but succeeded"),
+        (Err(err), Ok(_)) => panic!("creating reader should have succeeded but failed: {}", err),
+        (Err(reader_err), Err(expected_err)) => {
+            assert_eq!(format!("{}", reader_err), format!("{}", expected_err));
+        }
+        (Ok(mut reader), Ok(expected_read_result)) => {
+            let mut buf: Vec<u8> = vec![];
+            let read_result = reader.read_to_end(&mut buf).await;
+
+            match (read_result, expected_read_result) {
+                (Ok(_), Err(_)) => panic!("read_to_end should have failed but succeeded"),
+                (Err(err), Ok(_)) => {
+                    panic!("read_to_end should have succeeded but failed: {}", err)
+                }
+                (Err(read_err), Err(expected_read_err)) => {
+                    assert_eq!(read_err.kind(), expected_read_err);
+                }
+                (Ok(_n), Ok(expected_read_result)) => {
+                    assert_eq!(buf, expected_read_result.to_vec());
+                }
+            }
+        }
+    }
+}
+
+#[rstest]
+#[tokio::test]
+/// Check that the Reader does not allow starting a seek while another seek is running
+/// If this is not prevented, it might lead to futures piling up on the heap
+async fn seek_twice(
+    #[future] blob_service: Arc<dyn BlobService>,
+    #[future] directory_service: Arc<dyn DirectoryService>,
+) {
+    let mut reader = Reader::new(
+        CASTORE_NODE_COMPLICATED.clone(),
+        // don't put anything in the stores, as we don't actually do any requests.
+        blob_service.await,
+        directory_service.await,
+    )
+    .await
+    .expect("must succeed");
+
+    Pin::new(&mut reader)
+        .start_seek(io::SeekFrom::Start(1))
+        .expect("must succeed");
+    let seek_err = Pin::new(&mut reader)
+        .start_seek(io::SeekFrom::Start(2))
+        .expect_err("must fail");
+
+    assert_eq!(seek_err.kind(), io::ErrorKind::Other);
+    assert_eq!(seek_err.to_string(), "Already seeking".to_string());
+}
+
+#[rstest]
+#[tokio::test]
+async fn seek(
+    #[future] blob_service: Arc<dyn BlobService>,
+    #[future] directory_service: Arc<dyn DirectoryService>,
+) {
+    let mut reader = Reader::new(
+        CASTORE_NODE_HELLOWORLD.clone(),
+        // don't put anything in the stores, as we don't actually do any requests.
+        blob_service.await,
+        directory_service.await,
+    )
+    .await
+    .expect("must succeed");
+
+    let mut buf = [0u8; 10];
+
+    for position in [
+        io::SeekFrom::Start(0x65), // Just before the file contents
+        io::SeekFrom::Start(0x68), // Seek back the file contents
+        io::SeekFrom::Start(0x70), // Just before the end of the file contents
+    ] {
+        let n = reader.seek(position).await.expect("seek") as usize;
+        reader.read_exact(&mut buf).await.expect("read_exact");
+        assert_eq!(NAR_CONTENTS_HELLOWORLD[n..n + 10], buf);
+    }
+}