feat(hotkey): capture data part2 - declare coarse collector

src/server/hotkey_collector.cpp

@@ -19,18 +19,30 @@
 
 #include <dsn/dist/replication/replication_enums.h>
 #include <dsn/utility/smart_pointers.h>
+#include <dsn/utility/flags.h>
+#include <boost/functional/hash.hpp>
 #include "base/pegasus_key_schema.h"
 #include <dsn/dist/fmt_logging.h>
 
 namespace pegasus {
 namespace server {
 
+DSN_DEFINE_int32("pegasus.server",
+                 coarse_data_variance_threshold,
+                 3,
+                 "the threshold of variance calculate to find the outliers");
+
+DSN_DEFINE_int32("pegasus.server",
+                 data_capture_hash_bucket_num,
+                 37,
+                 "the number of data capture hash buckets");
+
 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>(this))
 {
 }
 
@@ -65,7 +77,26 @@ 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:
+        _internal_collector->analyse_data(_result);
+        if (_result.coarse_bucket_index != -1) {
+            // TODO: (Tangyanzhao) reset _internal_collector to hotkey_fine_data_collector
+            _state.store(hotkey_collector_state::FINE_DETECTING);
+        }
+        return;
+    default:
+        return;
+    }
+}
+
+/*static*/ int hotkey_collector::get_bucket_id(dsn::string_view data)
+{
+    size_t hash_value = boost::hash_range(data.begin(), data.end());
+    return static_cast<int>(hash_value % FLAGS_data_capture_hash_bucket_num);
+}
 
 void hotkey_collector::on_start_detect(dsn::replication::detect_hotkey_response &resp)
 {
@@ -88,7 +119,7 @@ void hotkey_collector::on_start_detect(dsn::replication::detect_hotkey_response
         dwarn_replica(hint);
         return;
     case hotkey_collector_state::STOPPED:
-        // TODO: (Tangyanzhao) start coarse detecting
+        _internal_collector.reset(new hotkey_coarse_data_collector(this));
         _state.store(hotkey_collector_state::COARSE_DETECTING);
         resp.err = dsn::ERR_OK;
         hint = fmt::format("starting to detect {} hotkey", dsn::enum_to_string(_hotkey_type));
@@ -113,5 +144,63 @@ void hotkey_collector::on_stop_detect(dsn::replication::detect_hotkey_response &
     ddebug_replica(hint);
 }
 
+hotkey_coarse_data_collector::hotkey_coarse_data_collector(replica_base *base)
+    : internal_collector_base(base), _hash_buckets(FLAGS_data_capture_hash_bucket_num)
+{
+    for (std::atomic<uint64_t> &bucket : _hash_buckets) {
+        bucket.store(0);
+    }
+}
+
+void hotkey_coarse_data_collector::capture_data(const dsn::blob &hash_key, uint64_t weight)
+{
+    _hash_buckets[hotkey_collector::get_bucket_id(hash_key)].fetch_add(weight);
+}
+
+void hotkey_coarse_data_collector::analyse_data(detect_hotkey_result &result)
+{
+    std::vector<uint64_t> buckets(FLAGS_data_capture_hash_bucket_num);
+    for (int i = 0; i < buckets.size(); i++) {
+        buckets[i] = _hash_buckets[i].load();
+        _hash_buckets[i].store(0);
+    }
+    result = internal_analysis_method(buckets, FLAGS_coarse_data_variance_threshold);
+}
+
+detect_hotkey_result
+hotkey_coarse_data_collector::internal_analysis_method(const std::vector<uint64_t> &captured_keys,
+                                                       int threshold)
+{
+
+    int data_size = captured_keys.size();
+    dassert(captured_keys.size() > 2, "data_capture_hash_bucket_num is too small");
+
+    double table_captured_key_sum = 0;
+    int hot_index = 0;
+    int hot_value = 0;
+    for (int i = 0; i < data_size; i++) {
+        table_captured_key_sum += captured_keys[i];
+        if (captured_keys[i] > hot_value) {
+            hot_index = i;
+            hot_value = captured_keys[i];
+        }
+    }
+    // TODO: (Tangyanzhao) increase a judgment of table_captured_key_sum
+    double captured_keys_avg_count =
+        (table_captured_key_sum - captured_keys[hot_index]) / (data_size - 1);
+    double standard_deviation = 0;
+    for (int i = 0; i < data_size; i++) {
+        if (i != hot_index) {
+            standard_deviation += pow((captured_keys[i] - captured_keys_avg_count), 2);
+        }
+    }
+    standard_deviation = sqrt(standard_deviation / (data_size - 2));
+    double hot_point = (hot_value - captured_keys_avg_count) / standard_deviation;
+    detect_hotkey_result result;
+    hot_point > threshold ? result.coarse_bucket_index = hot_index
+                          : result.coarse_bucket_index = -1;
+    return result;
+}
+
 } // namespace server
 } // namespace pegasus

src/server/hotkey_collector.h

@@ -26,6 +26,12 @@ namespace server {
 
 class internal_collector_base;
 
+struct detect_hotkey_result
+{
+    int coarse_bucket_index;
+    detect_hotkey_result() : coarse_bucket_index(-1) {}
+};
+
 //    hotkey_collector is responsible to find the hot keys after the partition
 //    was detected to be hot. The two types of hotkey, READ & WRITE, are detected
 //    separately.
@@ -75,28 +81,48 @@ class hotkey_collector : public dsn::replication::replica_base
     void analyse_data();
     void handle_rpc(const dsn::replication::detect_hotkey_request &req,
                     /*out*/ dsn::replication::detect_hotkey_response &resp);
+    static int get_bucket_id(dsn::string_view data);
 
 private:
     void on_start_detect(dsn::replication::detect_hotkey_response &resp);
     void on_stop_detect(dsn::replication::detect_hotkey_response &resp);
+
+    detect_hotkey_result _result;
     std::atomic<hotkey_collector_state> _state;
     const dsn::replication::hotkey_type::type _hotkey_type;
     std::shared_ptr<internal_collector_base> _internal_collector;
 };
 
-class internal_collector_base
+class internal_collector_base : public dsn::replication::replica_base
 {
 public:
+    explicit internal_collector_base(replica_base *base) : replica_base(base){};
     virtual void capture_data(const dsn::blob &hash_key, uint64_t weight) = 0;
-    virtual void analyse_data() = 0;
+    virtual void analyse_data(detect_hotkey_result &result) = 0;
 };
 
 // used in hotkey_collector_state::STOPPED and hotkey_collector_state::FINISHED, avoid null pointers
 class hotkey_empty_data_collector : public internal_collector_base
 {
 public:
-    void capture_data(const dsn::blob &hash_key, uint64_t size) {}
-    void analyse_data() {}
+    explicit hotkey_empty_data_collector(replica_base *base) : internal_collector_base(base){};
+    void capture_data(const dsn::blob &hash_key, uint64_t weight) override {}
+    void analyse_data(detect_hotkey_result &result) override {}
+};
+
+// TODO: (Tangyanzhao) add a unit test of hotkey_coarse_data_collector
+class hotkey_coarse_data_collector : public internal_collector_base
+{
+public:
+    explicit hotkey_coarse_data_collector(replica_base *base);
+    void capture_data(const dsn::blob &hash_key, uint64_t weight) override;
+    void analyse_data(detect_hotkey_result &result) override;
+
+private:
+    detect_hotkey_result internal_analysis_method(const std::vector<uint64_t> &captured_keys,
+                                                  int threshold);
+
+    std::vector<std::atomic<uint64_t>> _hash_buckets;
 };
 
 } // namespace server