mostly finish DeploymentIdLayer and StateChangeLayer

Author: jonaro00

auth/src/main.rs

@@ -1,7 +1,7 @@
 use std::io;
 
 use clap::Parser;
-use shuttle_common::backends::tracing::setup_tracing;
+use shuttle_common::{backends::tracing::setup_tracing, log::Backend};
 use sqlx::migrate::Migrator;
 use tracing::{info, trace};
 
@@ -15,7 +15,7 @@ async fn main() -> io::Result<()> {
 
     trace!(args = ?args, "parsed args");
 
-    setup_tracing(tracing_subscriber::registry(), "auth");
+    setup_tracing(tracing_subscriber::registry(), Backend::Auth);
 
     let db_path = args.state.join("authentication.sqlite");
 

common/src/backends/tracing.rs

@@ -20,9 +20,11 @@ use tracing::{debug_span, instrument::Instrumented, Instrument, Span, Subscriber
 use tracing_opentelemetry::OpenTelemetrySpanExt;
 use tracing_subscriber::{fmt, prelude::*, registry::LookupSpan, EnvFilter};
 
+use crate::log::Backend;
+
 const OTLP_ADDRESS: &str = "http://otel-collector:4317";
 
-pub fn setup_tracing<S>(subscriber: S, service_name: &str)
+pub fn setup_tracing<S>(subscriber: S, backend: Backend)
 where
     S: Subscriber + for<'a> LookupSpan<'a> + Send + Sync,
 {
@@ -46,7 +48,7 @@ where
         .with_trace_config(
             trace::config().with_resource(Resource::new(vec![KeyValue::new(
                 "service.name",
-                service_name.to_string(),
+                backend.to_string(),
             )])),
         )
         .install_batch(Tokio)
@@ -196,15 +198,17 @@ pub fn serde_json_map_to_key_value_list(
 /// Convert an [AnyValue] to a [serde_json::Value]
 pub fn from_any_value_to_serde_json_value(any_value: AnyValue) -> serde_json::Value {
     let Some(value) = any_value.value else {
-        return serde_json::Value::Null
+        return serde_json::Value::Null;
     };
 
     match value {
         any_value::Value::StringValue(s) => serde_json::Value::String(s),
         any_value::Value::BoolValue(b) => serde_json::Value::Bool(b),
         any_value::Value::IntValue(i) => serde_json::Value::Number(i.into()),
         any_value::Value::DoubleValue(f) => {
-            let Some(number) = serde_json::Number::from_f64(f) else {return serde_json::Value::Null};
+            let Some(number) = serde_json::Number::from_f64(f) else {
+                return serde_json::Value::Null;
+            };
             serde_json::Value::Number(number)
         }
         any_value::Value::ArrayValue(a) => {

common/src/lib.rs

@@ -12,7 +12,7 @@ pub type DeploymentId = Uuid;
 #[cfg(feature = "service")]
 pub mod log;
 #[cfg(feature = "service")]
-pub use log::Item as LogItem;
+pub use log::LogItem;
 #[cfg(feature = "models")]
 pub mod models;
 #[cfg(feature = "service")]

common/src/log.rs

@@ -1,43 +1,59 @@
 use chrono::{DateTime, Utc};
 #[cfg(feature = "display")]
-use crossterm::style::{StyledContent, Stylize};
+use crossterm::style::Stylize;
 use serde::{Deserialize, Serialize};
 use strum::EnumString;
+use tracing::{field::Visit, span, warn, Metadata, Subscriber};
+use tracing_subscriber::Layer;
 #[cfg(feature = "openapi")]
 use utoipa::ToSchema;
 use uuid::Uuid;
 
+/// Used to determine settings based on which backend crate does what
 #[derive(Clone, Debug, EnumString, Eq, PartialEq, Deserialize, Serialize)]
 #[cfg_attr(feature = "display", derive(strum::Display))]
 #[cfg_attr(feature = "openapi", derive(ToSchema))]
-pub enum InternalLogOrigin {
+pub enum Backend {
+    /// Is considered an error
     Unknown,
-    Deployer,
+
+    Auth,
     // Builder,
-    // ResourceRecorder,
+    Deployer,
+    Gateway,
+    Logger,
+    Provisioner,
+    ResourceRecorder,
 }
 
-impl Default for InternalLogOrigin {
+impl Default for Backend {
     fn default() -> Self {
         Self::Unknown
     }
 }
 
 #[derive(Clone, Debug, Deserialize, Serialize)]
 #[cfg_attr(feature = "openapi", derive(ToSchema))]
-#[cfg_attr(feature = "openapi", schema(as = shuttle_common::log::Item))]
-pub struct Item {
+#[cfg_attr(feature = "openapi", schema(as = shuttle_common::log::LogItem))]
+pub struct LogItem {
+    /// Deployment id
     #[cfg_attr(feature = "openapi", schema(value_type = KnownFormat::Uuid))]
     pub id: Uuid,
+
+    /// Internal service that produced this log
     #[cfg_attr(feature = "openapi", schema(value_type = shuttle_common::log::InternalLogOrigin))]
-    pub internal_origin: InternalLogOrigin,
+    pub internal_origin: Backend,
+
+    /// Time log was produced
     #[cfg_attr(feature = "openapi", schema(value_type = KnownFormat::DateTime))]
     pub timestamp: DateTime<Utc>,
+
+    /// The log line
     pub line: String,
 }
 
 #[cfg(feature = "display")]
-impl std::fmt::Display for Item {
+impl std::fmt::Display for LogItem {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let datetime: chrono::DateTime<chrono::Local> = DateTime::from(self.timestamp);
 
@@ -51,6 +67,110 @@ impl std::fmt::Display for Item {
     }
 }
 
+/// Records logs for the deployment progress
+pub trait LogRecorder: Clone + Send + 'static {
+    fn record(&self, log: LogItem);
+}
+
+/// Tracing subscriber layer which logs based on if the log
+/// is from a span that is tagged with a deployment id
+pub struct DeploymentLogLayer<R>
+where
+    R: LogRecorder + Send + Sync,
+{
+    pub recorder: R,
+    pub internal_service: Backend,
+}
+
+impl<R, S> Layer<S> for DeploymentLogLayer<R>
+where
+    S: Subscriber + for<'lookup> tracing_subscriber::registry::LookupSpan<'lookup>,
+    R: LogRecorder + Send + Sync + 'static,
+{
+    fn on_event(&self, event: &tracing::Event<'_>, ctx: tracing_subscriber::layer::Context<'_, S>) {
+        // We only care about events in some scope
+        let scope = if let Some(scope) = ctx.event_scope(event) {
+            scope
+        } else {
+            return;
+        };
+
+        // Find the outermost scope with the scope details containing the current state
+        for span in scope.from_root() {
+            let extensions = span.extensions();
+
+            if let Some(details) = extensions.get::<ScopeDetails>() {
+                self.recorder.record(LogItem {
+                    id: details.id,
+                    internal_origin: self.internal_service.clone(),
+                    timestamp: Utc::now(),
+                    line: "Test".into(),
+                });
+                break;
+            }
+        }
+    }
+    fn on_new_span(
+        &self,
+        attrs: &span::Attributes<'_>,
+        id: &span::Id,
+        ctx: tracing_subscriber::layer::Context<'_, S>,
+    ) {
+        // We only care about spans that change the state
+        if !DeploymentIdVisitor::is_valid(attrs.metadata()) {
+            return;
+        }
+        let mut visitor = DeploymentIdVisitor::default();
+        attrs.record(&mut visitor);
+        let details = visitor.details;
+
+        if details.id.is_nil() {
+            warn!("scope details does not have a valid id");
+            return;
+        }
+
+        // Safe to unwrap since this is the `on_new_span` method
+        let span = ctx.span(id).unwrap();
+        let mut extensions = span.extensions_mut();
+
+        self.recorder.record(LogItem {
+            id: details.id,
+            internal_origin: self.internal_service.clone(),
+            timestamp: Utc::now(),
+            line: "Test".into(),
+        });
+
+        extensions.insert::<ScopeDetails>(details);
+    }
+}
+
+#[derive(Debug, Default)]
+struct ScopeDetails {
+    id: Uuid,
+}
+/// To extract `id` field for scopes that have it
+#[derive(Default)]
+struct DeploymentIdVisitor {
+    details: ScopeDetails,
+}
+
+impl DeploymentIdVisitor {
+    /// Field containing the deployment identifier
+    const ID_IDENT: &'static str = "id";
+
+    fn is_valid(metadata: &Metadata) -> bool {
+        metadata.is_span() && metadata.fields().field(Self::ID_IDENT).is_some()
+    }
+}
+
+impl Visit for DeploymentIdVisitor {
+    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
+        if field.name() == Self::ID_IDENT {
+            self.details.id = Uuid::try_parse(&format!("{value:?}")).unwrap_or_default();
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -66,9 +186,9 @@ mod tests {
 
     #[test]
     fn test_timezone_formatting() {
-        let item = Item {
+        let item = LogItem {
             id: Uuid::new_v4(),
-            internal_origin: InternalLogOrigin::Deployer,
+            internal_origin: Backend::Deployer,
             timestamp: Utc::now(),
             line: r#"{"message": "Building"}"#.to_owned(),
         };

deployer/src/deployment/mod.rs

@@ -1,28 +1,27 @@
-pub mod deploy_layer;
-pub mod gateway_client;
-mod queue;
-mod run;
-
 use std::{path::PathBuf, sync::Arc};
 
 pub use queue::Queued;
 pub use run::{ActiveDeploymentsGetter, Built};
-use shuttle_common::storage_manager::ArtifactsStorageManager;
+use shuttle_common::{log::LogRecorder, storage_manager::ArtifactsStorageManager};
 use shuttle_proto::logger::logger_client::LoggerClient;
+use tokio::{
+    sync::{mpsc, Mutex},
+    task::JoinSet,
+};
 use tracing::{instrument, Span};
 use tracing_opentelemetry::OpenTelemetrySpanExt;
+use uuid::Uuid;
 
+pub mod gateway_client;
+mod queue;
+mod run;
+pub mod state_change_layer;
+
+use self::gateway_client::BuildQueueClient;
 use crate::{
     persistence::{DeploymentUpdater, ResourceManager, SecretGetter, SecretRecorder, State},
     RuntimeManager,
 };
-use tokio::{
-    sync::{mpsc, Mutex},
-    task::JoinSet,
-};
-use uuid::Uuid;
-
-use self::{deploy_layer::LogRecorder, gateway_client::BuildQueueClient};
 
 const QUEUE_BUFFER_SIZE: usize = 100;
 const RUN_BUFFER_SIZE: usize = 100;

deployer/src/deployment/queue.rs

@@ -11,8 +11,7 @@ use chrono::Utc;
 use crossbeam_channel::Sender;
 use flate2::read::GzDecoder;
 use opentelemetry::global;
-use shuttle_common::claims::Claim;
-use shuttle_service::builder::{build_workspace, BuiltService};
+use shuttle_common::log::LogRecorder;
 use tar::Archive;
 use tokio::fs;
 use tokio::task::JoinSet;
@@ -22,12 +21,17 @@ use tracing_opentelemetry::OpenTelemetrySpanExt;
 use ulid::Ulid;
 use uuid::Uuid;
 
-use super::deploy_layer::{Log, LogRecorder};
+use shuttle_common::{
+    claims::Claim,
+    storage_manager::{ArtifactsStorageManager, StorageManager},
+    LogItem,
+};
+use shuttle_service::builder::{build_workspace, BuiltService};
+
 use super::gateway_client::BuildQueueClient;
 use super::{Built, QueueReceiver, RunSender, State};
 use crate::error::{Error, Result, TestError};
 use crate::persistence::{DeploymentUpdater, SecretRecorder};
-use shuttle_common::storage_manager::{ArtifactsStorageManager, StorageManager};
 
 pub async fn task(
     mut recv: QueueReceiver,
@@ -192,18 +196,13 @@ impl Queued {
                 trace!(?message, "received cargo message");
                 // TODO: change these to `info!(...)` as [valuable] support increases.
                 // Currently it is not possible to turn these serde `message`s into a `valuable`, but once it is the passing down of `log_recorder` should be removed.
-                let log = match message {
-                    Message::TextLine(line) => Log {
-                        deployment_id: self.id,
-                        internal_origin: shuttle_common::log::InternalLogOrigin::Deployer, // will change to Builder
-                        tx_timestamp: Utc::now(),
-                        line,
-                    },
-                    message => Log {
-                        deployment_id: self.id,
-                        internal_origin: shuttle_common::log::InternalLogOrigin::Deployer, // will change to Builder
-                        tx_timestamp: Utc::now(),
-                        line: serde_json::to_string(&message).unwrap(),
+                let log = LogItem {
+                    id: self.id,
+                    internal_origin: shuttle_common::log::Backend::Deployer, // will change to Builder
+                    timestamp: Utc::now(),
+                    line: match message {
+                        Message::TextLine(line) => line,
+                        message => serde_json::to_string(&message).unwrap(),
                     },
                 };
                 log_recorder.record(log);