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/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