feat(replication): Aggregate data until threshold before replication (#5308)

Aggregate data until threshold is reached before replication. Threshold can be set with `replication_dispatch_threshold` flag and default is set to 1500. For stalled data there is
fiber that wakes up every `10ms` and sends data even if there is no threshold reached.

Signed-off-by: mkaruza <[email protected]>

Author: mkaruza

src/server/dflycmd.cc

@@ -658,7 +658,8 @@ void DflyCmd::StartStableSyncInThread(FlowInfo* flow, ExecutionState* exec_st, E
   DCHECK(shard);
   DCHECK(flow->conn);
 
-  flow->streamer.reset(new JournalStreamer(sf_->journal(), exec_st, JournalStreamer::SendLsn::YES));
+  flow->streamer.reset(
+      new JournalStreamer(sf_->journal(), exec_st, JournalStreamer::SendLsn::YES, true));
   flow->streamer->Start(flow->conn->socket());
 
   // Register cleanup.

src/server/journal/pending_buf.h

@@ -36,10 +36,9 @@ class PendingBuf {
     if (bufs_.back().buf.size() == Buf::kMaxBufSize) {
       bufs_.emplace_back();
     }
-    auto& fron_buf = bufs_.back();
-
-    fron_buf.mem_size += str.size();
-    fron_buf.buf.push_back(std::move(str));
+    auto& front_buf = bufs_.back();
+    front_buf.mem_size += str.size();
+    front_buf.buf.push_back(std::move(str));
   }
 
   // should be called to get the next buffer for sending
@@ -51,6 +50,10 @@ class PendingBuf {
     return bufs_.front();
   }
 
+  size_t FrontBufSize() const {
+    return bufs_.front().mem_size;
+  }
+
   // should be called when the buf from PrepareSendingBuf() method was sent
   void Pop() {
     DCHECK(bufs_.size() >= 2);

src/server/journal/streamer.cc

@@ -24,6 +24,9 @@ ABSL_FLAG(uint32_t, replication_stream_output_limit, 64_KB,
 ABSL_FLAG(uint32_t, migration_buckets_serialization_threshold, 100,
           "The Number of buckets to serialize on each iteration before yielding");
 
+ABSL_FLAG(uint32_t, replication_dispatch_threshold, 1500,
+          "Number of bytes to aggregate before replication");
+
 namespace dfly {
 using namespace util;
 using namespace journal;
@@ -36,13 +39,18 @@ iovec IoVec(io::Bytes src) {
 
 uint32_t replication_stream_output_limit_cached = 64_KB;
 uint32_t migration_buckets_serialization_threshold_cached = 100;
+uint32_t replication_dispatch_threshold = 1500;
+uint32_t stalled_writer_base_period_ms = 10;
 
 }  // namespace
 
-JournalStreamer::JournalStreamer(journal::Journal* journal, ExecutionState* cntx, SendLsn send_lsn)
-    : cntx_(cntx), journal_(journal), send_lsn_(send_lsn) {
+JournalStreamer::JournalStreamer(journal::Journal* journal, ExecutionState* cntx, SendLsn send_lsn,
+                                 bool is_stable_sync)
+    : cntx_(cntx), journal_(journal), is_stable_sync_(is_stable_sync), send_lsn_(send_lsn) {
   // cache the flag to avoid accessing it later.
   replication_stream_output_limit_cached = absl::GetFlag(FLAGS_replication_stream_output_limit);
+  replication_dispatch_threshold = absl::GetFlag(FLAGS_replication_dispatch_threshold);
+  last_async_write_time_ = fb2::ProactorBase::GetMonotonicTimeNs() / 1000000;
 }
 
 JournalStreamer::~JournalStreamer() {
@@ -75,34 +83,80 @@ void JournalStreamer::Start(util::FiberSocketBase* dest) {
   CHECK(dest_ == nullptr && dest != nullptr);
   dest_ = dest;
   journal_cb_id_ = journal_->RegisterOnChange(this);
+  StartStalledDataWriterFiber();
 }
 
 void JournalStreamer::Cancel() {
   VLOG(1) << "JournalStreamer::Cancel";
   waker_.notifyAll();
   journal_->UnregisterOnChange(journal_cb_id_);
-  if (!cntx_->IsError()) {
-    WaitForInflightToComplete();
-  }
+  StopStalledDataWriterFiber();
+  WaitForInflightToComplete();
 }
 
 size_t JournalStreamer::UsedBytes() const {
   return pending_buf_.Size();
 }
 
-void JournalStreamer::AsyncWrite() {
-  DCHECK(!pending_buf_.Empty());
+void JournalStreamer::Write(std::string str) {
+  DCHECK(!str.empty());
+  DVLOG(3) << "Writing " << str.size() << " bytes";
+
+  pending_buf_.Push(std::move(str));
+  AsyncWrite(false);
+}
+
+void JournalStreamer::StartStalledDataWriterFiber() {
+  if (is_stable_sync_ && !stalled_data_writer_.IsJoinable()) {
+    auto pb = fb2::ProactorBase::me();
+    std::chrono::milliseconds period_us(stalled_writer_base_period_ms);
+    stalled_data_writer_ = MakeFiber([this, index = pb->GetPoolIndex(), period_us]() mutable {
+      ThisFiber::SetName(absl::StrCat("fiber_periodic_journal_writer_", index));
+      this->StalledDataWriterFiber(period_us, &stalled_data_writer_done_);
+    });
+  }
+}
+
+void JournalStreamer::StalledDataWriterFiber(std::chrono::milliseconds period_ms,
+                                             util::fb2::Done* waiter) {
+  while (cntx_->IsRunning()) {
+    if (waiter->WaitFor(period_ms)) {
+      if (!cntx_->IsRunning()) {
+        return;
+      }
+    }
+
+    // We don't want to force async write to replicate if last data
+    // was written recent. Data needs to be stalled for period_ms duration.
+    if (!pending_buf_.Size() || in_flight_bytes_ > 0 ||
+        ((last_async_write_time_ + period_ms.count()) >
+         (fb2::ProactorBase::GetMonotonicTimeNs() / 1000000))) {
+      continue;
+    }
 
+    AsyncWrite(true);
+  }
+}
+
+void JournalStreamer::AsyncWrite(bool force_send) {
+  // Stable sync or RestoreStreamer replication can't write data until
+  // previous AsyncWriter finished.
   if (in_flight_bytes_ > 0) {
-    // We can not flush data while there are in flight requests because AsyncWrite
-    // is not atomic. Therefore, we just aggregate.
+    return;
+  }
+
+  // Writing in stable sync and outside of fiber needs to check
+  // threshold before writing data.
+  if (is_stable_sync_ && !force_send &&
+      pending_buf_.FrontBufSize() < replication_dispatch_threshold) {
     return;
   }
 
   const auto& cur_buf = pending_buf_.PrepareSendingBuf();
 
   in_flight_bytes_ = cur_buf.mem_size;
   total_sent_ += in_flight_bytes_;
+  last_async_write_time_ = fb2::ProactorBase::GetMonotonicTimeNs() / 1000000;
 
   const auto v_size = cur_buf.buf.size();
   absl::InlinedVector<iovec, 8> v(v_size);
@@ -112,18 +166,8 @@ void JournalStreamer::AsyncWrite() {
     v[i] = IoVec(io::Bytes(uptr, cur_buf.buf[i].size()));
   }
 
-  dest_->AsyncWrite(v.data(), v.size(), [this, len = in_flight_bytes_](std::error_code ec) {
-    OnCompletion(std::move(ec), len);
-  });
-}
-
-void JournalStreamer::Write(std::string str) {
-  DCHECK(!str.empty());
-  DVLOG(3) << "Writing " << str.size() << " bytes";
-
-  pending_buf_.Push(std::move(str));
-
-  AsyncWrite();
+  dest_->AsyncWrite(v.data(), v.size(),
+                    [this, len = in_flight_bytes_](std::error_code ec) { OnCompletion(ec, len); });
 }
 
 void JournalStreamer::OnCompletion(std::error_code ec, size_t len) {
@@ -136,7 +180,7 @@ void JournalStreamer::OnCompletion(std::error_code ec, size_t len) {
     if (ec) {
       cntx_->ReportError(ec);
     } else if (!pending_buf_.Empty()) {
-      AsyncWrite();
+      AsyncWrite(false);
     }
   }
 
@@ -176,13 +220,22 @@ void JournalStreamer::WaitForInflightToComplete() {
   }
 }
 
+void JournalStreamer::StopStalledDataWriterFiber() {
+  if (is_stable_sync_ && stalled_data_writer_.IsJoinable()) {
+    stalled_data_writer_done_.Notify();
+    if (stalled_data_writer_.IsJoinable()) {
+      stalled_data_writer_.Join();
+    }
+  }
+}
+
 bool JournalStreamer::IsStalled() const {
   return pending_buf_.Size() >= replication_stream_output_limit_cached;
 }
 
 RestoreStreamer::RestoreStreamer(DbSlice* slice, cluster::SlotSet slots, journal::Journal* journal,
                                  ExecutionState* cntx)
-    : JournalStreamer(journal, cntx, JournalStreamer::SendLsn::NO),
+    : JournalStreamer(journal, cntx, JournalStreamer::SendLsn::NO, false),
       db_slice_(slice),
       my_slots_(std::move(slots)) {
   DCHECK(slice != nullptr);

src/server/journal/streamer.h

@@ -21,7 +21,8 @@ namespace dfly {
 class JournalStreamer : public journal::JournalConsumerInterface {
  public:
   enum class SendLsn { NO = 0, YES = 1 };
-  JournalStreamer(journal::Journal* journal, ExecutionState* cntx, SendLsn send_lsn);
+  JournalStreamer(journal::Journal* journal, ExecutionState* cntx, SendLsn send_lsn,
+                  bool is_stable_sync);
   virtual ~JournalStreamer();
 
   // Self referential.
@@ -60,16 +61,26 @@ class JournalStreamer : public journal::JournalConsumerInterface {
   ExecutionState* cntx_;
 
  private:
-  void AsyncWrite();
+  void AsyncWrite(bool force_send);
   void OnCompletion(std::error_code ec, size_t len);
 
   bool IsStalled() const;
 
   journal::Journal* journal_;
 
+  util::fb2::Fiber stalled_data_writer_;
+  util::fb2::Done stalled_data_writer_done_;
+  void StartStalledDataWriterFiber();
+  void StopStalledDataWriterFiber();
+  void StalledDataWriterFiber(std::chrono::milliseconds period_ms, util::fb2::Done* waiter);
+
   PendingBuf pending_buf_;
 
+  // If we are replication in stable sync we can aggregate data before sending
+  bool is_stable_sync_;
   size_t in_flight_bytes_ = 0, total_sent_ = 0;
+  // Last time that send data in milliseconds
+  uint64_t last_async_write_time_ = 0;
   time_t last_lsn_time_ = 0;
   LSN last_lsn_writen_ = 0;
   util::fb2::EventCount waker_;

tests/dragonfly/cluster_test.py

@@ -835,7 +835,7 @@ async def test_cluster_replica_sets_non_owned_keys(df_factory: DflyInstanceFacto
         assert re.match(r"MOVED \d+ localhost:1111", e.value.args[0])
 
         await push_config(replica_config, [c_master_admin])
-        await asyncio.sleep(0.5)
+        await check_all_replicas_finished([c_replica], c_master)
         assert await c_master.execute_command("dbsize") == 0
         assert await c_replica.execute_command("dbsize") == 0
 
@@ -941,7 +941,7 @@ async def test_cluster_flush_slots_after_config_change(df_factory: DflyInstanceF
     """
     await push_config(config, [c_master_admin, c_replica_admin])
 
-    await asyncio.sleep(0.5)
+    await check_all_replicas_finished([c_replica], c_master)
 
     assert await c_master.execute_command("dbsize") == (100_000 - slot_0_size)
     assert await c_replica.execute_command("dbsize") == (100_000 - slot_0_size)
@@ -1146,6 +1146,10 @@ async def test_random_keys():
             assert await client.get(key) == "value"
 
     await test_random_keys()
+
+    for i in range(3):
+        await check_all_replicas_finished([c_replicas[i]], c_masters_admin[i])
+
     await asyncio.gather(*(wait_available_async(c) for c in c_replicas))
 
     # Make sure that getting a value from a replica works as well.

tests/dragonfly/replication_test.py

@@ -3016,6 +3016,8 @@ async def test_replica_snapshot_with_big_values_while_seeding(df_factory: DflyIn
     await seeder.stop(c_master)
     await stream_task
 
+    await check_all_replicas_finished([c_replica], c_master)
+
     # Check that everything is in sync
     hashes = await asyncio.gather(*(SeederV2.capture(c) for c in [c_master, c_replica]))
     assert len(set(hashes)) == 1