commit 7011500350a3b183c2a253d4d355992c8ec6eac9
author Smilencer <527646889@qq.com> Tue Nov 10 13:36:46 2020 +0800
committer GitHub <noreply@github.com> Tue Nov 10 13:36:46 2020 +0800
tree 9854542005c83664735f0c55d8ab39aba91aba7a
parent 27e9eefae1fc6563788cf694849a1657df0feab6

feat(hotkey): add thread_safe method to clear collector (#640)

src/server/hotkey_collector.cpp

diff --git a/src/server/hotkey_collector.cpp b/src/server/hotkey_collector.cpp
index 053d72c..c326e72 100644
--- a/src/server/hotkey_collector.cpp
+++ b/src/server/hotkey_collector.cpp
@@ -255,6 +255,13 @@
     }
 }
 
+void hotkey_coarse_data_collector::clear()
+{
+    for (int i = 0; i < FLAGS_hotkey_buckets_num; i++) {
+        _hash_buckets[i].store(0);
+    }
+}
+
 hotkey_fine_data_collector::hotkey_fine_data_collector(replica_base *base,
                                                        int target_bucket_index,
                                                        int max_queue_size)
@@ -332,5 +339,12 @@
     }
 }
 
+void hotkey_fine_data_collector::clear()
+{
+    std::pair<dsn::blob, uint64_t> key_weight_pair;
+    while (_capture_key_queue.try_dequeue(key_weight_pair)) {
+    }
+}
+
 } // namespace server
 } // namespace pegasus

src/server/hotkey_collector.h

diff --git a/src/server/hotkey_collector.h b/src/server/hotkey_collector.h
index 937d07d..9d6c3f1 100644
--- a/src/server/hotkey_collector.h
+++ b/src/server/hotkey_collector.h
@@ -102,6 +102,7 @@
     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(detect_hotkey_result &result) = 0;
+    virtual void clear() = 0;
 };
 
 // used in hotkey_collector_state::STOPPED and hotkey_collector_state::FINISHED, avoid null pointers
@@ -111,6 +112,7 @@
     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 {}
+    void clear() override {}
 };
 
 // TODO: (Tangyanzhao) add a unit test of hotkey_coarse_data_collector
@@ -120,6 +122,7 @@
     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;
+    void clear() override;
 
 private:
     hotkey_coarse_data_collector() = delete;
@@ -133,6 +136,7 @@
     hotkey_fine_data_collector(replica_base *base, int target_bucket_index, int max_queue_size);
     void capture_data(const dsn::blob &hash_key, uint64_t weight) override;
     void analyse_data(detect_hotkey_result &result) override;
+    void clear() override;
 
 private:
     hotkey_fine_data_collector() = delete;