common: add support for thread local sink overrides

source/common/common/logger.cc

@@ -27,8 +27,15 @@ void SinkDelegate::logWithStableName(absl::string_view, absl::string_view, absl:
 
 SinkDelegate::~SinkDelegate() {
   // The previous delegate should have never been set or should have been reset by now via
-  // restoreDelegate();
+  // restoreDelegate()/restoreTlsDelegate();
   assert(previous_delegate_ == nullptr);
+  assert(previous_tls_delegate_ == nullptr);
+}
+
+void SinkDelegate::setTlsDelegate() {
+  assert(previous_tls_delegate_ == nullptr);
+  previous_tls_delegate_ = log_sink_->tlsDelegate();
+  log_sink_->setTlsDelegate(this);
 }
 
 void SinkDelegate::setDelegate() {
@@ -38,6 +45,13 @@ void SinkDelegate::setDelegate() {
   log_sink_->setDelegate(this);
 }
 
+void SinkDelegate::restoreTlsDelegate() {
+  // Ensures stacked allocation of delegates.
+  assert(log_sink_->tlsDelegate() == this);
+  log_sink_->setTlsDelegate(previous_tls_delegate_);
+  previous_tls_delegate_ = nullptr;
+}
+
 void SinkDelegate::restoreDelegate() {
   // Ensures stacked allocation of delegates.
   assert(log_sink_->delegate() == this);
@@ -80,18 +94,27 @@ void DelegatingLogSink::log(const spdlog::details::log_msg& msg) {
   }
   lock.Release();
 
+  auto log_to_sink = [this, msg_view](SinkDelegate& sink) {
+    if (should_escape_) {
+      sink.log(escapeLogLine(msg_view));
+    } else {
+      sink.log(msg_view);
+    }
+  };
+  auto* tls_sink = tlsDelegate();
+  if (tls_sink != nullptr) {
+    log_to_sink(*tls_sink);
+    return;
+  }
+
   // Hold the sink mutex while performing the actual logging. This prevents the sink from being
   // swapped during an individual log event.
   // TODO(mattklein123): In production this lock will never be contended. In practice, thread
   // protection is really only needed in tests. It would be nice to figure out a test-only
   // mechanism for this that does not require extra locking that we don't explicitly need in the
   // prod code.
   absl::ReaderMutexLock sink_lock(&sink_mutex_);
-  if (should_escape_) {
-    sink_->log(escapeLogLine(msg_view));
-  } else {
-    sink_->log(msg_view);
-  }
+  log_to_sink(*sink_);
 }
 
 std::string DelegatingLogSink::escapeLogLine(absl::string_view msg_view) {

source/common/common/logger.h

@@ -115,21 +115,30 @@ class SinkDelegate : NonCopyable {
   virtual void flush() PURE;
 
 protected:
-  // Swap the current log sink delegate for this one. This should be called by the derived class
-  // constructor immediately before returning. This is required to match restoreDelegate(),
-  // otherwise it's possible for the previous delegate to get set in the base class constructor,
-  // the derived class constructor throws, and cleanup becomes broken.
+  // Swap the current thread local log sink delegate for this one. This should be called by the
+  // derived class constructor immediately before returning. This is required to match
+  // restoreTlsDelegate(), otherwise it's possible for the previous delegate to get set in the base
+  // class constructor, the derived class constructor throws, and cleanup becomes broken.
+  void setTlsDelegate();
+
+  // Swap the current *global* log sink delegate for this one. This behaves as setTlsDelegate, but
+  // operates on the global log sink instead of the thread local one.
   void setDelegate();
 
-  // Swap the current log sink (this) for the previous one. This should be called by the derived
-  // class destructor in the body. This is critical as otherwise it's possible for a log message
-  // to get routed to a partially destructed sink.
+  // Swap the current thread local log sink (this) for the previous one. This should be called by
+  // the derived class destructor in the body. This is critical as otherwise it's possible for a log
+  // message to get routed to a partially destructed sink.
+  void restoreTlsDelegate();
+
+  // Swap the current *global* log sink delegate for the previous one. This behaves as
+  // restoreTlsDelegate, but operates on the global sink instead of the thread local one.
   void restoreDelegate();
 
   SinkDelegate* previousDelegate() { return previous_delegate_; }
 
 private:
   SinkDelegate* previous_delegate_{nullptr};
+  SinkDelegate* previous_tls_delegate_{nullptr};
   DelegatingLogSinkSharedPtr log_sink_;
 };
 
@@ -165,12 +174,22 @@ class DelegatingLogSink : public spdlog::sinks::sink {
   template <class... Args>
   void logWithStableName(absl::string_view stable_name, absl::string_view level,
                          absl::string_view component, Args... msg) {
+    auto tls_sink = tlsDelegate();
+    if (tls_sink != nullptr) {
+      tls_sink->logWithStableName(stable_name, level, component, fmt::format(msg...));
+      return;
+    }
     absl::ReaderMutexLock sink_lock(&sink_mutex_);
     sink_->logWithStableName(stable_name, level, component, fmt::format(msg...));
   }
   // spdlog::sinks::sink
   void log(const spdlog::details::log_msg& msg) override;
   void flush() override {
+    auto* tls_sink = tlsDelegate();
+    if (tls_sink != nullptr) {
+      tls_sink->flush();
+      return;
+    }
     absl::ReaderMutexLock lock(&sink_mutex_);
     sink_->flush();
   }
@@ -221,6 +240,16 @@ class DelegatingLogSink : public spdlog::sinks::sink {
     return sink_;
   }
 
+  SinkDelegate** tlsSink() {
+    static thread_local SinkDelegate* tls_sink = nullptr;
+
+    return &tls_sink;
+  }
+
+  void setTlsDelegate(SinkDelegate* sink) { *tlsSink() = sink; }
+
+  SinkDelegate* tlsDelegate() { return *tlsSink(); }
+
   SinkDelegate* sink_ ABSL_GUARDED_BY(sink_mutex_){nullptr};
   absl::Mutex sink_mutex_;
   std::unique_ptr<StderrSinkDelegate> stderr_sink_; // Builtin sink to use as a last resort.

test/common/common/logger_test.cc

@@ -185,5 +185,49 @@ TEST_F(NamedLogTest, NamedLogsAreSentToSink) {
   ENVOY_LOG_EVENT_TO_LOGGER(Registry::getLog(Id::misc), debug, "misc_event", "log");
 }
 
+struct TlsLogSink : SinkDelegate {
+  TlsLogSink(DelegatingLogSinkSharedPtr log_sink) : SinkDelegate(log_sink) { setTlsDelegate(); }
+  ~TlsLogSink() override { restoreTlsDelegate(); }
+
+  MOCK_METHOD(void, log, (absl::string_view));
+  MOCK_METHOD(void, logWithStableName,
+              (absl::string_view, absl::string_view, absl::string_view, absl::string_view));
+  MOCK_METHOD(void, flush, ());
+};
+
+TEST(NamedLogtest, OverrideSink) {
+  NamedLogSink global_sink(Envoy::Logger::Registry::getSink());
+  testing::InSequence s;
+
+  {
+    TlsLogSink tls_sink(Envoy::Logger::Registry::getSink());
+
+    // Calls on the current thread goes to the TLS sink.
+    EXPECT_CALL(tls_sink, log(_));
+    ENVOY_LOG_MISC(info, "hello tls");
+
+    // Calls on other threads should use the global sink.
+    std::thread([&]() {
+      EXPECT_CALL(global_sink, log(_));
+      ENVOY_LOG_MISC(info, "hello global");
+    }).join();
+
+    // Sanity checking that we're still using the TLS sink.
+    EXPECT_CALL(tls_sink, log(_));
+    ENVOY_LOG_MISC(info, "hello tls");
+
+    // All the logging functions should be delegated to the TLS override.
+    EXPECT_CALL(tls_sink, flush());
+    Registry::getSink()->flush();
+
+    EXPECT_CALL(tls_sink, logWithStableName(_, _, _, _));
+    Registry::getSink()->logWithStableName("foo", "level", "bar", "msg");
+  }
+
+  // Now that the TLS sink is out of scope, log calls on this thread should use the global sink
+  // again.
+  EXPECT_CALL(global_sink, log(_));
+  ENVOY_LOG_MISC(info, "hello global 2");
+}
 } // namespace Logger
 } // namespace Envoy