Tracing: Allow extract span context data from request headers and inject new span context into it (#388)

RomanDzhabarov · web-flow · commit 21dbbafa4fb4 · 2017-01-30T14:59:24.000-08:00
diff --git a/include/envoy/http/header_map.h b/include/envoy/http/header_map.h
@@ -205,6 +205,7 @@ class HeaderEntry {
   HEADER_FUNC(Host)                                                                                \
   HEADER_FUNC(KeepAlive)                                                                           \
   HEADER_FUNC(Method)                                                                              \
+  HEADER_FUNC(OtSpanContext)                                                                       \
   HEADER_FUNC(Path)                                                                                \
   HEADER_FUNC(ProxyConnection)                                                                     \
   HEADER_FUNC(RequestId)                                                                           \
diff --git a/include/envoy/tracing/http_tracer.h b/include/envoy/tracing/http_tracer.h
@@ -36,7 +36,11 @@ class Driver {
 public:
   virtual ~Driver() {}
 
-  virtual SpanPtr startSpan(const std::string& operation_name, SystemTime start_time) PURE;
+  /**
+   * Start driver specific span.
+   */
+  virtual SpanPtr startSpan(Http::HeaderMap& request_headers, const std::string& operation_name,
+                            SystemTime start_time) PURE;
 };
 
 typedef std::unique_ptr<Driver> DriverPtr;
@@ -49,7 +53,7 @@ class HttpTracer {
 public:
   virtual ~HttpTracer() {}
 
-  virtual SpanPtr startSpan(const Config& config, const Http::HeaderMap& request_headers,
+  virtual SpanPtr startSpan(const Config& config, Http::HeaderMap& request_headers,
                             const Http::AccessLog::RequestInfo& request_info) PURE;
 };
 
diff --git a/source/common/http/headers.h b/source/common/http/headers.h
@@ -42,6 +42,7 @@ class Headers {
   const LowerCaseString KeepAlive{"keep-alive"};
   const LowerCaseString Location{"location"};
   const LowerCaseString Method{":method"};
+  const LowerCaseString OtSpanContext{"x-ot-span-context"};
   const LowerCaseString Path{":path"};
   const LowerCaseString ProxyConnection{"proxy-connection"};
   const LowerCaseString RequestId{"x-request-id"};
diff --git a/source/common/tracing/http_tracer_impl.cc b/source/common/tracing/http_tracer_impl.cc
@@ -135,9 +135,10 @@ void HttpTracerUtility::finalizeSpan(Span& active_span,
 HttpTracerImpl::HttpTracerImpl(DriverPtr&& driver, const LocalInfo::LocalInfo& local_info)
     : driver_(std::move(driver)), local_info_(local_info) {}
 
-SpanPtr HttpTracerImpl::startSpan(const Config& config, const Http::HeaderMap& request_headers,
+SpanPtr HttpTracerImpl::startSpan(const Config& config, Http::HeaderMap& request_headers,
                                   const Http::AccessLog::RequestInfo& request_info) {
-  SpanPtr active_span = driver_->startSpan(config.operationName(), request_info.startTime());
+  SpanPtr active_span =
+      driver_->startSpan(request_headers, config.operationName(), request_info.startTime());
   if (active_span) {
     HttpTracerUtility::populateSpan(*active_span, local_info_.nodeName(), request_headers,
                                     request_info);
@@ -244,11 +245,35 @@ LightStepDriver::LightStepDriver(const Json::Object& config,
   });
 }
 
-SpanPtr LightStepDriver::startSpan(const std::string& operation_name, SystemTime start_time) {
-  lightstep::Span span = tls_.getTyped<TlsLightStepTracer>(tls_slot_).tracer_.StartSpan(
-      operation_name, {lightstep::StartTimestamp(start_time)});
+SpanPtr LightStepDriver::startSpan(Http::HeaderMap& request_headers,
+                                   const std::string& operation_name, SystemTime start_time) {
+  lightstep::Tracer& tracer = tls_.getTyped<TlsLightStepTracer>(tls_slot_).tracer_;
+  LightStepSpanPtr active_span;
+
+  if (request_headers.OtSpanContext()) {
+    // Extract downstream context from HTTP carrier.
+    lightstep::envoy::CarrierStruct ctx;
+    ctx.ParseFromString(request_headers.OtSpanContext()->value().c_str());
+
+    lightstep::SpanContext parent_span_ctx = tracer.Extract(
+        lightstep::CarrierFormat::EnvoyProtoCarrier, lightstep::envoy::ProtoReader(ctx));
+    lightstep::Span ls_span =
+        tracer.StartSpan(operation_name, {lightstep::ChildOf(parent_span_ctx),
+                                          lightstep::StartTimestamp(start_time)});
+    active_span.reset(new LightStepSpan(ls_span));
+  } else {
+    lightstep::Span ls_span =
+        tracer.StartSpan(operation_name, {lightstep::StartTimestamp(start_time)});
+    active_span.reset(new LightStepSpan(ls_span));
+  }
+
+  // Inject newly created span context into HTTP carrier.
+  lightstep::envoy::CarrierStruct ctx;
+  tracer.Inject(active_span->context(), lightstep::CarrierFormat::EnvoyProtoCarrier,
+                lightstep::envoy::ProtoWriter(&ctx));
+  request_headers.insertOtSpanContext().value(ctx.SerializeAsString());
 
-  return SpanPtr{new LightStepSpan(span)};
+  return std::move(active_span);
 }
 
 void LightStepRecorder::onFailure(Http::AsyncClient::FailureReason) {
diff --git a/source/common/tracing/http_tracer_impl.h b/source/common/tracing/http_tracer_impl.h
@@ -9,6 +9,7 @@
 #include "common/http/header_map_impl.h"
 #include "common/json/json_loader.h"
 
+#include "lightstep/envoy.h"
 #include "lightstep/tracer.h"
 
 namespace Tracing {
@@ -67,8 +68,7 @@ class HttpTracerUtility {
 class HttpNullTracer : public HttpTracer {
 public:
   // Tracing::HttpTracer
-  SpanPtr startSpan(const Config&, const Http::HeaderMap&,
-                    const Http::AccessLog::RequestInfo&) override {
+  SpanPtr startSpan(const Config&, Http::HeaderMap&, const Http::AccessLog::RequestInfo&) override {
     return nullptr;
   }
 };
@@ -78,7 +78,7 @@ class HttpTracerImpl : public HttpTracer {
   HttpTracerImpl(DriverPtr&& driver, const LocalInfo::LocalInfo& local_info);
 
   // Tracing::HttpTracer
-  SpanPtr startSpan(const Config& config, const Http::HeaderMap& request_headers,
+  SpanPtr startSpan(const Config& config, Http::HeaderMap& request_headers,
                     const Http::AccessLog::RequestInfo& request_info) override;
 
 private:
@@ -94,10 +94,14 @@ class LightStepSpan : public Span {
   void finishSpan() override;
   void setTag(const std::string& name, const std::string& value) override;
 
+  lightstep::SpanContext context() { return span_.context(); }
+
 private:
   lightstep::Span span_;
 };
 
+typedef std::unique_ptr<LightStepSpan> LightStepSpanPtr;
+
 /**
  * LightStep (http://lightstep.com/) provides tracing capabilities, aggregation, visualization of
  * application trace data.
@@ -111,7 +115,8 @@ class LightStepDriver : public Driver {
                   std::unique_ptr<lightstep::TracerOptions> options);
 
   // Tracer::TracingDriver
-  SpanPtr startSpan(const std::string& operation_name, SystemTime start_time) override;
+  SpanPtr startSpan(Http::HeaderMap& request_headers, const std::string& operation_name,
+                    SystemTime start_time) override;
 
   Upstream::ClusterManager& clusterManager() { return cm_; }
   Upstream::ClusterInfoPtr cluster() { return cluster_; }
diff --git a/test/common/http/conn_manager_impl_test.cc b/test/common/http/conn_manager_impl_test.cc
@@ -243,8 +243,10 @@ TEST_F(HttpConnectionManagerImplTest, StartSpanOnlyHealthCheckRequest) {
   setup(false, "");
 
   NiceMock<Tracing::MockSpan>* span = new NiceMock<Tracing::MockSpan>();
+
   EXPECT_CALL(tracer_, startSpan_(_, _, _)).WillOnce(Return(span));
   EXPECT_CALL(*span, finishSpan()).Times(0);
+
   EXPECT_CALL(runtime_.snapshot_, featureEnabled("tracing.global_enabled", 100, _))
       .WillOnce(Return(true));
 
diff --git a/test/common/tracing/http_tracer_impl_test.cc b/test/common/tracing/http_tracer_impl_test.cc
@@ -352,9 +352,10 @@ class LightStepDriverTest : public Test {
   }
 
   const std::string operation_name_{"test"};
-  const Http::TestHeaderMapImpl request_headers_{
+  Http::TestHeaderMapImpl request_headers_{
       {":path", "/"}, {":method", "GET"}, {"x-request-id", "foo"}};
   const Http::TestHeaderMapImpl response_headers_{{":status", "500"}};
+  SystemTime start_time_;
 
   std::unique_ptr<LightStepDriver> driver_;
   NiceMock<Event::MockTimer>* timer_;
@@ -447,7 +448,7 @@ TEST(HttpTracerImplTest, BasicFunctionalityNullSpan) {
 
   HttpTracerImpl tracer(std::move(driver_ptr), local_info);
 
-  EXPECT_CALL(*driver, startSpan_(operation_name, time)).WillOnce(Return(nullptr));
+  EXPECT_CALL(*driver, startSpan_(_, operation_name, time)).WillOnce(Return(nullptr));
 
   tracer.startSpan(config, request_headers, request_info);
 }
@@ -471,7 +472,7 @@ TEST(HttpTracerImplTest, BasicFunctionalityNodeSet) {
   HttpTracerImpl tracer(std::move(driver_ptr), local_info);
 
   NiceMock<MockSpan>* span = new NiceMock<MockSpan>();
-  EXPECT_CALL(*driver, startSpan_(operation_name, time)).WillOnce(Return(span));
+  EXPECT_CALL(*driver, startSpan_(_, operation_name, time)).WillOnce(Return(span));
 
   EXPECT_CALL(*span, setTag(_, _)).Times(testing::AnyNumber());
   EXPECT_CALL(*span, setTag("node_id", "node_name"));
@@ -499,16 +500,17 @@ TEST_F(LightStepDriverTest, FlushSeveralSpans) {
             return &request;
           }));
 
-  SystemTime start_time;
   EXPECT_CALL(runtime_.snapshot_, getInteger("tracing.lightstep.min_flush_spans", 5))
       .Times(2)
       .WillRepeatedly(Return(2));
   EXPECT_CALL(runtime_.snapshot_, getInteger("tracing.lightstep.request_timeout", 5000U))
       .WillOnce(Return(5000U));
 
-  SpanPtr first_span = driver_->startSpan(operation_name_, start_time);
+  Http::TestHeaderMapImpl headers;
+  SpanPtr first_span = driver_->startSpan(headers, operation_name_, start_time_);
   first_span->finishSpan();
-  SpanPtr second_span = driver_->startSpan(operation_name_, start_time);
+
+  SpanPtr second_span = driver_->startSpan(headers, operation_name_, start_time_);
   second_span->finishSpan();
 
   Http::MessagePtr msg(new Http::ResponseMessageImpl(
@@ -539,11 +541,10 @@ TEST_F(LightStepDriverTest, FlushSpansTimer) {
   const Optional<std::chrono::milliseconds> timeout(std::chrono::seconds(5));
   EXPECT_CALL(cm_.async_client_, send_(_, _, timeout));
 
-  SystemTime start_time;
   EXPECT_CALL(runtime_.snapshot_, getInteger("tracing.lightstep.min_flush_spans", 5))
       .WillOnce(Return(5));
 
-  SpanPtr span = driver_->startSpan(operation_name_, start_time);
+  SpanPtr span = driver_->startSpan(request_headers_, operation_name_, start_time_);
   span->finishSpan();
 
   // Timer should be re-enabled.
@@ -579,15 +580,13 @@ TEST_F(LightStepDriverTest, FlushOneSpanGrpcFailure) {
 
             return &request;
           }));
-
-  SystemTime start_time;
-
   EXPECT_CALL(runtime_.snapshot_, getInteger("tracing.lightstep.min_flush_spans", 5))
       .WillOnce(Return(1));
   EXPECT_CALL(runtime_.snapshot_, getInteger("tracing.lightstep.request_timeout", 5000U))
       .WillOnce(Return(5000U));
 
-  SpanPtr span = driver_->startSpan(operation_name_, start_time);
+  Http::TestHeaderMapImpl headers;
+  SpanPtr span = driver_->startSpan(headers, operation_name_, start_time_);
   span->finishSpan();
 
   Http::MessagePtr msg(new Http::ResponseMessageImpl(
@@ -605,4 +604,32 @@ TEST_F(LightStepDriverTest, FlushOneSpanGrpcFailure) {
   EXPECT_EQ(1U, stats_.counter("tracing.lightstep.spans_sent").value());
 }
 
+TEST_F(LightStepDriverTest, SerializeAndDeserializeContext) {
+  setupValidDriver();
+
+  // Supply bogus context, that will be simply ignored.
+  const std::string invalid_context = "not valid context";
+  request_headers_.insertOtSpanContext().value(invalid_context);
+  driver_->startSpan(request_headers_, operation_name_, start_time_);
+
+  std::string injected_ctx = request_headers_.OtSpanContext()->value().c_str();
+  EXPECT_FALSE(injected_ctx.empty());
+
+  // Supply empty context.
+  request_headers_.removeOtSpanContext();
+  SpanPtr span = driver_->startSpan(request_headers_, operation_name_, start_time_);
+
+  injected_ctx = request_headers_.OtSpanContext()->value().c_str();
+  EXPECT_FALSE(injected_ctx.empty());
+
+  // Context can be parsed fine.
+  lightstep::envoy::CarrierStruct ctx;
+  ctx.ParseFromString(injected_ctx);
+
+  // Supply parent context, request_headers has properly populated x-ot-span-context.
+  SpanPtr span_with_parent = driver_->startSpan(request_headers_, operation_name_, start_time_);
+  injected_ctx = request_headers_.OtSpanContext()->value().c_str();
+  EXPECT_FALSE(injected_ctx.empty());
+}
+
 } // Tracing
diff --git a/test/mocks/tracing/mocks.h b/test/mocks/tracing/mocks.h
@@ -26,12 +26,12 @@ class MockHttpTracer : public HttpTracer {
   MockHttpTracer();
   ~MockHttpTracer();
 
-  SpanPtr startSpan(const Config& config, const Http::HeaderMap& request_headers,
+  SpanPtr startSpan(const Config& config, Http::HeaderMap& request_headers,
                     const Http::AccessLog::RequestInfo& request_info) override {
     return SpanPtr{startSpan_(config, request_headers, request_info)};
   }
 
-  MOCK_METHOD3(startSpan_, Span*(const Config& config, const Http::HeaderMap& request_headers,
+  MOCK_METHOD3(startSpan_, Span*(const Config& config, Http::HeaderMap& request_headers,
                                  const Http::AccessLog::RequestInfo& request_info));
 };
 
@@ -40,11 +40,13 @@ class MockDriver : public Driver {
   MockDriver();
   ~MockDriver();
 
-  SpanPtr startSpan(const std::string& operation_name, SystemTime start_time) override {
-    return SpanPtr{startSpan_(operation_name, start_time)};
+  SpanPtr startSpan(Http::HeaderMap& request_headers, const std::string& operation_name,
+                    SystemTime start_time) override {
+    return SpanPtr{startSpan_(request_headers, operation_name, start_time)};
   }
 
-  MOCK_METHOD2(startSpan_, Span*(const std::string& operation_name, SystemTime start_time));
+  MOCK_METHOD3(startSpan_, Span*(Http::HeaderMap& request_headers,
+                                 const std::string& operation_name, SystemTime start_time));
 };
 
 } // Tracing
