Fix resets of workflows causing a panic (#287)

Sushisource · web-flow · commit 12bd832fa9a2 · 2022-03-14T09:13:12.000-07:00
Problem was the internal run id was always being set to the original
run id, so when it changed because of a reset, lookup of the workflow
state failed
diff --git a/core/src/core_tests/activity_tasks.rs b/core/src/core_tests/activity_tasks.rs
@@ -8,12 +8,14 @@ use crate::{
     ActivityHeartbeat, MetricsContext, Worker, WorkerConfigBuilder,
 };
 use futures::FutureExt;
-use std::sync::Arc;
 use std::{
     cell::RefCell,
     collections::{hash_map::Entry, HashMap, VecDeque},
     rc::Rc,
-    sync::atomic::{AtomicUsize, Ordering},
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
     time::Duration,
 };
 use temporal_client::mocks::{mock_gateway, mock_manual_gateway};
diff --git a/core/src/workflow/machines/workflow_machines.rs b/core/src/workflow/machines/workflow_machines.rs
@@ -467,7 +467,6 @@ impl WorkflowMachines {
                     attrs,
                 )) = event.attributes
                 {
-                    self.run_id = attrs.original_execution_run_id.clone();
                     if let Some(st) = event.event_time {
                         let as_systime: SystemTime = st.try_into()?;
                         self.workflow_start_time = Some(as_systime);
diff --git a/core/src/workflow/workflow_tasks/mod.rs b/core/src/workflow/workflow_tasks/mod.rs
@@ -309,6 +309,7 @@ impl WorkflowTaskManager {
             task_token = %&work.task_token,
             history_length = %work.history.events.len(),
             attempt = %work.attempt,
+            run_id = %work.workflow_execution.run_id,
             "Applying new workflow task from server"
         );
         let task_start_time = Instant::now();
diff --git a/sdk-core-protos/src/lib.rs b/sdk-core-protos/src/lib.rs
@@ -260,8 +260,10 @@ pub mod coresdk {
             coresdk::{AsJsonPayloadExt, IntoPayloadsExt},
             temporal::api::common::v1::{Payload as ApiPayload, Payloads},
         };
-        use std::collections::HashMap;
-        use std::fmt::{Display, Formatter};
+        use std::{
+            collections::HashMap,
+            fmt::{Display, Formatter},
+        };
 
         impl<T> From<T> for Payload
         where
diff --git a/test-utils/src/lib.rs b/test-utils/src/lib.rs
@@ -11,7 +11,7 @@ use std::{
     convert::TryFrom, env, future::Future, net::SocketAddr, path::PathBuf, sync::Arc,
     time::Duration,
 };
-use temporal_client::WorkflowOptions;
+use temporal_client::{RetryGateway, ServerGateway, ServerGatewayApis, WorkflowOptions};
 use temporal_sdk::TestRustWorker;
 use temporal_sdk_core::{
     init_replay_worker, init_worker, replay::mock_gateway_from_history, telemetry_init,
@@ -29,6 +29,7 @@ use temporal_sdk_core_protos::{
     },
     temporal::api::history::v1::History,
 };
+use tokio::sync::OnceCell;
 use url::Url;
 
 pub const NAMESPACE: &str = "default";
@@ -77,7 +78,11 @@ pub struct CoreWfStarter {
     telemetry_options: TelemetryOptions,
     worker_config: WorkerConfig,
     wft_timeout: Option<Duration>,
-    initted_worker: Option<Arc<dyn Worker>>,
+    initted_worker: OnceCell<InitializedWorker>,
+}
+struct InitializedWorker {
+    worker: Arc<dyn Worker>,
+    client: Arc<RetryGateway<ServerGateway>>,
 }
 
 impl CoreWfStarter {
@@ -99,7 +104,7 @@ impl CoreWfStarter {
                 .build()
                 .unwrap(),
             wft_timeout: None,
-            initted_worker: None,
+            initted_worker: OnceCell::new(),
         }
     }
 
@@ -116,16 +121,11 @@ impl CoreWfStarter {
     }
 
     pub async fn get_worker(&mut self) -> Arc<dyn Worker> {
-        if self.initted_worker.is_none() {
-            telemetry_init(&self.telemetry_options).expect("Telemetry inits cleanly");
-            let gateway = get_integ_server_options()
-                .connect(self.worker_config.namespace.clone(), None)
-                .await
-                .expect("Must connect");
-            let worker = init_worker(self.worker_config.clone(), gateway);
-            self.initted_worker = Some(Arc::new(worker));
-        }
-        self.initted_worker.as_ref().unwrap().clone()
+        self.get_or_init().await.worker.clone()
+    }
+
+    pub async fn get_client(&mut self) -> Arc<RetryGateway<ServerGateway>> {
+        self.get_or_init().await.client.clone()
     }
 
     /// Start the workflow defined by the builder and return run id
@@ -137,13 +137,12 @@ impl CoreWfStarter {
     pub async fn start_wf_with_id(&self, workflow_id: String, mut opts: WorkflowOptions) -> String {
         opts.task_timeout = opts.task_timeout.or(self.wft_timeout);
         self.initted_worker
-            .as_ref()
+            .get()
             .expect(
-                "Core must be initted before starting a workflow.\
-                             Tests must call `get_core` first.",
+                "Worker must be initted before starting a workflow.\
+                             Tests must call `get_worker` first.",
             )
-            .as_ref()
-            .server_gateway()
+            .client
             .start_workflow(
                 vec![],
                 self.worker_config.task_queue.clone(),
@@ -218,6 +217,25 @@ impl CoreWfStarter {
         self.wft_timeout = Some(timeout);
         self
     }
+
+    async fn get_or_init(&mut self) -> &InitializedWorker {
+        self.initted_worker
+            .get_or_init(|| async {
+                telemetry_init(&self.telemetry_options).expect("Telemetry inits cleanly");
+                let gateway = Arc::new(
+                    get_integ_server_options()
+                        .connect(self.worker_config.namespace.clone(), None)
+                        .await
+                        .expect("Must connect"),
+                );
+                let worker = init_worker(self.worker_config.clone(), gateway.clone());
+                InitializedWorker {
+                    worker: Arc::new(worker),
+                    client: gateway,
+                }
+            })
+            .await
+    }
 }
 
 pub fn get_integ_server_options() -> ServerGatewayOptions {
diff --git a/tests/integ_tests/workflow_tests.rs b/tests/integ_tests/workflow_tests.rs
@@ -7,6 +7,7 @@ mod determinism;
 mod local_activities;
 mod patches;
 mod replay;
+mod resets;
 mod signals;
 mod stickyness;
 mod timers;
diff --git a/tests/integ_tests/workflow_tests/resets.rs b/tests/integ_tests/workflow_tests/resets.rs
@@ -0,0 +1,84 @@
+use futures::StreamExt;
+use std::{sync::Arc, time::Duration};
+use temporal_client::{ServerGatewayApis, WorkflowOptions, WorkflowService};
+use temporal_sdk::WfContext;
+use temporal_sdk_core_protos::temporal::api::{
+    common::v1::WorkflowExecution, workflowservice::v1::ResetWorkflowExecutionRequest,
+};
+use temporal_sdk_core_test_utils::{CoreWfStarter, NAMESPACE};
+use tokio::sync::Notify;
+
+const POST_RESET_SIG: &str = "post-reset";
+
+#[tokio::test]
+async fn reset_workflow() {
+    let wf_name = "reset_me_wf";
+    let mut starter = CoreWfStarter::new(wf_name);
+    let mut worker = starter.worker().await;
+    let notify = Arc::new(Notify::new());
+
+    let wf_notify = notify.clone();
+    worker.register_wf(wf_name.to_owned(), move |ctx: WfContext| {
+        let notify = wf_notify.clone();
+        async move {
+            // Make a couple workflow tasks
+            ctx.timer(Duration::from_secs(1)).await;
+            ctx.timer(Duration::from_secs(1)).await;
+            // Tell outer scope to send the reset
+            notify.notify_one();
+            let _ = ctx
+                .make_signal_channel(POST_RESET_SIG)
+                .next()
+                .await
+                .unwrap();
+            Ok(().into())
+        }
+    });
+
+    let run_id = worker
+        .submit_wf(
+            wf_name.to_owned(),
+            wf_name.to_owned(),
+            vec![],
+            WorkflowOptions::default(),
+        )
+        .await
+        .unwrap();
+
+    let client = starter.get_client().await;
+    let resetter_fut = async {
+        notify.notified().await;
+        // Do the reset
+        client
+            .get_client()
+            .raw_retry_client()
+            .reset_workflow_execution(ResetWorkflowExecutionRequest {
+                namespace: NAMESPACE.to_owned(),
+                workflow_execution: Some(WorkflowExecution {
+                    workflow_id: wf_name.to_owned(),
+                    run_id: run_id.clone(),
+                }),
+                // End of first WFT
+                workflow_task_finish_event_id: 4,
+                request_id: "test-req-id".to_owned(),
+                ..Default::default()
+            })
+            .await
+            .unwrap();
+
+        // Unblock the workflow by sending the signal. Run ID will have changed after reset so
+        // we use empty run id
+        client
+            .signal_workflow_execution(
+                wf_name.to_owned(),
+                "".to_owned(),
+                POST_RESET_SIG.to_owned(),
+                None,
+            )
+            .await
+            .unwrap();
+    };
+    let run_fut = worker.run_until_done();
+    let (_, rr) = tokio::join!(resetter_fut, run_fut);
+    rr.unwrap();
+}

Original file line number	Diff line number	Diff line change
`@@ -467,7 +467,6 @@ impl WorkflowMachines {`
`467`	`467`	`attrs,`
`468`	`468`	`)) = event.attributes`
`469`	`469`	`{`
`470`		`- self.run_id = attrs.original_execution_run_id.clone();`
`471`	`470`	`if let Some(st) = event.event_time {`
`472`	`471`	`let as_systime: SystemTime = st.try_into()?;`
`473`	`472`	`self.workflow_start_time = Some(as_systime);`