Add support for StrictSpans mode in Tracing

Goal: StrictSpans mode is meant to enforce the invariant that a parent span cannot
end while it has children running. If a parent span ends with active children, it is held open
and its end is triggered when the last child ends. This works even with a chain of spans, the last child
can trigger all parents up to the root to end.

Details:
1. Switch running_span_store_impl to key its map on SpanId instead of the pointer
2. Add FindSpanById method on RunningSpanStoreImpl
3. Add reference count called active_child_count_ on SpanImpl to keep track of active children
4. Add logic at span start and end time to update the reference count on parents,
5. Add tests

opencensus/trace/internal/running_span_store_impl.cc

@@ -34,35 +34,40 @@ namespace opencensus {
 namespace trace {
 namespace exporter {
 
-namespace {
-// Returns the memory address of the SpanImpl object, to be used as a key into
-// the map of spans.
-uintptr_t GetKey(const SpanImpl* span) {
-  return reinterpret_cast<uintptr_t>(span);
-}
-}  // namespace
-
 RunningSpanStoreImpl* RunningSpanStoreImpl::Get() {
   static RunningSpanStoreImpl* global_running_span_store =
       new RunningSpanStoreImpl;
   return global_running_span_store;
 }
 
-void RunningSpanStoreImpl::AddSpan(const std::shared_ptr<SpanImpl>& span) {
+void RunningSpanStoreImpl::AddSpan(const SpanId& id,
+                                   const std::shared_ptr<SpanImpl>& span) {
   absl::MutexLock l(&mu_);
-  spans_.insert({GetKey(span.get()), span});
+  spans_.insert({id, span});
 }
 
-bool RunningSpanStoreImpl::RemoveSpan(const std::shared_ptr<SpanImpl>& span) {
+bool RunningSpanStoreImpl::RemoveSpan(const SpanId& id) {
   absl::MutexLock l(&mu_);
-  auto iter = spans_.find(GetKey(span.get()));
+  auto iter = spans_.find(id);
   if (iter == spans_.end()) {
     return false;  // Not tracked.
   }
   spans_.erase(iter);
   return true;
 }
 
+bool RunningSpanStoreImpl::FindSpan(const SpanId& id,
+                                    std::shared_ptr<SpanImpl>* span) {
+  absl::MutexLock l(&mu_);
+  auto iter = spans_.find(id);
+  if (iter != spans_.end()) {
+    *span = iter->second;
+    return true;
+  } else {
+    return false;
+  }
+}
+
 RunningSpanStore::Summary RunningSpanStoreImpl::GetSummary() const {
   RunningSpanStore::Summary summary;
   absl::MutexLock l(&mu_);

opencensus/trace/internal/running_span_store_impl.h

@@ -38,11 +38,16 @@ class RunningSpanStoreImpl {
   static RunningSpanStoreImpl* Get();
 
   // Adds a new running Span.
-  void AddSpan(const std::shared_ptr<SpanImpl>& span) LOCKS_EXCLUDED(mu_);
+  void AddSpan(const SpanId& id, const std::shared_ptr<SpanImpl>& span)
+      LOCKS_EXCLUDED(mu_);
+
+  // Finds a Span with id SpanId, returns false if not found.
+  bool FindSpan(const SpanId& id, std::shared_ptr<SpanImpl>* span)
+      LOCKS_EXCLUDED(mu_);
 
   // Removes a Span that's no longer running. Returns true on success, false if
   // that Span was not being tracked.
-  bool RemoveSpan(const std::shared_ptr<SpanImpl>& span) LOCKS_EXCLUDED(mu_);
+  bool RemoveSpan(const SpanId& id) LOCKS_EXCLUDED(mu_);
 
   // Returns a summary of the data available in the RunningSpanStore.
   RunningSpanStore::Summary GetSummary() const LOCKS_EXCLUDED(mu_);
@@ -61,8 +66,12 @@ class RunningSpanStoreImpl {
 
   mutable absl::Mutex mu_;
 
-  // The key is the memory address of the underlying SpanImpl object.
-  std::unordered_map<uintptr_t, std::shared_ptr<SpanImpl>> spans_
+  // Necessary for using SpanId as the key in the map.
+  struct SpanIdHash {
+    std::size_t operator()(SpanId const& s) const noexcept { return s.hash(); }
+  };
+
+  std::unordered_map<SpanId, std::shared_ptr<SpanImpl>, SpanIdHash> spans_
       GUARDED_BY(mu_);
 };
 

opencensus/trace/internal/span.cc

@@ -69,13 +69,23 @@ class SpanGenerator {
     SpanId span_id = GenerateRandomSpanId();
     TraceId trace_id;
     SpanId parent_span_id;
-    TraceOptions trace_options;
+    TraceOptions trace_options = options.trace_options;
     if (parent_ctx == nullptr) {
       trace_id = GenerateRandomTraceId();
     } else {
       trace_id = parent_ctx->trace_id();
       parent_span_id = parent_ctx->span_id();
       trace_options = parent_ctx->trace_options();
+      if (trace_options.HasStrictSpans() && trace_options.IsSampled() &&
+          !has_remote_parent) {
+        std::shared_ptr<SpanImpl> span;
+        if (exporter::RunningSpanStoreImpl::Get()->FindSpan(parent_span_id,
+                                                            &span)) {
+          span->IncrementChildCount();
+        } else {
+          assert(false);
+        }
+      }
     }
     if (!trace_options.IsSampled()) {
       bool should_sample = false;
@@ -135,10 +145,13 @@ Span Span::StartSpanWithRemoteParent(absl::string_view name,
 Span::Span(const SpanContext& context, SpanImpl* impl)
     : context_(context), span_impl_(impl) {
   if (IsRecording()) {
-    exporter::RunningSpanStoreImpl::Get()->AddSpan(span_impl_);
+    exporter::RunningSpanStoreImpl::Get()->AddSpan(context.span_id(),
+                                                   span_impl_);
   }
 }
 
+void Span::IncrementChildCount() const { span_impl_->IncrementChildCount(); }
+
 void Span::AddAttribute(absl::string_view key,
                         AttributeValueRef attribute) const {
   if (IsRecording()) {
@@ -204,13 +217,37 @@ void Span::SetStatus(StatusCode canonical_code,
 
 void Span::End() const {
   if (IsRecording()) {
-    if (!span_impl_->End()) {
-      // The Span already ended, ignore this call.
+    EndSpanById(context_.span_id());
+  }
+}
+
+void Span::EndSpanById(const SpanId& id) {
+  std::shared_ptr<SpanImpl> span;
+  auto found = exporter::RunningSpanStoreImpl::Get()->FindSpan(id, &span);
+  assert(found && "Cannot call end span on an invalid ID");
+  if (!found) {
+    return;
+  }
+
+  if (!span->End()) {
+    return;
+  }
+  exporter::RunningSpanStoreImpl::Get()->RemoveSpan(id);
+  exporter::LocalSpanStoreImpl::Get()->AddSpan(span);
+  exporter::SpanExporterImpl::Get()->AddSpan(span);
+
+  if (span->context().trace_options().HasStrictSpans() &&
+      (!span->remote_parent() && span->parent_span_id().IsValid())) {
+    std::shared_ptr<SpanImpl> parent_span;
+
+    found = exporter::RunningSpanStoreImpl::Get()->FindSpan(
+        span->parent_span_id(), &parent_span);
+    assert(found && "Parent Span Id Not Found in Strict Spans Mode");
+
+    if (!found) {
       return;
     }
-    exporter::RunningSpanStoreImpl::Get()->RemoveSpan(span_impl_);
-    exporter::LocalSpanStoreImpl::Get()->AddSpan(span_impl_);
-    exporter::SpanExporterImpl::Get()->AddSpan(span_impl_);
+    parent_span->DecrementChildCount();
   }
 }
 

opencensus/trace/internal/span_id.cc

@@ -15,6 +15,7 @@
 #include "opencensus/trace/span_id.h"
 
 #include <cstring>
+#include <functional>
 #include <string>
 
 #include "absl/strings/escaping.h"
@@ -42,6 +43,15 @@ bool SpanId::IsValid() const {
   return tmp != 0;
 }
 
+std::size_t SpanId::hash() const {
+  // SpanIds are currently 64 bits. this hash function is only ever used
+  // internally within a given process, so this can be easily changed if we
+  // ever increase the size.
+  std::hash<uint64_t> hash_fn;
+  uint64_t val = *(reinterpret_cast<const uint64_t *>(rep_));
+  return hash_fn(val);
+}
+
 void SpanId::CopyTo(uint8_t *buf) const { memcpy(buf, rep_, kSize); }
 
 }  // namespace trace

opencensus/trace/internal/span_impl.cc

@@ -73,8 +73,6 @@ std::unordered_map<std::string, exporter::AttributeValue> CopyAttributes(
 }
 }  // namespace
 
-// SpanImpl::SpanImpl() : has_ended_(false), remote_parent_(false) {}
-
 SpanImpl::SpanImpl(const SpanContext& context, const TraceParams& trace_params,
                    absl::string_view name, const SpanId& parent_span_id,
                    bool remote_parent)
@@ -87,7 +85,9 @@ SpanImpl::SpanImpl(const SpanContext& context, const TraceParams& trace_params,
       links_(trace_params.max_links),
       attributes_(trace_params.max_attributes),
       has_ended_(false),
-      remote_parent_(remote_parent) {}
+      remote_parent_(remote_parent),
+      active_child_count_(0),
+      end_requested_(false) {}
 
 void SpanImpl::AddAttributes(AttributesRef attributes) {
   absl::MutexLock l(&mu_);
@@ -139,6 +139,10 @@ void SpanImpl::SetStatus(exporter::Status&& status) {
 
 bool SpanImpl::End() {
   absl::MutexLock l(&mu_);
+  if (context_.trace_options().HasStrictSpans() && active_child_count_ > 0) {
+    end_requested_ = true;
+    return false;
+  }
   if (has_ended_) {
     assert(false && "Invalid attempt to End() the same Span more than once.");
     // In non-debug builds, just ignore the second End().
@@ -154,6 +158,25 @@ bool SpanImpl::HasEnded() const {
   return has_ended_;
 }
 
+void SpanImpl::IncrementChildCount() const {
+  absl::MutexLock l(&mu_);
+  active_child_count_++;
+}
+
+void SpanImpl::DecrementChildCount() {
+  bool shouldEnd = false;
+  {
+    absl::MutexLock l(&mu_);
+    assert(active_child_count_ > 0);
+    active_child_count_--;
+    shouldEnd = (active_child_count_ == 0 && end_requested_);
+  }
+
+  if (shouldEnd) {
+    Span::EndSpanById(context_.span_id());
+  }
+}
+
 exporter::SpanData SpanImpl::ToSpanData() const {
   absl::MutexLock l(&mu_);
   // Make a deep copy of attributes.

opencensus/trace/internal/span_impl.h

@@ -102,6 +102,15 @@ class SpanImpl final {
 
   SpanId parent_span_id() const { return parent_span_id_; }
 
+  bool remote_parent() const { return remote_parent_; }
+
+  // Increments the counter of open child spans below this span.
+  void IncrementChildCount() const;
+
+  // Decrements the open child counter AND if the counter becomes
+  // zero then End() is called on the span.
+  void DecrementChildCount();
+
  private:
   friend class ::opencensus::trace::exporter::RunningSpanStoreImpl;
   friend class ::opencensus::trace::exporter::LocalSpanStoreImpl;
@@ -138,6 +147,11 @@ class SpanImpl final {
   bool has_ended_ GUARDED_BY(mu_);
   // True if the parent Span is in a different process.
   const bool remote_parent_;
+  // Counts the number of open children linked to this span.
+  mutable uint32_t active_child_count_ GUARDED_BY(mu_);
+  // Keeps track if End() was called on this span while it had open child
+  // spans and is waiting on them.
+  mutable bool end_requested_ GUARDED_BY(mu_);
 };
 
 }  // namespace trace

opencensus/trace/internal/span_test.cc

@@ -282,6 +282,71 @@ TEST(SpanTest, FullSpanTest) {
   expect_message_event(3, 4, 5, data.message_events().events()[1].event());
 }
 
+TEST(SpanTest, StrictSpanTest) {
+  AlwaysSampler sampler;
+  TraceOptions t = TraceOptions().WithStrictSpans(true);
+  StartSpanOptions opts = {&sampler, {}, t};
+  auto span = ::opencensus::trace::Span::StartSpan("MyRootSpan", nullptr, opts);
+  auto child = Span::StartSpan("child", &span);
+  EXPECT_TRUE(span.context().IsValid());
+
+  // Calling end on the parent span with an open child should fail/not cause it
+  // to end
+  span.End();
+  const exporter::SpanData data = SpanTestPeer::ToSpanData(&span);
+  EXPECT_EQ("MyRootSpan", data.name());
+  EXPECT_EQ(false, data.has_ended());
+
+  // Now end the child and see that the parent span is ended as well
+  child.End();
+  const exporter::SpanData data2 = SpanTestPeer::ToSpanData(&span);
+  EXPECT_EQ("MyRootSpan", data2.name());
+  EXPECT_EQ(true, data2.has_ended());
+
+  // Now, try ending child first and verify that that doesn't end the parent
+  // span
+  auto span2 = ::opencensus::trace::Span::StartSpan("root2", nullptr, opts);
+  auto child2 = Span::StartSpan("child2", &span2);
+  EXPECT_TRUE(span2.context().IsValid());
+
+  // Now end the child and see that the parent span is ended as well
+  child2.End();
+  const exporter::SpanData root2_pre = SpanTestPeer::ToSpanData(&span2);
+  EXPECT_EQ("root2", root2_pre.name());
+  EXPECT_EQ(false, root2_pre.has_ended());
+
+  // Calling end on the parent span with an open child should fail/not cause it
+  // to end
+  span2.End();
+  const exporter::SpanData root2_post = SpanTestPeer::ToSpanData(&span2);
+  EXPECT_EQ("root2", root2_post.name());
+  EXPECT_EQ(true, root2_post.has_ended());
+
+  // Now try with 3 tiers
+  auto grandparent =
+      ::opencensus::trace::Span::StartSpan("grandpa", nullptr, opts);
+  auto parent =
+      ::opencensus::trace::Span::StartSpan("parent", &grandparent, opts);
+  auto kid = Span::StartSpan("kid", &parent);
+
+  // End both grandparent and parent, neither should have ended
+  grandparent.End();
+  parent.End();
+  const exporter::SpanData grandparent_open_data =
+      SpanTestPeer::ToSpanData(&grandparent);
+  EXPECT_EQ(false, grandparent_open_data.has_ended());
+  const exporter::SpanData parent_open_data = SpanTestPeer::ToSpanData(&parent);
+  EXPECT_EQ(false, parent_open_data.has_ended());
+
+  kid.End();
+  const exporter::SpanData grandparent_closed_data =
+      SpanTestPeer::ToSpanData(&grandparent);
+  EXPECT_EQ(true, grandparent_closed_data.has_ended());
+  const exporter::SpanData parent_closed_data =
+      SpanTestPeer::ToSpanData(&parent);
+  EXPECT_EQ(true, parent_closed_data.has_ended());
+}
+
 TEST(SpanTest, ChildInheritsTraceOption) {
   constexpr uint8_t trace_id[] = {1, 2,  3,  4,  5,  6,  7,  8,
                                   9, 10, 11, 12, 13, 14, 15, 16};

opencensus/trace/internal/trace_options.cc

@@ -25,12 +25,15 @@ namespace trace {
 namespace {
 // One bit per option.
 constexpr uint8_t kIsSampled = 1;
+constexpr uint8_t kStrictSpans = 2;
 }  // namespace
 
 TraceOptions::TraceOptions(const uint8_t* buf) { memcpy(rep_, buf, kSize); }
 
 bool TraceOptions::IsSampled() const { return rep_[0] & kIsSampled; }
 
+bool TraceOptions::HasStrictSpans() const { return rep_[0] & kStrictSpans; }
+
 bool TraceOptions::operator==(const TraceOptions& that) const {
   return memcmp(rep_, that.rep_, kSize) == 0;
 }
@@ -49,5 +52,12 @@ TraceOptions TraceOptions::WithSampling(bool is_sampled) const {
   return TraceOptions(buf);
 }
 
+TraceOptions TraceOptions::WithStrictSpans(bool strict_spans) const {
+  uint8_t buf[kSize];
+  CopyTo(buf);
+  buf[0] = (buf[0] & ~kStrictSpans) | (strict_spans ? kStrictSpans : 0);
+  return TraceOptions(buf);
+}
+
 }  // namespace trace
 }  // namespace opencensus

opencensus/trace/internal/trace_options_test.cc

@@ -32,6 +32,10 @@ TEST(TraceOptionsTest, SetSampled) {
   opts = opts.WithSampling(false);
   EXPECT_EQ("00", opts.ToHex());
   EXPECT_FALSE(opts.IsSampled());
+  EXPECT_FALSE(opts.HasStrictSpans());
+  opts = opts.WithStrictSpans(true);
+  EXPECT_EQ("02", opts.ToHex());
+  EXPECT_TRUE(opts.HasStrictSpans());
 }
 
 TEST(TraceOptionsTest, Comparison) {