feat(tvix/nar-bridge): implement range request for NARs r/8737

With an implementation of AsyncRead + AsyncSeek, axum-range can answer
range requests.

We only use it if a range has been requested, as it uses more memory
than the linear variant.

Change-Id: I0072b0a09b328f3e932f14567a2caa3a49abcbf7
Reviewed-on: https://cl.tvl.fyi/c/depot/+/12509
Autosubmit: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI
Reviewed-by: raitobezarius <tvl@lahfa.xyz>
Reviewed-by: yuka <yuka@yuka.dev>

5 files changed, 352 insertions, 21 deletions

diff --git a/tvix/Cargo.lock b/tvix/Cargo.lock
index 22f4d194d43a..32674fbdfb91 100644
--- a/tvix/Cargo.lock
+++ b/tvix/Cargo.lock
@@ -362,6 +362,44 @@ dependencies = [
 ]
 
 [[package]]
+name = "axum-extra"
+version = "0.9.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0be6ea09c9b96cb5076af0de2e383bd2bc0c18f827cf1967bdd353e0b910d733"
+dependencies = [
+ "axum",
+ "axum-core",
+ "bytes",
+ "futures-util",
+ "headers",
+ "http",
+ "http-body",
+ "http-body-util",
+ "mime",
+ "pin-project-lite",
+ "serde",
+ "tower",
+ "tower-layer",
+ "tower-service",
+ "tracing",
+]
+
+[[package]]
+name = "axum-range"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b1c30398a7f716ebdd7f3c8a4f7a7a6df48a30e002007fd57b2a7a00fac864bd"
+dependencies = [
+ "axum",
+ "axum-extra",
+ "bytes",
+ "futures",
+ "http-body",
+ "pin-project",
+ "tokio",
+]
+
+[[package]]
 name = "backtrace"
 version = "0.3.69"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1501,6 +1539,30 @@ dependencies = [
 ]
 
 [[package]]
+name = "headers"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "322106e6bd0cba2d5ead589ddb8150a13d7c4217cf80d7c4f682ca994ccc6aa9"
+dependencies = [
+ "base64 0.21.7",
+ "bytes",
+ "headers-core",
+ "http",
+ "httpdate",
+ "mime",
+ "sha1",
+]
+
+[[package]]
+name = "headers-core"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "54b4a22553d4242c49fddb9ba998a99962b5cc6f22cb5a3482bec22522403ce4"
+dependencies = [
+ "http",
+]
+
+[[package]]
 name = "heck"
 version = "0.4.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -2133,6 +2195,8 @@ name = "nar-bridge"
 version = "0.1.0"
 dependencies = [
  "axum",
+ "axum-extra",
+ "axum-range",
  "bytes",
  "clap",
  "data-encoding",
diff --git a/tvix/Cargo.nix b/tvix/Cargo.nix
index 027aaba4845b..a1cb903b70a1 100644
--- a/tvix/Cargo.nix
+++ b/tvix/Cargo.nix
@@ -1338,6 +1338,173 @@ rec {
         };
         resolvedDefaultFeatures = [ "tracing" ];
       };
+      "axum-extra" = rec {
+        crateName = "axum-extra";
+        version = "0.9.3";
+        edition = "2021";
+        sha256 = "0cyp22wy0lykpmkikkr7z0c0rg6j7cw2xpphd83vav5rr44ymrhb";
+        libName = "axum_extra";
+        dependencies = [
+          {
+            name = "axum";
+            packageId = "axum";
+            usesDefaultFeatures = false;
+          }
+          {
+            name = "axum-core";
+            packageId = "axum-core";
+          }
+          {
+            name = "bytes";
+            packageId = "bytes";
+          }
+          {
+            name = "futures-util";
+            packageId = "futures-util";
+            usesDefaultFeatures = false;
+            features = [ "alloc" ];
+          }
+          {
+            name = "headers";
+            packageId = "headers";
+            optional = true;
+          }
+          {
+            name = "http";
+            packageId = "http";
+          }
+          {
+            name = "http-body";
+            packageId = "http-body";
+          }
+          {
+            name = "http-body-util";
+            packageId = "http-body-util";
+          }
+          {
+            name = "mime";
+            packageId = "mime";
+          }
+          {
+            name = "pin-project-lite";
+            packageId = "pin-project-lite";
+          }
+          {
+            name = "serde";
+            packageId = "serde";
+          }
+          {
+            name = "tower";
+            packageId = "tower";
+            usesDefaultFeatures = false;
+            features = [ "util" ];
+          }
+          {
+            name = "tower-layer";
+            packageId = "tower-layer";
+          }
+          {
+            name = "tower-service";
+            packageId = "tower-service";
+          }
+          {
+            name = "tracing";
+            packageId = "tracing";
+            optional = true;
+            usesDefaultFeatures = false;
+          }
+        ];
+        devDependencies = [
+          {
+            name = "axum";
+            packageId = "axum";
+          }
+          {
+            name = "serde";
+            packageId = "serde";
+            features = [ "derive" ];
+          }
+          {
+            name = "tower";
+            packageId = "tower";
+            features = [ "util" ];
+          }
+        ];
+        features = {
+          "async-read-body" = [ "dep:tokio-util" "tokio-util?/io" "dep:tokio" ];
+          "cookie" = [ "dep:cookie" ];
+          "cookie-key-expansion" = [ "cookie" "cookie?/key-expansion" ];
+          "cookie-private" = [ "cookie" "cookie?/private" ];
+          "cookie-signed" = [ "cookie" "cookie?/signed" ];
+          "default" = [ "tracing" ];
+          "erased-json" = [ "dep:serde_json" ];
+          "form" = [ "dep:serde_html_form" ];
+          "json-deserializer" = [ "dep:serde_json" "dep:serde_path_to_error" ];
+          "json-lines" = [ "dep:serde_json" "dep:tokio-util" "dep:tokio-stream" "tokio-util?/io" "tokio-stream?/io-util" "dep:tokio" ];
+          "multipart" = [ "dep:multer" ];
+          "protobuf" = [ "dep:prost" ];
+          "query" = [ "dep:serde_html_form" ];
+          "tracing" = [ "dep:tracing" "axum-core/tracing" ];
+          "typed-header" = [ "dep:headers" ];
+          "typed-routing" = [ "dep:axum-macros" "dep:percent-encoding" "dep:serde_html_form" "dep:form_urlencoded" ];
+        };
+        resolvedDefaultFeatures = [ "default" "tracing" "typed-header" ];
+      };
+      "axum-range" = rec {
+        crateName = "axum-range";
+        version = "0.4.0";
+        edition = "2021";
+        sha256 = "1gb4r3x00yiaggapy002w0q8mx3dg9x4z2iwgzfyn5pplyc07hxi";
+        libName = "axum_range";
+        dependencies = [
+          {
+            name = "axum";
+            packageId = "axum";
+            usesDefaultFeatures = false;
+          }
+          {
+            name = "axum-extra";
+            packageId = "axum-extra";
+            features = [ "typed-header" ];
+          }
+          {
+            name = "bytes";
+            packageId = "bytes";
+          }
+          {
+            name = "futures";
+            packageId = "futures";
+            usesDefaultFeatures = false;
+            features = [ "std" ];
+          }
+          {
+            name = "http-body";
+            packageId = "http-body";
+          }
+          {
+            name = "pin-project";
+            packageId = "pin-project";
+          }
+          {
+            name = "tokio";
+            packageId = "tokio";
+            features = [ "fs" "io-util" ];
+          }
+        ];
+        devDependencies = [
+          {
+            name = "axum";
+            packageId = "axum";
+            features = [ "macros" ];
+          }
+          {
+            name = "tokio";
+            packageId = "tokio";
+            features = [ "rt" "rt-multi-thread" "macros" ];
+          }
+        ];
+        features = { };
+      };
       "backtrace" = rec {
         crateName = "backtrace";
         version = "0.3.69";
@@ -1412,7 +1579,7 @@ rec {
           "default" = [ "std" ];
           "std" = [ "alloc" ];
         };
-        resolvedDefaultFeatures = [ "alloc" "std" ];
+        resolvedDefaultFeatures = [ "alloc" "default" "std" ];
       };
       "base64 0.22.1" = rec {
         crateName = "base64";
@@ -4734,6 +4901,63 @@ rec {
         };
         resolvedDefaultFeatures = [ "ahash" "allocator-api2" "default" "inline-more" "raw" ];
       };
+      "headers" = rec {
+        crateName = "headers";
+        version = "0.4.0";
+        edition = "2015";
+        sha256 = "1abari69kjl2yv2dg06g2x17qgd1a20xp7aqmmg2vfhcppk0c89j";
+        authors = [
+          "Sean McArthur <sean@seanmonstar.com>"
+        ];
+        dependencies = [
+          {
+            name = "base64";
+            packageId = "base64 0.21.7";
+          }
+          {
+            name = "bytes";
+            packageId = "bytes";
+          }
+          {
+            name = "headers-core";
+            packageId = "headers-core";
+          }
+          {
+            name = "http";
+            packageId = "http";
+          }
+          {
+            name = "httpdate";
+            packageId = "httpdate";
+          }
+          {
+            name = "mime";
+            packageId = "mime";
+          }
+          {
+            name = "sha1";
+            packageId = "sha1";
+          }
+        ];
+        features = { };
+      };
+      "headers-core" = rec {
+        crateName = "headers-core";
+        version = "0.3.0";
+        edition = "2015";
+        sha256 = "1r1w80i2bhmyh8s5mjr2dz6baqlrm6cak6yvzm4jq96lacjs5d2l";
+        libName = "headers_core";
+        authors = [
+          "Sean McArthur <sean@seanmonstar.com>"
+        ];
+        dependencies = [
+          {
+            name = "http";
+            packageId = "http";
+          }
+        ];
+
+      };
       "heck 0.4.1" = rec {
         crateName = "heck";
         version = "0.4.1";
@@ -6616,6 +6840,14 @@ rec {
             features = [ "http2" ];
           }
           {
+            name = "axum-extra";
+            packageId = "axum-extra";
+          }
+          {
+            name = "axum-range";
+            packageId = "axum-range";
+          }
+          {
             name = "bytes";
             packageId = "bytes";
           }
diff --git a/tvix/Cargo.toml b/tvix/Cargo.toml
index 16bf9c2dd99c..bcb8e6f7e78e 100644
--- a/tvix/Cargo.toml
+++ b/tvix/Cargo.toml
@@ -46,6 +46,8 @@ async-process = "2.2.4"
 async-stream = "0.3.5"
 async-tempfile = "0.4.0"
 axum = "0.7.5"
+axum-extra = "0.9.3"
+axum-range = "0.4.0"
 # https://github.com/liufuyang/bigtable_rs/pull/86
 bigtable_rs = { git = "https://github.com/liufuyang/bigtable_rs", rev = "1818355a5373a5bc2c84287e3a4e3807154ac8ef" }
 bitflags = "2.6.0"
diff --git a/tvix/nar-bridge/Cargo.toml b/tvix/nar-bridge/Cargo.toml
index 6ca0479a9a81..ac23b597311f 100644
--- a/tvix/nar-bridge/Cargo.toml
+++ b/tvix/nar-bridge/Cargo.toml
@@ -5,6 +5,8 @@ edition = "2021"
 
 [dependencies]
 axum = { workspace = true, features = ["http2"] }
+axum-extra = { workspace = true }
+axum-range = { workspace = true }
 tower = { workspace = true }
 tower-http = { workspace = true, features = ["trace"] }
 bytes = { workspace = true }
diff --git a/tvix/nar-bridge/src/nar.rs b/tvix/nar-bridge/src/nar.rs
index 2f6c0cc5a0ef..707c01e4bcda 100644
--- a/tvix/nar-bridge/src/nar.rs
+++ b/tvix/nar-bridge/src/nar.rs
@@ -1,7 +1,9 @@
-use axum::body::Body;
 use axum::extract::Query;
 use axum::http::StatusCode;
 use axum::response::Response;
+use axum::{body::Body, response::IntoResponse};
+use axum_extra::{headers::Range, TypedHeader};
+use axum_range::{KnownSize, Ranged};
 use bytes::Bytes;
 use data_encoding::BASE64URL_NOPAD;
 use futures::TryStreamExt;
@@ -22,6 +24,7 @@ pub(crate) struct GetNARParams {
 
 #[instrument(skip(blob_service, directory_service))]
 pub async fn get(
+    ranges: Option<TypedHeader<Range>>,
     axum::extract::Path(root_node_enc): axum::extract::Path<String>,
     axum::extract::Query(GetNARParams { nar_size }): Query<GetNARParams>,
     axum::extract::State(AppState {
@@ -29,7 +32,7 @@ pub async fn get(
         directory_service,
         ..
     }): axum::extract::State<AppState>,
-) -> Result<Response, StatusCode> {
+) -> Result<impl axum::response::IntoResponse, StatusCode> {
     use prost::Message;
     // b64decode the root node passed *by the user*
     let root_node_proto = BASE64URL_NOPAD
@@ -62,24 +65,52 @@ pub async fn get(
         return Err(StatusCode::BAD_REQUEST);
     }
 
-    let (w, r) = tokio::io::duplex(1024 * 8);
-
-    // spawn a task rendering the NAR to the client
-    tokio::spawn(async move {
-        if let Err(e) =
-            tvix_store::nar::write_nar(w, &root_node, blob_service, directory_service).await
-        {
-            warn!(err=%e, "failed to write out NAR");
-        }
-    });
-
-    Ok(Response::builder()
-        .status(StatusCode::OK)
-        .header("cache-control", "max-age=31536000, immutable")
-        .header("content-length", nar_size)
-        .header("content-type", nix_http::MIME_TYPE_NAR)
-        .body(Body::from_stream(ReaderStream::new(r)))
-        .unwrap())
+    Ok((
+        // headers
+        [
+            ("cache-control", "max-age=31536000, immutable"),
+            ("content-type", nix_http::MIME_TYPE_NAR),
+        ],
+        // If this is a range request, construct a seekable NAR reader
+        if let Some(TypedHeader(ranges)) = ranges {
+            let r =
+                tvix_store::nar::seekable::Reader::new(root_node, blob_service, directory_service)
+                    .await
+                    .map_err(|e| {
+                        warn!(err=%e, "failed to construct seekable nar reader");
+                        StatusCode::INTERNAL_SERVER_ERROR
+                    })?;
+
+            // ensure the user-supplied nar size was correct, no point returning data otherwise.
+            if r.stream_len() != nar_size {
+                warn!(
+                    actual_nar_size = r.stream_len(),
+                    supplied_nar_size = nar_size,
+                    "wrong nar size supplied"
+                );
+                return Err(StatusCode::BAD_REQUEST);
+            }
+            Ranged::new(Some(ranges), KnownSize::sized(r, nar_size)).into_response()
+        } else {
+            // use the non-seekable codepath if there's no range(s) requested,
+            // as it uses less memory.
+            let (w, r) = tokio::io::duplex(1024 * 8);
+
+            // spawn a task rendering the NAR to the client.
+            tokio::spawn(async move {
+                if let Err(e) =
+                    tvix_store::nar::write_nar(w, &root_node, blob_service, directory_service).await
+                {
+                    warn!(err=%e, "failed to write out NAR");
+                }
+            });
+
+            Response::builder()
+                .header("content-length", nar_size)
+                .body(Body::from_stream(ReaderStream::new(r)))
+                .unwrap()
+        },
+    ))
 }
 
 #[instrument(skip(blob_service, directory_service, request))]