Add trace submodule

Author: yoshuawuyts

Cargo.toml

@@ -39,6 +39,7 @@ url = "2.1.1"
 serde_json = "1.0.51"
 serde = { version = "1.0.106", features = ["derive"] }
 serde_urlencoded = "0.6.1"
+rand = "0.7.3"
 
 [dev-dependencies]
 http = "0.2.0"

src/headers/headers.rs

@@ -183,6 +183,18 @@ impl Debug for Headers {
     }
 }
 
+impl AsRef<Headers> for Headers {
+    fn as_ref(&self) -> &Headers {
+        self
+    }
+}
+
+impl AsMut<Headers> for Headers {
+    fn as_mut(&mut self) -> &mut Headers {
+        self
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/lib.rs

@@ -132,6 +132,7 @@ mod version;
 
 cfg_unstable! {
     pub mod upgrade;
+    pub mod trace;
 
     mod client;
     mod server;
@@ -165,10 +166,8 @@ pub use crate::url::Url;
 #[doc(inline)]
 pub use crate::cookies::Cookie;
 
-#[doc(inline)]
-pub mod trailers;
-
 pub mod security;
+pub mod trailers;
 
 #[cfg(feature = "hyperium_http")]
 mod hyperium_http;

src/trace/mod.rs

@@ -0,0 +1,268 @@
+//! Extract and inject [trace context](https://w3c.github.io/trace-context/) headers.
+//!
+//! ## Examples
+//!
+//! ```
+//! use http_types::trace::TraceContext;
+//!
+//! let mut res = http_types::Response::new(200);
+//!
+//! res.insert_header(
+//!     "traceparent",
+//!     "00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01"
+//! );
+//!
+//! let context = TraceContext::extract(&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);
+//! ```
+
+use rand::Rng;
+use std::fmt;
+
+use crate::Headers;
+
+/// A TraceContext object
+#[derive(Debug)]
+pub struct TraceContext {
+    id: u64,
+    version: u8,
+    trace_id: u128,
+    parent_id: Option<u64>,
+    flags: u8,
+}
+
+impl TraceContext {
+    /// Create and return TraceContext object based on `traceparent` HTTP header.
+    ///
+    /// ## Examples
+    /// ```
+    /// use http_types::trace::TraceContext;
+    ///
+    /// let mut res = http_types::Response::new(200);
+    /// res.insert_header(
+    ///   "traceparent",
+    ///   "00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01"
+    /// );
+    ///
+    /// let context = TraceContext::extract(&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);
+    /// ```
+    pub fn extract(headers: impl AsRef<Headers>) -> crate::Result<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 parts: Vec<&str> = traceparent.split('-').collect();
+
+        Ok(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,
+        }
+    }
+
+    /// Add the traceparent header to the http headers
+    ///
+    /// ## Examples
+    /// ```
+    /// use http_types::trace::TraceContext;
+    /// use http_types::{Request, Response, Url, Method};
+    ///
+    /// let mut req = Request::new(Method::Get, Url::parse("https://example.com").unwrap());
+    /// req.insert_header(
+    ///   "traceparent",
+    ///   "00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01"
+    /// );
+    ///
+    /// let parent = TraceContext::extract(&req).unwrap();
+    ///
+    /// let mut res = Response::new(200);
+    /// parent.inject(&mut res);
+    ///
+    /// let child = TraceContext::extract(&res).unwrap();
+    ///
+    /// assert_eq!(child.version(), parent.version());
+    /// assert_eq!(child.trace_id(), parent.trace_id());
+    /// assert_eq!(child.parent_id(), Some(parent.id()));
+    /// ```
+    pub fn inject(&self, mut headers: impl AsMut<Headers>) {
+        let headers = headers.as_mut();
+        headers.insert("traceparent", format!("{}", self));
+    }
+
+    /// Generate a child of the current TraceContext and return it.
+    ///
+    /// The child will have a new randomly genrated `id` and its `parent_id` will be set to the
+    /// `id` of this TraceContext.
+    pub fn child(&self) -> Self {
+        let mut rng = rand::thread_rng();
+
+        Self {
+            id: rng.gen(),
+            version: self.version,
+            trace_id: self.trace_id,
+            parent_id: Some(self.id),
+            flags: self.flags,
+        }
+    }
+
+    /// Return the id of the TraceContext.
+    pub fn id(&self) -> u64 {
+        self.id
+    }
+
+    /// Return the version of the TraceContext spec used.
+    ///
+    /// You probably don't need this.
+    pub fn version(&self) -> u8 {
+        self.version
+    }
+
+    /// Return the trace id of the TraceContext.
+    ///
+    /// All children will have the same `trace_id`.
+    pub fn trace_id(&self) -> u128 {
+        self.trace_id
+    }
+
+    /// Return the id of the parent TraceContext.
+    #[inline]
+    pub fn parent_id(&self) -> Option<u64> {
+        self.parent_id
+    }
+
+    /// Returns true if the trace is sampled
+    ///
+    /// ## Examples
+    ///
+    /// ```
+    /// 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();
+    /// assert_eq!(context.sampled(), true);
+    /// ```
+    pub fn sampled(&self) -> bool {
+        (self.flags & 0b00000001) == 1
+    }
+
+    /// Change sampled flag
+    ///
+    /// ## Examples
+    ///
+    /// ```
+    /// use http_types::trace::TraceContext;
+    ///
+    /// let mut context = TraceContext::new_root();
+    /// assert_eq!(context.sampled(), true);
+    /// context.set_sampled(false);
+    /// assert_eq!(context.sampled(), false);
+    /// ```
+    pub fn set_sampled(&mut self, sampled: bool) {
+        let x = sampled as u8;
+        self.flags ^= (x ^ self.flags) & (1 << 0);
+    }
+}
+
+impl fmt::Display for TraceContext {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(
+            f,
+            "{:02x}-{:032x}-{:016x}-{:02x}",
+            self.version, self.trace_id, self.id, self.flags
+        )
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn default() -> crate::Result<()> {
+        let mut headers = crate::Headers::new();
+        headers.insert("traceparent", "00-01-deadbeef-00");
+        let context = TraceContext::extract(&mut headers)?;
+        assert_eq!(context.version(), 0);
+        assert_eq!(context.trace_id(), 1);
+        assert_eq!(context.parent_id().unwrap(), 3735928559);
+        assert_eq!(context.flags, 0);
+        assert_eq!(context.sampled(), false);
+        Ok(())
+    }
+
+    #[test]
+    fn no_header() -> crate::Result<()> {
+        let mut headers = crate::Headers::new();
+        let context = TraceContext::extract(&mut headers)?;
+        assert_eq!(context.version(), 0);
+        assert_eq!(context.parent_id(), None);
+        assert_eq!(context.flags, 1);
+        assert_eq!(context.sampled(), true);
+        Ok(())
+    }
+
+    #[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)?;
+        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)?;
+        assert_eq!(context.sampled(), true);
+        Ok(())
+    }
+}