Free per-thread Context memory on thread exit. (#417)

* Add test for leaking Context memory.

Author: g-easy

opencensus/context/context.h

@@ -67,6 +67,7 @@ class Context {
 
   friend class ContextTestPeer;
   friend class WithContext;
+  friend class ContextWrapper;
   friend class ::opencensus::tags::ContextPeer;
   friend class ::opencensus::tags::WithTagMap;
   friend class ::opencensus::trace::ContextPeer;

opencensus/context/internal/context.cc

@@ -15,6 +15,7 @@
 #include "opencensus/context/context.h"
 
 #include <functional>
+#include <memory>
 #include <utility>
 
 #include "absl/strings/str_cat.h"
@@ -25,6 +26,20 @@
 namespace opencensus {
 namespace context {
 
+// Wrapper for per-thread Context, frees the Context on thread shutdown.
+class ContextWrapper {
+ public:
+  ContextWrapper() : ptr_(new Context) {}
+  Context* get() { return ptr_.get(); }
+
+ private:
+  std::unique_ptr<Context> ptr_;
+};
+
+namespace {
+thread_local ContextWrapper g_wrapper;
+}  // namespace
+
 Context::Context()
     : tags_(opencensus::tags::TagMap({})),
       span_(opencensus::trace::Span::BlankSpan()) {}
@@ -47,11 +62,7 @@ std::string Context::DebugString() const {
 }
 
 // static
-Context* Context::InternalMutableCurrent() {
-  static thread_local Context* thread_ctx = nullptr;
-  if (thread_ctx == nullptr) thread_ctx = new Context;
-  return thread_ctx;
-}
+Context* Context::InternalMutableCurrent() { return g_wrapper.get(); }
 
 void swap(Context& a, Context& b) {
   using std::swap;

opencensus/context/internal/context_test.cc

@@ -16,6 +16,7 @@
 
 #include <functional>
 #include <iostream>
+#include <thread>
 
 #include "gtest/gtest.h"
 #include "opencensus/tags/context_util.h"
@@ -87,6 +88,18 @@ TEST(ContextTest, WrapDoesNotLeak) {
   span.End();
 }
 
+TEST(ContextTest, ThreadDoesNotLeak) {
+  auto span = opencensus::trace::Span::StartSpan("MySpan");
+  std::thread t([span]() {
+    // Leak-sanitizer (part of ASAN) throws an error if this leaks.
+    opencensus::tags::WithTagMap wt(ExampleTagMap());
+    opencensus::trace::WithSpan ws(span);
+  });
+  t.join();
+  ExpectEmptyContext();
+  span.End();
+}
+
 TEST(ContextTest, WrappedFnIsCopiable) {
   std::function<void()> fn1, fn2;
   {