From b602520c0316cfd795105b3e10a9f0755550a74e Mon Sep 17 00:00:00 2001
From: mlehurau <max.lehuraux@analog.com>
Date: Wed, 28 Jul 2021 17:05:39 -0400
Subject: [PATCH] opentelemetry: Add extension to link spans

 ## Motivation

Discussed in #1121, the purpose of adding the `add_link` extension is to allow adding a link to a
propagated span and/or a closed span.
---
 tracing-opentelemetry/src/span_ext.rs | 56 ++++++++++++++++++++++++++-
 1 file changed, 55 insertions(+), 1 deletion(-)

diff --git a/tracing-opentelemetry/src/span_ext.rs b/tracing-opentelemetry/src/span_ext.rs
index 7140372ccf..45afd914db 100644
--- a/tracing-opentelemetry/src/span_ext.rs
+++ b/tracing-opentelemetry/src/span_ext.rs
@@ -1,5 +1,5 @@
 use crate::layer::WithContext;
-use opentelemetry::Context;
+use opentelemetry::{trace::TraceContextExt, Context};
 
 /// Utility functions to allow tracing [`Span`]s to accept and return
 /// [OpenTelemetry] [`Context`]s.
@@ -42,6 +42,40 @@ pub trait OpenTelemetrySpanExt {
     /// ```
     fn set_parent(&self, cx: Context);
 
+    /// Associates `self` with a given OpenTelemetry trace, using the provided
+    /// followed span [`Context`].
+    ///
+    /// [`Context`]: opentelemetry::Context
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use opentelemetry::{propagation::TextMapPropagator, trace::TraceContextExt};
+    /// use opentelemetry::sdk::propagation::TraceContextPropagator;
+    /// use tracing_opentelemetry::OpenTelemetrySpanExt;
+    /// use std::collections::HashMap;
+    /// use tracing::Span;
+    ///
+    /// // Example carrier, could be a framework header map that impls otel's `Extract`.
+    /// let mut carrier = HashMap::new();
+    ///
+    /// // Propagator can be swapped with b3 propagator, jaeger propagator, etc.
+    /// let propagator = TraceContextPropagator::new();
+    ///
+    /// // Extract otel parent context via the chosen propagator
+    /// let parent_context = propagator.extract(&carrier);
+    ///
+    /// // Generate a tracing span as usual
+    /// let app_root = tracing::span!(tracing::Level::INFO, "app_start");
+    ///
+    /// // Assign parent trace from external context
+    /// app_root.add_link(parent_context.clone());
+    ///
+    /// // Or if the current span has been created elsewhere:
+    /// Span::current().add_link(parent_context);
+    /// ```
+    fn add_link(&self, cx: Context);
+
     /// Extracts an OpenTelemetry [`Context`] from `self`.
     ///
     /// [`Context`]: opentelemetry::Context
@@ -86,6 +120,26 @@ impl OpenTelemetrySpanExt for tracing::Span {
         });
     }
 
+    fn add_link(&self, cx: Context) {
+        let mut cx = Some(cx);
+        self.with_subscriber(move |(id, subscriber)| {
+            if let Some(get_context) = subscriber.downcast_ref::<WithContext>() {
+                get_context.with_context(subscriber, id, move |builder, _tracer| {
+                    if let Some(cx) = cx.take() {
+                        let follows_context = cx.span().span_context().clone();
+                        let follows_link =
+                            opentelemetry::trace::Link::new(follows_context, Vec::new());
+                        if let Some(ref mut links) = builder.links {
+                            links.push(follows_link);
+                        } else {
+                            builder.links = Some(vec![follows_link]);
+                        }
+                    }
+                });
+            }
+        });
+    }
+
     fn context(&self) -> Context {
         let mut cx = None;
         self.with_subscriber(|(id, subscriber)| {