From 6cd697afe9c23974cb37d6436539a0a2ac3404e0 Mon Sep 17 00:00:00 2001
From: Bryan Garza <1396101+bryangarza@users.noreply.github.com>
Date: Wed, 18 May 2022 09:13:19 -0700
Subject: [PATCH] Add (optional) OpenTelemetry + Xray integration (#95) (#96)

This introduces all the necessary code to be able to send traces to
AWS Xray via `tracing-opentelemetry`. It can be optionally enabled
using the `xray` feature defined on this crate.

Also update the README.md with instructions on how to enable and use
this.
---
 .github/workflows/ci.yml |  6 +++-
 Cargo.toml               | 16 ++++++++--
 README.md                | 25 ++++++++++++++++
 src/bin/server.rs        | 65 ++++++++++++++++++++++++++++++++++++++--
 4 files changed, 106 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 513a5ca..85780c0 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -15,10 +15,14 @@ jobs:
     - uses: actions/checkout@v2
     - name: Build
       run: cargo build --verbose
+    - name: Build with OTel feature
+      run: cargo build --verbose --features otel
     - name: Run tests
       run: cargo test --verbose
+    - name: Run tests with OTel feature
+      run: cargo test --verbose --features otel
     - name: rustfmt
       uses: actions-rs/cargo@v1
       with:
         command: fmt
-        args: --all -- --check
+        args: --all -- --check
\ No newline at end of file
diff --git a/Cargo.toml b/Cargo.toml
index 7ed7c81..23103ef 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -24,13 +24,25 @@ path = "src/bin/server.rs"
 async-stream = "0.3.0"
 atoi = "0.3.2"
 bytes = "1"
+rand = "0.8.5"
 structopt = "0.3.14"
 tokio = { version = "1", features = ["full"] }
 tokio-stream = "0.1"
-tracing = "0.1.13"
+tracing = "0.1.34"
 tracing-futures = { version = "0.2.3" }
-tracing-subscriber = "0.2.2"
+tracing-subscriber = { version = "0.3.11", features = ["env-filter"] }
+# Implements the types defined in the OTel spec
+opentelemetry = { version = "0.17.0", optional = true }
+# Integration between the tracing crate and the opentelemetry crate
+tracing-opentelemetry = { version = "0.17.2", optional = true }
+# Provides a "propagator" to pass along an XrayId across services
+opentelemetry-aws = { version = "0.5.0", optional = true }
+# Allows you to send data to the OTel collector
+opentelemetry-otlp = { version = "0.10.0", optional = true }
 
 [dev-dependencies]
 # Enable test-utilities in dev mode only. This is mostly for tests.
 tokio = { version = "1", features = ["test-util"] }
+
+[features]
+otel = ["dep:opentelemetry", "dep:tracing-opentelemetry", "dep:opentelemetry-aws", "dep:opentelemetry-otlp"]
\ No newline at end of file
diff --git a/README.md b/README.md
index 0ec91df..e932766 100644
--- a/README.md
+++ b/README.md
@@ -59,6 +59,31 @@ cargo run --bin mini-redis-cli set foo bar
 cargo run --bin mini-redis-cli get foo
 ```
 
+## OpenTelemetry
+
+If you are running many instances of your application (which is usually the case
+when you are developing a cloud service, for example), you need a way to get all
+of your trace data out of your host and into a centralized place. There are many
+options here, such as Prometheus, Jaeger, DataDog, Honeycomb, AWS X-Ray etc.
+
+We leverage OpenTelemetry, because it's an open standard that allows for a
+single data format to be used for all the options mentioned above (and more).
+This eliminates the risk of vendor lock-in, since you can switch between
+providers if needed.
+
+### AWS X-Ray example
+
+To enable sending traces to X-Ray, use the `otel` feature:
+```
+RUST_LOG=debug cargo run --bin mini-redis-server --features otel
+```
+
+This will switch `tracing` to use `tracing-opentelemetry`. You will need to
+have a copy of AWSOtelCollector running on the same host.
+
+For demo purposes, you can follow the setup documented at
+https://github.com/aws-observability/aws-otel-collector/blob/main/docs/developers/docker-demo.md#run-a-single-aws-otel-collector-instance-in-docker
+
 ## Supported commands
 
 `mini-redis` currently supports the following commands.
diff --git a/src/bin/server.rs b/src/bin/server.rs
index 2f76ad7..5a9c783 100644
--- a/src/bin/server.rs
+++ b/src/bin/server.rs
@@ -12,11 +12,27 @@ use structopt::StructOpt;
 use tokio::net::TcpListener;
 use tokio::signal;
 
+#[cfg(feature = "otel")]
+// To be able to set the XrayPropagator
+use opentelemetry::global;
+#[cfg(feature = "otel")]
+// To configure certain options such as sampling rate
+use opentelemetry::sdk::trace as sdktrace;
+#[cfg(feature = "otel")]
+// For passing along the same XrayId across services
+use opentelemetry_aws::trace::XrayPropagator;
+#[cfg(feature = "otel")]
+// To be able to pass along the XrayId across services
+#[cfg(feature = "otel")]
+// The `Ext` traits are to allow the Registry to accept the
+// OpenTelemetry-specific types (such as `OpenTelemetryLayer`)
+use tracing_subscriber::{
+    fmt, layer::SubscriberExt, util::SubscriberInitExt, util::TryInitError, EnvFilter,
+};
+
 #[tokio::main]
 pub async fn main() -> mini_redis::Result<()> {
-    // enable logging
-    // see https://docs.rs/tracing for more info
-    tracing_subscriber::fmt::try_init()?;
+    set_up_logging()?;
 
     let cli = Cli::from_args();
     let port = cli.port.as_deref().unwrap_or(DEFAULT_PORT);
@@ -35,3 +51,46 @@ struct Cli {
     #[structopt(name = "port", long = "--port")]
     port: Option<String>,
 }
+
+#[cfg(not(feature = "otel"))]
+fn set_up_logging() -> mini_redis::Result<()> {
+    // See https://docs.rs/tracing for more info
+    tracing_subscriber::fmt::try_init()
+}
+
+#[cfg(feature = "otel")]
+fn set_up_logging() -> Result<(), TryInitError> {
+    // Set the global propagator to X-Ray propagator
+    // Note: If you need to pass the x-amzn-trace-id across services in the same trace,
+    // you will need this line. However, this requires additional code not pictured here.
+    // For a full example using hyper, see:
+    // https://github.com/open-telemetry/opentelemetry-rust/blob/main/examples/aws-xray/src/server.rs#L14-L26
+    global::set_text_map_propagator(XrayPropagator::default());
+
+    let tracer = opentelemetry_otlp::new_pipeline()
+        .tracing()
+        .with_exporter(opentelemetry_otlp::new_exporter().tonic())
+        .with_trace_config(
+            sdktrace::config()
+                .with_sampler(sdktrace::Sampler::AlwaysOn)
+                // Needed in order to convert the trace IDs into an Xray-compatible format
+                .with_id_generator(sdktrace::XrayIdGenerator::default()),
+        )
+        .install_simple()
+        .expect("Unable to initialize OtlpPipeline");
+
+    // Create a tracing layer with the configured tracer
+    let opentelemetry = tracing_opentelemetry::layer().with_tracer(tracer);
+
+    // Parse an `EnvFilter` configuration from the `RUST_LOG`
+    // environment variable.
+    let filter = EnvFilter::from_default_env();
+
+    // Use the tracing subscriber `Registry`, or any other subscriber
+    // that impls `LookupSpan`
+    tracing_subscriber::registry()
+        .with(opentelemetry)
+        .with(filter)
+        .with(fmt::Layer::default())
+        .try_init()
+}