chore: fixes - move yield points

Author: romange

helio

@@ -910,7 +910,7 @@ size_t SerializerBase::SerializedLen() const {
 io::Bytes SerializerBase::PrepareFlush(SerializerBase::FlushState flush_state) {
   size_t sz = mem_buf_.InputLen();
   if (sz == 0)
-    return mem_buf_.InputBuffer();
+    return {};
 
   bool is_last_chunk = flush_state == FlushState::kFlushEndEntry;
   VLOG(2) << "PrepareFlush:" << is_last_chunk << " " << number_of_chunks_;

src/server/rdb_save.cc

@@ -154,7 +154,7 @@ void SliceSnapshot::IterateBucketsFb(bool send_full_sync_cut) {
     stats_.keys_total += db_slice_->DbSize(db_indx);
   }
 
-  const uint64_t kCyclesPerYield = base::CycleClock::Frequency() >> 16;  // ~15usec.
+  const uint64_t kCyclesPerJiffy = base::CycleClock::Frequency() >> 16;  // ~15usec.
 
   for (DbIndex db_indx = 0; db_indx < db_array_.size(); ++db_indx) {
     if (!cntx_->IsRunning())
@@ -166,7 +166,6 @@ void SliceSnapshot::IterateBucketsFb(bool send_full_sync_cut) {
     PrimeTable* pt = &db_array_[db_indx]->prime;
     VLOG(1) << "Start traversing " << pt->size() << " items for index " << db_indx;
 
-    uint64_t start_cycles = base::CycleClock::Now();
     do {
       if (!cntx_->IsRunning()) {
         return;
@@ -175,13 +174,14 @@ void SliceSnapshot::IterateBucketsFb(bool send_full_sync_cut) {
       PrimeTable::Cursor next = pt->TraverseBuckets(
           cursor, [this, &db_indx](auto it) { return BucketSaveCb(db_indx, it); });
       cursor = next;
-      uint64_t now = base::CycleClock::Now();
-      if (now > start_cycles + kCyclesPerYield) {
-        ThisFiber::Yield();
-        start_cycles = base::CycleClock::Now();
-      }
 
-      PushSerialized(false);
+      // If we do not flush the data, and have not preempted,
+      // we may need to yield to other fibers to avoid grabbind the CPU for too long.
+      if (!PushSerialized(false)) {
+        if (ThisFiber::GetRunningTimeCycles() > kCyclesPerJiffy) {
+          ThisFiber::Yield();
+        }
+      }
     } while (cursor);
 
     DVLOG(2) << "after loop " << ThisFiber::GetName();
@@ -309,7 +309,6 @@ void SliceSnapshot::SerializeEntry(DbIndex db_indx, const PrimeKey& pk, const Pr
 
 size_t SliceSnapshot::FlushSerialized(SerializerBase::FlushState flush_state) {
   io::StringFile sfile;
-  uint64_t start = base::CycleClock::Now();
   error_code ec = serializer_->FlushToSink(&sfile, flush_state);
   CHECK(!ec);  // always succeeds
 
@@ -320,16 +319,8 @@ size_t SliceSnapshot::FlushSerialized(SerializerBase::FlushState flush_state) {
   uint64_t id = rec_id_++;
   DVLOG(2) << "Pushing " << id;
 
-  // FlushToSink can be quite slow for large values or due compression.
-  // In this case we sleep to avoid starvation of other fibers.
-  // Please note that we allocate the record id before we preempt in order to preserve ordering
-  // guarantees.
-  const uint64_t kSleepThreshold = base::CycleClock::Frequency() >> 15;  // ~30 usec
-  uint64_t elapsed = base::CycleClock::Now() - start;
-  if (elapsed > kSleepThreshold) {
-    // Balance the load between fibers.
-    ThisFiber::SleepFor(chrono::microseconds(200));
-  }
+  const uint64_t kLongRunThresholdCycles = base::CycleClock::Frequency() >> 15;  // ~30 usec
+  uint64_t running_cycles = ThisFiber::GetRunningTimeCycles();
 
   fb2::NoOpLock lk;
 
@@ -349,6 +340,14 @@ size_t SliceSnapshot::FlushSerialized(SerializerBase::FlushState flush_state) {
 
   VLOG(2) << "Pushed with Serialize() " << serialized;
 
+  // FlushToSink can be quite slow for large values or due compression.
+  // In this case we counter-balance CPU over-usage by forcing sleep.
+  // We measure running_cycles before the preemption points because they reset the counter.
+  if (running_cycles > kLongRunThresholdCycles) {
+    // Balance the load between fibers.
+    ThisFiber::SleepFor(chrono::microseconds(200));
+  }
+
   return serialized;
 }