feat(tvix/store): fix ctrl-c handling on mount command r/6557

This enables the tokio `signal` feature, and registers a ctrl_c signal
handler, which will use the unmount handle to unmount in case a ctrl-c
signal is received.

This avoids having disconnected mountpoints when Ctrl-C'ing a
`tvix-store mount` invocation.

In case the filesystem is unmounted externally (via `umount /path/to/
mountpoint`), the future is waiting for the signal is never resolved and
the task is stopped.

Change-Id: I149f705a6cb50188177f2a6c6a5fcd77218e2a3f
Reviewed-on: https://cl.tvl.fyi/c/depot/+/9218
Reviewed-by: tazjin <tazjin@tvl.su>
Autosubmit: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI

2 files changed, 25 insertions, 4 deletions

diff --git a/tvix/store/Cargo.toml b/tvix/store/Cargo.toml
index 0dcada31a61e..dece06be8fb4 100644
--- a/tvix/store/Cargo.toml
+++ b/tvix/store/Cargo.toml
@@ -17,7 +17,7 @@ sha2 = "0.10.6"
 sled = { version = "0.34.7", features = ["compression"] }
 thiserror = "1.0.38"
 tokio-stream = "0.1.14"
-tokio = { version = "1.28.0", features = ["rt-multi-thread", "net"] }
+tokio = { version = "1.28.0", features = ["net", "rt-multi-thread", "signal"] }
 tonic = "0.8.2"
 tracing = "0.1.37"
 tracing-subscriber = { version = "0.3.16", features = ["json"] }
diff --git a/tvix/store/src/bin/tvix-store.rs b/tvix/store/src/bin/tvix-store.rs
index d2a8927351f9..c11af5765bd9 100644
--- a/tvix/store/src/bin/tvix-store.rs
+++ b/tvix/store/src/bin/tvix-store.rs
@@ -264,16 +264,37 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 directory_service.clone(),
             )?;
 
-            tokio::task::spawn_blocking(move || {
+            let mut fuse_session = tokio::task::spawn_blocking(move || {
                 let f = FUSE::new(
                     blob_service,
                     directory_service,
                     path_info_service,
                     list_root,
                 );
-                fuser::mount2(f, &dest, &[])
+
+                fuser::Session::new(f, &dest, &[])
+            })
+            .await??;
+
+            // grab a handle to unmount the file system, and register a signal
+            // handler.
+            let mut fuse_unmounter = fuse_session.unmount_callable();
+            tokio::spawn(async move {
+                tokio::signal::ctrl_c().await.unwrap();
+                info!("interrupt received, unmounting…");
+                fuse_unmounter.unmount().unwrap();
+            });
+
+            // Start the fuse filesystem and wait for its completion, which
+            // happens when it's unmounted externally, or via the signal handler
+            // task.
+            tokio::task::spawn_blocking(move || -> io::Result<()> {
+                info!("mounting tvix-store on {:?}", fuse_session.mountpoint());
+                let res = fuse_session.run()?;
+                info!("unmount occured, terminating…");
+                Ok(res)
             })
-            .await??
+            .await??;
         }
     };
     Ok(())