feat(tvix/nar-bridge): support uploading NAR files r/8377

This ingests NAR files into the {Blob,Directory}Service, which are
already part of the AppState.

As we then need to correlate the root node to the uploaded PathInfo, we
need to keep a (short-lived) lookup table from NARHash to root node
around. We insert it into a `LruCache` after the NAR is uploaded, and
use `peek()` to do the lookup, which doesn't update the LRU list.

Change-Id: I48a4c6246bacf76559c5a4ccad2a0bc25c1b7900
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11986
Tested-by: BuildkiteCI
Reviewed-by: Brian Olsen <me@griff.name>

1 files changed, 20 insertions, 1 deletions

diff --git a/tvix/nar-bridge/src/lib.rs b/tvix/nar-bridge/src/lib.rs
index 5d9ec43c1cde..5f0e8c19d26a 100644
--- a/tvix/nar-bridge/src/lib.rs
+++ b/tvix/nar-bridge/src/lib.rs
@@ -1,18 +1,32 @@
-use axum::routing::head;
+use axum::routing::{head, put};
 use axum::{routing::get, Router};
+use lru::LruCache;
+use parking_lot::RwLock;
+use std::num::NonZeroUsize;
 use std::sync::Arc;
 use tvix_castore::blobservice::BlobService;
 use tvix_castore::directoryservice::DirectoryService;
+use tvix_castore::proto::node::Node;
 use tvix_store::pathinfoservice::PathInfoService;
 
 mod nar;
 mod narinfo;
 
+/// The capacity of the lookup table from NarHash to [Node].
+/// Should be bigger than the number of concurrent NAR upload.
+/// Cannot be [NonZeroUsize] here due to rust-analyzer going bananas.
+/// SAFETY: 1000 != 0
+const ROOT_NODES_CACHE_CAPACITY: usize = 1000;
+
 #[derive(Clone)]
 pub struct AppState {
     blob_service: Arc<dyn BlobService>,
     directory_service: Arc<dyn DirectoryService>,
     path_info_service: Arc<dyn PathInfoService>,
+
+    /// Lookup table from NarHash to [Node], necessary to populate the root_node
+    /// field of the PathInfo when processing the narinfo upload.
+    root_nodes: Arc<RwLock<LruCache<[u8; 32], Node>>>,
 }
 
 impl AppState {
@@ -25,6 +39,10 @@ impl AppState {
             blob_service,
             directory_service,
             path_info_service,
+            root_nodes: Arc::new(RwLock::new(LruCache::new({
+                // SAFETY: 1000 != 0
+                unsafe { NonZeroUsize::new_unchecked(ROOT_NODES_CACHE_CAPACITY) }
+            }))),
         }
     }
 }
@@ -32,6 +50,7 @@ impl AppState {
 pub fn gen_router(priority: u64) -> Router<AppState> {
     Router::new()
         .route("/", get(root))
+        .route("/nar/:nar_str", put(nar::put))
         .route("/nar/tvix-castore/:root_node_enc", get(nar::get))
         .route("/:narinfo_str", get(narinfo::get))
         .route("/:narinfo_str", head(narinfo::head))