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/dump_rtprio.cc b/aos/dump_rtprio.cc
index 5d9a0b4..5f4397f 100644
--- a/aos/dump_rtprio.cc
+++ b/aos/dump_rtprio.cc
@@ -7,17 +7,17 @@
 // exe,name,cpumask,policy,nice,priority,tid,pid,ppid,sid,cpu
 
 #include <sched.h>
-#include <stdlib.h>
-#include <stdio.h>
 #include <stdint.h>
-#include <sys/time.h>
+#include <stdio.h>
+#include <stdlib.h>
 #include <sys/resource.h>
+#include <sys/time.h>
 #include <unistd.h>
 
 #include <string>
 
-#include "aos/logging/logging.h"
 #include "aos/logging/implementations.h"
+#include "aos/logging/logging.h"
 #include "aos/time/time.h"
 
 namespace {
@@ -248,8 +248,6 @@
 
 int main() {
   ::aos::logging::Init();
-  ::aos::logging::SetImplementation(
-      new ::aos::logging::StreamLogImplementation(stdout));
 
   const int pid_max = find_pid_max();
   const cpu_set_t all_cpus = find_all_cpus();
diff --git a/aos/events/aos_logging.cc b/aos/events/aos_logging.cc
index d312925..5b6f72c 100644
--- a/aos/events/aos_logging.cc
+++ b/aos/events/aos_logging.cc
@@ -5,7 +5,7 @@
 void AosLogToFbs::Initialize(Sender<logging::LogMessageFbs> log_sender) {
   log_sender_ = std::move(log_sender);
   if (log_sender_) {
-    logging::RegisterCallbackImplementation(
+    implementation_ = std::make_shared<logging::CallbackLogImplementation>(
         [this](const logging::LogMessage &message_data) {
           aos::Sender<logging::LogMessageFbs>::Builder message =
               log_sender_.MakeBuilder();
@@ -23,8 +23,7 @@
           builder.add_name(name_str);
 
           message.Send(builder.Finish());
-        },
-        false);
+        });
   }
 }
 
diff --git a/aos/events/aos_logging.h b/aos/events/aos_logging.h
index b50569c..e99edb7 100644
--- a/aos/events/aos_logging.h
+++ b/aos/events/aos_logging.h
@@ -11,13 +11,14 @@
  public:
   AosLogToFbs() {}
 
-  // TODO(Tyler): Deregister logger on destruction to avoid memory leaks
-
   void Initialize(Sender<logging::LogMessageFbs> log_sender);
+  std::shared_ptr<logging::LogImplementation> implementation() const {
+    return implementation_;
+  }
 
  private:
   Sender<logging::LogMessageFbs> log_sender_;
-  logging::ScopedLogRestorer prev_logger_;
+  std::shared_ptr<logging::LogImplementation> implementation_;
 };
 
 }  // namespace aos
diff --git a/aos/events/event_loop.h b/aos/events/event_loop.h
index de397e2..1e7d25b 100644
--- a/aos/events/event_loop.h
+++ b/aos/events/event_loop.h
@@ -358,7 +358,7 @@
   // Returns the name of the underlying queue.
   const Channel *channel() const { return sender_->channel(); }
 
-  operator bool() { return sender_ ? true : false; }
+  operator bool() const { return sender_ ? true : false; }
 
   // Returns the time_points that the last message was sent at.
   aos::monotonic_clock::time_point monotonic_sent_time() const {
diff --git a/aos/events/shm_event_loop.cc b/aos/events/shm_event_loop.cc
index 29f01a4..406325d 100644
--- a/aos/events/shm_event_loop.cc
+++ b/aos/events/shm_event_loop.cc
@@ -964,9 +964,11 @@
   ReserveEvents();
 
   {
+    logging::ScopedLogRestorer prev_logger;
     AosLogToFbs aos_logger;
     if (!skip_logger_) {
       aos_logger.Initialize(MakeSender<logging::LogMessageFbs>("/aos"));
+      prev_logger.Swap(aos_logger.implementation());
     }
 
     aos::SetCurrentThreadName(name_.substr(0, 16));
diff --git a/aos/events/simulated_event_loop.cc b/aos/events/simulated_event_loop.cc
index a7301bf..c429eef 100644
--- a/aos/events/simulated_event_loop.cc
+++ b/aos/events/simulated_event_loop.cc
@@ -570,9 +570,8 @@
   void Setup() {
     MaybeScheduleTimingReports();
     if (!skip_logger_) {
-      logging::ScopedLogRestorer prev_logger;
       log_sender_.Initialize(MakeSender<logging::LogMessageFbs>("/aos"));
-      log_impl_ = logging::GetImplementation();
+      log_impl_ = log_sender_.implementation();
     }
   }
 
@@ -614,7 +613,7 @@
   const pid_t tid_;
 
   AosLogToFbs log_sender_;
-  logging::LogImplementation *log_impl_ = nullptr;
+  std::shared_ptr<logging::LogImplementation> log_impl_ = nullptr;
 };
 
 void SimulatedEventLoopFactory::set_send_delay(
@@ -720,8 +719,8 @@
   const monotonic_clock::time_point monotonic_now =
       simulated_event_loop_->monotonic_now();
   logging::ScopedLogRestorer prev_logger;
-  if (simulated_event_loop_->log_impl_ != nullptr) {
-    logging::SetImplementation(simulated_event_loop_->log_impl_);
+  if (simulated_event_loop_->log_impl_) {
+    prev_logger.Swap(simulated_event_loop_->log_impl_);
   }
   Context context = msgs_.front()->context;
 
@@ -837,8 +836,8 @@
   const ::aos::monotonic_clock::time_point monotonic_now =
       simulated_event_loop_->monotonic_now();
   logging::ScopedLogRestorer prev_logger;
-  if (simulated_event_loop_->log_impl_ != nullptr) {
-    logging::SetImplementation(simulated_event_loop_->log_impl_);
+  if (simulated_event_loop_->log_impl_) {
+    prev_logger.Swap(simulated_event_loop_->log_impl_);
   }
   if (token_ != scheduler_->InvalidToken()) {
     scheduler_->Deschedule(token_);
@@ -889,8 +888,8 @@
   monotonic_clock::time_point monotonic_now =
       simulated_event_loop_->monotonic_now();
   logging::ScopedLogRestorer prev_logger;
-  if (simulated_event_loop_->log_impl_ != nullptr) {
-    logging::SetImplementation(simulated_event_loop_->log_impl_);
+  if (simulated_event_loop_->log_impl_) {
+    prev_logger.Swap(simulated_event_loop_->log_impl_);
   }
   Call(
       [monotonic_now]() { return monotonic_now; },
diff --git a/aos/ipc_lib/eventfd_latency.cc b/aos/ipc_lib/eventfd_latency.cc
index a281a4c..9229f1a 100644
--- a/aos/ipc_lib/eventfd_latency.cc
+++ b/aos/ipc_lib/eventfd_latency.cc
@@ -161,8 +161,6 @@
   ::gflags::ParseCommandLineFlags(&argc, &argv, true);
 
   ::aos::logging::Init();
-  ::aos::logging::SetImplementation(
-      new ::aos::logging::StreamLogImplementation(stdout));
 
   return ::aos::Main(argc, argv);
 }
diff --git a/aos/ipc_lib/futex_latency.cc b/aos/ipc_lib/futex_latency.cc
index 60e7a70..59fa860 100644
--- a/aos/ipc_lib/futex_latency.cc
+++ b/aos/ipc_lib/futex_latency.cc
@@ -165,8 +165,6 @@
   ::gflags::ParseCommandLineFlags(&argc, &argv, true);
 
   ::aos::logging::Init();
-  ::aos::logging::SetImplementation(
-      new ::aos::logging::StreamLogImplementation(stdout));
 
   return ::aos::Main(argc, argv);
 }
diff --git a/aos/ipc_lib/ipc_comparison.cc b/aos/ipc_lib/ipc_comparison.cc
index 592327e..9191eae 100644
--- a/aos/ipc_lib/ipc_comparison.cc
+++ b/aos/ipc_lib/ipc_comparison.cc
@@ -902,8 +902,7 @@
   ::gflags::ParseCommandLineFlags(&argc, &argv, true);
 
   ::aos::InitNRT();
-  ::aos::logging::SetImplementation(
-      new ::aos::logging::StreamLogImplementation(stdout));
+  ::aos::logging::Init();
 
   return ::aos::Main(argc, argv);
 }
diff --git a/aos/ipc_lib/named_pipe_latency.cc b/aos/ipc_lib/named_pipe_latency.cc
index 6683f59..c3f5c5c 100644
--- a/aos/ipc_lib/named_pipe_latency.cc
+++ b/aos/ipc_lib/named_pipe_latency.cc
@@ -19,7 +19,8 @@
 // pipe.
 
 DEFINE_bool(sender, true, "If true, send signals to the other process.");
-DEFINE_string(fifo, "/dev/shm/aos/named_pipe_latency", "FIFO to use for the test.");
+DEFINE_string(fifo, "/dev/shm/aos/named_pipe_latency",
+              "FIFO to use for the test.");
 DEFINE_int32(seconds, 10, "Duration of the test to run");
 DEFINE_int32(
     latency_threshold, 1000,
@@ -36,7 +37,8 @@
 namespace aos {
 
 void SenderThread() {
-  int pipefd = open(FLAGS_fifo.c_str(), FD_CLOEXEC | O_NONBLOCK | O_WRONLY | O_NOATIME);
+  int pipefd =
+      open(FLAGS_fifo.c_str(), FD_CLOEXEC | O_NONBLOCK | O_WRONLY | O_NOATIME);
   const monotonic_clock::time_point end_time =
       monotonic_clock::now() + chrono::seconds(FLAGS_seconds);
   // Standard mersenne_twister_engine seeded with 0
@@ -75,7 +77,8 @@
 }
 
 void ReceiverThread() {
-  int pipefd = open(FLAGS_fifo.c_str(), O_CLOEXEC | O_NONBLOCK | O_RDONLY | O_NOATIME);
+  int pipefd =
+      open(FLAGS_fifo.c_str(), O_CLOEXEC | O_NONBLOCK | O_RDONLY | O_NOATIME);
   Tracing t;
   t.Start();
 
@@ -151,8 +154,7 @@
                 FLAGS_timer_priority);
   });
 
-  ::std::thread st(
-      []() { SenderThread(); });
+  ::std::thread st([]() { SenderThread(); });
 
   ReceiverThread();
   st.join();
@@ -167,8 +169,6 @@
   ::gflags::ParseCommandLineFlags(&argc, &argv, true);
 
   ::aos::logging::Init();
-  ::aos::logging::SetImplementation(
-      new ::aos::logging::StreamLogImplementation(stdout));
 
   return ::aos::Main(argc, argv);
 }
diff --git a/aos/ipc_lib/signal_stress.cc b/aos/ipc_lib/signal_stress.cc
index ffbfb06..ea9537f 100644
--- a/aos/ipc_lib/signal_stress.cc
+++ b/aos/ipc_lib/signal_stress.cc
@@ -192,8 +192,6 @@
   ::gflags::ParseCommandLineFlags(&argc, &argv, true);
 
   ::aos::logging::Init();
-  ::aos::logging::SetImplementation(
-      new ::aos::logging::StreamLogImplementation(stdout));
 
   return ::aos::Main(argc, argv);
 }
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
diff --git a/aos/mutex/mutex_test.cc b/aos/mutex/mutex_test.cc
index bd47ae0..07de9dc 100644
--- a/aos/mutex/mutex_test.cc
+++ b/aos/mutex/mutex_test.cc
@@ -71,7 +71,8 @@
   test_mutex_.Unlock();
   EXPECT_DEATH(
       {
-        logging::SetImplementation(new util::DeathTestLogImplementation());
+        logging::SetImplementation(
+            std::make_shared<util::DeathTestLogImplementation>());
         test_mutex_.Unlock();
       },
       ".*multiple unlock.*");
@@ -82,7 +83,8 @@
   logging::Init();
   EXPECT_DEATH(
       {
-        logging::SetImplementation(new util::DeathTestLogImplementation());
+        logging::SetImplementation(
+            std::make_shared<util::DeathTestLogImplementation>());
         test_mutex_.Unlock();
       },
       ".*multiple unlock.*");
@@ -92,7 +94,8 @@
 TEST_F(MutexDeathTest, RepeatLock) {
   EXPECT_DEATH(
       {
-        logging::SetImplementation(new util::DeathTestLogImplementation());
+        logging::SetImplementation(
+            std::make_shared<util::DeathTestLogImplementation>());
         ASSERT_FALSE(test_mutex_.Lock());
         ASSERT_FALSE(test_mutex_.Lock());
       },
@@ -107,9 +110,7 @@
       static_cast<Mutex *>(shm_malloc_aligned(sizeof(Mutex), alignof(Mutex)));
   new (mutex) Mutex();
 
-  std::thread thread([&]() {
-    ASSERT_FALSE(mutex->Lock());
-  });
+  std::thread thread([&]() { ASSERT_FALSE(mutex->Lock()); });
   thread.join();
   EXPECT_TRUE(mutex->Lock());
 
@@ -124,9 +125,7 @@
       static_cast<Mutex *>(shm_malloc_aligned(sizeof(Mutex), alignof(Mutex)));
   new (mutex) Mutex();
 
-  std::thread thread([&]() {
-    ASSERT_FALSE(mutex->Lock());
-  });
+  std::thread thread([&]() { ASSERT_FALSE(mutex->Lock()); });
   thread.join();
   EXPECT_EQ(Mutex::State::kOwnerDied, mutex->TryLock());
 
@@ -245,13 +244,12 @@
       static_cast<Mutex *>(shm_malloc_aligned(sizeof(Mutex), alignof(Mutex)));
   new (mutex) Mutex();
 
-  std::thread thread([&]() {
-    ASSERT_FALSE(mutex->Lock());
-  });
+  std::thread thread([&]() { ASSERT_FALSE(mutex->Lock()); });
   thread.join();
   EXPECT_DEATH(
       {
-        logging::SetImplementation(new util::DeathTestLogImplementation());
+        logging::SetImplementation(
+            std::make_shared<util::DeathTestLogImplementation>());
         MutexLocker locker(mutex);
       },
       ".*previous owner of mutex [^ ]+ died.*");
@@ -313,9 +311,7 @@
       static_cast<Mutex *>(shm_malloc_aligned(sizeof(Mutex), alignof(Mutex)));
   new (mutex) Mutex();
 
-  std::thread thread([&]() {
-    ASSERT_FALSE(mutex->Lock());
-  });
+  std::thread thread([&]() { ASSERT_FALSE(mutex->Lock()); });
   thread.join();
   {
     aos::IPCMutexLocker locker(mutex);
diff --git a/aos/stl_mutex/BUILD b/aos/stl_mutex/BUILD
index 06599e3..51229d4 100644
--- a/aos/stl_mutex/BUILD
+++ b/aos/stl_mutex/BUILD
@@ -7,7 +7,7 @@
     ],
     deps = [
         "//aos/ipc_lib:aos_sync",
-        "//aos/logging",
+        "@com_github_google_glog//:glog",
     ],
 )
 
@@ -20,6 +20,5 @@
         ":stl_mutex",
         "//aos:die",
         "//aos/testing:googletest",
-        "//aos/testing:test_logging",
     ],
 )
diff --git a/aos/stl_mutex/stl_mutex.h b/aos/stl_mutex/stl_mutex.h
index 4e14557..86a5988 100644
--- a/aos/stl_mutex/stl_mutex.h
+++ b/aos/stl_mutex/stl_mutex.h
@@ -4,9 +4,8 @@
 #include <mutex>
 
 #include "aos/ipc_lib/aos_sync.h"
-#include "aos/logging/logging.h"
-#include "aos/type_traits/type_traits.h"
 #include "aos/macros.h"
+#include "glog/logging.h"
 
 namespace aos {
 
@@ -31,7 +30,7 @@
         owner_died_ = true;
         break;
       default:
-        AOS_LOG(FATAL, "mutex_grab(%p) failed with %d\n", &native_handle_, ret);
+        LOG(FATAL) << "mutex_grab(" << &native_handle_ << ") failed: " << ret;
     }
   }
 
@@ -46,13 +45,13 @@
       case 4:
         return false;
       default:
-        AOS_LOG(FATAL, "mutex_trylock(%p) failed with %d\n", &native_handle_,
-                ret);
+        LOG(FATAL) << "mutex_trylock(" << &native_handle_
+                   << ") failed: " << ret;
     }
   }
 
   void unlock() {
-    AOS_CHECK(!owner_died_);
+    CHECK(!owner_died_);
     mutex_unlock(&native_handle_);
   }
 
@@ -84,20 +83,20 @@
 
   void lock() {
     if (mutex_islocked(mutex_.native_handle())) {
-      AOS_CHECK(!owner_died());
+      CHECK(!owner_died());
       ++recursive_locks_;
     } else {
       mutex_.lock();
       if (mutex_.owner_died()) {
         recursive_locks_ = 0;
       } else {
-        AOS_CHECK_EQ(0, recursive_locks_);
+        CHECK_EQ(0, recursive_locks_);
       }
     }
   }
   bool try_lock() {
     if (mutex_islocked(mutex_.native_handle())) {
-      AOS_CHECK(!owner_died());
+      CHECK(!owner_died());
       ++recursive_locks_;
       return true;
     } else {
@@ -105,7 +104,7 @@
         if (mutex_.owner_died()) {
           recursive_locks_ = 0;
         } else {
-          AOS_CHECK_EQ(0, recursive_locks_);
+          CHECK_EQ(0, recursive_locks_);
         }
         return true;
       } else {
diff --git a/aos/stl_mutex/stl_mutex_test.cc b/aos/stl_mutex/stl_mutex_test.cc
index 0fd052a..5e88c56 100644
--- a/aos/stl_mutex/stl_mutex_test.cc
+++ b/aos/stl_mutex/stl_mutex_test.cc
@@ -2,7 +2,6 @@
 
 #include "gtest/gtest.h"
 
-#include "aos/testing/test_logging.h"
 #include "aos/die.h"
 
 namespace aos {
@@ -10,10 +9,7 @@
 
 class StlMutexDeathTest : public ::testing::Test {
  protected:
-  void SetUp() override {
-    ::aos::testing::EnableTestLogging();
-    SetDieTestMode(true);
-  }
+  void SetUp() override { SetDieTestMode(true); }
 };
 
 typedef StlMutexDeathTest StlRecursiveMutexDeathTest;
@@ -58,7 +54,7 @@
   mutex.lock();
   ASSERT_TRUE(mutex.try_lock());
   mutex.unlock();
-	mutex.unlock();
+  mutex.unlock();
 }
 
 // Tests that unlocking an unlocked recursive mutex fails.
diff --git a/aos/testing/test_logging.cc b/aos/testing/test_logging.cc
index 3d97c05..db6246c 100644
--- a/aos/testing/test_logging.cc
+++ b/aos/testing/test_logging.cc
@@ -6,9 +6,9 @@
 
 #include "gtest/gtest.h"
 
+#include "absl/base/call_once.h"
 #include "aos/logging/implementations.h"
 #include "aos/mutex/mutex.h"
-#include "absl/base/call_once.h"
 
 using ::aos::logging::LogMessage;
 
@@ -38,9 +38,9 @@
 
   // This class has to be a singleton so that everybody can get access to the
   // same instance to read out the messages etc.
-  static TestLogImplementation *GetInstance() {
-    static absl::once_flag once;
-    absl::call_once(once, CreateInstance);
+  static std::shared_ptr<TestLogImplementation> GetInstance() {
+    static std::shared_ptr<TestLogImplementation> instance =
+        std::make_unique<TestLogImplementation>();
     return instance;
   }
 
@@ -72,19 +72,14 @@
 
   void PrintMessagesAsTheyComeIn() { print_as_messages_come_in_ = true; }
 
- private:
-  static TestLogImplementation *instance;
-  TestLogImplementation() {}
+  // Don't call these from outside this class.
   ~TestLogImplementation() {
     if (output_file_ != stdout) {
       fclose(output_file_);
     }
   }
 
-    static void CreateInstance() {
-      instance = new TestLogImplementation();
-  }
-
+ private:
   virtual void HandleMessage(const LogMessage &message) override {
     ::aos::MutexLocker locker(&messages_mutex_);
     if (message.level == FATAL || print_as_messages_come_in_) {
@@ -106,8 +101,6 @@
   static thread_local ::aos::monotonic_clock::time_point monotonic_now_;
 };
 
-TestLogImplementation *TestLogImplementation::instance;
-
 thread_local bool TestLogImplementation::mock_time_ = false;
 thread_local ::aos::monotonic_clock::time_point
     TestLogImplementation::monotonic_now_ = ::aos::monotonic_clock::min_time;
@@ -123,7 +116,7 @@
     }
   }
 
-  virtual void OnTestPartResult( const ::testing::TestPartResult &result) {
+  virtual void OnTestPartResult(const ::testing::TestPartResult &result) {
     if (result.failed()) {
       const char *failure_type = "unknown";
       switch (result.type()) {
@@ -136,11 +129,8 @@
         case ::testing::TestPartResult::Type::kSuccess:
           break;
       }
-      log_do(ERROR, "%s: %d: gtest %s failure\n%s\n",
-             result.file_name(),
-             result.line_number(),
-             failure_type,
-             result.message());
+      log_do(ERROR, "%s: %d: gtest %s failure\n%s\n", result.file_name(),
+             result.line_number(), failure_type, result.message());
     }
   }
 };
@@ -163,7 +153,7 @@
   absl::call_once(enable_test_logging_once, DoEnableTestLogging);
 }
 
-void SetLogFileName(const char* filename) {
+void SetLogFileName(const char *filename) {
   TestLogImplementation::GetInstance()->SetOutputFile(filename);
 }
 
@@ -174,9 +164,7 @@
 void MockTime(::aos::monotonic_clock::time_point monotonic_now) {
   TestLogImplementation::GetInstance()->MockTime(monotonic_now);
 }
-void UnMockTime() {
-  TestLogImplementation::GetInstance()->UnMockTime();
-}
+void UnMockTime() { TestLogImplementation::GetInstance()->UnMockTime(); }
 
 }  // namespace testing
 }  // namespace aos
diff --git a/aos/transaction/transaction_test.cc b/aos/transaction/transaction_test.cc
index 5d1fe90..52b297e 100644
--- a/aos/transaction/transaction_test.cc
+++ b/aos/transaction/transaction_test.cc
@@ -21,9 +21,7 @@
       i_ = i;
       test_ = test;
     }
-    void DoWork() {
-      test_->invoked_works()->push_back(i_);
-    }
+    void DoWork() { test_->invoked_works()->push_back(i_); }
 
     int i() const { return i_; }
 
@@ -37,9 +35,7 @@
   WorkStack<TestWork, 20> *work_stack() { return &work_stack_; }
 
   // Creates a TestWork with index i and adds it to work_stack().
-  void CreateWork(int i) {
-    work_stack_.AddWork(this, i);
-  }
+  void CreateWork(int i) { work_stack_.AddWork(this, i); }
 
  private:
   ::std::vector<int> created_works_, invoked_works_;
@@ -93,7 +89,8 @@
   logging::Init();
   EXPECT_DEATH(
       {
-        logging::SetImplementation(new util::DeathTestLogImplementation());
+        logging::SetImplementation(
+            std::make_shared<util::DeathTestLogImplementation>());
         for (int i = 0; i < 1000; ++i) {
           CreateWork(i);
         }
diff --git a/aos/vision/debug/debug_framework.cc b/aos/vision/debug/debug_framework.cc
index 47e0856..0b0d2c7 100644
--- a/aos/vision/debug/debug_framework.cc
+++ b/aos/vision/debug/debug_framework.cc
@@ -62,9 +62,8 @@
     }
 
     // Pass along the set exposure so that users can acceess it.
-    filter->InstallSetExposure([this](uint32_t abs_exp) {
-        this->SetExposure(abs_exp);
-    });
+    filter->InstallSetExposure(
+        [this](uint32_t abs_exp) { this->SetExposure(abs_exp); });
   }
 
   // This the first stage in the pipeline that takes
@@ -151,8 +150,6 @@
 void DebugFrameworkMain(int argc, char **argv, FilterHarness *filter,
                         CameraParams camera_params) {
   ::aos::logging::Init();
-  ::aos::logging::SetImplementation(
-      new ::aos::logging::StreamLogImplementation(stdout));
 
   gtk_init(&argc, &argv);
 
diff --git a/aos/vision/tools/camera_primer.cc b/aos/vision/tools/camera_primer.cc
index e1d5fc5..027f9dc 100644
--- a/aos/vision/tools/camera_primer.cc
+++ b/aos/vision/tools/camera_primer.cc
@@ -27,8 +27,6 @@
 // target_sender
 int main(int argc, char **argv) {
   ::aos::logging::Init();
-  ::aos::logging::SetImplementation(
-      new ::aos::logging::StreamLogImplementation(stdout));
 
   aos::vision::CameraParams params;
 
diff --git a/aos/vision/tools/jpeg_vision_test.cc b/aos/vision/tools/jpeg_vision_test.cc
index 4ba290c..738713b 100644
--- a/aos/vision/tools/jpeg_vision_test.cc
+++ b/aos/vision/tools/jpeg_vision_test.cc
@@ -107,13 +107,11 @@
   int dx = 0;
   int dy = 0;
 };
-}  // namespace aos
 }  // namespace vision
+}  // namespace aos
 
 int main(int argc, char *argv[]) {
   ::aos::logging::Init();
-  ::aos::logging::SetImplementation(
-      new ::aos::logging::StreamLogImplementation(stdout));
   aos::events::EpollLoop loop;
   gtk_init(&argc, &argv);