Fix possibly dropping final events in a WFT during cache miss (#537)

Author: Sushisource

core/src/worker/workflow/history_update.rs

@@ -176,11 +176,14 @@ impl HistoryPaginator {
             run_id: req.original_wft.work.execution.run_id.clone(),
             previous_wft_started_id: req.original_wft.work.update.previous_wft_started_id,
             wft_started_event_id: req.original_wft.work.update.wft_started_id,
-            id_of_last_event_in_last_extracted_update: None,
+            id_of_last_event_in_last_extracted_update: req
+                .original_wft
+                .paginator
+                .id_of_last_event_in_last_extracted_update,
             client,
             event_queue: Default::default(),
             next_page_token: NextPageToken::FetchFromStart,
-            final_events: vec![],
+            final_events: req.original_wft.work.update.events,
         };
         let first_update = paginator.extract_next_update().await?;
         req.original_wft.work.update = first_update;
@@ -244,7 +247,7 @@ impl HistoryPaginator {
     /// we have two, or until we are at the end of history.
     pub(crate) async fn extract_next_update(&mut self) -> Result<HistoryUpdate, tonic::Status> {
         loop {
-            let fetch_happened = !self.get_next_page().await?;
+            let no_next_page = !self.get_next_page().await?;
             let current_events = mem::take(&mut self.event_queue);
             let seen_enough_events = current_events
                 .back()
@@ -267,10 +270,7 @@ impl HistoryPaginator {
                 .id_of_last_event_in_last_extracted_update
                 .unwrap_or_default()
                 >= self.wft_started_event_id;
-            if current_events.is_empty()
-                && !fetch_happened
-                && already_sent_update_with_enough_events
-            {
+            if current_events.is_empty() && no_next_page && already_sent_update_with_enough_events {
                 // We must return an empty update which also says is contains the final WFT so we
                 // know we're done with replay.
                 return Ok(HistoryUpdate::from_events(
@@ -282,12 +282,15 @@ impl HistoryPaginator {
                 .0);
             }
 
-            if current_events.is_empty() || (fetch_happened && !seen_enough_events) {
+            if current_events.is_empty() || (no_next_page && !seen_enough_events) {
                 // If next page fetching happened, and we still ended up with no or insufficient
                 // events, something is wrong. We're expecting there to be more events to be able to
                 // extract this update, but server isn't giving us any. We have no choice except to
                 // give up and evict.
                 error!(
+                    current_events=?current_events,
+                    no_next_page,
+                    seen_enough_events,
                     "We expected to be able to fetch more events but server says there are none"
                 );
                 return Err(EMPTY_FETCH_ERR.clone());
@@ -319,13 +322,13 @@ impl HistoryPaginator {
         }
     }
 
-    /// Fetches the next page and adds it to the internal queue. Returns true if a fetch was
-    /// performed, false if there is no next page.
+    /// Fetches the next page and adds it to the internal queue.
+    /// Returns true if we still have a next page token after fetching.
     async fn get_next_page(&mut self) -> Result<bool, tonic::Status> {
         let history = loop {
             let npt = match mem::replace(&mut self.next_page_token, NextPageToken::Done) {
                 // If the last page token we got was empty, we're done.
-                NextPageToken::Done => return Ok(false),
+                NextPageToken::Done => break None,
                 NextPageToken::FetchFromStart => vec![],
                 NextPageToken::Next(v) => v,
             };
@@ -366,7 +369,7 @@ impl HistoryPaginator {
                 );
             }
         };
-        Ok(true)
+        Ok(!matches!(&self.next_page_token, NextPageToken::Done))
     }
 }
 
@@ -715,11 +718,21 @@ pub mod tests {
     use super::*;
     use crate::{
         replay::{HistoryInfo, TestHistoryBuilder},
-        test_help::canned_histories,
+        test_help::{canned_histories, mock_sdk_cfg, MockPollCfg, ResponseType},
         worker::client::mocks::mock_workflow_client,
     };
+    use futures::StreamExt;
     use futures_util::TryStreamExt;
-    use temporal_sdk_core_protos::temporal::api::workflowservice::v1::GetWorkflowExecutionHistoryResponse;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use temporal_client::WorkflowOptions;
+    use temporal_sdk::WfContext;
+    use temporal_sdk_core_protos::{
+        temporal::api::{
+            common::v1::WorkflowExecution, enums::v1::WorkflowTaskFailedCause,
+            workflowservice::v1::GetWorkflowExecutionHistoryResponse,
+        },
+        DEFAULT_WORKFLOW_TYPE,
+    };
 
     impl From<HistoryInfo> for HistoryUpdate {
         fn from(v: HistoryInfo) -> Self {
@@ -1252,4 +1265,123 @@ pub mod tests {
         let mut update = paginator.extract_next_update().await.unwrap();
         assert_matches!(update.take_next_wft_sequence(8), NextWFT::ReplayOver);
     }
+
+    #[tokio::test]
+    async fn weird_pagination_doesnt_drop_wft_events() {
+        crate::telemetry::test_telem_console();
+        let wf_id = "fakeid";
+        // 1: EVENT_TYPE_WORKFLOW_EXECUTION_STARTED
+        // 2: EVENT_TYPE_WORKFLOW_TASK_SCHEDULED
+        // 3: EVENT_TYPE_WORKFLOW_TASK_STARTED
+        // 4: EVENT_TYPE_WORKFLOW_TASK_COMPLETED
+        // empty page
+        // 5: EVENT_TYPE_WORKFLOW_EXECUTION_SIGNALED
+        // 6: EVENT_TYPE_WORKFLOW_TASK_SCHEDULED
+        // 7: EVENT_TYPE_WORKFLOW_TASK_STARTED
+        // 8: EVENT_TYPE_WORKFLOW_TASK_FAILED
+        // empty page
+        // 9: EVENT_TYPE_WORKFLOW_EXECUTION_SIGNALED
+        // 10: EVENT_TYPE_WORKFLOW_TASK_SCHEDULED
+        // 11: EVENT_TYPE_WORKFLOW_TASK_STARTED
+        // empty page
+        let mut t = TestHistoryBuilder::default();
+        t.add_by_type(EventType::WorkflowExecutionStarted);
+        t.add_full_wf_task();
+
+        t.add_we_signaled("hi", vec![]);
+        t.add_workflow_task_scheduled_and_started();
+        t.add_workflow_task_failed_with_failure(
+            WorkflowTaskFailedCause::UnhandledCommand,
+            Default::default(),
+        );
+
+        t.add_we_signaled("hi", vec![]);
+        t.add_workflow_task_scheduled_and_started();
+
+        let workflow_task = t.get_full_history_info().unwrap();
+        let mut wft_resp = workflow_task.as_poll_wft_response();
+        wft_resp.workflow_execution = Some(WorkflowExecution {
+            workflow_id: wf_id.to_string(),
+            run_id: t.get_orig_run_id().to_string(),
+        });
+        // Just 9/10/11 in WFT
+        wft_resp.history.as_mut().unwrap().events.drain(0..8);
+
+        let mut resp_1: GetWorkflowExecutionHistoryResponse =
+            t.get_full_history_info().unwrap().into();
+        resp_1.next_page_token = vec![1];
+        resp_1.history.as_mut().unwrap().events.truncate(4);
+
+        let mut mock_client = mock_workflow_client();
+        mock_client
+            .expect_get_workflow_execution_history()
+            .returning(move |_, _, _| Ok(resp_1.clone()))
+            .times(1);
+        mock_client
+            .expect_get_workflow_execution_history()
+            .returning(move |_, _, _| {
+                Ok(GetWorkflowExecutionHistoryResponse {
+                    history: Some(History { events: vec![] }),
+                    raw_history: vec![],
+                    next_page_token: vec![2],
+                    archived: false,
+                })
+            })
+            .times(1);
+        let mut resp_2: GetWorkflowExecutionHistoryResponse =
+            t.get_full_history_info().unwrap().into();
+        resp_2.next_page_token = vec![3];
+        resp_2.history.as_mut().unwrap().events.drain(0..4);
+        resp_2.history.as_mut().unwrap().events.truncate(4);
+        mock_client
+            .expect_get_workflow_execution_history()
+            .returning(move |_, _, _| Ok(resp_2.clone()))
+            .times(1);
+        mock_client
+            .expect_get_workflow_execution_history()
+            .returning(move |_, _, _| {
+                Ok(GetWorkflowExecutionHistoryResponse {
+                    history: Some(History { events: vec![] }),
+                    raw_history: vec![],
+                    next_page_token: vec![],
+                    archived: false,
+                })
+            })
+            .times(1);
+
+        let wf_type = DEFAULT_WORKFLOW_TYPE;
+        let mh =
+            MockPollCfg::from_resp_batches(wf_id, t, [ResponseType::Raw(wft_resp)], mock_client);
+        let mut worker = mock_sdk_cfg(mh, |cfg| {
+            cfg.max_cached_workflows = 2;
+            cfg.ignore_evicts_on_shutdown = false;
+        });
+
+        let sig_ctr = Arc::new(AtomicUsize::new(0));
+        let sig_ctr_clone = sig_ctr.clone();
+        worker.register_wf(wf_type.to_owned(), move |ctx: WfContext| {
+            let sig_ctr_clone = sig_ctr_clone.clone();
+            async move {
+                let mut sigchan = ctx.make_signal_channel("hi");
+                while sigchan.next().await.is_some() {
+                    if sig_ctr_clone.fetch_add(1, Ordering::AcqRel) == 1 {
+                        break;
+                    }
+                }
+                Ok(().into())
+            }
+        });
+
+        worker
+            .submit_wf(
+                wf_id.to_owned(),
+                wf_type.to_owned(),
+                vec![],
+                WorkflowOptions::default(),
+            )
+            .await
+            .unwrap();
+        worker.run_until_done().await.unwrap();
+        assert_eq!(sig_ctr.load(Ordering::Acquire), 2);
+    }
 }

sdk/src/workflow_context.rs

@@ -11,7 +11,7 @@ use crate::{
     SignalExternalWfResult, TimerResult, UnblockEvent, Unblockable,
 };
 use crossbeam::channel::{Receiver, Sender};
-use futures::{task::Context, FutureExt, Stream};
+use futures::{task::Context, FutureExt, Stream, StreamExt};
 use parking_lot::RwLock;
 use std::{
     collections::HashMap,
@@ -398,6 +398,14 @@ impl DrainableSignalStream {
         }
         signals
     }
+
+    pub fn drain_ready(&mut self) -> Vec<SignalData> {
+        let mut signals = vec![];
+        while let Some(s) = self.0.next().now_or_never().flatten() {
+            signals.push(s);
+        }
+        signals
+    }
 }
 
 impl Stream for DrainableSignalStream {

sdk/src/workflow_context/options.rs

@@ -272,7 +272,7 @@ impl Signal {
 }
 
 /// Data contained within a signal
-#[derive(Default)]
+#[derive(Default, Debug)]
 pub struct SignalData {
     /// The arguments the signal will receive
     pub input: Vec<Payload>,