Add rate limiting

Author: ls4154

core/client.h

@@ -15,37 +15,43 @@
 #include "db.h"
 #include "core_workload.h"
 #include "utils/countdown_latch.h"
+#include "utils/rate_limit.h"
 #include "utils/utils.h"
 
 namespace ycsbc {
 
 inline int ClientThread(ycsbc::DB *db, ycsbc::CoreWorkload *wl, const int num_ops, bool is_loading,
-                        bool init_db, bool cleanup_db, CountDownLatch *latch) {
-    try {
-        if (init_db) {
-            db->Init();
-        }
-
-        int ops = 0;
-        for (int i = 0; i < num_ops; ++i) {
-            if (is_loading) {
-                wl->DoInsert(*db);
-            } else {
-                wl->DoTransaction(*db);
-            }
-            ops++;
-        }
-
-        if (cleanup_db) {
-            db->Cleanup();
-        }
-
-        latch->CountDown();
-        return ops;
-    } catch(const utils::Exception& e) {
-        std::cerr<<"Caught exception: "<<e.what()<<std::endl;
-        exit(1);
+                        bool init_db, bool cleanup_db, utils::CountDownLatch *latch, utils::RateLimiter *rlim) {
+
+  try {
+    if (init_db) {
+      db->Init();
+    }
+
+    int ops = 0;
+    for (int i = 0; i < num_ops; ++i) {
+      if (rlim) {
+        rlim->Consume(1);
+      }
+
+      if (is_loading) {
+        wl->DoInsert(*db);
+      } else {
+        wl->DoTransaction(*db);
+      }
+      ops++;
     }
+
+    if (cleanup_db) {
+      db->Cleanup();
+    }
+
+    latch->CountDown();
+    return ops;
+  } catch (const utils::Exception &e) {
+    std::cerr << "Caught exception: " << e.what() << std::endl;
+    exit(1);
+  }
 }
 
 } // ycsbc

core/ycsbc.cc

@@ -22,14 +22,14 @@
 #include "db_factory.h"
 #include "measurements.h"
 #include "utils/countdown_latch.h"
+#include "utils/rate_limit.h"
 #include "utils/timer.h"
 #include "utils/utils.h"
 
 void UsageMessage(const char *command);
 bool StrStartWith(const char *str, const char *pre);
 void ParseCommandLine(int argc, const char *argv[], ycsbc::utils::Properties &props);
 
-void StatusThread(ycsbc::Measurements *measurements, CountDownLatch *latch, int interval) {
 void StatusThread(ycsbc::Measurements *measurements, ycsbc::utils::CountDownLatch *latch, int interval) {
   using namespace std::chrono;
   time_point<system_clock> start = system_clock::now();
@@ -51,6 +51,39 @@ void StatusThread(ycsbc::Measurements *measurements, ycsbc::utils::CountDownLatc
   };
 }
 
+void RateLimitThread(std::string rate_file, std::vector<ycsbc::utils::RateLimiter *> rate_limiters,
+                     ycsbc::utils::CountDownLatch *latch) {
+  std::ifstream ifs;
+  ifs.open(rate_file);
+
+  if (!ifs.is_open()) {
+    ycsbc::utils::Exception("failed to open: " + rate_file);
+  }
+
+  int64_t num_threads = rate_limiters.size();
+
+  int64_t last_time = 0;
+  while (!ifs.eof()) {
+    int64_t next_time;
+    int64_t next_rate;
+    ifs >> next_time >> next_rate;
+
+    if (next_time <= last_time) {
+      ycsbc::utils::Exception("invalid rate file");
+    }
+
+    bool done = latch->AwaitFor(next_time - last_time);
+    if (done) {
+      break;
+    }
+    last_time = next_time;
+
+    for (auto x : rate_limiters) {
+      x->SetRate(next_rate / num_threads);
+    }
+  }
+}
+
 int main(const int argc, const char *argv[]) {
   ycsbc::utils::Properties props;
   ParseCommandLine(argc, argv, props);
@@ -83,6 +116,7 @@ int main(const int argc, const char *argv[]) {
   ycsbc::CoreWorkload wl;
   wl.Init(props);
 
+  // print status periodically
   const bool show_status = (props.GetProperty("status", "false") == "true");
   const int status_interval = std::stoi(props.GetProperty("status.interval", "10"));
 
@@ -105,8 +139,9 @@ int main(const int argc, const char *argv[]) {
       if (i < total_ops % num_threads) {
         thread_ops++;
       }
+
       client_threads.emplace_back(std::async(std::launch::async, ycsbc::ClientThread, dbs[i], &wl,
-                                             thread_ops, true, true, !do_transaction, &latch));
+                                             thread_ops, true, true, !do_transaction, &latch, nullptr));
     }
     assert((int)client_threads.size() == num_threads);
 
@@ -129,8 +164,14 @@ int main(const int argc, const char *argv[]) {
   measurements->Reset();
   std::this_thread::sleep_for(std::chrono::seconds(stoi(props.GetProperty("sleepafterload", "0"))));
 
+
   // transaction phase
   if (do_transaction) {
+    // initial ops per second, unlimited if <= 0
+    const int64_t ops_limit = std::stoi(props.GetProperty("limit.ops", "0"));
+    // rate file path for dynamic rate limiting, format "time_stamp_sec new_ops_per_second" per line
+    std::string rate_file = props.GetProperty("limit.file", "");
+
     const int total_ops = stoi(props[ycsbc::CoreWorkload::OPERATION_COUNT_PROPERTY]);
 
     ycsbc::utils::CountDownLatch latch(num_threads);
@@ -143,14 +184,27 @@ int main(const int argc, const char *argv[]) {
                                  measurements, &latch, status_interval);
     }
     std::vector<std::future<int>> client_threads;
+    std::vector<ycsbc::utils::RateLimiter *> rate_limiters;
     for (int i = 0; i < num_threads; ++i) {
       int thread_ops = total_ops / num_threads;
       if (i < total_ops % num_threads) {
         thread_ops++;
       }
+      ycsbc::utils::RateLimiter *rlim = nullptr;
+      if (ops_limit > 0 || rate_file != "") {
+        int64_t per_thread_ops = ops_limit / num_threads;
+        rlim = new ycsbc::utils::RateLimiter(per_thread_ops, per_thread_ops);
+      }
+      rate_limiters.push_back(rlim);
       client_threads.emplace_back(std::async(std::launch::async, ycsbc::ClientThread, dbs[i], &wl,
-                                             thread_ops, false, !do_load, true,  &latch));
+                                             thread_ops, false, !do_load, true, &latch, rlim));
     }
+
+    std::future<void> rlim_future;
+    if (rate_file != "") {
+      rlim_future = std::async(std::launch::async, RateLimitThread, rate_file, rate_limiters, &latch);
+    }
+
     assert((int)client_threads.size() == num_threads);
 
     int sum = 0;

utils/rate_limit.h

@@ -0,0 +1,80 @@
+//
+//  rate_limit.h
+//  YCSB-cpp
+//
+//  Copyright (c) 2023 Youngjae Lee <ls4154.lee@gmail.com>.
+//
+
+#ifndef YCSB_C_RATE_LIMIT_H_
+#define YCSB_C_RATE_LIMIT_H_
+
+#include <algorithm>
+#include <chrono>
+#include <cstdio>
+#include <mutex>
+#include <ratio>
+#include <thread>
+
+namespace ycsbc {
+
+namespace utils {
+
+// Token bucket rate limiter for single client
+class RateLimiter {
+ public:
+  RateLimiter(int64_t r, int64_t b) : r_(r * TOKEN_PRECISION), b_(b * TOKEN_PRECISION), tokens_(0), last_(Clock::now()) {}
+
+  inline void Consume(int64_t n) {
+    std::unique_lock<std::mutex> lock(mutex_);
+
+    if (r_ <= 0) {
+      return;
+    }
+
+    // refill tokens
+    auto now = Clock::now();
+    auto diff = std::chrono::duration_cast<Duration>(now - last_);
+    tokens_ = std::min(b_, tokens_ + diff.count() * r_ / 1000000000);
+    last_ = now;
+
+    // check tokens
+    tokens_ -= n * TOKEN_PRECISION;
+
+    // sleep
+    if (tokens_ < 0) {
+      lock.unlock();
+      int64_t wait_time = -tokens_ * 1000000000 / r_;
+      std::this_thread::sleep_for(std::chrono::nanoseconds(wait_time));
+    }
+  }
+
+  inline void SetRate(int64_t r) {
+    std::lock_guard<std::mutex> lock(mutex_);
+
+    // refill tokens
+    auto now = Clock::now();
+    auto diff = std::chrono::duration_cast<Duration>(now - last_);
+    tokens_ = std::min(b_, tokens_ + diff.count() * r_ * TOKEN_PRECISION / 1000000000);
+    last_ = now;
+
+    // set rate
+    r_ = r * TOKEN_PRECISION;
+  }
+
+ private:
+  using Clock = std::chrono::steady_clock;
+  using Duration = std::chrono::nanoseconds;
+  static constexpr int64_t TOKEN_PRECISION = 10000;
+
+  std::mutex mutex_;
+  int64_t r_;
+  int64_t b_;
+  int64_t tokens_;
+  Clock::time_point last_;
+};
+
+} // utils
+
+} // ycsbc
+
+#endif // YCSB_C_RATE_LIMIT_H_