[next] refactor: tracing

cargo-shuttle/src/lib.rs

@@ -472,6 +472,8 @@ impl Shuttle {
                     .expect("failed to find cargo home dir")
                     .join("bin/shuttle-next");
 
+                println!("Installing shuttle runtime. This can take a while...");
+
                 if cfg!(debug_assertions) {
                     // Canonicalized path to shuttle-runtime for dev to work on windows
                     let path = std::fs::canonicalize(format!("{MANIFEST_DIR}/../runtime"))

codegen/src/shuttle_main/mod.rs

@@ -258,10 +258,10 @@ impl ToTokens for Loader {
                         .or_else(|_| shuttle_runtime::tracing_subscriber::EnvFilter::try_new("INFO"))
                         .unwrap();
 
-                shuttle_runtime::tracing_subscriber::registry()
+                let _guard = shuttle_runtime::tracing_subscriber::registry()
                     .with(filter_layer)
                     .with(logger)
-                    .init();
+                    .set_default(); // Scope our runtime logger to this thread scope only
 
                 #vars
                 #(let #fn_inputs = #fn_inputs_builder::new()#fn_inputs_builder_options.build(&mut #factory_ident).await.context(format!("failed to provision {}", stringify!(#fn_inputs_builder)))?;)*
@@ -317,10 +317,10 @@ mod tests {
                         .or_else(|_| shuttle_runtime::tracing_subscriber::EnvFilter::try_new("INFO"))
                         .unwrap();
 
-                shuttle_runtime::tracing_subscriber::registry()
+                let _guard = shuttle_runtime::tracing_subscriber::registry()
                     .with(filter_layer)
                     .with(logger)
-                    .init();
+                    .set_default();
 
                 simple().await
             }
@@ -398,10 +398,10 @@ mod tests {
                         .or_else(|_| shuttle_runtime::tracing_subscriber::EnvFilter::try_new("INFO"))
                         .unwrap();
 
-                shuttle_runtime::tracing_subscriber::registry()
+                let _guard = shuttle_runtime::tracing_subscriber::registry()
                     .with(filter_layer)
                     .with(logger)
-                    .init();
+                    .set_default();
 
                 let pool = shuttle_shared_db::Postgres::new().build(&mut factory).await.context(format!("failed to provision {}", stringify!(shuttle_shared_db::Postgres)))?;
                 let redis = shuttle_shared_db::Redis::new().build(&mut factory).await.context(format!("failed to provision {}", stringify!(shuttle_shared_db::Redis)))?;
@@ -513,10 +513,10 @@ mod tests {
                         .or_else(|_| shuttle_runtime::tracing_subscriber::EnvFilter::try_new("INFO"))
                         .unwrap();
 
-                shuttle_runtime::tracing_subscriber::registry()
+                let _guard = shuttle_runtime::tracing_subscriber::registry()
                     .with(filter_layer)
                     .with(logger)
-                    .init();
+                    .set_default();
 
                 let vars = std::collections::HashMap::from_iter(factory.get_secrets().await?.into_iter().map(|(key, value)| (format!("secrets.{}", key), value)));
                 let pool = shuttle_shared_db::Postgres::new().size(&shuttle_runtime::strfmt("10Gb", &vars)?).public(false).build(&mut factory).await.context(format!("failed to provision {}", stringify!(shuttle_shared_db::Postgres)))?;

runtime/src/alpha/mod.rs

@@ -15,7 +15,7 @@ use core::future::Future;
 use shuttle_common::{
     backends::{
         auth::{AuthPublicKey, JwtAuthenticationLayer},
-        tracing::ExtractPropagationLayer,
+        tracing::{setup_tracing, ExtractPropagationLayer},
     },
     claims::{Claim, ClaimLayer, InjectPropagationLayer},
     resource,
@@ -54,6 +54,8 @@ pub async fn start(loader: impl Loader<ProvisionerFactory> + Send + 'static) {
     let args = Args::parse();
     let addr = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), args.port);
 
+    setup_tracing(tracing_subscriber::registry(), "shuttle-alpha");
+
     let provisioner_address = args.provisioner_address;
     let mut server_builder = Server::builder()
         .http2_keepalive_interval(Some(Duration::from_secs(60)))

runtime/src/bin/shuttle-next.rs

@@ -4,33 +4,25 @@ use std::{
 };
 
 use clap::Parser;
+use shuttle_common::backends::tracing::{setup_tracing, ExtractPropagationLayer};
 use shuttle_proto::runtime::runtime_server::RuntimeServer;
 use shuttle_runtime::{AxumWasm, NextArgs};
 use tonic::transport::Server;
 use tracing::trace;
-use tracing_subscriber::{fmt, prelude::*, EnvFilter};
 
 #[tokio::main(flavor = "multi_thread")]
 async fn main() {
     let args = NextArgs::parse();
 
-    // TODO: replace with tracing helper from main branch
-    let fmt_layer = fmt::layer();
-    let filter_layer = EnvFilter::try_from_default_env()
-        .or_else(|_| EnvFilter::try_new("info"))
-        .unwrap();
-
-    tracing_subscriber::registry()
-        .with(filter_layer)
-        .with(fmt_layer)
-        .init();
+    setup_tracing(tracing_subscriber::registry(), "shuttle-next");
 
     trace!(args = ?args, "parsed args");
 
     let addr = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), args.port);
 
-    let mut server_builder =
-        Server::builder().http2_keepalive_interval(Some(Duration::from_secs(60)));
+    let mut server_builder = Server::builder()
+        .http2_keepalive_interval(Some(Duration::from_secs(60)))
+        .layer(ExtractPropagationLayer);
 
     let axum = AxumWasm::default();
     let svc = RuntimeServer::new(axum);