From d4696317e378067cf0bb3d84ab582f9873b992f0 Mon Sep 17 00:00:00 2001
From: Doodle <13706157+critical27@users.noreply.github.com>
Date: Mon, 6 Sep 2021 20:10:03 +0800
Subject: [PATCH] skip the wait in rate limiter in CI

---
 src/kvstore/CMakeLists.txt     |  1 +
 src/kvstore/RateLimiter.cpp    | 11 +++++++++++
 src/kvstore/RateLimiter.h      |  8 ++++++--
 tests/common/nebula_service.py |  1 +
 4 files changed, 19 insertions(+), 2 deletions(-)
 create mode 100644 src/kvstore/RateLimiter.cpp

diff --git a/src/kvstore/CMakeLists.txt b/src/kvstore/CMakeLists.txt
index b3524f6a9c5..ffcc13401ab 100644
--- a/src/kvstore/CMakeLists.txt
+++ b/src/kvstore/CMakeLists.txt
@@ -8,6 +8,7 @@ nebula_add_library(
     RocksEngineConfig.cpp
     LogEncoder.cpp
     NebulaSnapshotManager.cpp
+    RateLimiter.cpp
     plugins/elasticsearch/ESListener.cpp
 )
 
diff --git a/src/kvstore/RateLimiter.cpp b/src/kvstore/RateLimiter.cpp
new file mode 100644
index 00000000000..6419f456d7b
--- /dev/null
+++ b/src/kvstore/RateLimiter.cpp
@@ -0,0 +1,11 @@
+/* Copyright (c) 2021 vesoft inc. All rights reserved.
+ *
+ * This source code is licensed under Apache 2.0 License,
+ * attached with Common Clause Condition 1.0, found in the LICENSES directory.
+ */
+
+#include "kvstore/RateLimiter.h"
+
+DEFINE_bool(skip_wait_in_rate_limiter,
+            false,
+            "skip the waiting of first second in rate limiter in CI");
diff --git a/src/kvstore/RateLimiter.h b/src/kvstore/RateLimiter.h
index 96eb5266973..eb87512a52f 100644
--- a/src/kvstore/RateLimiter.h
+++ b/src/kvstore/RateLimiter.h
@@ -12,6 +12,8 @@
 #include "common/thrift/ThriftTypes.h"
 #include "common/time/WallClock.h"
 
+DECLARE_bool(skip_wait_in_rate_limiter);
+
 namespace nebula {
 namespace kvstore {
 
@@ -27,7 +29,8 @@ class RateLimiter {
     DCHECK(buckets_.find({spaceId, partId}) == buckets_.end());
     // token will be available after 1 second, to prevent speed spike at the beginning
     auto now = time::WallClock::fastNowInSec();
-    folly::TokenBucket bucket(rate_, burstSize_, static_cast<double>(now + 1));
+    int64_t waitInSec = FLAGS_skip_wait_in_rate_limiter ? 0 : 1;
+    folly::TokenBucket bucket(rate_, burstSize_, static_cast<double>(now + waitInSec));
     buckets_.emplace(std::make_pair(spaceId, partId), std::move(bucket));
   }
 
@@ -40,7 +43,8 @@ class RateLimiter {
   // Caller must make sure the **the parition has been add, and won't be removed during consume.**
   // Snaphot and rebuild index follow this principle by design.
   void consume(GraphSpaceID spaceId, PartitionID partId, size_t toConsume) {
-    DCHECK(buckets_.find({spaceId, partId}) != buckets_.end());
+    // todo(doodle): enable this DCHECK later
+    // DCHECK(buckets_.find({spaceId, partId}) != buckets_.end());
     auto iter = buckets_.find({spaceId, partId});
     if (iter != buckets_.end()) {
       if (toConsume > burstSize_) {
diff --git a/tests/common/nebula_service.py b/tests/common/nebula_service.py
index 86d3f6b36f4..fd7531ca12e 100644
--- a/tests/common/nebula_service.py
+++ b/tests/common/nebula_service.py
@@ -77,6 +77,7 @@ def _format_nebula_command(self, name, meta_port, ports, debug_log=True):
         if name == 'storaged':
             params.append('--local_config=false')
             params.append('--raft_heartbeat_interval_secs=30')
+            params.append('--skip_wait_in_rate_limiter=true')
         if debug_log:
             params.append('--v=4')
         param_format = " ".join(params)