feat(hotkey): build a fundamental framework of hotkey detection

src/server/capacity_unit_calculator.cpp

@@ -3,13 +3,21 @@
 // can be found in the LICENSE file in the root directory of this source tree.
 
 #include "capacity_unit_calculator.h"
+
 #include <dsn/utility/config_api.h>
 #include <rocksdb/status.h>
+#include "hotkey_collector.h"
 
 namespace pegasus {
 namespace server {
 
-capacity_unit_calculator::capacity_unit_calculator(replica_base *r) : replica_base(r)
+capacity_unit_calculator::capacity_unit_calculator(
+    replica_base *r,
+    std::shared_ptr<hotkey_collector> read_hotkey_collector,
+    std::shared_ptr<hotkey_collector> write_hotkey_collector)
+    : replica_base(r),
+      _read_hotkey_collector(read_hotkey_collector),
+      _write_hotkey_collector(write_hotkey_collector)
 {
     _read_capacity_unit_size =
         dsn_config_get_value_uint64("pegasus.server",
@@ -97,6 +105,8 @@ void capacity_unit_calculator::add_get_cu(int32_t status,
         return;
     }
 
+    _read_hotkey_collector->capture_raw_key(key, key.size() + value.size());
+
     if (status == rocksdb::Status::kNotFound) {
         add_read_cu(1);
         return;

src/server/capacity_unit_calculator.h

@@ -11,10 +11,14 @@
 namespace pegasus {
 namespace server {
 
+class hotkey_collector;
+
 class capacity_unit_calculator : public dsn::replication::replica_base
 {
 public:
-    explicit capacity_unit_calculator(replica_base *r);
+    capacity_unit_calculator(replica_base *r,
+                             std::shared_ptr<hotkey_collector> read_hotkey_collector,
+                             std::shared_ptr<hotkey_collector> write_hotkey_collector);
 
     void add_get_cu(int32_t status, const dsn::blob &key, const dsn::blob &value);
     void add_multi_get_cu(int32_t status,
@@ -70,6 +74,9 @@ class capacity_unit_calculator : public dsn::replication::replica_base
     ::dsn::perf_counter_wrapper _pfc_multi_put_bytes;
     ::dsn::perf_counter_wrapper _pfc_check_and_set_bytes;
     ::dsn::perf_counter_wrapper _pfc_check_and_mutate_bytes;
+
+    std::shared_ptr<hotkey_collector> _read_hotkey_collector;
+    std::shared_ptr<hotkey_collector> _write_hotkey_collector;
 };
 
 } // namespace server

src/server/hotkey_collector.cpp

@@ -0,0 +1,31 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "hotkey_collector.h"
+
+namespace pegasus {
+namespace server {
+
+void hotkey_collector::handle_operation(const dsn::replication::detect_hotkey_request &req,
+                                        dsn::replication::detect_hotkey_response &resp)
+{
+}
+
+void hotkey_collector::capture_raw_key(const dsn::blob &raw_key, uint64_t size) {}
+
+} // namespace server
+} // namespace pegasus

src/server/hotkey_collector.h

@@ -0,0 +1,73 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <dsn/utility/string_view.h>
+#include <dsn/dist/replication/replication_types.h>
+
+namespace pegasus {
+namespace server {
+
+//    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.
+//
+//    +--------------------+  +----------------------------------------------------+
+//    |   Replcia server   |  | Hotkey collector                                   |
+//    |                    |  | +-----------------------------------------------+  |
+//    | +----------------+ |  | | Corase capture                                |  |
+//    | |                | |--> |                 +----------+                  |  |
+//    | |  RPC received  | || | |                 |   Data   |                  |  |
+//    | |                | || | |                 +-----+----+                  |  |
+//    | +-------+--------+ || | |                       |                       |  |
+//    |         |          || | |  +---------------+----v--+-------+---------+  |  |
+//    |         v          || | |  |       |Hot    |       |       |         |  |  |
+//    | +-------+--------+ || | |  |Bucket |Bucket |Bucket |Bucket |Bucket   |  |  |
+//    | |   Replication  | || | |  +-----------+-----------------------------+  |  |
+//    | | (only on the   | || | |              |                                |  |
+//    | |  write path))  | || | +--------------|--------------------------------+  |
+//    | +-------+--------+ || |             +--v---+                               |
+//    |         |          || |             | Data |                               |
+//    |         v          || |             +------+                               |
+//    | +-------+--------+ || |          +-----|-------+-------------+             |
+//    | |                | || |   +------|-------------|-------------|---------+   |
+//    | |  Capture data  ---| |   | Fine |capture      |             |         |   |
+//    | |                | |  |   |      |             |             |         |   |
+//    | +-------+--------+ |  |   | +----v----+   +----v----+   +----v----+    |   |
+//    |         |          |  |   | |  queue  |   |  queue  |   |  queue  |    |   |
+//    |         v          |  |   | +----+----+   +----+----+   +----+----+    |   |
+//    | +-------+--------+ |  |   |      |             |             |         |   |
+//    | |                | |  |   | +----v-------------v-------------v------+  |   |
+//    | |   Place data   | |  |   | |             Analsis pool              |  |   |
+//    | |   to the disk  | |  |   | +-----------------|---------------------+  |   |
+//    | |                | |  |   +-------------------|------------------------+   |
+//    | +----------------+ |  |                       v                            |
+//    |                    |  |                     Hotkey                         |
+//    +--------------------+  +----------------------------------------------------+
+
+class hotkey_collector
+{
+public:
+    // size: the cu size of raw_key/hash_key calculated by `capacity_unit_calculator`
+    void capture_raw_key(const dsn::blob &raw_key, uint64_t size);
+    void handle_operation(const dsn::replication::detect_hotkey_request &req,
+                          dsn::replication::detect_hotkey_response &resp);
+};
+
+} // namespace server
+} // namespace pegasus

src/server/pegasus_server_impl.cpp

@@ -22,6 +22,7 @@
 #include "capacity_unit_calculator.h"
 #include "pegasus_server_write.h"
 #include "meta_store.h"
+#include "hotkey_collector.h"
 
 using namespace dsn::literals::chrono_literals;
 
@@ -1485,7 +1486,8 @@ ::dsn::error_code pegasus_server_impl::start(int argc, char **argv)
     });
 
     // initialize cu calculator and write service after server being initialized.
-    _cu_calculator = dsn::make_unique<capacity_unit_calculator>(this);
+    _cu_calculator = dsn::make_unique<capacity_unit_calculator>(
+        this, _read_hotkey_collector, _write_hotkey_collector);
     _server_write = dsn::make_unique<pegasus_server_write>(this, _verbose_log);
 
     return ::dsn::ERR_OK;
@@ -2798,5 +2800,28 @@ void pegasus_server_impl::set_ingestion_status(dsn::replication::ingestion_statu
     _ingestion_status = status;
 }
 
+void pegasus_server_impl::on_detect_hotkey(const dsn::replication::detect_hotkey_request &req,
+                                           dsn::replication::detect_hotkey_response &resp)
+{
+
+    if (dsn_unlikely(req.action != dsn::replication::detect_action::START &&
+                     req.action != dsn::replication::detect_action::STOP)) {
+        resp.err = dsn::ERR_INVALID_PARAMETERS;
+        resp.__set_err_hint("invalid detect_action");
+        return;
+    }
+
+    if (dsn_unlikely(req.type != dsn::replication::hotkey_type::READ &&
+                     req.type != dsn::replication::hotkey_type::WRITE)) {
+        resp.err = dsn::ERR_INVALID_PARAMETERS;
+        resp.__set_err_hint("invalid hotkey_type");
+        return;
+    }
+
+    auto collector = req.type == dsn::replication::hotkey_type::READ ? _read_hotkey_collector
+                                                                     : _write_hotkey_collector;
+    collector->handle_operation(req, resp);
+}
+
 } // namespace server
 } // namespace pegasus

src/server/pegasus_server_impl.h

@@ -28,6 +28,7 @@ namespace server {
 class meta_store;
 class capacity_unit_calculator;
 class pegasus_server_write;
+class hotkey_collector;
 
 class pegasus_server_impl : public pegasus_read_service
 {
@@ -318,6 +319,9 @@ class pegasus_server_impl : public pegasus_read_service
 
     ::dsn::error_code flush_all_family_columns(bool wait);
 
+    void on_detect_hotkey(const dsn::replication::detect_hotkey_request &req,
+                          dsn::replication::detect_hotkey_response &resp) override;
+
 private:
     static const std::chrono::seconds kServerStatUpdateTimeSec;
     static const std::string COMPRESSION_HEADER;
@@ -382,6 +386,9 @@ class pegasus_server_impl : public pegasus_read_service
 
     dsn::task_tracker _tracker;
 
+    std::shared_ptr<hotkey_collector> _read_hotkey_collector;
+    std::shared_ptr<hotkey_collector> _write_hotkey_collector;
+
     // perf counters
     ::dsn::perf_counter_wrapper _pfc_get_qps;
     ::dsn::perf_counter_wrapper _pfc_multi_get_qps;

src/server/test/CMakeLists.txt

@@ -10,6 +10,7 @@ set(MY_PROJ_SRC "../pegasus_server_impl.cpp"
                 "../pegasus_mutation_duplicator.cpp"
                 "../hotspot_partition_calculator.cpp"
                 "../meta_store.cpp"
+                "../hotkey_collector.cpp"
 )
 
 set(MY_SRC_SEARCH_MODE "GLOB")

src/server/test/capacity_unit_calculator_test.cpp

@@ -6,6 +6,7 @@
 #include "server/capacity_unit_calculator.h"
 
 #include <dsn/dist/replication/replica_base.h>
+#include "server/hotkey_collector.h"
 
 namespace pegasus {
 namespace server {
@@ -26,7 +27,8 @@ class mock_capacity_unit_calculator : public capacity_unit_calculator
     }
 
     explicit mock_capacity_unit_calculator(dsn::replication::replica_base *r)
-        : capacity_unit_calculator(r)
+        : capacity_unit_calculator(
+              r, std::make_shared<hotkey_collector>(), std::make_shared<hotkey_collector>())
     {
     }