Easily specify a custom Formatter

Summary:
StandardLogHandlerFactory has a good amount of logic in it for handling
options. It would be nice to reuse it when specifying a custom Formatter.

This also exposes the WriterFactory from StreamHandlerFactory, so that a user
can re-use the code to create the default LogWriter when using a custom
Formatter.

Reviewed By: yfeldblum

Differential Revision: D16074437

fbshipit-source-id: d79947e2fe93b8920b83294438a88c4dd871986b

folly/logging/StandardLogHandlerFactory.cpp

@@ -106,7 +106,14 @@ std::shared_ptr<StandardLogHandler> StandardLogHandlerFactory::createHandler(
     throw std::invalid_argument(
         to<string>("unknown log formatter type \"", *formatterType, "\""));
   }
+  return createHandler(type, writerFactory, formatterFactory.get(), options);
+}
 
+std::shared_ptr<StandardLogHandler> StandardLogHandlerFactory::createHandler(
+    StringPiece type,
+    WriterFactory* writerFactory,
+    FormatterFactory* formatterFactory,
+    const Options& options) {
   Optional<LogLevel> logLevel;
   Optional<LogLevel> syncLevel;
 

folly/logging/StandardLogHandlerFactory.h

@@ -66,6 +66,12 @@ class StandardLogHandlerFactory {
       StringPiece type,
       WriterFactory* writerFactory,
       const Options& options);
+
+  static std::shared_ptr<StandardLogHandler> createHandler(
+      StringPiece type,
+      WriterFactory* writerFactory,
+      FormatterFactory* formatterFactory,
+      const Options& options);
 };
 
 } // namespace folly

folly/logging/StreamHandlerFactory.cpp

@@ -22,40 +22,32 @@
 
 namespace folly {
 
-class StreamHandlerFactory::WriterFactory
-    : public StandardLogHandlerFactory::WriterFactory {
- public:
-  bool processOption(StringPiece name, StringPiece value) override {
-    if (name == "stream") {
-      stream_ = value.str();
-      return true;
-    }
-    return fileWriterFactory_.processOption(name, value);
+bool StreamHandlerFactory::WriterFactory::processOption(
+    StringPiece name,
+    StringPiece value) {
+  if (name == "stream") {
+    stream_ = value.str();
+    return true;
   }
+  return fileWriterFactory_.processOption(name, value);
+}
 
-  std::shared_ptr<LogWriter> createWriter() override {
-    // Get the output file to use
-    File outputFile;
-    if (stream_.empty()) {
-      throw std::invalid_argument(
-          "no stream name specified for stream handler");
-    } else if (stream_ == "stderr") {
-      outputFile = File{STDERR_FILENO, /* ownsFd */ false};
-    } else if (stream_ == "stdout") {
-      outputFile = File{STDOUT_FILENO, /* ownsFd */ false};
-    } else {
-      throw std::invalid_argument(to<std::string>(
-          "unknown stream \"",
-          stream_,
-          "\": expected one of stdout or stderr"));
-    }
-
-    return fileWriterFactory_.createWriter(std::move(outputFile));
+std::shared_ptr<LogWriter> StreamHandlerFactory::WriterFactory::createWriter() {
+  // Get the output file to use
+  File outputFile;
+  if (stream_.empty()) {
+    throw std::invalid_argument("no stream name specified for stream handler");
+  } else if (stream_ == "stderr") {
+    outputFile = File{STDERR_FILENO, /* ownsFd */ false};
+  } else if (stream_ == "stdout") {
+    outputFile = File{STDOUT_FILENO, /* ownsFd */ false};
+  } else {
+    throw std::invalid_argument(to<std::string>(
+        "unknown stream \"", stream_, "\": expected one of stdout or stderr"));
   }
 
-  std::string stream_;
-  FileWriterFactory fileWriterFactory_;
-};
+  return fileWriterFactory_.createWriter(std::move(outputFile));
+}
 
 std::shared_ptr<LogHandler> StreamHandlerFactory::createHandler(
     const Options& options) {

folly/logging/StreamHandlerFactory.h

@@ -15,7 +15,9 @@
  */
 #pragma once
 
+#include <folly/logging/FileWriterFactory.h>
 #include <folly/logging/LogHandlerFactory.h>
+#include <folly/logging/StandardLogHandlerFactory.h>
 
 namespace folly {
 
@@ -37,8 +39,15 @@ class StreamHandlerFactory : public LogHandlerFactory {
 
   std::shared_ptr<LogHandler> createHandler(const Options& options) override;
 
- private:
-  class WriterFactory;
+  class WriterFactory : public StandardLogHandlerFactory::WriterFactory {
+   public:
+    bool processOption(StringPiece name, StringPiece value) override;
+    std::shared_ptr<LogWriter> createWriter() override;
+
+   private:
+    std::string stream_;
+    FileWriterFactory fileWriterFactory_;
+  };
 };
 
 } // namespace folly

folly/logging/test/FileHandlerFactoryTest.cpp

@@ -357,3 +357,47 @@ TEST(StreamHandlerFactory, errors) {
         "unknown option \"foo\"");
   }
 }
+
+TEST(StreamHandlerFactory, writerFactory) {
+  StreamHandlerFactory::WriterFactory factory;
+  factory.processOption("stream", "stderr");
+  {
+    auto writer = factory.createWriter();
+    checkAsyncWriter(
+        writer.get(), STDERR_FILENO, AsyncFileWriter::kDefaultMaxBufferSize);
+  }
+
+  // For duplicate option, prefer the latter
+  factory.processOption("stream", "stdout");
+  {
+    auto writer = factory.createWriter();
+    checkAsyncWriter(
+        writer.get(), STDOUT_FILENO, AsyncFileWriter::kDefaultMaxBufferSize);
+  }
+}
+
+TEST(StreamHandlerFactory, writerFactoryError) {
+  {
+    StreamHandlerFactory::WriterFactory factory;
+    factory.processOption("stream", "nope");
+    EXPECT_THROW_RE(
+        factory.createWriter(),
+        std::invalid_argument,
+        "^unknown stream \"nope\": expected one of stdout or stderr$");
+  }
+  {
+    StreamHandlerFactory::WriterFactory factory;
+    EXPECT_THROW_RE(
+        factory.createWriter(),
+        std::invalid_argument,
+        "^no stream name specified for stream handler$");
+  }
+  {
+    StreamHandlerFactory::WriterFactory factory;
+    factory.processOption("stream", "");
+    EXPECT_THROW_RE(
+        factory.createWriter(),
+        std::invalid_argument,
+        "^no stream name specified for stream handler$");
+  }
+}

folly/logging/test/StandardLogHandlerFactoryTest.cpp

@@ -16,28 +16,40 @@
 
 #include <folly/test/TestUtils.h>
 #include <queue>
+#include <vector>
 
 #include <folly/logging/Init.h>
 #include <folly/logging/LogConfigParser.h>
+#include <folly/logging/LogFormatter.h>
 #include <folly/logging/LogHandlerFactory.h>
 #include <folly/logging/LogWriter.h>
 #include <folly/logging/LoggerDB.h>
 #include <folly/logging/StandardLogHandler.h>
 #include <folly/logging/StandardLogHandlerFactory.h>
 #include <folly/logging/xlog.h>
+#include <folly/portability/GMock.h>
 
 namespace folly {
 
+class TestLogFormatter : public LogFormatter {
+ public:
+  std::string formatMessage(
+      const LogMessage& logMessage,
+      const LogCategory* /* handlerCategory */) override {
+    return "Test Formatter! " + logMessage.getMessage();
+  }
+};
+
 class TestLogWriter : public LogWriter {
  public:
-  void writeMessage(folly::StringPiece /* buffer */, uint32_t /* flags */ = 0)
+  void writeMessage(folly::StringPiece buffer, uint32_t /* flags */ = 0)
       override {
-    message_count++;
+    messages.push_back(buffer.str());
   }
 
   void flush() override {}
 
-  int message_count{0};
+  std::vector<std::string> messages;
 
   bool ttyOutput() const override {
     return false;
@@ -46,8 +58,10 @@ class TestLogWriter : public LogWriter {
 
 class TestHandlerFactory : public LogHandlerFactory {
  public:
-  explicit TestHandlerFactory(const std::shared_ptr<TestLogWriter> writer)
-      : writer_(writer) {}
+  explicit TestHandlerFactory(
+      const std::shared_ptr<TestLogWriter> writer,
+      const std::shared_ptr<TestLogFormatter> formatter = nullptr)
+      : writer_(writer), formatter_(formatter) {}
 
   StringPiece getType() const override {
     return "test";
@@ -55,12 +69,18 @@ class TestHandlerFactory : public LogHandlerFactory {
 
   std::shared_ptr<LogHandler> createHandler(const Options& options) override {
     TestWriterFactory writerFactory{writer_};
+    if (formatter_ == nullptr) {
+      return StandardLogHandlerFactory::createHandler(
+          getType(), &writerFactory, options);
+    }
+    TestFormatterFactory formatterFactory{formatter_};
     return StandardLogHandlerFactory::createHandler(
-        getType(), &writerFactory, options);
+        getType(), &writerFactory, &formatterFactory, options);
   }
 
  private:
   std::shared_ptr<TestLogWriter> writer_;
+  std::shared_ptr<TestLogFormatter> formatter_;
   class TestWriterFactory : public StandardLogHandlerFactory::WriterFactory {
    public:
     explicit TestWriterFactory(std::shared_ptr<TestLogWriter> writer)
@@ -78,6 +98,26 @@ class TestHandlerFactory : public LogHandlerFactory {
    private:
     std::shared_ptr<TestLogWriter> writer_;
   };
+
+  class TestFormatterFactory
+      : public StandardLogHandlerFactory::FormatterFactory {
+   public:
+    explicit TestFormatterFactory(std::shared_ptr<LogFormatter> formatter)
+        : formatter_(formatter) {}
+
+    bool processOption(StringPiece /* name */, StringPiece /* value */)
+        override {
+      return false;
+    }
+
+    std::shared_ptr<LogFormatter> createFormatter(
+        const std::shared_ptr<LogWriter>& /* logWriter */) override {
+      return formatter_;
+    }
+
+   private:
+    std::shared_ptr<LogFormatter> formatter_;
+  };
 };
 
 } // namespace folly
@@ -107,7 +147,7 @@ TEST_F(StandardLogHandlerFactoryTest, LogLevelTest) {
   FB_LOG(logger, WARN) << "WARN 1";
   FB_LOG(logger, WARN) << "WARN 2";
 
-  EXPECT_EQ(writer->message_count, 2);
+  EXPECT_EQ(writer->messages.size(), 2);
 }
 
 TEST_F(StandardLogHandlerFactoryTest, LogLevelReverseTest) {
@@ -120,7 +160,7 @@ TEST_F(StandardLogHandlerFactoryTest, LogLevelReverseTest) {
   FB_LOG(logger, DBG2) << "DBG2";
   FB_LOG(logger, WARN) << "WARN 1";
 
-  EXPECT_EQ(writer->message_count, 1);
+  EXPECT_EQ(writer->messages.size(), 1);
 }
 
 TEST_F(StandardLogHandlerFactoryTest, MultipleLoggerTest) {
@@ -138,5 +178,28 @@ TEST_F(StandardLogHandlerFactoryTest, MultipleLoggerTest) {
   FB_LOG(logger, DBG3) << "DBG3"; // only "default" logs this
   FB_LOG(logger, WARN) << "WARN 1"; // both log handlers log this
 
-  EXPECT_EQ(writer->message_count, 3);
+  EXPECT_EQ(writer->messages.size(), 3);
+}
+
+TEST_F(StandardLogHandlerFactoryTest, CustomFormatterTest) {
+  Logger logger{&db, "test"};
+  db.resetConfig(parseLogConfig("test=WARN:default; default=test"));
+
+  // Log once with default formatter
+  FB_LOG(logger, WARN) << "pollution secretary bean";
+
+  // Switch to test formatter and log again
+  db.registerHandlerFactory(
+      std::make_unique<TestHandlerFactory>(
+          writer, std::make_shared<TestLogFormatter>()),
+      true);
+  db.resetConfig(parseLogConfig("test=WARN:default; default=test"));
+
+  FB_LOG(logger, WARN) << "ethereal potato kick";
+
+  EXPECT_EQ(writer->messages.size(), 2);
+  EXPECT_THAT(
+      writer->messages[0],
+      testing::MatchesRegex("^.+pollution secretary bean.+$"));
+  EXPECT_EQ(writer->messages[1], "Test Formatter! ethereal potato kick");
 }