test: e2e testing for OpenTelemetry Tracing

src/environmentd/Cargo.toml

@@ -66,6 +66,7 @@ rand = "0.8.5"
 reqwest = { version = "0.11.13", features = ["json"] }
 rlimit = "0.8.3"
 sentry = { version = "0.29.1", optional = true }
+sentry-tracing = "0.29.1"
 serde = { version = "1.0.152", features = ["derive"] }
 serde_json = "1.0.89"
 shell-words = "1.1.0"
@@ -79,6 +80,7 @@ tokio-postgres = { git = "https://github.com/MaterializeInc/rust-postgres" }
 tokio-stream = { version = "0.1.11", features = ["net"] }
 tower-http = { version = "0.3.5", features = ["cors"] }
 tracing = "0.1.37"
+tracing-core = "0.1.30"
 tracing-opentelemetry = { git = "https://github.com/MaterializeInc/tracing.git", branch = "v0.1.x" }
 tracing-subscriber = "0.3.16"
 tungstenite = { version = "0.18.0", features = ["native-tls"] }

src/environmentd/tests/sql.rs

@@ -2274,3 +2274,37 @@ fn test_isolation_level_notice() {
         })
         .unwrap();
 }
+
+#[test]
+fn test_emit_tracing_notice() {
+    let config = util::Config::default().with_enable_tracing(true);
+    let server = util::start_server(config).unwrap();
+
+    let (tx, mut rx) = futures::channel::mpsc::unbounded();
+
+    let mut client = server
+        .pg_config()
+        .notice_callback(move |notice| {
+            tx.unbounded_send(notice).unwrap();
+        })
+        .connect(postgres::NoTls)
+        .unwrap();
+
+    client
+        .execute("SET emit_trace_id_notice = true", &[])
+        .unwrap();
+    let _row = client.query_one("SELECT 1;", &[]).unwrap();
+
+    let tracing_re = Regex::new("trace id: (.*)").unwrap();
+    match rx.try_next() {
+        Ok(Some(msg)) => {
+            // assert the NOTICE we recieved contained a trace_id
+            let captures = tracing_re.captures(msg.message()).expect("no matches?");
+            let trace_id = captures.get(1).expect("trace_id not captured?").as_str();
+
+            assert!(trace_id.is_ascii());
+            assert_eq!(trace_id.len(), 32);
+        }
+        x => panic!("failed to read message from channel, {:?}", x),
+    }
+}

src/environmentd/tests/util.rs

@@ -103,14 +103,15 @@ use mz_ore::metrics::MetricsRegistry;
 use mz_ore::now::{EpochMillis, NowFn, SYSTEM_TIME};
 use mz_ore::retry::Retry;
 use mz_ore::task;
-use mz_ore::tracing::TracingHandle;
+use mz_ore::tracing::{TracingHandle, TracingConfig, StderrLogConfig, StderrLogFormat, OpenTelemetryConfig, TracingGuard};
 use mz_persist_client::cache::PersistClientCache;
 use mz_persist_client::cfg::PersistConfig;
 use mz_persist_client::PersistLocation;
 use mz_secrets::SecretsController;
 use mz_sql::catalog::EnvironmentId;
 use mz_stash::StashFactory;
 use mz_storage_client::types::connections::ConnectionContext;
+use tracing_subscriber::filter::Targets;
 
 pub static KAFKA_ADDRS: Lazy<String> =
     Lazy::new(|| env::var("KAFKA_ADDRS").unwrap_or_else(|_| "localhost:9092".into()));
@@ -128,6 +129,7 @@ pub struct Config {
     default_cluster_replica_size: String,
     builtin_cluster_replica_size: String,
     propagate_crashes: bool,
+    enable_tracing: bool,
 }
 
 impl Default for Config {
@@ -144,6 +146,7 @@ impl Default for Config {
             default_cluster_replica_size: "1".to_string(),
             builtin_cluster_replica_size: "1".to_string(),
             propagate_crashes: false,
+            enable_tracing: false,
         }
     }
 }
@@ -216,6 +219,11 @@ impl Config {
         self.propagate_crashes = propagate_crashes;
         self
     }
+
+    pub fn with_enable_tracing(mut self, enable_tracing: bool) -> Self {
+        self.enable_tracing = enable_tracing;
+        self
+    }
 }
 
 pub fn start_server(config: Config) -> Result<Server, anyhow::Error> {
@@ -277,6 +285,33 @@ pub fn start_server(config: Config) -> Result<Server, anyhow::Error> {
     let postgres_factory = StashFactory::new(&metrics_registry);
     let secrets_controller = Arc::clone(&orchestrator);
     let connection_context = ConnectionContext::for_tests(orchestrator.reader());
+    let (tracing_handle, tracing_guard) = if config.enable_tracing {
+        let config = TracingConfig::<fn(&tracing::Metadata) -> sentry_tracing::EventFilter> {
+            service_name: "environmentd",
+            stderr_log: StderrLogConfig {
+                format: StderrLogFormat::Json,
+                filter: Targets::default(),
+            },
+            opentelemetry: Some(OpenTelemetryConfig {
+                endpoint: "http://fake_address_for_testing:8080".to_string(),
+                headers: http::HeaderMap::new(),
+                filter: Targets::default().with_default(tracing_core::Level::DEBUG),
+                resource: opentelemetry::sdk::resource::Resource::default(),
+                start_enabled: true,
+            }),
+            #[cfg(feature = "tokio-console")]
+            tokio_console: None,
+            sentry: None,
+            build_version: &mz_environmentd::BUILD_INFO.version,
+            build_sha: &mz_environmentd::BUILD_INFO.sha,
+            build_time: &mz_environmentd::BUILD_INFO.time,
+        };
+        let (tracing_handle, tracing_guard) = runtime.block_on(mz_ore::tracing::configure(config))?;
+        (tracing_handle, Some(tracing_guard))
+    } else {
+        (TracingHandle::disabled(), None)
+    };
+
     let inner = runtime.block_on(mz_environmentd::serve(mz_environmentd::Config {
         adapter_stash_url,
         controller: ControllerConfig {
@@ -315,7 +350,7 @@ pub fn start_server(config: Config) -> Result<Server, anyhow::Error> {
         bootstrap_system_parameters: Default::default(),
         availability_zones: Default::default(),
         connection_context,
-        tracing_handle: TracingHandle::disabled(),
+        tracing_handle,
         storage_usage_collection_interval: config.storage_usage_collection_interval,
         segment_api_key: None,
         egress_ips: vec![],
@@ -330,6 +365,7 @@ pub fn start_server(config: Config) -> Result<Server, anyhow::Error> {
         runtime,
         metrics_registry,
         _temp_dir: temp_dir,
+        _tracing_guard: tracing_guard,
     };
     Ok(server)
 }
@@ -339,6 +375,7 @@ pub struct Server {
     pub runtime: Arc<Runtime>,
     _temp_dir: Option<TempDir>,
     pub metrics_registry: MetricsRegistry,
+    pub _tracing_guard: Option<TracingGuard>,
 }
 
 impl Server {

src/ore/src/tracing.rs

@@ -197,6 +197,7 @@ impl std::fmt::Debug for TracingHandle {
 /// A guard for the tracing infrastructure configured with [`configure`].
 ///
 /// This guard should be kept alive for the lifetime of the program.
+#[must_use = "Must hold for the lifetime of the program, otherwise tracing will be shutdown"]
 pub struct TracingGuard {
     _sentry_guard: Option<sentry::ClientInitGuard>,
 }