feat(tvix/nix-compat): worker protocol operation parser r/7784

Change-Id: I7776635b17c44534223603d28cf59c7eebd976e0
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11229
Reviewed-by: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI

7 files changed, 206 insertions, 11 deletions

diff --git a/tvix/Cargo.lock b/tvix/Cargo.lock
index 332418cdf7f5..797cdd5d54dd 100644
--- a/tvix/Cargo.lock
+++ b/tvix/Cargo.lock
@@ -874,6 +874,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c34f04666d835ff5d62e058c3995147c06f42fe86ff053337632bca83e42702d"
 
 [[package]]
+name = "enum-primitive-derive"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ba7795da175654fe16979af73f81f26a8ea27638d8d9823d317016888a63dc4c"
+dependencies = [
+ "num-traits",
+ "quote 1.0.35",
+ "syn 2.0.48",
+]
+
+[[package]]
 name = "equivalent"
 version = "1.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1863,12 +1874,14 @@ dependencies = [
  "data-encoding",
  "ed25519",
  "ed25519-dalek",
+ "enum-primitive-derive",
  "futures",
  "futures-util",
  "glob",
  "hex-literal",
  "lazy_static",
  "nom",
+ "num-traits",
  "pretty_assertions",
  "serde",
  "serde_json",
@@ -1912,9 +1925,9 @@ dependencies = [
 
 [[package]]
 name = "num-traits"
-version = "0.2.17"
+version = "0.2.18"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "39e3200413f237f41ab11ad6d161bc7239c84dcb631773ccd7de3dfe4b5c267c"
+checksum = "da0df0e5185db44f69b44f26786fe401b6c293d1907744beaa7fa62b2e5a517a"
 dependencies = [
  "autocfg",
  "libm",
diff --git a/tvix/Cargo.nix b/tvix/Cargo.nix
index 0f0268bfe737..001b63d526fb 100644
--- a/tvix/Cargo.nix
+++ b/tvix/Cargo.nix
@@ -2692,6 +2692,32 @@ rec {
         ];
 
       };
+      "enum-primitive-derive" = rec {
+        crateName = "enum-primitive-derive";
+        version = "0.3.0";
+        edition = "2018";
+        sha256 = "0k6wcf58h5kh64yq5nfq71va53kaya0kzxwsjwbgwm2n2zd9axxs";
+        procMacro = true;
+        authors = [
+          "Doug Goldstein <cardoe@cardoe.com>"
+        ];
+        dependencies = [
+          {
+            name = "num-traits";
+            packageId = "num-traits";
+            usesDefaultFeatures = false;
+          }
+          {
+            name = "quote";
+            packageId = "quote 1.0.35";
+          }
+          {
+            name = "syn";
+            packageId = "syn 2.0.48";
+          }
+        ];
+
+      };
       "equivalent" = rec {
         crateName = "equivalent";
         version = "1.0.1";
@@ -5722,6 +5748,10 @@ rec {
             packageId = "ed25519-dalek";
           }
           {
+            name = "enum-primitive-derive";
+            packageId = "enum-primitive-derive";
+          }
+          {
             name = "futures-util";
             packageId = "futures-util";
             optional = true;
@@ -5736,6 +5766,10 @@ rec {
             packageId = "nom";
           }
           {
+            name = "num-traits";
+            packageId = "num-traits";
+          }
+          {
             name = "serde";
             packageId = "serde";
             features = [ "derive" ];
@@ -5886,9 +5920,9 @@ rec {
       };
       "num-traits" = rec {
         crateName = "num-traits";
-        version = "0.2.17";
+        version = "0.2.18";
         edition = "2018";
-        sha256 = "0z16bi5zwgfysz6765v3rd6whfbjpihx3mhsn4dg8dzj2c221qrr";
+        sha256 = "0yjib8p2p9kzmaz48xwhs69w5dh1wipph9jgnillzd2x33jz03fs";
         authors = [
           "The Rust Project Developers"
         ];
diff --git a/tvix/nix-compat/Cargo.toml b/tvix/nix-compat/Cargo.toml
index 181eb9428915..6e3df0485db1 100644
--- a/tvix/nix-compat/Cargo.toml
+++ b/tvix/nix-compat/Cargo.toml
@@ -14,9 +14,11 @@ bstr = { version = "1.6.0", features = ["alloc", "unicode", "serde"] }
 data-encoding = "2.3.3"
 ed25519 = "2.2.3"
 ed25519-dalek = "2.1.0"
+enum-primitive-derive = "0.3.0"
 futures-util = { version = "0.3.30", features = ["io"], optional = true }
 glob = "0.3.0"
 nom = "7.1.3"
+num-traits = "0.2.18"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 sha2 = "0.10.6"
diff --git a/tvix/nix-compat/src/wire/mod.rs b/tvix/nix-compat/src/wire/mod.rs
index 9444ebbcfe59..cd7b87aacd85 100644
--- a/tvix/nix-compat/src/wire/mod.rs
+++ b/tvix/nix-compat/src/wire/mod.rs
@@ -6,3 +6,6 @@ pub mod bytes;
 
 #[cfg(feature = "async")]
 pub mod primitive;
+
+#[cfg(feature = "async")]
+pub mod worker_protocol;
diff --git a/tvix/nix-compat/src/wire/worker_protocol.rs b/tvix/nix-compat/src/wire/worker_protocol.rs
new file mode 100644
index 000000000000..82f227c72362
--- /dev/null
+++ b/tvix/nix-compat/src/wire/worker_protocol.rs
@@ -0,0 +1,83 @@
+use std::io::{Error, ErrorKind};
+
+use enum_primitive_derive::Primitive;
+use num_traits::{FromPrimitive, ToPrimitive};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+use crate::wire::primitive;
+
+/// Worker Operation
+///
+/// These operations are encoded as unsigned 64 bits before being sent
+/// to the wire. See the [read_op] and
+/// [write_op] operations to serialize/deserialize the
+/// operation on the wire.
+///
+/// Note: for now, we're using the Nix 2.20 operation description. The
+/// operations marked as obsolete are obsolete for Nix 2.20, not
+/// necessarily for Nix 2.3. We'll revisit this later on.
+#[derive(Debug, PartialEq, Primitive)]
+pub enum Operation {
+    IsValidPath = 1,
+    HasSubstitutes = 3,
+    QueryPathHash = 4,   // obsolete
+    QueryReferences = 5, // obsolete
+    QueryReferrers = 6,
+    AddToStore = 7,
+    AddTextToStore = 8, // obsolete since 1.25, Nix 3.0. Use WorkerProto::Op::AddToStore
+    BuildPaths = 9,
+    EnsurePath = 10,
+    AddTempRoot = 11,
+    AddIndirectRoot = 12,
+    SyncWithGC = 13,
+    FindRoots = 14,
+    ExportPath = 16,   // obsolete
+    QueryDeriver = 18, // obsolete
+    SetOptions = 19,
+    CollectGarbage = 20,
+    QuerySubstitutablePathInfo = 21,
+    QueryDerivationOutputs = 22, // obsolete
+    QueryAllValidPaths = 23,
+    QueryFailedPaths = 24,
+    ClearFailedPaths = 25,
+    QueryPathInfo = 26,
+    ImportPaths = 27,                // obsolete
+    QueryDerivationOutputNames = 28, // obsolete
+    QueryPathFromHashPart = 29,
+    QuerySubstitutablePathInfos = 30,
+    QueryValidPaths = 31,
+    QuerySubstitutablePaths = 32,
+    QueryValidDerivers = 33,
+    OptimiseStore = 34,
+    VerifyStore = 35,
+    BuildDerivation = 36,
+    AddSignatures = 37,
+    NarFromPath = 38,
+    AddToStoreNar = 39,
+    QueryMissing = 40,
+    QueryDerivationOutputMap = 41,
+    RegisterDrvOutput = 42,
+    QueryRealisation = 43,
+    AddMultipleToStore = 44,
+    AddBuildLog = 45,
+    BuildPathsWithResults = 46,
+    AddPermRoot = 47,
+}
+
+/// Read a worker [Operation] from the wire.
+pub async fn read_op<R: AsyncReadExt + Unpin>(r: &mut R) -> std::io::Result<Operation> {
+    let op_number = primitive::read_u64(r).await?;
+    Operation::from_u64(op_number).ok_or(Error::new(
+        ErrorKind::Other,
+        format!("Invalid OP number {}", op_number),
+    ))
+}
+
+/// Write a worker [Operation] to the wire.
+pub async fn write_op<W: AsyncWriteExt + Unpin>(w: &mut W, op: &Operation) -> std::io::Result<()> {
+    let op = Operation::to_u64(op).ok_or(Error::new(
+        ErrorKind::Other,
+        format!("Can't convert the OP {:?} to u64", op),
+    ))?;
+    w.write_u64(op).await
+}
diff --git a/users/picnoir/tvix-daemon/Cargo.lock b/users/picnoir/tvix-daemon/Cargo.lock
index c53c72c09ee3..635c54fe4a8a 100644
--- a/users/picnoir/tvix-daemon/Cargo.lock
+++ b/users/picnoir/tvix-daemon/Cargo.lock
@@ -656,6 +656,17 @@ dependencies = [
 ]
 
 [[package]]
+name = "enum-primitive-derive"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ba7795da175654fe16979af73f81f26a8ea27638d8d9823d317016888a63dc4c"
+dependencies = [
+ "num-traits",
+ "quote",
+ "syn 2.0.52",
+]
+
+[[package]]
 name = "equivalent"
 version = "1.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1393,9 +1404,11 @@ dependencies = [
  "data-encoding",
  "ed25519",
  "ed25519-dalek",
+ "enum-primitive-derive",
  "futures-util",
  "glob",
  "nom",
+ "num-traits",
  "serde",
  "serde_json",
  "sha2",
diff --git a/users/picnoir/tvix-daemon/Cargo.nix b/users/picnoir/tvix-daemon/Cargo.nix
index 77fd020f9045..f284ba334238 100644
--- a/users/picnoir/tvix-daemon/Cargo.nix
+++ b/users/picnoir/tvix-daemon/Cargo.nix
@@ -1,5 +1,5 @@
-# This file was @generated by crate2nix 0.12.0 with the command:
-#   "generate"
+# This file was @generated by crate2nix 0.13.0 with the command:
+#   "generate" "--all-features"
 # See https://github.com/kolloch/crate2nix for more info.
 
 { nixpkgs ? <nixpkgs>
@@ -2039,6 +2039,32 @@ rec {
         };
         resolvedDefaultFeatures = [ "alloc" "default" ];
       };
+      "enum-primitive-derive" = rec {
+        crateName = "enum-primitive-derive";
+        version = "0.3.0";
+        edition = "2018";
+        sha256 = "0k6wcf58h5kh64yq5nfq71va53kaya0kzxwsjwbgwm2n2zd9axxs";
+        procMacro = true;
+        authors = [
+          "Doug Goldstein <cardoe@cardoe.com>"
+        ];
+        dependencies = [
+          {
+            name = "num-traits";
+            packageId = "num-traits";
+            usesDefaultFeatures = false;
+          }
+          {
+            name = "quote";
+            packageId = "quote";
+          }
+          {
+            name = "syn";
+            packageId = "syn 2.0.52";
+          }
+        ];
+
+      };
       "equivalent" = rec {
         crateName = "equivalent";
         version = "1.0.1";
@@ -4266,6 +4292,10 @@ rec {
             packageId = "ed25519-dalek";
           }
           {
+            name = "enum-primitive-derive";
+            packageId = "enum-primitive-derive";
+          }
+          {
             name = "futures-util";
             packageId = "futures-util";
             optional = true;
@@ -4280,6 +4310,10 @@ rec {
             packageId = "nom";
           }
           {
+            name = "num-traits";
+            packageId = "num-traits";
+          }
+          {
             name = "serde";
             packageId = "serde";
             features = [ "derive" ];
@@ -4388,7 +4422,7 @@ rec {
           "default" = [ "std" ];
           "libm" = [ "dep:libm" ];
         };
-        resolvedDefaultFeatures = [ "std" ];
+        resolvedDefaultFeatures = [ "default" "std" ];
       };
       "num_cpus" = rec {
         crateName = "num_cpus";
@@ -9182,7 +9216,7 @@ rec {
           {
             name = "object_store";
             packageId = "object_store";
-            features = [ "aws" "azure" "gcp" "http" ];
+            features = [ "http" ];
           }
           {
             name = "parking_lot";
@@ -9260,12 +9294,13 @@ rec {
           }
         ];
         features = {
+          "cloud" = [ "object_store/aws" "object_store/azure" "object_store/gcp" ];
           "fs" = [ "dep:libc" "dep:fuse-backend-rs" ];
           "fuse" = [ "fs" ];
           "tonic-reflection" = [ "dep:tonic-reflection" ];
           "virtiofs" = [ "fs" "dep:vhost" "dep:vhost-user-backend" "dep:virtio-queue" "dep:vm-memory" "dep:vmm-sys-util" "dep:virtio-bindings" "fuse-backend-rs?/vhost-user-fs" "fuse-backend-rs?/virtiofs" ];
         };
-        resolvedDefaultFeatures = [ "default" "fs" "fuse" "tonic-reflection" ];
+        resolvedDefaultFeatures = [ "cloud" "default" "fs" "fuse" "tonic-reflection" ];
       };
       "tvix-daemon" = rec {
         crateName = "tvix-daemon";
@@ -9508,13 +9543,14 @@ rec {
           }
         ];
         features = {
-          "default" = [ "fuse" "otlp" "tonic-reflection" ];
+          "cloud" = [ "tvix-castore/cloud" ];
+          "default" = [ "cloud" "fuse" "otlp" "tonic-reflection" ];
           "fuse" = [ "tvix-castore/fuse" ];
           "otlp" = [ "dep:opentelemetry" "dep:opentelemetry-otlp" "dep:opentelemetry_sdk" ];
           "tonic-reflection" = [ "dep:tonic-reflection" "tvix-castore/tonic-reflection" ];
           "virtiofs" = [ "tvix-castore/virtiofs" ];
         };
-        resolvedDefaultFeatures = [ "default" "fuse" "otlp" "tonic-reflection" ];
+        resolvedDefaultFeatures = [ "cloud" "default" "fuse" "otlp" "tonic-reflection" ];
       };
       "typenum" = rec {
         crateName = "typenum";
@@ -11665,6 +11701,7 @@ rec {
               (
                 _: {
                   buildTests = true;
+                  release = false;
                 }
               );
             # If the user hasn't set any pre/post commands, we don't want to
@@ -11689,6 +11726,16 @@ rec {
             # recreate a file hierarchy as when running tests with cargo
 
             # the source for test data
+            # It's necessary to locate the source in $NIX_BUILD_TOP/source/
+            # instead of $NIX_BUILD_TOP/
+            # because we compiled those test binaries in the former and not the latter.
+            # So all paths will expect source tree to be there and not in the build top directly.
+            # For example: $NIX_BUILD_TOP := /build in general, if you ask yourself.
+            # TODO(raitobezarius): I believe there could be more edge cases if `crate.sourceRoot`
+            # do exist but it's very hard to reason about them, so let's wait until the first bug report.
+            mkdir -p source/
+            cd source/
+
             ${pkgs.buildPackages.xorg.lndir}/bin/lndir ${crate.src}
 
             # build outputs