1 files changed, 110 insertions, 0 deletions
diff --git a/src/rocksdb/db/periodic_task_scheduler.h b/src/rocksdb/db/periodic_task_scheduler.h
new file mode 100644
index 000000000..f45b80c4d
--- /dev/null
+++ b/src/rocksdb/db/periodic_task_scheduler.h
@@ -0,0 +1,110 @@
+//  Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#pragma once
+
+#ifndef ROCKSDB_LITE
+
+#include "util/timer.h"
+
+namespace ROCKSDB_NAMESPACE {
+class SystemClock;
+
+using PeriodicTaskFunc = std::function<void()>;
+
+constexpr uint64_t kInvalidPeriodSec = 0;
+
+// List of task types
+enum class PeriodicTaskType : uint8_t {
+  kDumpStats = 0,
+  kPersistStats,
+  kFlushInfoLog,
+  kRecordSeqnoTime,
+  kMax,
+};
+
+// PeriodicTaskScheduler contains the periodic task scheduled from the DB
+// instance. It's used to schedule/unschedule DumpStats(), PersistStats(),
+// FlushInfoLog(), etc. Each type of the task can only have one instance,
+// re-register the same task type would only update the repeat period.
+//
+// Internally, it uses a global single threaded timer object to run the periodic
+// task functions. Timer thread will always be started since the info log
+// flushing cannot be disabled.
+class PeriodicTaskScheduler {
+ public:
+  explicit PeriodicTaskScheduler() = default;
+
+  PeriodicTaskScheduler(const PeriodicTaskScheduler&) = delete;
+  PeriodicTaskScheduler(PeriodicTaskScheduler&&) = delete;
+  PeriodicTaskScheduler& operator=(const PeriodicTaskScheduler&) = delete;
+  PeriodicTaskScheduler& operator=(PeriodicTaskScheduler&&) = delete;
+
+  // Register a task with its default repeat period
+  Status Register(PeriodicTaskType task_type, const PeriodicTaskFunc& fn);
+
+  // Register a task with specified repeat period. 0 is an invalid argument
+  // (kInvalidPeriodSec). To stop the task, please use Unregister() specifically
+  Status Register(PeriodicTaskType task_type, const PeriodicTaskFunc& fn,
+                  uint64_t repeat_period_seconds);
+
+  // Unregister the task
+  Status Unregister(PeriodicTaskType task_type);
+
+#ifndef NDEBUG
+  // Override the timer for the unittest
+  void TEST_OverrideTimer(SystemClock* clock);
+
+  // Call Timer TEST_WaitForRun() which wait until Timer starting waiting.
+  void TEST_WaitForRun(const std::function<void()>& callback) const {
+    if (timer_ != nullptr) {
+      timer_->TEST_WaitForRun(callback);
+    }
+  }
+
+  // Get global valid task number in the Timer
+  size_t TEST_GetValidTaskNum() const {
+    if (timer_ != nullptr) {
+      return timer_->TEST_GetPendingTaskNum();
+    }
+    return 0;
+  }
+
+  // If it has the specified task type registered
+  bool TEST_HasTask(PeriodicTaskType task_type) const {
+    auto it = tasks_map_.find(task_type);
+    return it != tasks_map_.end();
+  }
+#endif  // NDEBUG
+
+ private:
+  // default global Timer instance
+  static Timer* Default();
+
+  // Internal structure to store task information
+  struct TaskInfo {
+    TaskInfo(std::string _name, uint64_t _repeat_every_sec)
+        : name(std::move(_name)), repeat_every_sec(_repeat_every_sec) {}
+    std::string name;
+    uint64_t repeat_every_sec;
+  };
+
+  // Internal tasks map
+  std::map<PeriodicTaskType, TaskInfo> tasks_map_;
+
+  // Global timer pointer, which doesn't support synchronous add/cancel tasks
+  // so having a global `timer_mutex` for add/cancel task.
+  Timer* timer_ = Default();
+
+  // Global task id, protected by the global `timer_mutex`
+  inline static uint64_t id_;
+
+  static constexpr uint64_t kMicrosInSecond = 1000U * 1000U;
+};
+
+}  // namespace ROCKSDB_NAMESPACE
+
+#endif  // ROCKSDB_LITE