Skip to content

Commit

Permalink
Support the new "b3 single" header format
Browse files Browse the repository at this point in the history 
Closes #21
  • Loading branch information
@sfackler
sfackler committed Aug 31, 2018
1 parent 4cdc1b3 commit 8f1fd85
Showing 1 changed file with 232 additions and 1 deletion.
233 changes: 232 additions & 1 deletion http-zipkin/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ extern crate http;
extern crate zipkin;

use http::header::{HeaderMap, HeaderValue};
use std::fmt::Write;
use std::str::FromStr;
use zipkin::{SamplingFlags, TraceContext};

Expand All @@ -28,6 +29,22 @@ const X_B3_FLAGS: &'static str = "X-B3-Flags";
const X_B3_TRACEID: &'static str = "X-B3-TraceId";
const X_B3_PARENTSPANID: &'static str = "X-B3-ParentSpanId";
const X_B3_SPANID: &'static str = "X-B3-SpanId";
const B3: &'static str = "b3";

/// Serializes sampling flags into the `b3` HTTP header.
///
/// This form is more compact than the old `X-B3-` set of headers, but some implementations may not support it.
pub fn set_sampling_flags_single(flags: SamplingFlags, headers: &mut HeaderMap) {
if flags.debug() {
headers.insert(B3, HeaderValue::from_static("d"));
} else if flags.sampled() == Some(true) {
headers.insert(B3, HeaderValue::from_static("1"));
} else if flags.sampled() == Some(false) {
headers.insert(B3, HeaderValue::from_static("0"));
} else {
headers.remove(B3);
}
}

/// Serializes sampling flags into a set of HTTP headers.
pub fn set_sampling_flags(flags: SamplingFlags, headers: &mut HeaderMap) {
Expand All @@ -52,6 +69,29 @@ pub fn set_sampling_flags(flags: SamplingFlags, headers: &mut HeaderMap) {

/// Deserializes sampling flags from a set of HTTP headers.
pub fn get_sampling_flags(headers: &HeaderMap) -> SamplingFlags {
match headers.get(B3) {
Some(value) => get_sampling_flags_single(value),
None => get_sampling_flags_multi(headers),
}
}

fn get_sampling_flags_single(value: &HeaderValue) -> SamplingFlags {
let mut builder = SamplingFlags::builder();

if value == "d" {
builder.debug(true);
} else if value == "1" {
builder.sampled(true);
} else if value == "0" {
builder.sampled(false);
} else if let Some(context) = get_trace_context_single(value) {
return context.sampling_flags();
}

builder.build()
}

fn get_sampling_flags_multi(headers: &HeaderMap) -> SamplingFlags {
let mut builder = SamplingFlags::builder();

if let Some(flags) = headers.get(X_B3_FLAGS) {
Expand All @@ -69,6 +109,25 @@ pub fn get_sampling_flags(headers: &HeaderMap) -> SamplingFlags {
builder.build()
}

/// Serializes a trace context into the `b3` header.
///
/// This form is more compact than the old `X-B3-` set of headers, but some implementations may not support it.
pub fn set_trace_context_single(context: TraceContext, headers: &mut HeaderMap) {
let mut value = String::new();
write!(value, "{}-{}", context.trace_id(), context.span_id()).unwrap();
if context.debug() {
value.push_str("-d");
} else if context.sampled() == Some(true) {
value.push_str("-1");
} else if context.sampled() == Some(false) {
value.push_str("-0");
}
if let Some(parent_id) = context.parent_id() {
write!(value, "-{}", parent_id).unwrap();
}
headers.insert(B3, HeaderValue::from_str(&value).unwrap());
}

/// Serializes a trace context into a set of HTTP headers.
pub fn set_trace_context(context: TraceContext, headers: &mut HeaderMap) {
set_sampling_flags(context.sampling_flags(), headers);
Expand Down Expand Up @@ -96,14 +155,55 @@ pub fn set_trace_context(context: TraceContext, headers: &mut HeaderMap) {

/// Deserializes a trace context from a set of HTTP headers.
pub fn get_trace_context(headers: &HeaderMap) -> Option<TraceContext> {
match headers.get(B3) {
Some(value) => get_trace_context_single(value),
None => get_trace_context_multi(headers),
}
}

fn get_trace_context_single(value: &HeaderValue) -> Option<TraceContext> {
let mut parts = value.to_str().ok()?.split('-');

let trace_id = parts.next()?.parse().ok()?;
let span_id = parts.next()?.parse().ok()?;

let mut builder = TraceContext::builder();
builder.trace_id(trace_id).span_id(span_id);

let maybe_sampling = match parts.next() {
Some(next) => next,
None => return Some(builder.build()),
};

let parent_id = if maybe_sampling == "d" {
builder.debug(true);
parts.next()
} else if maybe_sampling == "1" {
builder.sampled(true);
parts.next()
} else if maybe_sampling == "0" {
builder.sampled(false);
parts.next()
} else {
Some(maybe_sampling)
};

if let Some(parent_id) = parent_id {
builder.parent_id(parent_id.parse().ok()?);
}

Some(builder.build())
}

fn get_trace_context_multi(headers: &HeaderMap) -> Option<TraceContext> {
let trace_id = parse_header(headers, X_B3_TRACEID)?;
let span_id = parse_header(headers, X_B3_SPANID)?;

let mut builder = TraceContext::builder();
builder
.trace_id(trace_id)
.span_id(span_id)
.sampling_flags(get_sampling_flags(headers));
.sampling_flags(get_sampling_flags_multi(headers));

if let Some(parent_id) = parse_header(headers, X_B3_PARENTSPANID) {
builder.parent_id(parent_id);
Expand Down Expand Up @@ -138,6 +238,18 @@ mod test {
assert_eq!(get_sampling_flags(&headers), flags);
}

#[test]
fn flags_empty_single() {
let mut headers = HeaderMap::new();
let flags = SamplingFlags::builder().build();
set_sampling_flags_single(flags, &mut headers);

let expected_headers = HeaderMap::new();
assert_eq!(headers, expected_headers);

assert_eq!(get_sampling_flags(&headers), flags);
}

#[test]
fn flags_debug() {
let mut headers = HeaderMap::new();
Expand All @@ -151,6 +263,19 @@ mod test {
assert_eq!(get_sampling_flags(&headers), flags);
}

#[test]
fn flags_debug_single() {
let mut headers = HeaderMap::new();
let flags = SamplingFlags::builder().debug(true).build();
set_sampling_flags_single(flags, &mut headers);

let mut expected_headers = HeaderMap::new();
expected_headers.insert("b3", HeaderValue::from_static("d"));
assert_eq!(headers, expected_headers);

assert_eq!(get_sampling_flags(&headers), flags);
}

#[test]
fn flags_sampled() {
let mut headers = HeaderMap::new();
Expand All @@ -164,6 +289,19 @@ mod test {
assert_eq!(get_sampling_flags(&headers), flags);
}

#[test]
fn flags_sampled_single() {
let mut headers = HeaderMap::new();
let flags = SamplingFlags::builder().sampled(true).build();
set_sampling_flags_single(flags, &mut headers);

let mut expected_headers = HeaderMap::new();
expected_headers.insert("b3", HeaderValue::from_static("1"));
assert_eq!(headers, expected_headers);

assert_eq!(get_sampling_flags(&headers), flags);
}

#[test]
fn flags_unsampled() {
let mut headers = HeaderMap::new();
Expand All @@ -177,6 +315,19 @@ mod test {
assert_eq!(get_sampling_flags(&headers), flags);
}

#[test]
fn flags_unsampled_single() {
let mut headers = HeaderMap::new();
let flags = SamplingFlags::builder().sampled(false).build();
set_sampling_flags_single(flags, &mut headers);

let mut expected_headers = HeaderMap::new();
expected_headers.insert("b3", HeaderValue::from_static("0"));
assert_eq!(headers, expected_headers);

assert_eq!(get_sampling_flags(&headers), flags);
}

#[test]
fn trace_context() {
let mut headers = HeaderMap::new();
Expand All @@ -200,4 +351,84 @@ mod test {

assert_eq!(get_trace_context(&headers), Some(context));
}

#[test]
fn trace_context_single() {
let mut headers = HeaderMap::new();
let context = TraceContext::builder()
.trace_id([0, 1, 2, 3, 4, 5, 6, 7].into())
.parent_id([1, 2, 3, 4, 5, 6, 7, 8].into())
.span_id([2, 3, 4, 5, 6, 7, 8, 9].into())
.sampled(true)
.build();
set_trace_context_single(context, &mut headers);

let mut expected_headers = HeaderMap::new();
expected_headers.insert(
"b3",
HeaderValue::from_static("0001020304050607-0203040506070809-1-0102030405060708"),
);
assert_eq!(headers, expected_headers);

assert_eq!(get_trace_context(&headers), Some(context));
}

#[test]
fn trace_context_unsampled_single() {
let mut headers = HeaderMap::new();
let context = TraceContext::builder()
.trace_id([0, 1, 2, 3, 4, 5, 6, 7].into())
.parent_id([1, 2, 3, 4, 5, 6, 7, 8].into())
.span_id([2, 3, 4, 5, 6, 7, 8, 9].into())
.build();
set_trace_context_single(context, &mut headers);

let mut expected_headers = HeaderMap::new();
expected_headers.insert(
"b3",
HeaderValue::from_static("0001020304050607-0203040506070809-0102030405060708"),
);
assert_eq!(headers, expected_headers);

assert_eq!(get_trace_context(&headers), Some(context));
}

#[test]
fn trace_context_parentless_single() {
let mut headers = HeaderMap::new();
let context = TraceContext::builder()
.trace_id([0, 1, 2, 3, 4, 5, 6, 7].into())
.span_id([2, 3, 4, 5, 6, 7, 8, 9].into())
.sampled(true)
.build();
set_trace_context_single(context, &mut headers);

let mut expected_headers = HeaderMap::new();
expected_headers.insert(
"b3",
HeaderValue::from_static("0001020304050607-0203040506070809-1"),
);
assert_eq!(headers, expected_headers);

assert_eq!(get_trace_context(&headers), Some(context));
}

#[test]
fn trace_context_minimal_single() {
let mut headers = HeaderMap::new();
let context = TraceContext::builder()
.trace_id([0, 1, 2, 3, 4, 5, 6, 7].into())
.span_id([2, 3, 4, 5, 6, 7, 8, 9].into())
.build();
set_trace_context_single(context, &mut headers);

let mut expected_headers = HeaderMap::new();
expected_headers.insert(
"b3",
HeaderValue::from_static("0001020304050607-0203040506070809"),
);
assert_eq!(headers, expected_headers);

assert_eq!(get_trace_context(&headers), Some(context));
}
}

0 comments on commit 8f1fd85

Please sign in to comment.