mlog: Add async I/O and sampling support to MLogger infrastructure

moxygen/mlog/CMakeLists.txt

@@ -50,3 +50,13 @@ moxygen_add_library(moxygen_mlog_file_mlogger_factory
     moxygen_mlog_file_mlogger
     moxygen_mlog_mlogger_factory
 )
+
+moxygen_add_library(moxygen_mlog_sampling_mlogger_factory
+  SRCS
+    SamplingMLoggerFactory.cpp
+  EXPORTED_DEPS
+    moxygen_mlog_mlogger_factory
+    Folly::folly_random
+)
+
+add_subdirectory(test)

moxygen/mlog/FileMLogger.cpp

@@ -5,19 +5,56 @@
  */
 
 #include "moxygen/mlog/FileMLogger.h"
+#include <folly/Random.h>
 #include <folly/json/json.h>
+#include <folly/logging/xlog.h>
 #include <fstream>
 
 namespace moxygen {
 
+std::optional<std::string> FileMLogger::derivePath() const {
+  if (dir_) {
+    if (!dcid_) {
+      XLOG(WARN) << "Skipping mlog output: dcid must be set when dir is configured";
+      return std::nullopt;
+    }
+    const auto dcidHex = dcid_->hex();
+    if (dcidHex.empty()) {
+      XLOG(WARN)
+          << "Skipping mlog output: dcid must have a non-empty hex representation when dir is configured";
+      return std::nullopt;
+    }
+    return *dir_ + "/" + dcidHex + ".mlog";
+  }
+  return path_;
+}
+
 void FileMLogger::outputLogs() {
-  std::ofstream fileObj(path_);
+  auto path = derivePath();
+  if (!path) {
+    return;
+  }
+
+  // Pre-format logs on the calling thread (where formatLog() is safe to
+  // call), then execute the write lambda either inline or via the executor.
+  std::vector<std::string> serialized;
+  serialized.reserve(logs_.size());
   for (const auto& log : logs_) {
-    auto obj = formatLog(log);
-    std::string jsonLog = folly::toPrettyJson(obj);
-    fileObj << jsonLog << std::endl;
+    serialized.push_back(folly::toPrettyJson(formatLog(log)));
+  }
+
+  auto write = [p = std::move(*path), lines = std::move(serialized)]() {
+    std::ofstream fileObj(p);
+    for (const auto& line : lines) {
+      fileObj << line << '\n';
+    }
+  };
+
+  if (writeExecutor_) {
+    writeExecutor_->add(std::move(write));
+  } else {
+    write();
   }
-  fileObj.close();
 }
 
 } // namespace moxygen

moxygen/mlog/FileMLogger.h

@@ -6,6 +6,10 @@
 
 #pragma once
 
+#include <folly/Executor.h>
+#include <memory>
+#include <optional>
+#include <string>
 #include "moxygen/mlog/MLogger.h"
 
 namespace moxygen {
@@ -14,17 +18,50 @@ namespace moxygen {
  * FileMLogger is a concrete MLogger implementation that outputs logs to a file.
  * It inherits from the abstract MLogger base class and implements outputLogs()
  * to write formatted JSON logs to the path specified via setPath().
+ *
+ * If a write executor is set via setWriteExecutor(), outputLogs() will
+ * pre-format logs on the calling thread and schedule the file write
+ * asynchronously on that executor.
+ *
+ * If a directory is set via setDir(), setPath() will prepend it automatically.
  */
 class FileMLogger : public MLogger {
  public:
   explicit FileMLogger(VantagePoint vantagePoint) : MLogger(vantagePoint) {}
   ~FileMLogger() override = default;
 
   /**
-   * Outputs all accumulated logs to the file specified by path_.
+   * Set an executor for asynchronous file writes.
+   */
+  void setWriteExecutor(std::shared_ptr<folly::Executor> executor) {
+    writeExecutor_ = std::move(executor);
+  }
+
+  /**
+   * Set an output directory
+   */
+  void setDir(std::string dir) {
+    dir_ = std::move(dir);
+  }
+
+  /**
+   * Outputs all accumulated logs to the file specified by path_ or
+   * auto-derived path (if dir is set).
    * Logs are formatted as pretty-printed JSON.
    */
   void outputLogs() override;
+
+ private:
+  /**
+     * Derives the output file path. If dir is set, dcid with a non-empty hex
+     * representation is required and the path is {dir}/{dcid_hex}.mlog.
+     * Otherwise returns the explicit path_. Returns nullopt when logging should
+     * be skipped.
+   */
+    std::optional<std::string> derivePath() const;
+
+  std::shared_ptr<folly::Executor> writeExecutor_;
+  std::optional<std::string> dir_;
 };
 
 } // namespace moxygen

moxygen/mlog/FileMLoggerFactory.h

@@ -6,27 +6,55 @@
 
 #pragma once
 
+#include <folly/Executor.h>
 #include <moxygen/mlog/FileMLogger.h>
+#include <memory>
+#include <optional>
 #include <string>
 #include "moxygen/mlog/MLoggerFactory.h"
 
 namespace moxygen {
 
-// Creates an MLogger that writes to a file
+// Creates FileMLogger instances per session.
+// Optionally forwards a write executor and/or output directory to each created
+// logger.
 class FileMLoggerFactory : public MLoggerFactory {
  public:
-  FileMLoggerFactory(const std::string& path, VantagePoint vantagePoint)
-      : path_(path), vantagePoint_(vantagePoint) {}
+  explicit FileMLoggerFactory(VantagePoint vantagePoint = VantagePoint::SERVER)
+      : vantagePoint_(vantagePoint) {}
+
+  explicit FileMLoggerFactory(
+      std::string path,
+      VantagePoint vantagePoint = VantagePoint::SERVER)
+      : vantagePoint_(vantagePoint), path_(std::move(path)) {}
 
   std::shared_ptr<MLogger> createMLogger() override {
     auto logger = std::make_shared<FileMLogger>(vantagePoint_);
-    logger->setPath(path_);
+    if (writeExecutor_) {
+      logger->setWriteExecutor(writeExecutor_);
+    }
+    if (path_) {
+      logger->setPath(*path_);
+    }
+    if (dir_) {
+      logger->setDir(*dir_);
+    }
     return logger;
   }
 
+  void setWriteExecutor(std::shared_ptr<folly::Executor> executor) {
+    writeExecutor_ = std::move(executor);
+  }
+
+  void setDir(std::string dir) {
+    dir_ = std::move(dir);
+  }
+
  private:
-  std::string path_;
   VantagePoint vantagePoint_;
+  std::shared_ptr<folly::Executor> writeExecutor_;
+  std::optional<std::string> path_;
+  std::optional<std::string> dir_;
 };
 
 } // namespace moxygen

moxygen/mlog/SamplingMLoggerFactory.cpp

@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * This source code is licensed under the Apache 2.0 license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include "moxygen/mlog/SamplingMLoggerFactory.h"
+
+namespace moxygen {
+
+SamplingMLoggerFactory::SamplingMLoggerFactory(
+    std::shared_ptr<MLoggerFactory> inner,
+    float sampleRate)
+    : inner_(std::move(inner)), sampleRate_(sampleRate) {
+  // Validate and normalize sampleRate to valid range [0.0, 1.0]
+  if (sampleRate_ <= 0.0f) {
+    sampleRate_ = 0.0f;
+    bucketSize_ = 0; // Sentinel: no sampling
+  } else if (sampleRate_ >= 1.0f) {
+    sampleRate_ = 1.0f;
+    bucketSize_ = 1; // All sessions logged
+  } else {
+    // bucketSize = how many sessions per one logged session, e.g. 0.01 -> 100
+    bucketSize_ = static_cast<uint32_t>(1.0f / sampleRate_);
+    if (bucketSize_ == 0) {
+      bucketSize_ = 1; // guard against rounding to zero for very high rates
+    }
+  }
+}
+
+std::shared_ptr<MLogger> SamplingMLoggerFactory::createMLogger() {
+  if (bucketSize_ == 0) {
+    // No sampling: sampleRate was <= 0
+    return nullptr;
+  }
+  if (bucketSize_ == 1) {
+    // Log all: sampleRate was >= 1
+    return inner_->createMLogger();
+  }
+  // folly::Random::oneIn() is thread-safe and bucketSize_ is always > 1 here.
+  if (folly::Random::oneIn(bucketSize_)) {
+    return inner_->createMLogger();
+  }
+  return nullptr;
+}
+
+} // namespace moxygen

moxygen/mlog/SamplingMLoggerFactory.h

@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * This source code is licensed under the Apache 2.0 license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <folly/Random.h>
+#include <memory>
+#include <moxygen/mlog/MLoggerFactory.h>
+
+namespace moxygen {
+
+// Wraps a MLoggerFactory and applies probabilistic sampling.
+// For each createMLogger() call, returns a logger with probability sampleRate,
+// or nullptr otherwise (meaning the session is not logged).
+//
+// THREAD SAFETY: Uses folly::Random::oneIn() which relies on ThreadLocalPRNG
+// and is safe for concurrent calls. 
+class SamplingMLoggerFactory : public MLoggerFactory {
+ public:
+  SamplingMLoggerFactory(
+      std::shared_ptr<MLoggerFactory> inner,
+      float sampleRate);
+
+  std::shared_ptr<MLogger> createMLogger() override;
+
+ private:
+  std::shared_ptr<MLoggerFactory> inner_;
+  float sampleRate_;
+  uint32_t bucketSize_; // ceil(1 / sampleRate_)
+};
+
+} // namespace moxygen

moxygen/mlog/test/CMakeLists.txt

@@ -0,0 +1,22 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This source code is licensed under the Apache 2.0 license found in the
+# LICENSE file in the root directory of this source tree.
+
+if(NOT BUILD_TESTS)
+    return()
+endif()
+
+moxygen_add_test(TARGET MLoggerTests
+  PREFIX "MLogger"
+  SOURCES
+    FileMLoggerTest.cpp
+    FileMLoggerFactoryTest.cpp
+    SamplingMLoggerFactoryTest.cpp
+  DEPENDS
+    moxygen::moxygen_mlog_file_mlogger
+    moxygen::moxygen_mlog_file_mlogger_factory
+    moxygen::moxygen_mlog_sampling_mlogger_factory
+    Folly::folly_executors_manual_executor
+    mvfst::mvfst_codec_types
+    GTest::gtest_main
+)