feat(hotkey): collector can be terminated by timeout

src/server/hotkey_collector.cpp

@@ -21,16 +21,24 @@
 #include <dsn/utility/smart_pointers.h>
 #include "base/pegasus_key_schema.h"
 #include <dsn/dist/fmt_logging.h>
+#include <dsn/utility/flags.h>
 
 namespace pegasus {
 namespace server {
 
+DSN_DEFINE_int32(
+    "pegasus.server",
+    max_seconds_to_detect_hotkey,
+    150,
+    "the max time (in seconds) allowed to capture hotkey, will stop if hotkey's not found");
+
 hotkey_collector::hotkey_collector(dsn::replication::hotkey_type::type hotkey_type,
                                    dsn::replication::replica_base *r_base)
     : replica_base(r_base),
       _state(hotkey_collector_state::STOPPED),
       _hotkey_type(hotkey_type),
-      _internal_collector(std::make_shared<hotkey_empty_data_collector>())
+      _internal_collector(std::make_shared<hotkey_empty_data_collector>()),
+      _collector_start_time(0)
 {
 }
 
@@ -65,7 +73,17 @@ void hotkey_collector::capture_hash_key(const dsn::blob &hash_key, int64_t weigh
     _internal_collector->capture_data(hash_key, weight);
 }
 
-void hotkey_collector::analyse_data() { _internal_collector->analyse_data(); }
+void hotkey_collector::analyse_data()
+{
+    switch (_state.load()) {
+    case hotkey_collector_state::COARSE_DETECTING:
+        terminate_if_timeout();
+        _internal_collector->analyse_data();
+        return;
+    default:
+        return;
+    }
+}
 
 void hotkey_collector::on_start_detect(dsn::replication::detect_hotkey_response &resp)
 {
@@ -82,12 +100,13 @@ void hotkey_collector::on_start_detect(dsn::replication::detect_hotkey_response
         return;
     case hotkey_collector_state::FINISHED:
         resp.err = dsn::ERR_INVALID_STATE;
-        hint = fmt::format(
-            "{} hotkey result has been found, you can send a stop rpc to restart hotkey detection",
-            dsn::enum_to_string(_hotkey_type));
+        hint = fmt::format("{} hotkey result has been found, you can send a stop rpc to "
+                           "restart hotkey detection",
+                           dsn::enum_to_string(_hotkey_type));
         dwarn_replica(hint);
         return;
     case hotkey_collector_state::STOPPED:
+        _collector_start_time = dsn_now_s();
         // TODO: (Tangyanzhao) start coarse detecting
         _state.store(hotkey_collector_state::COARSE_DETECTING);
         resp.err = dsn::ERR_OK;
@@ -105,13 +124,31 @@ void hotkey_collector::on_start_detect(dsn::replication::detect_hotkey_response
 
 void hotkey_collector::on_stop_detect(dsn::replication::detect_hotkey_response &resp)
 {
-    _state.store(hotkey_collector_state::STOPPED);
-    _internal_collector.reset();
+    terminate();
     resp.err = dsn::ERR_OK;
     std::string hint =
         fmt::format("{} hotkey stopped, cache cleared", dsn::enum_to_string(_hotkey_type));
     ddebug_replica(hint);
 }
 
+void hotkey_collector::terminate()
+{
+    _state.store(hotkey_collector_state::STOPPED);
+    _internal_collector.reset();
+    _collector_start_time = 0;
+}
+
+void hotkey_collector::terminate_if_timeout()
+{
+    if (_collector_start_time == 0) {
+        return;
+    }
+    if (dsn_now_s() >= _collector_start_time + FLAGS_max_seconds_to_detect_hotkey) {
+        ddebug_replica("hotkey collector work time is exhausted but no hotkey has been found");
+        terminate();
+        return;
+    }
+};
+
 } // namespace server
 } // namespace pegasus

src/server/hotkey_collector.h

@@ -79,9 +79,13 @@ class hotkey_collector : public dsn::replication::replica_base
 private:
     void on_start_detect(dsn::replication::detect_hotkey_response &resp);
     void on_stop_detect(dsn::replication::detect_hotkey_response &resp);
+    void terminate();
+    void terminate_if_timeout();
+
     std::atomic<hotkey_collector_state> _state;
     const dsn::replication::hotkey_type::type _hotkey_type;
     std::shared_ptr<internal_collector_base> _internal_collector;
+    uint64_t _collector_start_time;
 };
 
 class internal_collector_base

src/server/test/capacity_unit_calculator_test.cpp

@@ -30,8 +30,8 @@ class mock_capacity_unit_calculator : public capacity_unit_calculator
     explicit mock_capacity_unit_calculator(dsn::replication::replica_base *r)
         : capacity_unit_calculator(
               r,
-              std::make_shared<hotkey_collector>(dsn::replication::hotkey_type::READ, this),
-              std::make_shared<hotkey_collector>(dsn::replication::hotkey_type::WRITE, this))
+              std::make_shared<hotkey_collector>(dsn::replication::hotkey_type::READ, r),
+              std::make_shared<hotkey_collector>(dsn::replication::hotkey_type::WRITE, r))
     {
     }