Fix possible issue with queries not always being handled before next WFT

Author: Sushisource

core/src/worker/workflow/mod.rs

@@ -14,7 +14,6 @@ mod workflow_stream;
 pub(crate) use driven_workflow::DrivenWorkflow;
 pub(crate) use history_update::HistoryUpdate;
 
-use crate::protosext::ValidPollWFTQResponse;
 use crate::{
     MetricsContext,
     abstractions::{
@@ -23,7 +22,9 @@ use crate::{
     },
     internal_flags::InternalFlags,
     pollers::TrackedPermittedTqResp,
-    protosext::{legacy_query_failure, protocol_messages::IncomingProtocolMessage},
+    protosext::{
+        ValidPollWFTQResponse, legacy_query_failure, protocol_messages::IncomingProtocolMessage,
+    },
     telemetry::{TelemetryInstance, VecDisplayer, set_trace_subscriber_for_current_thread},
     worker::{
         LocalActRequest, LocalActivityExecutionResult, LocalActivityResolution,
@@ -49,7 +50,6 @@ use std::{
     collections::VecDeque,
     fmt::Debug,
     future::Future,
-    mem,
     ops::DerefMut,
     rc::Rc,
     result,
@@ -1099,8 +1099,7 @@ struct BufferedTasks {
     /// current one has been processed).
     query_only_tasks: VecDeque<PermittedWFT>,
     /// These are query-only tasks for the *buffered* wft, if any. They will all be discarded if
-    /// a buffered wft is replaced before being handled. They move to `query_only_tasks` once the
-    /// buffered task is taken.
+    /// a buffered wft is replaced before being handled.
     query_only_tasks_for_buffered: VecDeque<PermittedWFT>,
 }
 
@@ -1135,9 +1134,13 @@ impl BufferedTasks {
         if let Some(q) = self.query_only_tasks.pop_front() {
             return Some(q);
         }
-        if let Some(t) = self.wft.take() {
-            self.query_only_tasks = mem::take(&mut self.query_only_tasks_for_buffered);
-            return Some(t);
+        if self.wft.is_some() {
+            if let Some(q) = self.query_only_tasks_for_buffered.pop_front() {
+                return Some(q);
+            }
+            if let Some(t) = self.wft.take() {
+                return Some(t);
+            }
         }
         None
     }