From 0c5ad949148b178c13b0cb8f2ed3fe274e9cf522 Mon Sep 17 00:00:00 2001
From: Florian Klink <flokli@flokli.de>
Date: Sat, 9 Nov 2024 16:00:23 +0000
Subject: feat(tvix/nar-bridge): make root_nodes_cache_capacity configurable

Allow this to be overridden via the CLI, also getting rid of the unsafe
because this doesn't need to be const anymore.

Change-Id: I5e51b52e42522a21f59ef69628b464477c0764d1
Reviewed-on: https://cl.tvl.fyi/c/depot/+/12753
Tested-by: BuildkiteCI
Autosubmit: flokli <flokli@flokli.de>
Reviewed-by: Ilan Joselevich <personal@ilanjoselevich.com>
---
 tvix/nar-bridge/src/bin/nar-bridge.rs | 13 ++++++++++++-
 tvix/nar-bridge/src/lib.rs            | 12 ++----------
 2 files changed, 14 insertions(+), 11 deletions(-)

(limited to 'tvix/nar-bridge/src')

diff --git a/tvix/nar-bridge/src/bin/nar-bridge.rs b/tvix/nar-bridge/src/bin/nar-bridge.rs
index 48eba0ac44ac..d9fce9dff322 100644
--- a/tvix/nar-bridge/src/bin/nar-bridge.rs
+++ b/tvix/nar-bridge/src/bin/nar-bridge.rs
@@ -1,6 +1,7 @@
 use clap::Parser;
 use mimalloc::MiMalloc;
 use nar_bridge::AppState;
+use std::num::NonZeroUsize;
 use tower::ServiceBuilder;
 use tower_http::trace::{DefaultMakeSpan, TraceLayer};
 use tracing::info;
@@ -25,6 +26,11 @@ struct Cli {
     #[clap(flatten)]
     listen_args: tokio_listener::ListenerAddressLFlag,
 
+    /// The capacity of the lookup table from NarHash to [Node].
+    /// Should be bigger than the number of concurrent NAR uploads.
+    #[arg(long, env, default_value_t = NonZeroUsize::new(1000).unwrap())]
+    root_nodes_cache_capacity: NonZeroUsize,
+
     #[cfg(feature = "otlp")]
     /// Whether to configure OTLP. Set --otlp=false to disable.
     #[arg(long, default_missing_value = "true", default_value = "true", num_args(0..=1), require_equals(true), action(clap::ArgAction::Set))]
@@ -51,7 +57,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     let (blob_service, directory_service, path_info_service, _nar_calculation_service) =
         tvix_store::utils::construct_services(cli.service_addrs).await?;
 
-    let state = AppState::new(blob_service, directory_service, path_info_service);
+    let state = AppState::new(
+        blob_service,
+        directory_service,
+        path_info_service,
+        cli.root_nodes_cache_capacity,
+    );
 
     let app = nar_bridge::gen_router(cli.priority)
         .layer(
diff --git a/tvix/nar-bridge/src/lib.rs b/tvix/nar-bridge/src/lib.rs
index 9246a334c7f6..db926e8cede4 100644
--- a/tvix/nar-bridge/src/lib.rs
+++ b/tvix/nar-bridge/src/lib.rs
@@ -15,12 +15,6 @@ 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>,
@@ -37,15 +31,13 @@ impl AppState {
         blob_service: Arc<dyn BlobService>,
         directory_service: Arc<dyn DirectoryService>,
         path_info_service: Arc<dyn PathInfoService>,
+        root_nodes_cache_capacity: NonZeroUsize,
     ) -> Self {
         Self {
             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) }
-            }))),
+            root_nodes: Arc::new(RwLock::new(LruCache::new(root_nodes_cache_capacity))),
         }
     }
 }
-- 
cgit 1.4.1