Add an option for WithSpan() to End() the Span. (#327)

Fixes #323.

Author: g-easy

opencensus/trace/internal/with_span.cc

@@ -26,14 +26,15 @@ using ::opencensus::trace::Span;
 namespace opencensus {
 namespace trace {
 
-WithSpan::WithSpan(const Span& span, bool cond)
+WithSpan::WithSpan(const Span& span, bool cond, bool end_span)
     : swapped_span_(span)
 #ifndef NDEBUG
       ,
       original_context_(Context::InternalMutableCurrent())
 #endif
       ,
-      cond_(cond) {
+      cond_(cond),
+      end_span_(end_span) {
   ConditionalSwap();
 }
 
@@ -43,6 +44,9 @@ WithSpan::~WithSpan() {
          "WithSpan must be destructed on the same thread as it was "
          "constructed.");
 #endif
+  if (cond_ && end_span_) {
+    Context::InternalMutableCurrent()->span_.End();
+  }
   ConditionalSwap();
 }
 

opencensus/trace/internal/with_span_test.cc

@@ -107,6 +107,17 @@ TEST(WithSpanTest, DisabledViaConditional) {
   span1.End();
 }
 
+TEST(WithSpanTest, EndSpan) {
+  auto span = opencensus::trace::Span::StartSpan("MySpan");
+  ExpectNoSpan();
+  {
+    opencensus::trace::WithSpan ws(span, /*cond=*/true, /*end_span=*/true);
+    EXPECT_EQ(span.context(), ContextTestPeer::CurrentCtx());
+  }
+  ExpectNoSpan();
+  // TODO: Check End() was called.
+}
+
 #ifndef NDEBUG
 TEST(WithSpanDeathTest, DestructorOnWrongThread) {
   auto span = opencensus::trace::Span::StartSpan("MySpan");

opencensus/trace/with_span.h

@@ -23,7 +23,8 @@ namespace trace {
 
 // WithSpan is a scoped object that sets the current Span to the given one,
 // until the WithSpan object is destroyed. If the condition is false, it doesn't
-// do anything.
+// do anything. If the condition is true and end_span is true, it calls End() on
+// the Span when it falls out of scope.
 //
 // Because WithSpan changes the current (thread local) context, NEVER allocate a
 // WithSpan in one thread and deallocate in another. A simple way to ensure this
@@ -36,7 +37,7 @@ namespace trace {
 // }
 class WithSpan {
  public:
-  explicit WithSpan(const Span& span, bool cond = true);
+  explicit WithSpan(const Span& span, bool cond = true, bool end_span = false);
   ~WithSpan();
 
   // No Span&& constructor because it encourages "consuming" the Span with a
@@ -57,6 +58,7 @@ class WithSpan {
   const ::opencensus::context::Context* original_context_;
 #endif
   const bool cond_;
+  const bool end_span_;
 };
 
 }  // namespace trace