From da9da9c9e77e8fa4dfd6f7e060dfde34ac89ba17 Mon Sep 17 00:00:00 2001
From: hjiang <hjiang@anyscale.com>
Date: Tue, 17 Dec 2024 17:30:14 +0000
Subject: [PATCH] customized rotation filename

---
 include/spdlog/sinks/rotating_file_sink-inl.h | 21 ++++++++++++---
 include/spdlog/sinks/rotating_file_sink.h     | 12 +++++++++
 tests/test_file_logging.cpp                   | 27 +++++++++++++++++++
 3 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/include/spdlog/sinks/rotating_file_sink-inl.h b/include/spdlog/sinks/rotating_file_sink-inl.h
index 420bafb0..785a5ec3 100644
--- a/include/spdlog/sinks/rotating_file_sink-inl.h
+++ b/include/spdlog/sinks/rotating_file_sink-inl.h
@@ -13,6 +13,7 @@
 #include <spdlog/details/null_mutex.h>
 #include <spdlog/fmt/fmt.h>
 
+#include <cassert>
 #include <cerrno>
 #include <chrono>
 #include <ctime>
@@ -41,7 +42,7 @@ SPDLOG_INLINE rotating_file_sink<Mutex>::rotating_file_sink(
     if (max_files > 200000) {
         throw_spdlog_ex("rotating sink constructor: max_files arg cannot exceed 200000");
     }
-    file_helper_.open(calc_filename(base_filename_, 0));
+    file_helper_.open(get_filename_for_rotation_(base_filename_, 0));
     current_size_ = file_helper_.size();  // expensive. called only once
     if (rotate_on_open && current_size_ > 0) {
         rotate_();
@@ -49,6 +50,20 @@ SPDLOG_INLINE rotating_file_sink<Mutex>::rotating_file_sink(
     }
 }
 
+template <typename Mutex>
+SPDLOG_INLINE void rotating_file_sink<Mutex>::set_rotate_filename_format(std::function<filename_t(const filename_t &filename, std::size_t index)> rotation_file_format) {
+    assert(!rotation_file_format_);
+    rotation_file_format_ = std::move(rotation_file_format);
+}
+
+template <typename Mutex>
+SPDLOG_INLINE filename_t rotating_file_sink<Mutex>::get_filename_for_rotation_(const filename_t &filename, std::size_t index) {
+    if (rotation_file_format_) {
+      return rotation_file_format_(filename, index);
+    }
+    return calc_filename(filename, index);
+}
+
 // calc filename according to index and file extension if exists.
 // e.g. calc_filename("logs/mylog.txt, 3) => "logs/mylog.3.txt".
 template <typename Mutex>
@@ -112,11 +127,11 @@ SPDLOG_INLINE void rotating_file_sink<Mutex>::rotate_() {
 
     file_helper_.close();
     for (auto i = max_files_; i > 0; --i) {
-        filename_t src = calc_filename(base_filename_, i - 1);
+        filename_t src = get_filename_for_rotation_(base_filename_, i - 1);
         if (!path_exists(src)) {
             continue;
         }
-        filename_t target = calc_filename(base_filename_, i);
+        filename_t target = get_filename_for_rotation_(base_filename_, i);
 
         if (!rename_file_(src, target)) {
             // if failed try again after a small delay.
diff --git a/include/spdlog/sinks/rotating_file_sink.h b/include/spdlog/sinks/rotating_file_sink.h
index 42bd3760..75af0757 100644
--- a/include/spdlog/sinks/rotating_file_sink.h
+++ b/include/spdlog/sinks/rotating_file_sink.h
@@ -9,6 +9,7 @@
 #include <spdlog/sinks/base_sink.h>
 
 #include <chrono>
+#include <functional>
 #include <mutex>
 #include <string>
 
@@ -26,10 +27,17 @@ public:
                        std::size_t max_files,
                        bool rotate_on_open = false,
                        const file_event_handlers &event_handlers = {});
+    // Default function to get rotation filename by base filename and rotation file index.
     static filename_t calc_filename(const filename_t &filename, std::size_t index);
     filename_t filename();
     void rotate_now();
 
+    // Set the file format for rotation files.
+    // NOTE:
+    // 1. The format function is supposed to be called only once, otherwise check failure.
+    // 2. If [index] is 0, [filename] is expected to return.
+    void set_rotate_filename_format(std::function<filename_t(const filename_t &filename, std::size_t index)> rotation_file_format);
+
 protected:
     void sink_it_(const details::log_msg &msg) override;
     void flush_() override;
@@ -46,11 +54,15 @@ private:
     // return true on success, false otherwise.
     bool rename_file_(const filename_t &src_filename, const filename_t &target_filename);
 
+    // A wrapper around rotation filename format function(s).
+    filename_t get_filename_for_rotation_(const filename_t &filename, std::size_t index);
+
     filename_t base_filename_;
     std::size_t max_size_;
     std::size_t max_files_;
     std::size_t current_size_;
     details::file_helper file_helper_;
+    std::function<filename_t(const filename_t &filename, std::size_t index)> rotation_file_format_;
 };
 
 using rotating_file_sink_mt = rotating_file_sink<std::mutex>;
diff --git a/tests/test_file_logging.cpp b/tests/test_file_logging.cpp
index 71d06a6d..7a9ffd75 100644
--- a/tests/test_file_logging.cpp
+++ b/tests/test_file_logging.cpp
@@ -141,3 +141,30 @@ TEST_CASE("rotating_file_logger4", "[rotating_logger]") {
     REQUIRE(get_filesize(ROTATING_LOG) > 0);
     REQUIRE(get_filesize(ROTATING_LOG ".1") > 0);
 }
+
+// Test customized rotation filename.
+TEST_CASE("rotating_file_logger5", "[rotating_logger]") {
+    prepare_logdir();
+    size_t max_size = 1024 * 10;
+    spdlog::filename_t basename = SPDLOG_FILENAME_T(ROTATING_LOG);
+    auto sink = std::make_shared<spdlog::sinks::rotating_file_sink_st>(basename, max_size, 2);
+    sink->set_rotate_filename_format([](const spdlog::filename_t &filename, std::size_t index) {
+        if (index == 0u) {
+            return filename;
+        }
+        const auto old_fname = spdlog::sinks::rotating_file_sink_st::calc_filename(filename, index);
+        return old_fname + ".test_suffix";
+    });
+    auto logger = std::make_shared<spdlog::logger>("rotating_sink_logger", sink);
+
+    logger->info("Test message - pre-rotation");
+    logger->flush();
+
+    sink->rotate_now();
+
+    logger->info("Test message - post-rotation");
+    logger->flush();
+
+    REQUIRE(get_filesize(ROTATING_LOG) > 0);
+    REQUIRE(get_filesize(ROTATING_LOG ".1.test_suffix") > 0);
+}