Skip to content

Commit

Permalink
Stackdriver - adding support for Google X-Cloud-Trace-Context propaga…
Browse files Browse the repository at this point in the history
…tion format
  • Loading branch information
ivan-brko committed Jan 17, 2024
1 parent 3827119 commit 0e397d5
Show file tree
Hide file tree
Showing 5 changed files with 306 additions and 0 deletions.
1 change: 1 addition & 0 deletions opentelemetry-stackdriver/CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# Changelog

## vNext
- Added support for `GoogleTraceContextPropagator`

WARNING The current version relies on features only in upstream git version. This should be modified before releasing.

Expand Down
7 changes: 7 additions & 0 deletions opentelemetry-stackdriver/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,9 @@ prost-types = "0.11.1"
thiserror = "1.0.30"
tonic = { version = "0.9.0", features = ["gzip", "tls", "transport"] }
yup-oauth2 = { version = "8.1.0", optional = true }
once_cell = { version = "1.19", optional = true }
# we don't need unicode support
regex = { version = "1.10", default-features = false, features = ["std", "perf"], optional = true }

# Futures
futures-core = "0.3"
Expand All @@ -38,6 +41,7 @@ default = ["yup-authorizer", "tls-native-roots"]
yup-authorizer = ["hyper-rustls", "yup-oauth2"]
tls-native-roots = ["tonic/tls-roots"]
tls-webpki-roots = ["tonic/tls-webpki-roots"]
propagator = ["once_cell", "regex"]

[dev-dependencies]
reqwest = "0.11.9"
Expand All @@ -46,3 +50,6 @@ tokio = "1"
tonic-build = "0.9.0"
walkdir = "2.3.2"
futures-util = { version = "0.3", default-features = false }
# TODO: Replace with opentelemetry version before release
opentelemetry = { git = "https://github.com/open-telemetry/opentelemetry-rust.git", rev = "6175c3", features = ["testing"] }

8 changes: 8 additions & 0 deletions opentelemetry-stackdriver/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -5,3 +5,11 @@
This crate provides an `opentelemetry` exporter for use with Google StackDriver trace. It uses gRPC to send tracing spans.

It is licensed under the Apache 2.0 license. Contributions are welcome.

### Propagator
Feature flag `propagator` will enable the `GoogleTraceContextPropagator` which implements the `TextMapPropagator` trait for Google `X-Cloud-Trace-Context` Trace Context format.

Example usage:
```rust
opentelemetry::global::set_text_map_propagator(GoogleTraceContextPropagator::new());
```
287 changes: 287 additions & 0 deletions opentelemetry-stackdriver/src/google_trace_context_propagator.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,287 @@
use once_cell::sync::Lazy;
use opentelemetry::propagation::text_map_propagator::FieldIter;
use opentelemetry::propagation::{Extractor, Injector, TextMapPropagator};
use opentelemetry::trace::{SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState};
use opentelemetry::Context;
use regex::Regex;

/// Propagates span context in the Google Cloud Trace format,
/// using the __X-Cloud-Trace-Context__ header.
///
/// See https://cloud.google.com/trace/docs/setup/#force-trace for details on the format.
#[derive(Clone, Debug, Default)]
pub struct GoogleTraceContextPropagator {
_private: (),
}

// https://cloud.google.com/trace/docs/setup/#force-trace
// documentation for the structure of X-Cloud-Trace-Context header is not very detailed

// this implementation is based on the official GCP golang library implementation:
// https://github.com/GoogleCloudPlatform/opentelemetry-operations-go/blob/main/propagator/propagator.go
// the regex they use: "^(?P<trace_id>[0-9a-f]{32})/(?P<span_id>[0-9]{1,20})(;o=(?P<trace_flags>[0-9]))?$"
// - trace id is 32 hex characters, mandatory
// - span id is 1-20 decimal characters, mandatory
// - trace flags is optional, 0 to 9 (0 - not sampled, missing or any other number - sampled)

const CLOUD_TRACE_CONTEXT_HEADER: &str = "X-Cloud-Trace-Context";
const GOOGLE_PROPAGATION_HEADER_VALUE_REGEX_STR: &str =
r"^(?P<trace_id>[0-9a-f]{32})/(?P<span_id>[0-9]{1,20})(;o=(?P<trace_flags>[0-9]))?$";

static TRACE_CONTEXT_HEADER_FIELDS: Lazy<[String; 1]> =
Lazy::new(|| [CLOUD_TRACE_CONTEXT_HEADER.to_owned()]);

static GOOGLE_PROPAGATION_HEADER_VALUE_REGEX: Lazy<Option<Regex>> =
Lazy::new(|| Regex::new(GOOGLE_PROPAGATION_HEADER_VALUE_REGEX_STR).ok());

impl GoogleTraceContextPropagator {
/// Create a new `GoogleTraceContextPropagator`.
pub fn new() -> Self {
GoogleTraceContextPropagator { _private: () }
}

fn extract_span_context(&self, extractor: &dyn Extractor) -> Result<SpanContext, ()> {
let Some(header_value) = extractor.get(CLOUD_TRACE_CONTEXT_HEADER).map(|v| v.trim()) else {
return Err(());
};

let Some(regex) = GOOGLE_PROPAGATION_HEADER_VALUE_REGEX.as_ref() else {
return Err(());
};

// we could do a quick check if the header value matches the regex here to avoid more expensive regex parsing in the next step if it doesn't match
// but the assumption is that the header value will almost always be valid, so that would add an unnecessary check in majority of cases

if let Some(caps) = regex.captures(header_value) {
let trace_id_hex = caps.name("trace_id").map(|m| m.as_str());
let span_id_dec = caps.name("span_id").map(|m| m.as_str());

// the request is sampled by default, it's not sampled only if explicitly set to 0
let trace_flags = caps.name("trace_flags").map_or(TraceFlags::SAMPLED, |m| {
if m.as_str() == "0" {
TraceFlags::NOT_SAMPLED
} else {
TraceFlags::SAMPLED
}
});

match (trace_id_hex, span_id_dec) {
(Some(trace_id_hex), Some(span_id_dec)) => {
Self::construct_span_context(trace_flags, trace_id_hex, span_id_dec)
}
_ => Err(()),
}
} else {
Err(())
}
}

fn construct_span_context(
trace_flags: TraceFlags,
trace_id_hex: &str,
span_id_dec: &str,
) -> Result<SpanContext, ()> {
let trace_id = TraceId::from_hex(trace_id_hex).map_err(|_| ())?;

let span_id = span_id_dec.parse::<u64>()
.map(|v| SpanId::from_bytes(v.to_be_bytes())) // we can create SPAN ID only from bytes or hex string
.map_err(|_| ())?;

let span_context = SpanContext::new(trace_id, span_id, trace_flags, true, TraceState::NONE);

// Ensure span is valid
if !span_context.is_valid() {
return Err(());
}

Ok(span_context)
}
}

impl TextMapPropagator for GoogleTraceContextPropagator {
fn inject_context(&self, cx: &Context, injector: &mut dyn Injector) {
let span = cx.span();
let span_context = span.span_context();
let sampled_flag = span_context.trace_flags().to_u8();
if span_context.is_valid() {
let header_value = format!(
"{:032x}/{};o={}",
span_context.trace_id(),
// at the moment we can only get span id as bytes
u64::from_be_bytes(span_context.span_id().to_bytes()),
sampled_flag
);
injector.set(CLOUD_TRACE_CONTEXT_HEADER, header_value);
}
}

fn extract_with_context(&self, cx: &Context, extractor: &dyn Extractor) -> Context {
self.extract_span_context(extractor)
.map(|sc| cx.with_remote_span_context(sc))
.unwrap_or_else(|_| cx.clone())
}

fn fields(&self) -> FieldIter<'_> {
FieldIter::new(TRACE_CONTEXT_HEADER_FIELDS.as_ref())
}
}

#[cfg(test)]
mod tests {
use super::*;
use opentelemetry::testing::trace::TestSpan;
use opentelemetry::trace::TraceState;
use std::collections::HashMap;

#[test]
fn test_google_propagation_header_value_regex_str_valid() {
// Try to create a Regex from the string
let regex_result = Regex::new(GOOGLE_PROPAGATION_HEADER_VALUE_REGEX_STR);

// Assert that the Regex was created successfully
assert!(
regex_result.is_ok(),
"Failed to create Regex from GOOGLE_PROPAGATION_HEADER_VALUE_REGEX_STR"
);

// If the Regex was created successfully, validate it against a known valid string
let regex = regex_result.unwrap();
let valid_string = "105445aa7843bc8bf206b12000100000/1;o=1";
assert!(
regex.is_match(valid_string),
"Regex does not match known valid string"
);
}

#[test]
fn test_extract_span_context_valid() {
let propagator = GoogleTraceContextPropagator::new();
let mut headers = HashMap::new();
headers.insert(
// hashmap implementation of Extractor trait uses lowercase keys
CLOUD_TRACE_CONTEXT_HEADER.to_string().to_lowercase(),
"105445aa7843bc8bf206b12000100000/1;o=1".to_string(),
);

let span_context = propagator.extract_span_context(&headers).unwrap();
assert_eq!(
format!("{:x}", span_context.trace_id()),
"105445aa7843bc8bf206b12000100000"
);
assert_eq!(u64::from_be_bytes(span_context.span_id().to_bytes()), 1);
assert!(span_context.is_sampled());
}

#[test]
fn test_extract_span_context_valid_without_options() {
let propagator = GoogleTraceContextPropagator::new();
let mut headers = HashMap::new();
headers.insert(
// hashmap implementation of Extractor trait uses lowercase keys
CLOUD_TRACE_CONTEXT_HEADER.to_string().to_lowercase(),
"105445aa7843bc8bf206b12000100000/1".to_string(),
);

let span_context = propagator.extract_span_context(&headers).unwrap();
assert_eq!(
format!("{:x}", span_context.trace_id()),
"105445aa7843bc8bf206b12000100000"
);
assert_eq!(u64::from_be_bytes(span_context.span_id().to_bytes()), 1);
assert!(span_context.is_sampled());
}

#[test]
fn test_extract_span_context_valid_not_sampled() {
let propagator = GoogleTraceContextPropagator::new();
let mut headers = HashMap::new();
headers.insert(
// hashmap implementation of Extractor trait uses lowercase keys
CLOUD_TRACE_CONTEXT_HEADER.to_string().to_lowercase(),
"105445aa7843bc8bf206b12000100000/1;o=0".to_string(),
);

let span_context = propagator.extract_span_context(&headers).unwrap();
assert_eq!(
format!("{:x}", span_context.trace_id()),
"105445aa7843bc8bf206b12000100000"
);
assert_eq!(u64::from_be_bytes(span_context.span_id().to_bytes()), 1);
assert!(!span_context.is_sampled());
}

#[test]
fn test_extract_span_context_invalid() {
let propagator = GoogleTraceContextPropagator::new();
let headers = HashMap::new();

assert!(propagator.extract_span_context(&headers).is_err());
}

#[test]
fn test_inject_context_valid() {
let propagator = GoogleTraceContextPropagator::new();
let mut headers = HashMap::new();
let span = TestSpan(SpanContext::new(
TraceId::from_hex("105445aa7843bc8bf206b12000100000").unwrap(),
SpanId::from_hex("0000000000000001").unwrap(),
TraceFlags::SAMPLED,
true,
TraceState::default(),
));
let cx = Context::current_with_span(span);

propagator.inject_context(&cx, &mut headers);
assert_eq!(
// hashmap implementation of Extractor trait uses lowercase keys
headers.get(CLOUD_TRACE_CONTEXT_HEADER.to_lowercase().as_str()),
Some(&"105445aa7843bc8bf206b12000100000/1;o=1".to_string())
);
}

#[test]
fn test_extract_with_context_valid() {
let propagator = GoogleTraceContextPropagator::new();
let mut headers = HashMap::new();
headers.insert(
CLOUD_TRACE_CONTEXT_HEADER.to_string().to_lowercase(),
"105445aa7843bc8bf206b12000100000/1;o=1".to_string(),
);
let cx = Context::current();

let new_cx = propagator.extract_with_context(&cx, &headers);
assert!(new_cx.span().span_context().is_valid());
}

#[test]
fn test_extract_with_context_invalid_trace_id() {
let propagator = GoogleTraceContextPropagator::new();
let mut headers = HashMap::new();
// Insert a trace ID with less than 32 characters
headers.insert(
CLOUD_TRACE_CONTEXT_HEADER.to_string().to_lowercase(),
"105445aa7843bc8b/1;o=1".to_string(), // This trace ID is shorter than 32 characters
);
let cx = Context::current();

let new_cx = propagator.extract_with_context(&cx, &headers);
// Assert that the span context is not valid
assert!(!new_cx.span().span_context().is_valid());
}

#[test]
fn test_extract_with_context_invalid_span_id() {
let propagator = GoogleTraceContextPropagator::new();
let mut headers = HashMap::new();
// Insert a trace ID with less than 32 characters
headers.insert(
CLOUD_TRACE_CONTEXT_HEADER.to_string().to_lowercase(),
"105445aa7843bc8b/1abc;o=1".to_string(), // This span id is not decimal
);
let cx = Context::current();

let new_cx = propagator.extract_with_context(&cx, &headers);
// Assert that the span context is not valid
assert!(!new_cx.span().span_context().is_valid());
}
}
3 changes: 3 additions & 0 deletions opentelemetry-stackdriver/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,9 @@ use yup_oauth2::authenticator::Authenticator;
#[allow(clippy::derive_partial_eq_without_eq)] // tonic doesn't derive Eq for generated types
pub mod proto;

#[cfg(feature = "propagator")]
pub mod google_trace_context_propagator;

const HTTP_HOST: &str = "http.host";
const HTTP_PATH: &str = "http.path";
const HTTP_USER_AGENT: &str = "http.user_agent";
Expand Down

0 comments on commit 0e397d5

Please sign in to comment.