Merge pull request #204 from http-rs/standardize-trace

Standardize trace

Author: yoshuawuyts

src/headers/constants.rs

@@ -143,6 +143,9 @@ pub const SERVER_TIMING: HeaderName = HeaderName::from_lowercase_str("server-tim
 ///  The `Te` Header
 pub const TE: HeaderName = HeaderName::from_lowercase_str("te");
 
+///  The `Traceparent` Header
+pub const TRACEPARENT: HeaderName = HeaderName::from_lowercase_str("traceparent");
+
 ///  The `Trailer` Header
 pub const TRAILER: HeaderName = HeaderName::from_lowercase_str("trailer");
 

src/trace/trace_context.rs

@@ -1,13 +1,16 @@
 use rand::Rng;
 use std::fmt;
 
-use crate::Headers;
+use crate::headers::{HeaderName, HeaderValue, Headers, TRACEPARENT};
+use crate::Status;
 
-/// Extract and inject [Trace-Context](https://w3c.github.io/trace-context/) headers.
+/// Extract and apply [Trace-Context](https://w3c.github.io/trace-context/) headers.
 ///
-/// ## Examples
+/// # Examples
 ///
 /// ```
+/// # fn main() -> http_types::Result<()> {
+/// #
 /// use http_types::trace::TraceContext;
 ///
 /// let mut res = http_types::Response::new(200);
@@ -17,14 +20,16 @@ use crate::Headers;
 ///     "00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01"
 /// );
 ///
-/// let context = TraceContext::extract(&res).unwrap();
+/// let context = TraceContext::from_headers(&res)?.unwrap();
 ///
 /// let trace_id = u128::from_str_radix("0af7651916cd43dd8448eb211c80319c", 16);
 /// let parent_id = u64::from_str_radix("00f067aa0ba902b7", 16);
 ///
 /// assert_eq!(context.trace_id(), trace_id.unwrap());
 /// assert_eq!(context.parent_id(), parent_id.ok());
 /// assert_eq!(context.sampled(), true);
+/// #
+/// # Ok(()) }
 /// ```
 #[derive(Debug)]
 pub struct TraceContext {
@@ -36,10 +41,44 @@ pub struct TraceContext {
 }
 
 impl TraceContext {
+    /// Generate a new TraceContect object without a parent.
+    ///
+    /// By default root TraceContext objects are sampled.
+    /// To mark it unsampled, call `context.set_sampled(false)`.
+    ///
+    /// # Examples
+    /// ```
+    /// use http_types::trace::TraceContext;
+    ///
+    /// let context = TraceContext::new();
+    ///
+    /// assert_eq!(context.parent_id(), None);
+    /// assert_eq!(context.sampled(), true);
+    /// ```
+    pub fn new() -> Self {
+        let mut rng = rand::thread_rng();
+
+        Self {
+            id: rng.gen(),
+            version: 0,
+            trace_id: rng.gen(),
+            parent_id: None,
+            flags: 1,
+        }
+    }
+
     /// Create and return TraceContext object based on `traceparent` HTTP header.
     ///
-    /// ## Examples
+    /// # Errors
+    ///
+    /// This function may error if the header is malformed. An error with a
+    /// status code of `400: Bad Request` will be generated.
+    ///
+    /// # Examples
+    ///
     /// ```
+    /// # fn main() -> http_types::Result<()> {
+    /// #
     /// use http_types::trace::TraceContext;
     ///
     /// let mut res = http_types::Response::new(200);
@@ -48,65 +87,43 @@ impl TraceContext {
     ///   "00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01"
     /// );
     ///
-    /// let context = TraceContext::extract(&res).unwrap();
+    /// let context = TraceContext::from_headers(&res)?.unwrap();
     ///
     /// let trace_id = u128::from_str_radix("0af7651916cd43dd8448eb211c80319c", 16);
     /// let parent_id = u64::from_str_radix("00f067aa0ba902b7", 16);
     ///
     /// assert_eq!(context.trace_id(), trace_id.unwrap());
     /// assert_eq!(context.parent_id(), parent_id.ok());
     /// assert_eq!(context.sampled(), true);
+    /// #
+    /// # Ok(()) }
     /// ```
-    pub fn extract(headers: impl AsRef<Headers>) -> crate::Result<Self> {
+    pub fn from_headers(headers: impl AsRef<Headers>) -> crate::Result<Option<Self>> {
         let headers = headers.as_ref();
         let mut rng = rand::thread_rng();
 
-        let traceparent = match headers.get("traceparent") {
-            Some(header) => header.as_str(),
-            None => return Ok(Self::new_root()),
+        let traceparent = match headers.get(TRACEPARENT) {
+            Some(header) => header,
+            None => return Ok(None),
         };
+        let parts: Vec<&str> = traceparent.as_str().split('-').collect();
 
-        let parts: Vec<&str> = traceparent.split('-').collect();
-
-        Ok(Self {
+        Ok(Some(Self {
             id: rng.gen(),
             version: u8::from_str_radix(parts[0], 16)?,
-            trace_id: u128::from_str_radix(parts[1], 16)?,
-            parent_id: Some(u64::from_str_radix(parts[2], 16)?),
-            flags: u8::from_str_radix(parts[3], 16)?,
-        })
-    }
-
-    /// Generate a new TraceContect object without a parent.
-    ///
-    /// By default root TraceContext objects are sampled.
-    /// To mark it unsampled, call `context.set_sampled(false)`.
-    ///
-    /// ## Examples
-    /// ```
-    /// use http_types::trace::TraceContext;
-    ///
-    /// let context = TraceContext::new_root();
-    ///
-    /// assert_eq!(context.parent_id(), None);
-    /// assert_eq!(context.sampled(), true);
-    /// ```
-    pub fn new_root() -> Self {
-        let mut rng = rand::thread_rng();
-
-        Self {
-            id: rng.gen(),
-            version: 0,
-            trace_id: rng.gen(),
-            parent_id: None,
-            flags: 1,
-        }
+            trace_id: u128::from_str_radix(parts[1], 16).status(400)?,
+            parent_id: Some(u64::from_str_radix(parts[2], 16).status(400)?),
+            flags: u8::from_str_radix(parts[3], 16).status(400)?,
+        }))
     }
 
     /// Add the traceparent header to the http headers
     ///
-    /// ## Examples
+    /// # Examples
+    ///
     /// ```
+    /// # fn main() -> http_types::Result<()> {
+    /// #
     /// use http_types::trace::TraceContext;
     /// use http_types::{Request, Response, Url, Method};
     ///
@@ -116,20 +133,33 @@ impl TraceContext {
     ///   "00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01"
     /// );
     ///
-    /// let parent = TraceContext::extract(&req).unwrap();
+    /// let parent = TraceContext::from_headers(&req)?.unwrap();
     ///
     /// let mut res = Response::new(200);
-    /// parent.inject(&mut res);
+    /// parent.apply(&mut res);
     ///
-    /// let child = TraceContext::extract(&res).unwrap();
+    /// let child = TraceContext::from_headers(&res)?.unwrap();
     ///
     /// assert_eq!(child.version(), parent.version());
     /// assert_eq!(child.trace_id(), parent.trace_id());
     /// assert_eq!(child.parent_id(), Some(parent.id()));
+    /// #
+    /// # Ok(()) }
     /// ```
-    pub fn inject(&self, mut headers: impl AsMut<Headers>) {
+    pub fn apply(&self, mut headers: impl AsMut<Headers>) {
         let headers = headers.as_mut();
-        headers.insert("traceparent", format!("{}", self));
+        headers.insert(TRACEPARENT, self.value());
+    }
+
+    /// Get the `HeaderName`.
+    pub fn name(&self) -> HeaderName {
+        TRACEPARENT
+    }
+
+    /// Get the `HeaderValue`.
+    pub fn value(&self) -> HeaderValue {
+        let output = format!("{}", self);
+        unsafe { HeaderValue::from_bytes_unchecked(output.into()) }
     }
 
     /// Generate a child of the current TraceContext and return it.
@@ -175,29 +205,33 @@ impl TraceContext {
 
     /// Returns true if the trace is sampled
     ///
-    /// ## Examples
+    /// # Examples
     ///
     /// ```
+    /// # fn main() -> http_types::Result<()> {
+    /// #
     /// use http_types::trace::TraceContext;
     /// use http_types::Response;
     ///
     /// let mut res = Response::new(200);
     /// res.insert_header("traceparent", "00-00000000000000000000000000000001-0000000000000002-01");
-    /// let context = TraceContext::extract(&res).unwrap();
+    /// let context = TraceContext::from_headers(&res)?.unwrap();
     /// assert_eq!(context.sampled(), true);
+    /// #
+    /// # Ok(()) }
     /// ```
     pub fn sampled(&self) -> bool {
         (self.flags & 0b00000001) == 1
     }
 
     /// Change sampled flag
     ///
-    /// ## Examples
+    /// # Examples
     ///
     /// ```
     /// use http_types::trace::TraceContext;
     ///
-    /// let mut context = TraceContext::new_root();
+    /// let mut context = TraceContext::new();
     /// assert_eq!(context.sampled(), true);
     /// context.set_sampled(false);
     /// assert_eq!(context.sampled(), false);
@@ -225,8 +259,8 @@ mod test {
     #[test]
     fn default() -> crate::Result<()> {
         let mut headers = crate::Headers::new();
-        headers.insert("traceparent", "00-01-deadbeef-00");
-        let context = TraceContext::extract(&mut headers)?;
+        headers.insert(TRACEPARENT, "00-01-deadbeef-00");
+        let context = TraceContext::from_headers(&mut headers)?.unwrap();
         assert_eq!(context.version(), 0);
         assert_eq!(context.trace_id(), 1);
         assert_eq!(context.parent_id().unwrap(), 3735928559);
@@ -237,8 +271,7 @@ mod test {
 
     #[test]
     fn no_header() -> crate::Result<()> {
-        let mut headers = crate::Headers::new();
-        let context = TraceContext::extract(&mut headers)?;
+        let context = TraceContext::new();
         assert_eq!(context.version(), 0);
         assert_eq!(context.parent_id(), None);
         assert_eq!(context.flags, 1);
@@ -249,17 +282,17 @@ mod test {
     #[test]
     fn not_sampled() -> crate::Result<()> {
         let mut headers = crate::Headers::new();
-        headers.insert("traceparent", "00-01-02-00");
-        let context = TraceContext::extract(&mut headers)?;
+        headers.insert(TRACEPARENT, "00-01-02-00");
+        let context = TraceContext::from_headers(&mut headers)?.unwrap();
         assert_eq!(context.sampled(), false);
         Ok(())
     }
 
     #[test]
     fn sampled() -> crate::Result<()> {
         let mut headers = crate::Headers::new();
-        headers.insert("traceparent", "00-01-02-01");
-        let context = TraceContext::extract(&mut headers)?;
+        headers.insert(TRACEPARENT, "00-01-02-01");
+        let context = TraceContext::from_headers(&mut headers)?.unwrap();
         assert_eq!(context.sampled(), true);
         Ok(())
     }