Make WithContext conditional. (#207)

While here, add an example to the docs.

Author: g-easy

opencensus/context/internal/with_context.cc

@@ -22,26 +22,26 @@
 namespace opencensus {
 namespace context {
 
-WithContext::WithContext(const Context& ctx)
-    : swapped_context_(ctx)
+WithContext::WithContext(const Context& ctx, bool cond)
+    : swapped_context_(cond ? ctx : Context())
 #ifndef NDEBUG
       ,
       original_context_(Context::InternalMutableCurrent())
 #endif
-{
-  using std::swap;
-  swap(*Context::InternalMutableCurrent(), swapped_context_);
+      ,
+      cond_(cond) {
+  ConditionalSwap();
 }
 
-WithContext::WithContext(Context&& ctx)
-    : swapped_context_(std::move(ctx))
+WithContext::WithContext(Context&& ctx, bool cond)
+    : swapped_context_(cond ? std::move(ctx) : Context())
 #ifndef NDEBUG
       ,
       original_context_(Context::InternalMutableCurrent())
 #endif
-{
-  using std::swap;
-  swap(*Context::InternalMutableCurrent(), swapped_context_);
+      ,
+      cond_(cond) {
+  ConditionalSwap();
 }
 
 WithContext::~WithContext() {
@@ -50,8 +50,14 @@ WithContext::~WithContext() {
          "WithContext must be destructed on the same thread as it was "
          "constructed.");
 #endif
-  using std::swap;
-  swap(*Context::InternalMutableCurrent(), swapped_context_);
+  ConditionalSwap();
+}
+
+void WithContext::ConditionalSwap() {
+  if (cond_) {
+    using std::swap;
+    swap(*Context::InternalMutableCurrent(), swapped_context_);
+  }
 }
 
 }  // namespace context

opencensus/context/internal/with_context_test.cc

@@ -35,6 +35,13 @@ TEST(WithContextTest, WithContextMovable) {
   opencensus::context::WithContext wc(std::move(ctx));
 }
 
+TEST(WithContextTest, WithContextConditional) {
+  opencensus::context::Context ctx = opencensus::context::Context::Current();
+  opencensus::context::WithContext wc1(ctx, true);
+  opencensus::context::WithContext wc2(ctx, false);
+  // TODO: Test swaps.
+}
+
 #ifndef NDEBUG
 TEST(WithContextDeathTest, DestructorOnWrongThread) {
   opencensus::context::Context ctx = opencensus::context::Context::Current();

opencensus/context/with_context.h

@@ -21,15 +21,22 @@ namespace opencensus {
 namespace context {
 
 // WithContext is a scoped object that sets the current Context to the given
-// one, until the WithContext object is destroyed.
+// one, until the WithContext object is destroyed. If the condition is false, it
+// doesn't do anything.
 //
-// Because it changes the current (thread local) context, NEVER allocate a
-// WithContext in one thread and deallocate in another. A simple way to ensure
-// this is to only ever stack-allocate it.
+// Because WithContext changes the current (thread local) context, NEVER
+// allocate a WithContext in one thread and deallocate in another. A simple way
+// to ensure this is to only ever stack-allocate it.
+//
+// Example usage:
+// {
+//   WithContext wc(op.ctx_);
+//   // Do work.
+// }
 class WithContext {
  public:
-  explicit WithContext(const Context& ctx);
-  explicit WithContext(Context&& ctx);
+  explicit WithContext(const Context& ctx, bool cond = true);
+  explicit WithContext(Context&& ctx, bool cond = true);
   ~WithContext();
 
  private:
@@ -39,10 +46,13 @@ class WithContext {
   WithContext& operator=(const WithContext&) = delete;
   WithContext& operator=(WithContext&&) = delete;
 
+  void ConditionalSwap();
+
   Context swapped_context_;
 #ifndef NDEBUG
   const Context* original_context_;
 #endif
+  const bool cond_;
 };
 
 }  // namespace context