Add support for adding context string callback.

Summary: For program running in different general user contexts, e.g. services, multithreading, it's very useful to extend log prefixes to contain the domain specific context strings in debugging. This diff adds one interface function to push one callback to the logger. The callbacks will be executed when the log message data is constructed, and the strings returned by callbacks will be appended to the tail of log entry prefixes.

Reviewed By: simpkins

Differential Revision: D21291284

fbshipit-source-id: e141bf1c907f6730ea741d0dc06975add93d489c

folly/logging/CustomLogFormatter.cpp

@@ -83,7 +83,8 @@ struct FormatKeys {
  *
  * TODO: Support including thread names and thread context info.
  */
-constexpr std::array<FormatKeys, 11> formatKeys{{
+constexpr std::array<FormatKeys, 12> formatKeys{{
+    FormatKeys(/* key */ "CTX", /*    argIndex  */ 11),
     FormatKeys(/* key */ "D", /*      argIndex  */ 2, /* width */ 2),
     FormatKeys(/* key */ "FILE", /*   argIndex  */ 8),
     FormatKeys(/* key */ "FUN", /*    argIndex  */ 9),
@@ -263,6 +264,7 @@ std::string CustomLogFormatter::formatMessage(
         basename,
         message.getFunctionName(),
         message.getLineNumber(),
+        message.getContextString(),
         // NOTE: THE FOLLOWING ARGUMENTS ALWAYS NEED TO BE THE LAST 3:
         message.getMessage(),
         // If colored logs are enabled, the singleLineLogFormat_ will contain
@@ -286,7 +288,8 @@ std::string CustomLogFormatter::formatMessage(
         message.getThreadID(),
         basename,
         message.getFunctionName(),
-        message.getLineNumber());
+        message.getLineNumber(),
+        message.getContextString());
 
     // Estimate header length. If this still isn't long enough the string will
     // grow as necessary, so the code will still be correct, but just slightly

folly/logging/GlogStyleFormatter.cpp

@@ -61,7 +61,7 @@ std::string GlogStyleFormatter::formatMessage(
 
   auto basename = message.getFileBaseName();
   auto headerFormatter = folly::format(
-      "{}{:02d}{:02d} {:02d}:{:02d}:{:02d}.{:06d} {:5d} {}:{}] ",
+      "{}{:02d}{:02d} {:02d}:{:02d}:{:02d}.{:06d} {:5d} {}:{}{}] ",
       getGlogLevelName(message.getLevel())[0],
       ltime.tm_mon + 1,
       ltime.tm_mday,
@@ -71,7 +71,8 @@ std::string GlogStyleFormatter::formatMessage(
       usecs.count(),
       message.getThreadID(),
       basename,
-      message.getLineNumber());
+      message.getLineNumber(),
+      message.getContextString());
 
   // TODO: Support including thread names and thread context info.
 

folly/logging/LogMessage.cpp

@@ -16,10 +16,18 @@
 
 #include <folly/logging/LogMessage.h>
 
+#include <folly/logging/LogCategory.h>
+#include <folly/logging/LoggerDB.h>
 #include <folly/system/ThreadId.h>
 
 using std::chrono::system_clock;
 
+namespace {
+std::string getContextStringFromCategory(const folly::LogCategory* category) {
+  return category->getDB()->getContextString();
+}
+} // namespace
+
 namespace folly {
 
 LogMessage::LogMessage(
@@ -36,6 +44,7 @@ LogMessage::LogMessage(
       filename_{filename},
       lineNumber_{lineNumber},
       functionName_{functionName},
+      contextString_{getContextStringFromCategory(category_)},
       rawMessage_{std::move(msg)} {
   sanitizeMessage();
 }
@@ -55,6 +64,7 @@ LogMessage::LogMessage(
       filename_{filename},
       lineNumber_{lineNumber},
       functionName_{functionName},
+      contextString_{getContextStringFromCategory(category_)},
       rawMessage_{std::move(msg)} {
   sanitizeMessage();
 }

folly/logging/LogMessage.h

@@ -123,6 +123,10 @@ class LogMessage {
     return numNewlines_;
   }
 
+  const std::string& getContextString() const {
+    return contextString_;
+  }
+
  private:
   void sanitizeMessage();
 
@@ -155,6 +159,14 @@ class LogMessage {
    */
   size_t numNewlines_{0};
 
+  /**
+   * contextString_ contains user defined context information.
+   *
+   * This can be customized by adding new callback through
+   * addLogMessageContextCallback().
+   */
+  std::string contextString_;
+
   /**
    * rawMessage_ contains the original message.
    *

folly/logging/LoggerDB.cpp

@@ -599,6 +599,87 @@ LogCategory* LoggerDB::xlogInitCategory(
   return category;
 }
 
+class LoggerDB::ContextCallbackList::CallbacksObj {
+  using StorageBlock = std::array<ContextCallback, 16>;
+
+ public:
+  CallbacksObj() : end_(block_.begin()) {}
+
+  template <typename F>
+  void forEach(F&& f) const {
+    auto end = end_.load(std::memory_order_acquire);
+
+    for (auto it = block_.begin(); it != end; it = std::next(it)) {
+      f(*it);
+    }
+  }
+
+  /**
+   * Callback list is implemented as an unsynchronized array, so an atomic
+   * end flag is used for list readers to get a synchronized view of the
+   * list with concurrent writers, protecting the underlying array.
+   * There can be also race condition between list writers, because the end
+   * flag is firstly tested before written, which should be serialized with
+   * another global mutex to prevent TOCTOU bug.
+   */
+  void push(ContextCallback callback) {
+    auto end = end_.load(std::memory_order_relaxed);
+    if (end == block_.end()) {
+      folly::throw_exception(std::length_error(
+          "Exceeding limit for the number of pushed context callbacks."));
+    }
+    *end = std::move(callback);
+    end_.store(std::next(end), std::memory_order_release);
+  }
+
+ private:
+  StorageBlock block_;
+  std::atomic<StorageBlock::iterator> end_;
+};
+
+LoggerDB::ContextCallbackList::~ContextCallbackList() {
+  auto callback = callbacks_.load(std::memory_order_relaxed);
+  if (callback != nullptr) {
+    delete callback;
+  }
+}
+
+void LoggerDB::ContextCallbackList::addCallback(ContextCallback callback) {
+  std::lock_guard<std::mutex> g(writeMutex_);
+  auto callbacks = callbacks_.load(std::memory_order_relaxed);
+  if (!callbacks) {
+    callbacks = new CallbacksObj();
+    callbacks_.store(callbacks, std::memory_order_relaxed);
+  }
+  callbacks->push(std::move(callback));
+}
+
+std::string LoggerDB::ContextCallbackList::getContextString() const {
+  auto callbacks = callbacks_.load(std::memory_order_relaxed);
+  if (callbacks == nullptr) {
+    return {};
+  }
+
+  std::string ret;
+  callbacks->forEach([&](const auto& callback) {
+    ret += ' ';
+    try {
+      ret += callback();
+    } catch (const std::exception& e) {
+      folly::toAppend("[error:", folly::exceptionStr(e), "]", &ret);
+    };
+  });
+  return ret;
+}
+
+void LoggerDB::addContextCallback(ContextCallback callback) {
+  contextCallbacks_.addCallback(std::move(callback));
+}
+
+std::string LoggerDB::getContextString() const {
+  return contextCallbacks_.getContextString();
+}
+
 std::atomic<LoggerDB::InternalWarningHandler> LoggerDB::warningHandler_;
 
 void LoggerDB::internalWarningImpl(

folly/logging/LoggerDB.h

@@ -41,6 +41,8 @@ enum class LogLevel : uint32_t;
  * LoggerDB stores the set of LogCategory objects.
  */
 class LoggerDB {
+  using ContextCallback = folly::Function<std::string() const>;
+
  public:
   /**
    * Get the main LoggerDB singleton.
@@ -189,6 +191,23 @@ class LoggerDB {
    */
   explicit LoggerDB(TestConstructorArg);
 
+  /**
+   * Add a new context string callback to the list.
+   *
+   * The callbacks will be invoked during the construction of log messages,
+   * and returned strings will be appended in order to the tail of log
+   * log entry prefixes with space prepended to each item.
+   */
+  void addContextCallback(ContextCallback);
+
+  /**
+   * Return a context string to be appended after default log prefixes.
+   *
+   * The context string is cutomized through adding context callbacks to
+   * LoggerDB objects.
+   */
+  std::string getContextString() const;
+
   /**
    * internalWarning() is used to report a problem when something goes wrong
    * internally in the logging library.
@@ -241,6 +260,18 @@ class LoggerDB {
     HandlerMap handlers;
   };
 
+  class ContextCallbackList {
+   public:
+    void addCallback(ContextCallback);
+    std::string getContextString() const;
+    ~ContextCallbackList();
+
+   private:
+    class CallbacksObj;
+    std::atomic<CallbacksObj*> callbacks_{nullptr};
+    std::mutex writeMutex_;
+  };
+
   // Forbidden copy constructor and assignment operator
   LoggerDB(LoggerDB const&) = delete;
   LoggerDB& operator=(LoggerDB const&) = delete;
@@ -298,6 +329,13 @@ class LoggerDB {
    */
   folly::Synchronized<HandlerInfo> handlerInfo_;
 
+  /**
+   * Callbacks returning context strings.
+   *
+   * Exceptions from the callbacks are catched and reflected in corresponding
+   * position in log entries.
+   */
+  ContextCallbackList contextCallbacks_;
   static std::atomic<InternalWarningHandler> warningHandler_;
 };
 

folly/logging/test/LogMessageTest.cpp

@@ -37,10 +37,13 @@ using namespace folly;
     EXPECT_EQ(static_cast<int>(hasNewlines), checkMsg.containsNewlines());    \
     EXPECT_EQ(__FILE__, checkMsg.getFileName());                              \
     EXPECT_EQ(__LINE__, checkMsg.getLineNumber());                            \
+    EXPECT_EQ(" context string", checkMsg.getContextString());                \
   }
 
 TEST(LogMessage, sanitize) {
   LoggerDB db{LoggerDB::TESTING};
+  db.addContextCallback([]() { return "context"; });
+  db.addContextCallback([]() { return "string"; });
   Logger logger{&db, "test"};
   auto* category = logger.getCategory();