Initial OpenTelemetry tracing support

crates/cli/src/main.rs

@@ -23,7 +23,6 @@ use std::path::PathBuf;
 use anyhow::Context;
 use clap::Clap;
 use mas_config::ConfigurationSection;
-use tracing::trace;
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Registry};
 
 use self::{
@@ -35,6 +34,7 @@ mod config;
 mod database;
 mod manage;
 mod server;
+mod telemetry;
 mod templates;
 
 #[derive(Clap, Debug)]
@@ -91,27 +91,50 @@ impl RootCommand {
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
+    // We're splitting the "fallible" part of main in another function to have a
+    // chance to shutdown the telemetry exporters regardless of if there was an
+    // error or not
+    let res = try_main().await;
+    telemetry::shutdown();
+    res
+}
+
+async fn try_main() -> anyhow::Result<()> {
     // Load environment variables from .env files
-    if let Err(e) = dotenv::dotenv() {
+    // We keep the path to log it afterwards
+    let dotenv_path: Option<PathBuf> = dotenv::dotenv()
+        .map(Some)
         // Display the error if it is something other than the .env file not existing
-        if !e.not_found() {
-            return Err(e).context("could not load .env file");
-        }
-    }
+        .or_else(|e| if e.not_found() { Ok(None) } else { Err(e) })?;
 
     // Setup logging & tracing
+    let (tracer, _meter) = telemetry::setup()?;
+    let telemetry_layer = tracing_opentelemetry::layer().with_tracer(tracer);
+
+    // This writes logs to stderr
     let fmt_layer = tracing_subscriber::fmt::layer().with_writer(std::io::stderr);
     let filter_layer = EnvFilter::try_from_default_env().or_else(|_| EnvFilter::try_new("info"))?;
 
-    let subscriber = Registry::default().with(filter_layer).with(fmt_layer);
+    let subscriber = Registry::default()
+        .with(telemetry_layer)
+        .with(filter_layer)
+        .with(fmt_layer);
     subscriber
         .try_init()
         .context("could not initialize logging")?;
 
+    // Now that logging is set up, we can log stuff, like if the .env file was
+    // loaded or not
+    if let Some(path) = dotenv_path {
+        tracing::info!(?path, "Loaded environment variables from file");
+    }
+
     // Parse the CLI arguments
     let opts = RootCommand::parse();
 
     // And run the command
-    trace!(?opts, "Running command");
-    opts.run().await
+    tracing::trace!(?opts, "Running command");
+    opts.run().await?;
+
+    Ok(())
 }

crates/cli/src/telemetry.rs

@@ -0,0 +1,88 @@
+// Copyright 2021 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::time::Duration;
+
+use futures::stream::{Stream, StreamExt};
+use opentelemetry::{
+    global,
+    sdk::{
+        self,
+        metrics::{self, PushController},
+        trace::{self, Tracer},
+        Resource,
+    },
+};
+use opentelemetry_semantic_conventions as semcov;
+
+pub fn setup() -> anyhow::Result<(Tracer, PushController)> {
+    global::set_error_handler(|e| tracing::error!("{}", e))?;
+
+    Ok((tracer()?, meter()?))
+}
+
+pub fn shutdown() {
+    global::shutdown_tracer_provider();
+}
+
+fn tracer() -> anyhow::Result<Tracer> {
+    let exporter = opentelemetry_otlp::new_exporter().tonic();
+
+    let tracer = opentelemetry_otlp::new_pipeline()
+        .tracing()
+        .with_exporter(exporter)
+        .with_trace_config(trace_config())
+        .install_batch(opentelemetry::runtime::Tokio)?;
+
+    Ok(tracer)
+}
+
+fn interval(duration: Duration) -> impl Stream<Item = tokio::time::Instant> {
+    // Skip first immediate tick from tokio
+    opentelemetry::util::tokio_interval_stream(duration).skip(1)
+}
+
+fn meter() -> anyhow::Result<PushController> {
+    let exporter = opentelemetry_otlp::new_exporter().tonic();
+
+    let meter = opentelemetry_otlp::new_pipeline()
+        .metrics(tokio::spawn, interval)
+        .with_exporter(exporter)
+        .with_aggregator_selector(metrics::selectors::simple::Selector::Exact)
+        .build()?;
+
+    Ok(meter)
+}
+
+fn trace_config() -> trace::Config {
+    trace::config().with_resource(resource())
+}
+
+fn resource() -> Resource {
+    let resource = Resource::new(vec![
+        semcov::resource::SERVICE_NAME.string(env!("CARGO_PKG_NAME")),
+        semcov::resource::SERVICE_VERSION.string(env!("CARGO_PKG_VERSION")),
+    ]);
+
+    let detected = Resource::from_detectors(
+        Duration::from_secs(5),
+        vec![
+            Box::new(sdk::resource::EnvResourceDetector::new()),
+            Box::new(sdk::resource::OsResourceDetector),
+            Box::new(sdk::resource::ProcessResourceDetector),
+        ],
+    );
+
+    resource.merge(&detected)
+}