From a506cf846edca75b93e5457aca51c568378201be Mon Sep 17 00:00:00 2001
From: easy <g-easy@users.noreply.github.com>
Date: Thu, 20 Jun 2019 22:42:20 +1000
Subject: [PATCH] Add an option for WithSpan() to End() the Span. (#327)

Fixes #323.
---
 opencensus/trace/internal/with_span.cc      |  8 ++++++--
 opencensus/trace/internal/with_span_test.cc | 11 +++++++++++
 opencensus/trace/with_span.h                |  6 ++++--
 3 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/opencensus/trace/internal/with_span.cc b/opencensus/trace/internal/with_span.cc
index 76660e21..67fa9baa 100644
--- a/opencensus/trace/internal/with_span.cc
+++ b/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();
 }
 
diff --git a/opencensus/trace/internal/with_span_test.cc b/opencensus/trace/internal/with_span_test.cc
index 03ed037c..0ca8121b 100644
--- a/opencensus/trace/internal/with_span_test.cc
+++ b/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");
diff --git a/opencensus/trace/with_span.h b/opencensus/trace/with_span.h
index 230e3168..7cac6b56 100644
--- a/opencensus/trace/with_span.h
+++ b/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