chore(wp): conditional

Signed-off-by: Gil Levkovich <69595609+glevkovich@users.noreply.github.com>

Author: glevkovich

src/facade/dragonfly_connection.cc

@@ -846,8 +846,9 @@ pair<string, string> Connection::GetClientInfoBeforeAfterTid() const {
   }
   string after;
   absl::StrAppend(&after, " irqmatch=", int(cpu == my_cpu_id));
-  if (dispatch_q_.size()) {
-    absl::StrAppend(&after, " pipeline=", dispatch_q_.size());
+  size_t total_pending = dispatch_q_.size() + pending_pipeline_cmd_cnt_;
+  if (total_pending > 0) {
+    absl::StrAppend(&after, " pipeline=", total_pending);
     absl::StrAppend(&after, " pbuf=", pending_pipeline_bytes_);
   }
   absl::StrAppend(&after, " age=", now - creation_time_, " idle=", now - last_interaction_);
@@ -1389,13 +1390,12 @@ bool Connection::ShouldEndAsyncFiber(const MessageHandle& msg) {
     return false;
   }
 
-  if (dispatch_q_.empty()) {
+  if (dispatch_q_.empty() && (parsed_head_ == nullptr) && (parsed_to_execute_ == nullptr)) {
     // Migration requests means we should terminate this function (and allow the fiber to
     // join), so that we can re-launch the fiber in the new thread.
     // We intentionally return and not break in order to keep the connection open.
     return true;
   }
-
   // There shouldn't be any other migration requests in the queue, but it's worth checking
   // as otherwise it would lead to an endless loop.
   bool has_migration_req =
@@ -1607,8 +1607,15 @@ void Connection::AsyncFiber() {
     bool subscriber_over_limit =
         stats_->dispatch_queue_subscriber_bytes >= qbp.publish_buffer_limit;
 
+    // Detect if we are blocked by a migration request while having pending data.
+    // If so, we must skip the admin queue temporarily to drain the data pipeline.
+    bool blocked_migration =
+        !dispatch_q_.empty() &&
+        std::holds_alternative<MigrationRequestMessage>(dispatch_q_.front().handle) &&
+        (parsed_head_ != nullptr);
+
     // 1. Prioritize Control/Admin messages (dispatch_q_) over the data pipeline.
-    if (!dispatch_q_.empty()) {
+    if (!dispatch_q_.empty() && !blocked_migration) {
       MessageHandle msg = std::move(dispatch_q_.front());
       dispatch_q_.pop_front();
 
@@ -1637,7 +1644,7 @@ void Connection::AsyncFiber() {
     }
 
     // 2. Handle Data Pipeline (Intrusive List) - only runs when dispatch_q_ is empty.
-    DCHECK(dispatch_q_.empty());
+    DCHECK(dispatch_q_.empty() || blocked_migration);
     if (parsed_head_) {
       // Ensure pointers are synchronized before deciding how to execute.
       DCHECK_EQ(parsed_head_, parsed_to_execute_) << "Pointers diverged before dispatch decision.";
@@ -1651,16 +1658,9 @@ void Connection::AsyncFiber() {
         if (reply_builder_->GetError())
           break;
       } else {
-        auto replies_recorded_before{reply_builder_->RepliesRecorded()};
+        // Execute the batch. (This puts data in the buffer but does NOT flush).
         if (!ExecuteRedisBatch())
           break;  // Fatal error (e.g. socket closed or builder error)
-
-        // If we drained the pipeline but produced no new output in this batch
-        // (e.g. filtered PubSub messages), the ReplyBuilder might still be holding buffered
-        // data from previous commands. We must flush now to ensure that data is sent.
-        if (!parsed_head_ && (replies_recorded_before == reply_builder_->RepliesRecorded())) {
-          reply_builder_->Flush();
-        }
       }
     }
 
@@ -1679,26 +1679,36 @@ void Connection::AsyncFiber() {
 }
 
 bool Connection::ExecuteRedisBatch() {
+  size_t processed_in_this_loop{};
+  static constexpr size_t kYieldInterval{32};
+  uint64_t start_cycles{CycleClock::Now()};
+
   while (parsed_to_execute_) {
     // During the loop, new admin commands may arrive in the dispatch_q_. They should be
     // prioritized.
     if (!dispatch_q_.empty()) {
       return true;
     }
 
-    // Yield if we have been running too long to prevent starvation.
-    if (ThisFiber::GetRunningTimeCycles() > max_busy_read_cycles_cached) {
+    // Optimization: Check the yield budget only every kYieldInterval commands.
+    // We use && to short-circuit and avoid calling the expensive CycleClock::Now()
+    // on every iteration.
+    if ((++processed_in_this_loop > kYieldInterval) &&
+        (CycleClock::Now() - start_cycles > max_busy_read_cycles_cached)) {
       ThisFiber::Yield();
+      processed_in_this_loop = 0;
+      start_cycles = CycleClock::Now();
     }
 
     ++local_stats_.cmds;
     last_interaction_ = time(nullptr);
 
     auto* cmd{parsed_to_execute_};
-    service_->DispatchCommand(ParsedArgs{*cmd}, cmd, facade::AsyncPreference::ONLY_ASYNC);
     parsed_to_execute_ = cmd->next;
     DCHECK_EQ(cmd, parsed_head_);
     parsed_head_ = parsed_to_execute_;
+
+    service_->DispatchCommand(ParsedArgs{*cmd}, cmd, facade::AsyncPreference::ONLY_ASYNC);
     ReleaseParsedCommand(cmd, true);
 
     if (reply_builder_->GetError()) {
@@ -1709,8 +1719,10 @@ bool Connection::ExecuteRedisBatch() {
   // pipeline queue is empty - flush manually since we are not sure if more commands are coming.
   if (parsed_head_ == nullptr) {
     parsed_tail_ = nullptr;
-    // NOTE: we do not flush here to allow for more commands to batch and let the
-    // caller decide when to flush.
+    parsed_to_execute_ = nullptr;
+    // We finished the pipeline, so we MUST send whatever data is pending
+    // to prevent clients from hanging.
+    reply_builder_->Flush();
   }
 
   DCHECK_EQ(parsed_head_, parsed_to_execute_) << "ExecuteRedisBatch: Pointers diverged.";
@@ -1753,7 +1765,17 @@ bool Connection::Migrate(util::fb2::ProactorBase* dest) {
   CHECK(!cc_->async_dispatch);
   CHECK_EQ(cc_->subscriptions, 0);  // are bound to thread local caches
   CHECK_EQ(self_.use_count(), 1u);  // references cache our thread and backpressure
-                                    //
+  //
+  // Migration is only safe when both admin queue and pipeline are empty.
+  const bool has_pending_cmds = !dispatch_q_.empty() || (parsed_head_ != nullptr);
+  // IoLoop v1 uses the async fiber; v2 runs inline. In both cases we must be idle.
+  DCHECK(dispatch_q_.empty()) << DebugInfo();
+  DCHECK(parsed_head_ == nullptr);
+  if (has_pending_cmds) {
+    LOG(ERROR) << "Migration skipped (" << (ioloop_v2_ ? "io_v2" : "io_v1")
+               << "): pending work remains " << DebugInfo();
+    return false;
+  }
   if (ioloop_v2_ && socket_ && socket_->IsOpen()) {
     socket_->ResetOnRecvHook();
   }