agent: add tracing + graceful shutdown

agent/Cargo.lock

@@ -433,6 +433,16 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "nu-ansi-term"
+version = "0.46.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84"
+dependencies = [
+ "overload",
+ "winapi",
+]
+
 [[package]]
 name = "num_cpus"
 version = "1.15.0"
@@ -449,6 +459,12 @@ version = "1.18.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "dd8b5dd2ae5ed71462c540258bedcb51965123ad7e7ccf4b9a8cafaa4a63576d"
 
+[[package]]
+name = "overload"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39"
+
 [[package]]
 name = "percent-encoding"
 version = "2.3.0"
@@ -587,6 +603,8 @@ dependencies = [
  "tokio-stream",
  "tonic",
  "tonic-build",
+ "tracing",
+ "tracing-subscriber",
 ]
 
 [[package]]
@@ -678,6 +696,15 @@ version = "1.0.164"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9e8c8cf938e98f769bc164923b06dce91cea1751522f46f8466461af04c9027d"
 
+[[package]]
+name = "sharded-slab"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "900fba806f70c630b0a382d0d825e17a0f19fcd059a2ade1ff237bcddf446b31"
+dependencies = [
+ "lazy_static",
+]
+
 [[package]]
 name = "signal-hook-registry"
 version = "1.4.1"
@@ -696,6 +723,12 @@ dependencies = [
  "autocfg",
 ]
 
+[[package]]
+name = "smallvec"
+version = "1.10.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a507befe795404456341dfab10cef66ead4c041f62b8b11bbb92bffe5d0953e0"
+
 [[package]]
 name = "socket2"
 version = "0.4.9"
@@ -762,6 +795,16 @@ dependencies = [
  "windows-sys",
 ]
 
+[[package]]
+name = "thread_local"
+version = "1.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3fdd6f064ccff2d6567adcb3873ca630700f00b5ad3f060c25b5dcfd9a4ce152"
+dependencies = [
+ "cfg-if",
+ "once_cell",
+]
+
 [[package]]
 name = "tokio"
 version = "1.28.2"
@@ -929,6 +972,32 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0955b8137a1df6f1a2e9a37d8a6656291ff0297c1a97c24e0d8425fe2312f79a"
 dependencies = [
  "once_cell",
+ "valuable",
+]
+
+[[package]]
+name = "tracing-log"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "78ddad33d2d10b1ed7eb9d1f518a5674713876e97e5bb9b7345a7984fbb4f922"
+dependencies = [
+ "lazy_static",
+ "log",
+ "tracing-core",
+]
+
+[[package]]
+name = "tracing-subscriber"
+version = "0.3.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "30a651bc37f915e81f087d86e62a18eec5f79550c7faff886f7090b4ea757c77"
+dependencies = [
+ "nu-ansi-term",
+ "sharded-slab",
+ "smallvec",
+ "thread_local",
+ "tracing-core",
+ "tracing-log",
 ]
 
 [[package]]
@@ -943,6 +1012,12 @@ version = "1.0.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b15811caf2415fb889178633e7724bad2509101cde276048e013b9def5e51fa0"
 
+[[package]]
+name = "valuable"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d"
+
 [[package]]
 name = "want"
 version = "0.3.0"

agent/Cargo.toml

@@ -7,9 +7,11 @@ edition = "2021"
 anyhow = "1.0.71"
 prost = "0.11.9"
 sysinfo = { version = "0.29.2", default-features = false }
-tokio = { version = "1.28.2", features = ["rt-multi-thread", "macros", "io-util", "process"] }
+tokio = { version = "1.28.2", features = ["rt-multi-thread", "macros", "io-util", "process", "signal"] }
 tokio-stream = { version = "0.1.14", features = ["net"] }
 tonic = "0.9.2"
+tracing = "0.1.37"
+tracing-subscriber = "0.3.17"
 
 [build-dependencies]
 tonic-build = "0.9.2"

agent/src/bin/prymn_agent.rs

@@ -1,8 +1,28 @@
 use prymn_agent::new_server;
+use tokio::sync::oneshot;
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
-    new_server().serve("0.0.0.0:5043".parse()?).await?;
+    tracing_subscriber::fmt().init();
+
+    let (shutdown_tx, shutdown_rx) = oneshot::channel();
+
+    tokio::spawn(async {
+        tokio::signal::ctrl_c()
+            .await
+            .expect("could not listen to ctrl_c signal");
+
+        shutdown_tx.send(()).expect("bug: channel closed");
+    });
+
+    let addr = "0.0.0.0:5012".parse()?;
+    tracing::info!("Serving new agent at address: {}", addr);
+
+    new_server()
+        .serve_with_shutdown(addr, async {
+            let _ = shutdown_rx.await;
+        })
+        .await?;
 
     Ok(())
 }