Actually manage memory in the old-style AOS logging

LeakSanitizer should be happy with it now. It's also still just as
thread-safe.

Change-Id: Id09a0349657cf4f719267b053f0ea3d8ec366256
diff --git a/aos/logging/BUILD b/aos/logging/BUILD
index a499d68..6686a3c 100644
--- a/aos/logging/BUILD
+++ b/aos/logging/BUILD
@@ -5,20 +5,28 @@
     name = "logging",
     srcs = [
         "context.cc",
+        "implementations.cc",
         "interface.cc",
     ],
     hdrs = [
         "context.h",
+        "implementations.h",
         "interface.h",
         "logging.h",
     ],
     visibility = ["//visibility:public"],
     deps = [
+        ":printf_formats",
         ":sizes",
         "//aos:complex_thread_local",
         "//aos:die",
         "//aos:macros",
         "//aos/libc:aos_strerror",
+        "//aos/mutex",
+        "//aos/stl_mutex",
+        "//aos/time",
+        "//aos/type_traits",
+        "@com_google_absl//absl/base",
     ],
 )
 
@@ -66,26 +74,9 @@
 
 cc_library(
     name = "implementations",
-    srcs = [
-        "implementations.cc",
-    ],
-    hdrs = [
-        "implementations.h",
-    ],
-    linkopts = [
-        "-lpthread",
-    ],
     visibility = ["//visibility:public"],
     deps = [
         ":logging",
-        ":printf_formats",
-        ":sizes",
-        "//aos:die",
-        "//aos:macros",
-        "//aos/mutex",
-        "//aos/time",
-        "//aos/type_traits",
-        "@com_google_absl//absl/base",
     ],
 )
 
diff --git a/aos/logging/context.cc b/aos/logging/context.cc
index 84e4e80..2c689ed 100644
--- a/aos/logging/context.cc
+++ b/aos/logging/context.cc
@@ -21,6 +21,7 @@
 
 #include "aos/complex_thread_local.h"
 #include "aos/die.h"
+#include "aos/logging/implementations.h"
 
 namespace aos {
 namespace logging {
@@ -71,10 +72,7 @@
 
 }  // namespace
 
-::std::atomic<LogImplementation *> global_top_implementation(NULL);
-
-Context::Context()
-    : implementation(global_top_implementation.load()), sequence(0) {
+Context::Context() : implementation(GetImplementation()), sequence(0) {
   cork_data.Reset();
 }
 
@@ -105,6 +103,11 @@
 
 void Context::Delete() { delete_current_context = true; }
 
+void Context::DeleteNow() {
+  my_context.Clear();
+  delete_current_context = false;
+}
+
 }  // namespace internal
 }  // namespace logging
 }  // namespace aos
diff --git a/aos/logging/context.h b/aos/logging/context.h
index b2abc0d..2e8250e 100644
--- a/aos/logging/context.h
+++ b/aos/logging/context.h
@@ -2,9 +2,10 @@
 #define AOS_LOGGING_CONTEXT_H_
 
 #include <inttypes.h>
+#include <limits.h>
 #include <stddef.h>
 #include <sys/types.h>
-#include <limits.h>
+#include <memory>
 
 #include <atomic>
 
@@ -19,8 +20,6 @@
 // goes.
 namespace internal {
 
-extern ::std::atomic<LogImplementation *> global_top_implementation;
-
 // An separate instance of this class is accessible from each task/thread.
 // NOTE: It will get deleted in the child of a fork.
 //
@@ -43,9 +42,11 @@
   // still work to clean up any state.
   static void Delete();
 
+  static void DeleteNow();
+
   // Which one to log to right now.
   // Will be NULL if there is no logging implementation to use right now.
-  LogImplementation *implementation;
+  std::shared_ptr<LogImplementation> implementation;
 
   // A name representing this task/(process and thread).
   char name[LOG_MESSAGE_NAME_LEN];
diff --git a/aos/logging/implementations.cc b/aos/logging/implementations.cc
index 094ecbb..096127a 100644
--- a/aos/logging/implementations.cc
+++ b/aos/logging/implementations.cc
@@ -5,10 +5,12 @@
 
 #include <algorithm>
 #include <chrono>
+#include <mutex>
 
 #include "absl/base/call_once.h"
 #include "aos/die.h"
 #include "aos/logging/printf_formats.h"
+#include "aos/stl_mutex/stl_mutex.h"
 #include "aos/time/time.h"
 
 namespace aos {
@@ -17,50 +19,15 @@
 
 namespace chrono = ::std::chrono;
 
-// The root LogImplementation. It only logs to stderr/stdout.
-// Some of the things specified in the LogImplementation documentation doesn't
-// apply here (mostly the parts about being able to use AOS_LOG) because this is
-// the root one.
-class RootLogImplementation : public LogImplementation {
- protected:
-  virtual ::aos::monotonic_clock::time_point monotonic_now() const {
-    return ::aos::monotonic_clock::now();
-  }
-
- private:
-  __attribute__((format(GOOD_PRINTF_FORMAT_TYPE, 3, 0))) void DoLog(
-      log_level level, const char *format, va_list ap) override {
-    LogMessage message;
-    internal::FillInMessage(level, monotonic_now(), format, ap, &message);
-    internal::PrintMessage(stderr, message);
+struct GlobalState {
+  std::shared_ptr<LogImplementation> implementation;
+  aos::stl_mutex lock;
+  static GlobalState *Get() {
+    static GlobalState r;
+    return &r;
   }
 };
 
-RootLogImplementation *root_implementation = nullptr;
-
-void SetGlobalImplementation(LogImplementation *implementation) {
-  if (root_implementation == nullptr) {
-    fputs("Somebody didn't call logging::Init()!\n", stderr);
-    abort();
-  }
-
-  internal::Context *context = internal::Context::Get();
-
-  context->implementation = implementation;
-  internal::global_top_implementation.store(implementation);
-}
-
-void NewContext() { internal::Context::Delete(); }
-
-void DoInit() {
-  SetGlobalImplementation(root_implementation = new RootLogImplementation());
-
-  if (pthread_atfork(NULL /*prepare*/, NULL /*parent*/, NewContext /*child*/) !=
-      0) {
-    AOS_LOG(FATAL, "pthread_atfork(NULL, NULL, %p) failed\n", NewContext);
-  }
-}
-
 }  // namespace
 namespace internal {
 namespace {
@@ -127,58 +94,56 @@
   internal::PrintMessage(stream_, message);
 }
 
-void SetImplementation(LogImplementation *implementation, bool update_global) {
-  internal::Context *context = internal::Context::Get();
+void SetImplementation(std::shared_ptr<LogImplementation> implementation) {
+  Init();
+  GlobalState *const global = GlobalState::Get();
+  std::unique_lock<aos::stl_mutex> locker(global->lock);
+  global->implementation = std::move(implementation);
+}
 
-  context->implementation = implementation;
-  if (update_global) {
-    SetGlobalImplementation(implementation);
+std::shared_ptr<LogImplementation> SwapImplementation(
+    std::shared_ptr<LogImplementation> implementation) {
+  std::shared_ptr<LogImplementation> result;
+  {
+    GlobalState *const global = GlobalState::Get();
+    std::unique_lock<aos::stl_mutex> locker(global->lock);
+    result = std::move(global->implementation);
+    global->implementation = std::move(implementation);
   }
+  Cleanup();
+  return result;
 }
 
-LogImplementation *SwapImplementation(LogImplementation *implementation) {
-  internal::Context *context = internal::Context::Get();
-
-  LogImplementation *old = context->implementation;
-  context->implementation = implementation;
-
-  return old;
+std::shared_ptr<LogImplementation> GetImplementation() {
+  GlobalState *const global = GlobalState::Get();
+  std::unique_lock<aos::stl_mutex> locker(global->lock);
+  CHECK(global->implementation);
+  return global->implementation;
 }
 
-LogImplementation *GetImplementation() {
-  return internal::Context::Get()->implementation;
-}
+namespace {
 
-void Init() {
-  static absl::once_flag once;
-  absl::call_once(once, DoInit);
-}
+struct DoInit {
+  DoInit() {
+    GlobalState *const global = GlobalState::Get();
+    std::unique_lock<aos::stl_mutex> locker(global->lock);
+    CHECK(!global->implementation);
+    global->implementation = std::make_shared<StreamLogImplementation>(stdout);
+  }
+};
+
+}  // namespace
+
+void Init() { static DoInit do_init; }
 
 void Load() { internal::Context::Get(); }
 
 void Cleanup() { internal::Context::Delete(); }
 
-namespace {
-
-class CallbackLogImplementation : public HandleMessageLogImplementation {
- public:
-  CallbackLogImplementation(
-      const ::std::function<void(const LogMessage &)> &callback)
-      : callback_(callback) {}
-
- private:
-  void HandleMessage(const LogMessage &message) override { callback_(message); }
-
-  ::std::function<void(const LogMessage &)> callback_;
-};
-
-}  // namespace
-
 void RegisterCallbackImplementation(
-    const ::std::function<void(const LogMessage &)> &callback,
-    bool update_global) {
+    const ::std::function<void(const LogMessage &)> &callback) {
   Init();
-  SetImplementation(new CallbackLogImplementation(callback), update_global);
+  SetImplementation(std::make_shared<CallbackLogImplementation>(callback));
 }
 
 }  // namespace logging
diff --git a/aos/logging/implementations.h b/aos/logging/implementations.h
index 0e17165..e8a2213 100644
--- a/aos/logging/implementations.h
+++ b/aos/logging/implementations.h
@@ -11,6 +11,7 @@
 
 #include <atomic>
 #include <functional>
+#include <memory>
 #include <string>
 
 #include "aos/logging/context.h"
@@ -118,31 +119,31 @@
   FILE *const stream_;
 };
 
-// Adds another implementation to the stack of implementations in this
-// task/thread.
-// Any tasks/threads created after this call will also use this implementation.
-// The cutoff is when the state in a given task/thread is created (either lazily
-// when needed or by calling Load()).
-// The logging system takes ownership of implementation. It will delete it if
-// necessary, so it must be created with new.
-// TODO: Log implementations are never deleted. Need means to safely deregister.
-void SetImplementation(LogImplementation *implementation,
-                       bool update_global = true);
+std::shared_ptr<LogImplementation> GetImplementation();
 
-// Updates the log implementation for the current thread, returning the current
-// implementation.
-LogImplementation *SwapImplementation(LogImplementation *implementation);
+// Sets the current implementation.
+void SetImplementation(std::shared_ptr<LogImplementation> implementation);
 
-LogImplementation *GetImplementation();
+// Updates the log implementation, returning the current implementation.
+std::shared_ptr<LogImplementation> SwapImplementation(
+    std::shared_ptr<LogImplementation> implementation);
 
 // Must be called at least once per process/load before anything else is
 // called. This function is safe to call multiple times from multiple
 // tasks/threads.
 void Init();
 
-// Forces all of the state that is usually lazily created when first needed to
-// be created when called. Cleanup() will delete it.
-void Load();
+class CallbackLogImplementation : public HandleMessageLogImplementation {
+ public:
+  CallbackLogImplementation(
+      const ::std::function<void(const LogMessage &)> &callback)
+      : callback_(callback) {}
+
+ private:
+  void HandleMessage(const LogMessage &message) override { callback_(message); }
+
+  ::std::function<void(const LogMessage &)> callback_;
+};
 
 // Resets all information in this task/thread to its initial state.
 // NOTE: This is not the opposite of Init(). The state that this deletes is
@@ -150,17 +151,25 @@
 void Cleanup();
 
 void RegisterCallbackImplementation(
-    const ::std::function<void(const LogMessage &)> &callback,
-    bool update_global = true);
+    const ::std::function<void(const LogMessage &)> &callback);
 
 class ScopedLogRestorer {
  public:
-  ScopedLogRestorer() { prev_impl_ = GetImplementation(); }
+  ScopedLogRestorer() = default;
 
-  ~ScopedLogRestorer() { SetImplementation(prev_impl_); }
+  ~ScopedLogRestorer() {
+    if (prev_impl_) {
+      SetImplementation(std::move(prev_impl_));
+    }
+    Cleanup();
+  }
+
+  void Swap(std::shared_ptr<LogImplementation> new_impl) {
+    prev_impl_ = SwapImplementation(std::move(new_impl));
+  }
 
  private:
-  LogImplementation *prev_impl_;
+  std::shared_ptr<LogImplementation> prev_impl_;
 };
 
 // This is where all of the code that is only used by actual LogImplementations
diff --git a/aos/logging/implementations_test.cc b/aos/logging/implementations_test.cc
index 272214a..84a1bed 100644
--- a/aos/logging/implementations_test.cc
+++ b/aos/logging/implementations_test.cc
@@ -41,11 +41,8 @@
  public:
   const LogMessage &message() { return message_; }
   bool used() { return used_; }
-  void reset_used() { used_ = false; }
 
-  TestLogImplementation() : used_(false) {}
-
-  bool used_;
+  bool used_ = false;
 };
 class LoggingTest : public ::testing::Test {
  protected:
@@ -90,21 +87,19 @@
 
  private:
   void SetUp() override {
-    static bool first = true;
-    if (first) {
-      first = false;
-
-      Init();
-      SetImplementation(log_implementation = new TestLogImplementation());
-    }
-
-    log_implementation->reset_used();
+    log_implementation = std::make_shared<TestLogImplementation>();
+    SetImplementation(log_implementation);
   }
-  void TearDown() override { Cleanup(); }
+  void TearDown() override {
+    SetImplementation(nullptr);
+    Cleanup();
+    internal::Context::DeleteNow();
+    CHECK_EQ(log_implementation.use_count(), 1);
+    log_implementation.reset();
+  }
 
-  static TestLogImplementation *log_implementation;
+  std::shared_ptr<TestLogImplementation> log_implementation;
 };
-TestLogImplementation *LoggingTest::log_implementation(NULL);
 typedef LoggingTest LoggingDeathTest;
 
 // Tests both basic logging functionality and that the test setup works
@@ -245,16 +240,18 @@
 }
 
 TEST(ScopedLogRestorerTest, RestoreTest) {
-  LogImplementation *curr_impl = GetImplementation();
+  SetImplementation(std::make_shared<StreamLogImplementation>(stdout));
+  LogImplementation *curr_impl = GetImplementation().get();
 
   {
     ScopedLogRestorer log_restorer;
 
-    logging::RegisterCallbackImplementation([] (const LogMessage&) {});
-    ASSERT_NE(curr_impl, GetImplementation());
+    log_restorer.Swap(
+        std::make_shared<CallbackLogImplementation>([](const LogMessage &) {}));
+    ASSERT_NE(curr_impl, GetImplementation().get());
   }
 
-  ASSERT_EQ(curr_impl, GetImplementation());
+  ASSERT_EQ(curr_impl, GetImplementation().get());
 }
 
 }  // namespace testing
diff --git a/aos/logging/interface.cc b/aos/logging/interface.cc
index e6e9203..6a201aa 100644
--- a/aos/logging/interface.cc
+++ b/aos/logging/interface.cc
@@ -35,11 +35,12 @@
     ::std::function<void(LogImplementation *)> function) {
   Context *context = Context::Get();
 
-  LogImplementation *const implementation = context->implementation;
+  const std::shared_ptr<LogImplementation> implementation =
+      context->implementation;
   if (implementation == NULL) {
     Die("no logging implementation to use\n");
   }
-  function(implementation);
+  function(implementation.get());
 }
 
 }  // namespace internal