From 830fdda8d47895dbbe145faaebfba27a1ad32289 Mon Sep 17 00:00:00 2001
From: Florian Klink <flokli@flokli.de>
Date: Sun, 30 Jun 2024 22:29:11 +0300
Subject: feat(tvix/store): use tokio-listener for tvix-store daemon command

This allows binding on unix sockets, as well as systemd socket
activation.

Change-Id: Icf648c4fd0895468c52607deb6397b8b5928102b
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11901
Autosubmit: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI
Reviewed-by: Connor Brewster <cbrewster@hey.com>
---
 tvix/Cargo.lock                  |  1 +
 tvix/Cargo.nix                   | 16 ++++++++++++++--
 tvix/store/Cargo.toml            |  2 +-
 tvix/store/src/bin/tvix-store.rs | 31 +++++++++++++++----------------
 4 files changed, 31 insertions(+), 19 deletions(-)

(limited to 'tvix')

diff --git a/tvix/Cargo.lock b/tvix/Cargo.lock
index d61dd20927a6..e758b359760d 100644
--- a/tvix/Cargo.lock
+++ b/tvix/Cargo.lock
@@ -3831,6 +3831,7 @@ version = "0.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4134661e12ec11c6276be73544a43144a357b08dfab5c41fd226e15b5bc9a6b2"
 dependencies = [
+ "clap",
  "document-features",
  "futures-core",
  "futures-util",
diff --git a/tvix/Cargo.nix b/tvix/Cargo.nix
index 5a2a3e2b5142..321ed98d4484 100644
--- a/tvix/Cargo.nix
+++ b/tvix/Cargo.nix
@@ -11722,6 +11722,13 @@ rec {
         edition = "2021";
         sha256 = "1cm6r5dmpq96s8gw9dgsinq5g8s466j48dg7dckwc4gc28g6cd21";
         dependencies = [
+          {
+            name = "clap";
+            packageId = "clap";
+            optional = true;
+            usesDefaultFeatures = false;
+            features = [ "derive" "std" ];
+          }
           {
             name = "document-features";
             packageId = "document-features";
@@ -11776,6 +11783,11 @@ rec {
           }
         ];
         devDependencies = [
+          {
+            name = "clap";
+            packageId = "clap";
+            features = [ "help" ];
+          }
           {
             name = "tokio";
             packageId = "tokio";
@@ -11801,7 +11813,7 @@ rec {
           "unix_path_tools" = [ "nix" ];
           "user_facing_default" = [ "inetd" "unix" "unix_path_tools" "sd_listen" "socket_options" ];
         };
-        resolvedDefaultFeatures = [ "default" "inetd" "nix" "sd_listen" "socket2" "socket_options" "tokio-util" "tonic011" "unix" "unix_path_tools" "user_facing_default" ];
+        resolvedDefaultFeatures = [ "clap" "default" "inetd" "multi-listener" "nix" "sd_listen" "socket2" "socket_options" "tokio-util" "tonic011" "unix" "unix_path_tools" "user_facing_default" ];
       };
       "tokio-macros" = rec {
         crateName = "tokio-macros";
@@ -14252,7 +14264,7 @@ rec {
           {
             name = "tokio-listener";
             packageId = "tokio-listener";
-            features = [ "tonic011" ];
+            features = [ "clap" "multi-listener" "sd_listen" "tonic011" ];
           }
           {
             name = "tokio-stream";
diff --git a/tvix/store/Cargo.toml b/tvix/store/Cargo.toml
index 98760589cc78..e8fbaf02c594 100644
--- a/tvix/store/Cargo.toml
+++ b/tvix/store/Cargo.toml
@@ -26,7 +26,7 @@ sha2 = "0.10.6"
 sled = { version = "0.34.7" }
 thiserror = "1.0.38"
 tokio = { version = "1.32.0", features = ["fs", "macros", "net", "rt", "rt-multi-thread", "signal"] }
-tokio-listener = { version = "0.4.1", features = [ "tonic011" ] }
+tokio-listener = { version = "0.4.2", features = [ "clap", "multi-listener", "sd_listen", "tonic011" ] }
 tokio-stream = { version = "0.1.14", features = ["fs"] }
 tokio-util = { version = "0.7.9", features = ["io", "io-util", "compat"] }
 tonic = { version = "0.11.0", features = ["tls", "tls-roots"] }
diff --git a/tvix/store/src/bin/tvix-store.rs b/tvix/store/src/bin/tvix-store.rs
index 657ce06720ee..04b4c7a6aab0 100644
--- a/tvix/store/src/bin/tvix-store.rs
+++ b/tvix/store/src/bin/tvix-store.rs
@@ -9,9 +9,6 @@ use serde::Deserialize;
 use serde::Serialize;
 use std::path::PathBuf;
 use std::sync::Arc;
-use tokio_listener::Listener;
-use tokio_listener::SystemOptions;
-use tokio_listener::UserOptions;
 use tonic::transport::Server;
 use tower::ServiceBuilder;
 use tower_http::trace::{DefaultMakeSpan, TraceLayer};
@@ -66,8 +63,9 @@ struct Cli {
 enum Commands {
     /// Runs the tvix-store daemon.
     Daemon {
-        #[arg(long, short = 'l')]
-        listen_address: Option<String>,
+        /// The address to listen on.
+        #[clap(flatten)]
+        listen_args: tokio_listener::ListenerAddressLFlag,
 
         #[arg(
             long,
@@ -198,7 +196,7 @@ fn default_threads() -> usize {
 async fn run_cli(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
     match cli.command {
         Commands::Daemon {
-            listen_address,
+            listen_args,
             blob_service_addr,
             directory_service_addr,
             path_info_service_addr,
@@ -212,11 +210,6 @@ async fn run_cli(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
                 )
                 .await?;
 
-            let listen_address = listen_address
-                .unwrap_or_else(|| "[::]:8000".to_string())
-                .parse()
-                .unwrap();
-
             let mut server = Server::builder().layer(
                 ServiceBuilder::new()
                     .layer(
@@ -251,15 +244,21 @@ async fn run_cli(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
                 router = router.add_service(reflection_svc);
             }
 
-            info!(listen_address=%listen_address, "starting daemon");
+            let listen_address = &listen_args.listen_address.unwrap_or_else(|| {
+                "[::]:8000"
+                    .parse()
+                    .expect("invalid fallback listen address")
+            });
 
-            let listener = Listener::bind(
-                &listen_address,
-                &SystemOptions::default(),
-                &UserOptions::default(),
+            let listener = tokio_listener::Listener::bind(
+                listen_address,
+                &Default::default(),
+                &listen_args.listener_options,
             )
             .await?;
 
+            info!(listen_address=%listen_address, "starting daemon");
+
             router.serve_with_incoming(listener).await?;
         }
         Commands::Import {
-- 
cgit 1.4.1