Allow to migrate cluster slot by raw write batch

kvrocks.conf

@@ -536,6 +536,16 @@ compaction-checker-range 0-7
 # rename-command KEYS ""
 
 ################################ MIGRATE #####################################
+# Slot migration supports two ways:
+# - redis-command: Migrate data by redis serialization protocol(RESP).
+# - raw-key-value: Migrate the raw key value data of the storage engine directly.
+#                  This way eliminates the overhead of converting to the redis
+#                  command, reduces resource consumption, improves migration
+#                  efficiency, and can implement a finer rate limit.
+#
+# Default: redis-command
+migrate-type redis-command
+
 # If the network bandwidth is completely consumed by the migration task,
 # it will affect the availability of kvrocks. To avoid this situation,
 # migrate-speed is adopted to limit the migrating speed.
@@ -562,6 +572,18 @@ migrate-pipeline-size 16
 # Default: 10000
 migrate-sequence-gap 10000
 
+# The raw-key-value migration way uses batch for migration. This option sets the batch size
+# for each migration.
+#
+# Default: 16kb
+migrate-batch-size-kb 16
+
+# Rate limit for migration based on raw-key-value, representing the maximum number of data
+# that can be migrated per second. 0 means no limit.
+#
+# Default: 16M
+migrate-batch-rate-limit-mb 16
+
 ################################ ROCKSDB #####################################
 
 # Specify the capacity of column family block cache. A larger block cache

src/cluster/batch_sender.cc

@@ -0,0 +1,129 @@
+/*
+ * 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 "batch_sender.h"
+
+#include "io_util.h"
+#include "server/redis_reply.h"
+#include "time_util.h"
+
+Status BatchSender::Put(rocksdb::ColumnFamilyHandle *cf, const rocksdb::Slice &key, const rocksdb::Slice &value) {
+  // If the data is too large to fit in one batch, it needs to be split into multiple batches.
+  // To cover this case, we append the log data when first add metadata.
+  if (pending_entries_ == 0 && !prefix_logdata_.empty()) {
+    auto s = PutLogData(prefix_logdata_);
+    if (!s.IsOK()) {
+      return s;
+    }
+  }
+  auto s = write_batch_.Put(cf, key, value);
+  if (!s.ok()) {
+    return {Status::NotOK, fmt::format("failed to put key value to migration batch, {}", s.ToString())};
+  }
+
+  pending_entries_++;
+  entries_num_++;
+  return Status::OK();
+}
+
+Status BatchSender::Delete(rocksdb::ColumnFamilyHandle *cf, const rocksdb::Slice &key) {
+  auto s = write_batch_.Delete(cf, key);
+  if (!s.ok()) {
+    return {Status::NotOK, fmt::format("failed to delete key from migration batch, {}", s.ToString())};
+  }
+  pending_entries_++;
+  entries_num_++;
+  return Status::OK();
+}
+
+Status BatchSender::PutLogData(const rocksdb::Slice &blob) {
+  auto s = write_batch_.PutLogData(blob);
+  if (!s.ok()) {
+    return {Status::NotOK, fmt::format("failed to put log data to migration batch, {}", s.ToString())};
+  }
+  pending_entries_++;
+  entries_num_++;
+  return Status::OK();
+}
+
+void BatchSender::SetPrefixLogData(const std::string &prefix_logdata) { prefix_logdata_ = prefix_logdata; }
+
+Status BatchSender::Send() {
+  if (pending_entries_ == 0) {
+    return Status::OK();
+  }
+
+  // rate limit
+  if (bytes_per_sec_ > 0) {
+    auto single_burst = rate_limiter_->GetSingleBurstBytes();
+    auto left = static_cast<int64_t>(write_batch_.GetDataSize());
+    while (left > 0) {
+      auto request_size = std::min(left, single_burst);
+      rate_limiter_->Request(request_size, rocksdb::Env::IOPriority::IO_HIGH, nullptr);
+      left -= request_size;
+    }
+  }
+
+  auto s = sendApplyBatchCmd(dst_fd_, write_batch_);
+  if (!s.IsOK()) {
+    return s.Prefixed("failed to send APPLYBATCH command");
+  }
+
+  sent_bytes_ += write_batch_.GetDataSize();
+  sent_batches_num_++;
+  pending_entries_ = 0;
+  write_batch_.Clear();
+  return Status::OK();
+}
+
+Status BatchSender::sendApplyBatchCmd(int fd, const rocksdb::WriteBatch &write_batch) {
+  if (fd <= 0) {
+    return {Status::NotOK, "invalid fd"};
+  }
+
+  GET_OR_RET(util::SockSend(fd, redis::ArrayOfBulkStrings({"APPLYBATCH", write_batch.Data()})));
+
+  std::string line = GET_OR_RET(util::SockReadLine(fd));
+
+  if (line.compare(0, 1, "-") == 0) {
+    return {Status::NotOK, line};
+  }
+
+  return Status::OK();
+}
+
+void BatchSender::SetBytesPerSecond(size_t bytes_per_sec) {
+  if (bytes_per_sec_ == bytes_per_sec) {
+    return;
+  }
+  bytes_per_sec_ = bytes_per_sec;
+  if (bytes_per_sec > 0) {
+    rate_limiter_->SetBytesPerSecond(static_cast<int64_t>(bytes_per_sec));
+  }
+}
+
+double BatchSender::GetRate(uint64_t since) const {
+  auto t = util::GetTimeStampMS();
+  if (t <= since) {
+    return 0;
+  }
+
+  return ((static_cast<double>(sent_bytes_) / 1024.0) / (static_cast<double>(t - since) / 1000.0));
+}

src/cluster/batch_sender.h

@@ -0,0 +1,71 @@
+/*
+ * 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 <rocksdb/rate_limiter.h>
+#include <rocksdb/write_batch.h>
+
+#include "status.h"
+
+class BatchSender {
+ public:
+  BatchSender() = default;
+  BatchSender(int fd, size_t max_bytes, size_t bytes_per_sec)
+      : dst_fd_(fd),
+        max_bytes_(max_bytes),
+        bytes_per_sec_(bytes_per_sec),
+        rate_limiter_(std::unique_ptr<rocksdb::RateLimiter>(
+            rocksdb::NewGenericRateLimiter(static_cast<int64_t>(bytes_per_sec_)))) {}
+
+  ~BatchSender() = default;
+
+  Status Put(rocksdb::ColumnFamilyHandle *cf, const rocksdb::Slice &key, const rocksdb::Slice &value);
+  Status Delete(rocksdb::ColumnFamilyHandle *cf, const rocksdb::Slice &key);
+  Status PutLogData(const rocksdb::Slice &blob);
+  void SetPrefixLogData(const std::string &prefix_logdata);
+  Status Send();
+
+  void SetMaxBytes(size_t max_bytes) {
+    if (max_bytes_ != max_bytes) max_bytes_ = max_bytes;
+  }
+  bool IsFull() const { return write_batch_.GetDataSize() >= max_bytes_; }
+  uint64_t GetSentBytes() const { return sent_bytes_; }
+  uint32_t GetSentBatchesNum() const { return sent_batches_num_; }
+  uint32_t GetEntriesNum() const { return entries_num_; }
+  void SetBytesPerSecond(size_t bytes_per_sec);
+  double GetRate(uint64_t since) const;
+
+ private:
+  static Status sendApplyBatchCmd(int fd, const rocksdb::WriteBatch &write_batch);
+
+  rocksdb::WriteBatch write_batch_{};
+  std::string prefix_logdata_{};
+  uint64_t sent_bytes_ = 0;
+  uint32_t sent_batches_num_ = 0;
+  uint32_t entries_num_ = 0;
+  uint32_t pending_entries_ = 0;
+
+  int dst_fd_;
+  size_t max_bytes_;
+
+  size_t bytes_per_sec_ = 0;  // 0 means no limit
+  std::unique_ptr<rocksdb::RateLimiter> rate_limiter_;
+};