Sampling

.gitignore

@@ -32,3 +32,5 @@
 *.app
 
 bazel-*
+
+.build/

zipkin/include/zipkin/span_context.h

@@ -58,6 +58,8 @@ class SpanContext {
       : trace_id_{trace_id}, id_{id}, parent_id_{parent_id}, flags_{flags},
         is_initialized_{true} {}
 
+  bool isSampled() const { return flags_ & zipkin::sampled_flag; }
+
   /**
    * @return the span id as an integer
    */

zipkin/include/zipkin/zipkin_core_types.h

@@ -331,6 +331,9 @@ class Span : public ZipkinBase {
   Span()
       : trace_id_(0), name_(), id_(0), debug_(false), monotonic_start_time_(0),
         tracer_(nullptr) {}
+
+  void setSampled(const bool val) { sampled_ = val; }
+  bool isSampled() const { return sampled_; }
 
   /**
    * Sets the span's trace id attribute.
@@ -566,6 +569,7 @@ class Span : public ZipkinBase {
   std::string name_;
   uint64_t id_;
   Optional<TraceId> parent_id_;
+  bool sampled_;
   bool debug_;
   std::vector<Annotation> annotations_;
   std::vector<BinaryAnnotation> binary_annotations_;

zipkin/src/span_context.cc

@@ -8,6 +8,11 @@ SpanContext::SpanContext(const Span &span) {
   trace_id_ = span.traceId();
   id_ = span.id();
   parent_id_ = span.isSetParentId() ? span.parentId() : 0;
+
+  if (span.isSampled()) {
+    flags_ |= static_cast<unsigned char>(zipkin::sampled_flag);
+    flags_ |= static_cast<unsigned char>(zipkin::sampling_set_flag);
+  }
 
   for (const Annotation &annotation : span.annotations()) {
     if (annotation.value() == ZipkinCoreConstants::get().CLIENT_RECV) {

zipkin/src/zipkin_core_types.cc

@@ -185,7 +185,9 @@ const std::string Span::toJson() {
 
 void Span::finish() {
   if (auto t = tracer()) {
-    t->reportSpan(std::move(*this));
+    if (this->isSampled()) {
+      t->reportSpan(std::move(*this));
+    }
   }
 }
 

zipkin_opentracing/CMakeLists.txt

@@ -8,7 +8,8 @@ set(ZIPKIN_OPENTRACING_SRCS src/utility.cc
                             src/propagation.cc
                             src/dynamic_load.cc
                             src/tracer_factory.cc
-                            src/opentracing.cc)
+                            src/opentracing.cc
+                            src/sampling.cc)
 
 if (BUILD_SHARED_LIBS)               
   add_library(zipkin_opentracing SHARED ${ZIPKIN_OPENTRACING_SRCS})

zipkin_opentracing/include/zipkin/opentracing.h

@@ -8,6 +8,7 @@ struct ZipkinOtTracerOptions {
   uint32_t collector_port = 9411;
   SteadyClock::duration reporting_period = DEFAULT_REPORTING_PERIOD;
   size_t max_buffered_spans = DEFAULT_SPAN_BUFFER_SIZE;
+  double sample_rate = 1.0;
 
   std::string service_name;
   IpAddress service_address;
@@ -19,4 +20,5 @@ makeZipkinOtTracer(const ZipkinOtTracerOptions &options);
 std::shared_ptr<opentracing::Tracer>
 makeZipkinOtTracer(const ZipkinOtTracerOptions &options,
                    std::unique_ptr<Reporter> &&reporter);
+
 } // namespace zipkin

zipkin_opentracing/src/opentracing.cc

@@ -1,11 +1,14 @@
 #include <zipkin/opentracing.h>
+#include <opentracing/util.h>
 
+#include "sampling.h"
 #include "propagation.h"
 #include "utility.h"
 #include <atomic>
 #include <cstring>
 #include <mutex>
 #include <unordered_map>
+#include <random>
 #include <zipkin/tracer.h>
 #include <zipkin/utility.h>
 #include <zipkin/zipkin_core_types.h>
@@ -99,6 +102,14 @@ class OtSpanContext : public ot::SpanContext {
     return injectSpanContext(writer, span_context_, baggage_);
   }
 
+  bool isSampled() const {
+    return span_context_.isSampled();
+  }
+
+  bool isValid() const {
+    return span_context_.id() != 0 && !span_context_.trace_id().empty();
+  }
+
 private:
   zipkin::SpanContext span_context_;
   mutable std::mutex baggage_mutex_;
@@ -156,7 +167,7 @@ class OtSpan : public ot::Span {
           parent_span_context->baggage_mutex_};
       auto baggage = parent_span_context->baggage_;
       span_context_ =
-          OtSpanContext{zipkin::SpanContext{*span_}, std::move(baggage)};
+          OtSpanContext{zipkin::SpanContext{*span_}, std::move(baggage)}; 
     } else {
       span_context_ = OtSpanContext{zipkin::SpanContext{*span_}};
     }
@@ -271,17 +282,27 @@ class OtSpan : public ot::Span {
 class OtTracer : public ot::Tracer,
                  public std::enable_shared_from_this<OtTracer> {
 public:
-  explicit OtTracer(TracerPtr &&tracer) : tracer_{std::move(tracer)} {}
+  explicit OtTracer(TracerPtr &&tracer) : tracer_{std::move(tracer)}, sampler_{std::move(new ProbabilisticSampler(1.0))} {}
+  explicit OtTracer(TracerPtr &&tracer, SamplerPtr &&sampler) : tracer_{std::move(tracer)}, sampler_{std::move(sampler)} {}
 
   std::unique_ptr<ot::Span>
   StartSpanWithOptions(string_view operation_name,
                        const ot::StartSpanOptions &options) const
       noexcept override {
+
     // Create the core zipkin span.
     SpanPtr span{new zipkin::Span{}};
     span->setName(operation_name);
     span->setTracer(tracer_.get());
 
+    auto parent = findSpanContext(options.references);
+
+    if (parent && parent->isValid()) {
+      span->setSampled(parent->isSampled());
+    } else {
+      span->setSampled(sampler_->ShouldSample());
+    }
+
     Endpoint endpoint{tracer_->serviceName(), tracer_->address()};
 
     // Add a binary annotation for the serviceName.
@@ -329,6 +350,7 @@ class OtTracer : public ot::Tracer,
 
 private:
   TracerPtr tracer_;
+  SamplerPtr sampler_;
 
   template <class Carrier>
   expected<void> InjectImpl(const ot::SpanContext &sc, Carrier &writer) const
@@ -368,7 +390,8 @@ makeZipkinOtTracer(const ZipkinOtTracerOptions &options,
                    std::unique_ptr<Reporter> &&reporter) {
   TracerPtr tracer{new Tracer{options.service_name, options.service_address}};
   tracer->setReporter(std::move(reporter));
-  return std::shared_ptr<ot::Tracer>{new OtTracer{std::move(tracer)}};
+  SamplerPtr sampler{new ProbabilisticSampler{options.sample_rate}};
+  return std::shared_ptr<ot::Tracer>{new OtTracer{std::move(tracer), std::move(sampler)}};
 }
 
 std::shared_ptr<ot::Tracer>

zipkin_opentracing/src/propagation.cc

@@ -64,7 +64,7 @@ injectSpanContext(const opentracing::TextMapWriter &carrier,
   if (!result) {
     return result;
   }
-  result = carrier.Set(zipkin_sampled, "1");
+  result = carrier.Set(zipkin_sampled, span_context.isSampled() ? "1" : "0");
   if (!result) {
     return result;
   }
@@ -105,7 +105,7 @@ extractSpanContext(const opentracing::TextMapReader &carrier,
   TraceId trace_id;
   Optional<TraceId> parent_id;
   uint64_t span_id;
-  flags_t flags = 0;
+  flags_t flags = 0;  
   auto result = carrier.ForeachKey(
       [&](ot::string_view key, ot::string_view value) -> ot::expected<void> {
         if (keyCompare(key, zipkin_trace_id)) {
@@ -135,7 +135,10 @@ extractSpanContext(const opentracing::TextMapReader &carrier,
           if (!parseBool(value, sampled)) {
             return ot::make_unexpected(ot::span_context_corrupted_error);
           }
-          flags |= sampled_flag;
+          if (sampled) {
+            flags |= sampled_flag;
+          }
+          flags |= sampling_set_flag;
         } else if (keyCompare(key, zipkin_parent_span_id)) {
           parent_id = Hex::hexToTraceId(value);
           if (!parent_id.valid()) {

zipkin_opentracing/src/sampling.cc

@@ -0,0 +1,11 @@
+#include "sampling.h"
+#include <random>
+
+namespace zipkin {
+  bool
+  ProbabilisticSampler::ShouldSample() {
+    static thread_local std::mt19937 rng(std::random_device{}());
+    std::bernoulli_distribution dist(sample_rate_);
+    return dist(rng);
+  }
+} // namespace zipkin

zipkin_opentracing/src/sampling.h

@@ -0,0 +1,22 @@
+#pragma once
+#include <memory>
+#include <algorithm>
+
+namespace zipkin {
+  class Sampler {
+    public:
+      virtual ~Sampler() = default;
+      virtual bool ShouldSample() = 0;
+  };
+
+  class ProbabilisticSampler : public Sampler {
+    public:
+      ProbabilisticSampler(double sample_rate) : sample_rate_(std::max(0.0, std::min(sample_rate, 1.0))) {};
+      bool ShouldSample() override;
+
+    private:
+      double sample_rate_;
+  };
+
+  typedef std::unique_ptr<Sampler> SamplerPtr;
+} // namespace zipkin

zipkin_opentracing/src/tracer_factory.cc

@@ -37,6 +37,11 @@ const char *const configuration_schema = R"(
       "description":
         "The maximum number of spans to buffer before sending them to the collector",
       "minimum": 1
+    },
+    "sample_rate": {
+      "type": "float",
+      "minimum": 0.0,
+      "maxiumum": 1.0
     }
   }
 }
@@ -91,6 +96,9 @@ OtTracerFactory::MakeTracer(const char *configuration,
   if (document.HasMember("max_buffered_spans")) {
     options.max_buffered_spans = document["max_buffered_spans"].GetInt();
   }
+  if (document.HasMember("sample_rate")) {
+    options.sample_rate = document["sample_rate"].GetDouble();
+  }
   return makeZipkinOtTracer(options);
 } catch (const std::bad_alloc &) {
   return opentracing::make_unexpected(

zipkin_opentracing/test/ot_tracer_test.cc

@@ -3,6 +3,7 @@
 #include <algorithm>
 #include <opentracing/noop.h>
 #include <zipkin/opentracing.h>
+#include "../src/sampling.h"
 
 #define CATCH_CONFIG_MAIN
 #include <zipkin/catch/catch.hpp>
@@ -55,6 +56,42 @@ TEST_CASE("ot_tracer") {
     CHECK(hasTag(span, "xyz", true));
   }
 
+  SECTION("Uses sampling rate to determine whether to sample a span.") {
+    auto r = new InMemoryReporter();
+
+    ZipkinOtTracerOptions options;
+    options.sample_rate = 0.0;
+    auto t = makeZipkinOtTracer(options, std::unique_ptr<Reporter>(r));
+
+    auto span_a = t->StartSpan("a");
+    CHECK(span_a);
+    span_a->Finish();
+
+    CHECK(r->spans().empty());
+  }
+
+  SECTION("Propagates sampling decision to child span") {
+    ZipkinOtTracerOptions no_sampling;
+    no_sampling.sample_rate = 0.0;
+    auto r1 = new InMemoryReporter();
+    auto no_sampling_tracer = makeZipkinOtTracer(no_sampling, std::unique_ptr<Reporter>(r1));
+
+    ZipkinOtTracerOptions always_sample;
+    always_sample.sample_rate = 1.0;
+    auto r2 = new InMemoryReporter();
+    auto sampling_tracer = makeZipkinOtTracer(always_sample, std::unique_ptr<Reporter>(r2));
+
+    auto span_a = no_sampling_tracer->StartSpan("a");
+    CHECK(span_a);
+    span_a->Finish();
+    auto span_b = sampling_tracer->StartSpan("b", {ChildOf(&span_a->context())});
+    CHECK(span_b);
+    span_b->Finish();
+
+    CHECK(r1->spans().empty());
+    CHECK(r2->spans().empty());
+  }
+
   SECTION("You can set a single child-of reference when starting a span.") {
     auto span_a = tracer->StartSpan("a");
     CHECK(span_a);