diff --git a/crates/sdk-core/CHANGELOG.md b/crates/sdk-core/CHANGELOG.md index 183eb44691..18eddd03db 100644 --- a/crates/sdk-core/CHANGELOG.md +++ b/crates/sdk-core/CHANGELOG.md @@ -33,6 +33,13 @@ relevant information. ## Unreleased +### Added +* A worker can take a `WorkerExtension` whose `WorkflowTaskInterceptor` runs before and after + each workflow task completion Core sends, including the local activity heartbeat completion + and the completion of output buffered before an eviction. An error from `before_complete`, or + a hook that misses the task's deadline, fails the workflow task. Core waits at most five + seconds on `after_complete`. The replay worker never calls the hooks. + ### Security * Ephemeral server downloads now use `zip` 9.0 for ZIP extraction, replacing the 8.x dependency [yanked upstream for security issues](https://github.com/zip-rs/zip2/issues/1075#issuecomment-6209381974). diff --git a/crates/sdk-core/src/core_tests/mod.rs b/crates/sdk-core/src/core_tests/mod.rs index 6dd4479bb7..881d35d7c1 100644 --- a/crates/sdk-core/src/core_tests/mod.rs +++ b/crates/sdk-core/src/core_tests/mod.rs @@ -5,6 +5,7 @@ mod replay_flag; mod updates; mod workers; mod workflow_cancels; +mod workflow_task_interceptor; mod workflow_tasks; use crate::{ diff --git a/crates/sdk-core/src/core_tests/workflow_task_interceptor.rs b/crates/sdk-core/src/core_tests/workflow_task_interceptor.rs new file mode 100644 index 0000000000..1e2b38be0d --- /dev/null +++ b/crates/sdk-core/src/core_tests/workflow_task_interceptor.rs @@ -0,0 +1,628 @@ +//! The hooks a [WorkflowTaskInterceptor] gets around each workflow task completion. + +use crate::{ + WorkerExtension, WorkflowTaskCompletionInfo, WorkflowTaskFailedInfo, + WorkflowTaskInterceptError, WorkflowTaskInterceptor, init_replay_worker, + replay::{HistoryForReplay, ReplayWorkerInput, TestHistoryBuilder, canned_histories}, + test_help::{ + MockPollCfg, PollWFTRespExt, ResponseType, WorkerExt, WorkerTestHelpers, + build_mock_pollers, hist_to_poll_resp, mock_worker, schedule_local_activity_cmd, + start_timer_cmd, test_worker_cfg, + }, + worker::client::mocks::mock_worker_client, +}; +use futures_util::stream; +use std::{ + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; +use temporalio_common::protos::{ + coresdk::{ + ActivityTaskCompletion, + activity_result::ActivityExecutionResult, + activity_task::activity_task, + workflow_activation::{WorkflowActivation, workflow_activation_job}, + workflow_commands::{ + ActivityCancellationType, CompleteWorkflowExecution, UpdateResponse, + update_response::Response, workflow_command, + }, + workflow_completion::WorkflowActivationCompletion, + }, + temporal::api::{ + enums::v1::EventType, workflowservice::v1::RespondWorkflowTaskCompletedResponse, + }, +}; + +/// What the hooks and the mock server saw, in order. +#[derive(Default)] +struct Recorder { + log: Mutex>, + infos: Mutex>, + /// `before_complete` fails for the tasks with these started event ids. + fail_before: Vec, + /// `before_complete` waits this long first. + before_delay: Option, + /// `after_complete` never returns for the tasks with these started event ids. + hang_after: Vec, +} + +impl Recorder { + fn push(&self, entry: impl Into) { + self.log.lock().unwrap().push(entry.into()); + } + + fn log(&self) -> Vec { + self.log.lock().unwrap().clone() + } + + fn infos(&self) -> Vec { + self.infos.lock().unwrap().clone() + } +} + +#[async_trait::async_trait] +impl WorkflowTaskInterceptor for Recorder { + async fn before_complete( + &self, + task: &WorkflowTaskCompletionInfo, + ) -> Result<(), WorkflowTaskInterceptError> { + self.push(format!("before {}", task.started_event_id)); + self.infos.lock().unwrap().push(task.clone()); + if let Some(delay) = self.before_delay { + tokio::time::sleep(delay).await; + } + if self.fail_before.contains(&task.started_event_id) { + return Err(WorkflowTaskInterceptError::new("store down")); + } + Ok(()) + } + + async fn after_complete(&self, task: &WorkflowTaskCompletionInfo, accepted: bool) { + self.push(format!("after {} {accepted}", task.started_event_id)); + if self.hang_after.contains(&task.started_event_id) { + futures_util::future::pending::<()>().await; + } + } + + fn on_task_failed(&self, task: &WorkflowTaskFailedInfo) { + self.push(format!("failed {} {}", task.attempt, task.reported)); + } + + fn on_evicted(&self, _run_id: &str) { + self.push("evicted"); + } +} + +struct Extension(Arc); + +impl WorkerExtension for Extension { + fn name(&self) -> &str { + "recorder" + } + + fn workflow_task_interceptor(&self) -> Option> { + Some(self.0.clone()) + } +} + +/// Builds a worker with the recorder as its interceptor. The mock server logs each completion it +/// gets, and answers with `answer`. +fn worker_with( + mut cfg: MockPollCfg, + recorder: &Arc, + mut answer: impl FnMut(usize) -> Result + + Send + + 'static, +) -> crate::Worker { + // An empty poll stream would start shutdown before the run is evicted. + cfg.make_poll_stream_interminable = true; + let log = recorder.clone(); + let mut completions = 0; + cfg.completion_mock_fn = Some(Box::new(move |c| { + completions += 1; + log.push(format!( + "sent commands={} force_new={}", + c.commands.len(), + c.force_create_new_workflow_task + )); + answer(completions) + })); + let mut mock = build_mock_pollers(cfg); + let extension: Arc = Arc::new(Extension(recorder.clone())); + mock.worker_cfg(|wc| { + wc.max_cached_workflows = 1; + wc.extension = Some(extension); + }); + mock_worker(mock) +} + +fn accept(_: usize) -> Result { + Ok(Default::default()) +} + +async fn complete( + worker: &crate::Worker, + activation: &WorkflowActivation, + commands: Vec, +) { + worker + .complete_workflow_activation(WorkflowActivationCompletion::from_cmds( + activation.run_id.clone(), + commands, + )) + .await + .unwrap(); +} + +/// Completes the run, then answers the eviction that follows. +async fn complete_and_evict(worker: &crate::Worker, run_id: &str) { + worker.complete_execution(run_id).await; + let eviction = worker.poll_workflow_activation().await.unwrap(); + assert_matches!( + eviction.jobs[0].variant, + Some(workflow_activation_job::Variant::RemoveFromCache(_)) + ); + worker + .complete_workflow_activation(WorkflowActivationCompletion::empty(eviction.run_id)) + .await + .unwrap(); +} + +fn timer_then_complete() -> TestHistoryBuilder { + let mut t = TestHistoryBuilder::default(); + t.add_wfe_started_with_wft_timeout(Duration::from_secs(10)); + t.add_full_wf_task(); + let timer_started = t.add_by_type(EventType::TimerStarted); + t.add_timer_fired(timer_started, "1".to_string()); + t.add_workflow_task_scheduled_and_started(); + t +} + +#[tokio::test] +async fn hooks_run_before_and_after_each_completion() { + let recorder = Arc::new(Recorder::default()); + let worker = worker_with( + MockPollCfg::from_resp_batches( + "fakeid", + timer_then_complete(), + [1, 2], + mock_worker_client(), + ), + &recorder, + accept, + ); + let first = worker.poll_workflow_activation().await.unwrap(); + let received = Instant::now(); + complete( + &worker, + &first, + vec![start_timer_cmd(1, Duration::from_secs(1))], + ) + .await; + let second = worker.poll_workflow_activation().await.unwrap(); + complete_and_evict(&worker, &second.run_id).await; + worker.drain_pollers_and_shutdown().await; + + assert_eq!( + recorder.log(), + [ + "before 3", + "sent commands=1 force_new=false", + "after 3 true", + "before 8", + "sent commands=1 force_new=false", + "after 8 true", + "evicted", + ] + ); + let infos = recorder.infos(); + assert_eq!(infos[0].run_id, first.run_id); + assert_eq!(infos[0].workflow_id, "fakeid"); + assert_eq!( + infos[0].namespace, + test_worker_cfg().build().unwrap().namespace + ); + assert_eq!(infos[0].attempt, 1); + assert!(infos[0].started_time.is_some()); + assert_eq!(infos[0].closes_run, None); + assert_eq!(infos[1].closes_run, Some(crate::RunCloseKind::Completed)); + // The hook gets nine of the ten seconds of the task timeout, from the task's receipt. + let left = infos[0].deadline - received; + assert!(left <= Duration::from_secs(9), "{left:?}"); + assert!(left > Duration::from_secs(8), "{left:?}"); + assert!(infos[1].deadline > infos[0].deadline); +} + +#[tokio::test] +async fn an_error_fails_the_task_and_evicts_the_run() { + let recorder = Arc::new(Recorder { + fail_before: vec![3], + ..Default::default() + }); + let mut cfg = + MockPollCfg::from_resp_batches("fakeid", timer_then_complete(), [1], mock_worker_client()); + cfg.num_expected_fails = 1; + cfg.num_expected_completions = Some(0.into()); + let worker = worker_with(cfg, &recorder, accept); + + let first = worker.poll_workflow_activation().await.unwrap(); + complete( + &worker, + &first, + vec![start_timer_cmd(1, Duration::from_secs(1))], + ) + .await; + let eviction = worker.poll_workflow_activation().await.unwrap(); + assert_matches!( + eviction.jobs[0].variant, + Some(workflow_activation_job::Variant::RemoveFromCache(_)) + ); + worker + .complete_workflow_activation(WorkflowActivationCompletion::empty(eviction.run_id)) + .await + .unwrap(); + worker.drain_pollers_and_shutdown().await; + + assert_eq!(recorder.log(), ["before 3", "failed 1 true", "evicted"]); +} + +#[tokio::test] +async fn a_hook_that_misses_the_deadline_fails_the_task() { + let recorder = Arc::new(Recorder { + before_delay: Some(Duration::from_secs(2)), + ..Default::default() + }); + let mut t = TestHistoryBuilder::default(); + t.add_wfe_started_with_wft_timeout(Duration::from_secs(1)); + t.add_workflow_task_scheduled_and_started(); + let mut cfg = MockPollCfg::from_resp_batches("fakeid", t, [1], mock_worker_client()); + cfg.num_expected_fails = 1; + cfg.num_expected_completions = Some(0.into()); + let worker = worker_with(cfg, &recorder, accept); + + let first = worker.poll_workflow_activation().await.unwrap(); + let started = Instant::now(); + complete( + &worker, + &first, + vec![start_timer_cmd(1, Duration::from_secs(1))], + ) + .await; + assert!(started.elapsed() < Duration::from_secs(2)); + let eviction = worker.poll_workflow_activation().await.unwrap(); + worker + .complete_workflow_activation(WorkflowActivationCompletion::empty(eviction.run_id)) + .await + .unwrap(); + worker.drain_pollers_and_shutdown().await; + + assert_eq!(recorder.log(), ["before 3", "failed 1 true", "evicted"]); +} + +#[tokio::test] +async fn a_refused_completion_is_not_accepted() { + let recorder = Arc::new(Recorder::default()); + let mut cfg = + MockPollCfg::from_resp_batches("fakeid", timer_then_complete(), [1], mock_worker_client()); + cfg.num_expected_completions = Some(1.into()); + let worker = worker_with(cfg, &recorder, |_| { + Err(tonic::Status::invalid_argument("UnhandledCommand")) + }); + + let first = worker.poll_workflow_activation().await.unwrap(); + complete( + &worker, + &first, + vec![start_timer_cmd(1, Duration::from_secs(1))], + ) + .await; + let eviction = worker.poll_workflow_activation().await.unwrap(); + worker + .complete_workflow_activation(WorkflowActivationCompletion::empty(eviction.run_id)) + .await + .unwrap(); + worker.drain_pollers_and_shutdown().await; + + assert_eq!( + recorder.log(), + [ + "before 3", + "sent commands=1 force_new=false", + "after 3 false", + "evicted" + ] + ); +} + +#[tokio::test] +async fn a_dropped_speculative_task_is_not_accepted() { + let wfid = "fakeid"; + let mut t = TestHistoryBuilder::default(); + t.add_by_type(EventType::WorkflowExecutionStarted); + t.add_full_wf_task(); + let timer_started = t.add_by_type(EventType::TimerStarted); + let mut spec_task_hist = t.clone(); + spec_task_hist.add_workflow_task_scheduled_and_started(); + let mut real_hist = t.clone(); + real_hist.add_timer_fired(timer_started, "1".to_string()); + real_hist.add_workflow_task_scheduled_and_started(); + + let update_id = "upd-1"; + let mut speculative_task = hist_to_poll_resp(&spec_task_hist, wfid, ResponseType::OneTask(2)); + speculative_task.add_update_request(update_id, 1); + let cfg = MockPollCfg::from_resp_batches( + wfid, + real_hist, + [ + ResponseType::ToTaskNum(1), + speculative_task.into(), + ResponseType::OneTask(2), + ], + mock_worker_client(), + ); + let recorder = Arc::new(Recorder::default()); + let worker = worker_with(cfg, &recorder, |n| { + let mut r = RespondWorkflowTaskCompletedResponse::default(); + if n == 2 { + // The rejection drops the speculative task from History. + r.reset_history_event_id = 3; + } + Ok(r) + }); + + let first = worker.poll_workflow_activation().await.unwrap(); + complete( + &worker, + &first, + vec![start_timer_cmd(1, Duration::from_secs(1))], + ) + .await; + let update = worker.poll_workflow_activation().await.unwrap(); + complete( + &worker, + &update, + vec![ + UpdateResponse { + protocol_instance_id: update_id.to_string(), + response: Some(Response::Rejected(Default::default())), + } + .into(), + ], + ) + .await; + let fired = worker.poll_workflow_activation().await.unwrap(); + complete_and_evict(&worker, &fired.run_id).await; + worker.drain_pollers_and_shutdown().await; + + assert_eq!( + recorder.log(), + [ + "before 3", + "sent commands=1 force_new=false", + "after 3 true", + "before 7", + "sent commands=0 force_new=false", + "after 7 false", + "before 8", + "sent commands=1 force_new=false", + "after 8 true", + "evicted", + ] + ); +} + +#[tokio::test] +async fn the_local_activity_heartbeat_completion_runs_the_hooks() { + let mut t = TestHistoryBuilder::default(); + t.add_wfe_started_with_wft_timeout(Duration::from_secs(1)); + t.add_full_wf_task(); + t.add_workflow_task_scheduled_and_started(); + let cfg = MockPollCfg::from_resp_batches("fakeid", t, [1, 2], mock_worker_client()); + let recorder = Arc::new(Recorder::default()); + let worker = worker_with(cfg, &recorder, accept); + + let first = worker.poll_workflow_activation().await.unwrap(); + complete( + &worker, + &first, + vec![schedule_local_activity_cmd( + 1, + "1", + ActivityCancellationType::TryCancel, + Duration::from_secs(60), + )], + ) + .await; + let la = worker.poll_activity_task().await.unwrap(); + assert_matches!(&la.variant, Some(activity_task::Variant::Start(s)) if s.is_local); + // Core heartbeats the task at 80% of its one-second timeout, with no lang activation. + let resolved = { + let wf_poll = worker.poll_workflow_activation(); + tokio::pin!(wf_poll); + tokio::select! { + act = &mut wf_poll => panic!("Unexpected activation {act:?}"), + _ = tokio::time::sleep(Duration::from_millis(1500)) => {} + } + assert_eq!( + recorder.log(), + ["before 3", "sent commands=0 force_new=true", "after 3 true"] + ); + worker + .complete_activity_task(ActivityTaskCompletion { + task_token: la.task_token, + result: Some(ActivityExecutionResult::ok(vec![1].into())), + }) + .await + .unwrap(); + wf_poll.await.unwrap() + }; + complete_and_evict(&worker, &resolved.run_id).await; + worker.drain_pollers_and_shutdown().await; + + assert_eq!( + recorder.log(), + [ + "before 3", + "sent commands=0 force_new=true", + "after 3 true", + "before 6", + "sent commands=2 force_new=false", + "after 6 true", + "evicted", + ] + ); +} + +#[tokio::test] +async fn an_eviction_with_output_runs_the_hooks_before_on_evicted() { + let mut t = TestHistoryBuilder::default(); + // Keep the LA heartbeat from firing mid-test. + t.add_wfe_started_with_wft_timeout(Duration::from_secs(60)); + t.add_workflow_task_scheduled_and_started(); + let cfg = MockPollCfg::from_resp_batches("fakeid", t, [1], mock_worker_client()); + let recorder = Arc::new(Recorder::default()); + let worker = worker_with(cfg, &recorder, accept); + + let first = worker.poll_workflow_activation().await.unwrap(); + let run_id = first.run_id.clone(); + let la_cmd = |seq: u32| { + schedule_local_activity_cmd( + seq, + &seq.to_string(), + ActivityCancellationType::WaitCancellationCompleted, + Duration::from_secs(60), + ) + }; + complete(&worker, &first, vec![la_cmd(1), la_cmd(2)]).await; + let mut tokens = vec![]; + for _ in 0..2 { + let task = worker.poll_activity_task().await.unwrap(); + let Some(activity_task::Variant::Start(start)) = &task.variant else { + panic!("Expected LA start, got {task:?}"); + }; + tokens.push((start.activity_id.clone(), task.task_token)); + } + tokens.sort(); + let (_, first_token) = tokens.remove(0); + worker + .complete_activity_task(ActivityTaskCompletion { + task_token: first_token, + result: Some(ActivityExecutionResult::ok(vec![1].into())), + }) + .await + .unwrap(); + // The marker of the first LA is output that waits on the second one. + let resolved = worker.poll_workflow_activation().await.unwrap(); + complete(&worker, &resolved, vec![]).await; + assert!(recorder.log().is_empty()); + + worker.request_workflow_eviction(&run_id); + let cancel = worker.poll_activity_task().await.unwrap(); + assert_matches!(cancel.variant, Some(activity_task::Variant::Cancel(_))); + worker + .complete_activity_task(ActivityTaskCompletion { + task_token: cancel.task_token, + result: Some(ActivityExecutionResult::cancel_from_details(None)), + }) + .await + .unwrap(); + let eviction = worker.poll_workflow_activation().await.unwrap(); + assert_matches!( + eviction.jobs[0].variant, + Some(workflow_activation_job::Variant::RemoveFromCache(_)) + ); + worker + .complete_workflow_activation(WorkflowActivationCompletion::empty(eviction.run_id)) + .await + .unwrap(); + worker.drain_pollers_and_shutdown().await; + + assert_eq!( + recorder.log(), + [ + "before 3", + "sent commands=1 force_new=true", + "after 3 true", + "evicted" + ] + ); +} + +#[tokio::test] +async fn the_replay_worker_never_calls_the_hooks() { + let recorder = Arc::new(Recorder::default()); + let extension: Arc = Arc::new(Extension(recorder.clone())); + let mut cfg = test_worker_cfg().build().unwrap(); + cfg.extension = Some(extension); + let worker = init_replay_worker(ReplayWorkerInput::new( + cfg, + stream::iter([HistoryForReplay::from(canned_histories::single_timer("1"))]), + )) + .unwrap(); + + let first = worker.poll_workflow_activation().await.unwrap(); + complete( + &worker, + &first, + vec![start_timer_cmd(1, Duration::from_secs(1))], + ) + .await; + let fired = worker.poll_workflow_activation().await.unwrap(); + complete( + &worker, + &fired, + vec![CompleteWorkflowExecution { result: None }.into()], + ) + .await; + while let Ok(activation) = worker.poll_workflow_activation().await { + worker + .complete_workflow_activation(WorkflowActivationCompletion::empty(activation.run_id)) + .await + .unwrap(); + } + worker.shutdown().await; + + assert!(recorder.log().is_empty()); +} + +#[tokio::test] +async fn a_hung_after_complete_does_not_stall_the_run() { + let recorder = Arc::new(Recorder { + hang_after: vec![3], + ..Default::default() + }); + let worker = worker_with( + MockPollCfg::from_resp_batches( + "fakeid", + timer_then_complete(), + [1, 2], + mock_worker_client(), + ), + &recorder, + accept, + ); + + let first = worker.poll_workflow_activation().await.unwrap(); + complete( + &worker, + &first, + vec![start_timer_cmd(1, Duration::from_secs(1))], + ) + .await; + let second = worker.poll_workflow_activation().await.unwrap(); + complete_and_evict(&worker, &second.run_id).await; + worker.drain_pollers_and_shutdown().await; + + assert_eq!( + recorder.log(), + [ + "before 3", + "sent commands=1 force_new=false", + "after 3 true", + "before 8", + "sent commands=1 force_new=false", + "after 8 true", + "evicted", + ] + ); +} diff --git a/crates/sdk-core/src/lib.rs b/crates/sdk-core/src/lib.rs index cf0f68c112..3e2b6e27e9 100644 --- a/crates/sdk-core/src/lib.rs +++ b/crates/sdk-core/src/lib.rs @@ -50,6 +50,10 @@ pub use worker::{ WorkerConfigBuilder, WorkerTuner, WorkerValidationError, WorkerVersioningStrategy, WorkflowErrorType, WorkflowSlotKind, }; +pub use worker::{ + RunCloseKind, WorkerExtension, WorkflowTaskCompletionInfo, WorkflowTaskFailedInfo, + WorkflowTaskInterceptError, WorkflowTaskInterceptor, +}; use crate::{ replay::{HistoryForReplay, ReplayWorkerInput}, diff --git a/crates/sdk-core/src/protosext/mod.rs b/crates/sdk-core/src/protosext/mod.rs index 506b953c65..0544a42eba 100644 --- a/crates/sdk-core/src/protosext/mod.rs +++ b/crates/sdk-core/src/protosext/mod.rs @@ -12,7 +12,7 @@ use std::{ collections::HashMap, convert::TryFrom, fmt::{Debug, Display, Formatter}, - time::{Duration, SystemTime}, + time::{Duration, Instant, SystemTime}, }; use temporalio_common::protos::{ constants::{LOCAL_ACTIVITY_MARKER_NAME, PATCH_MARKER_NAME}, @@ -57,6 +57,10 @@ pub(crate) struct ValidPollWFTQResponse { pub(crate) attempt: u32, pub(crate) previous_started_event_id: i64, pub(crate) started_event_id: i64, + /// When the server started the task. From the `started_time` field. + pub(crate) started_time: Option, + /// When the worker received the task, before it waits on a busy run or a history fetch. + pub(crate) received: Instant, /// If this is present, `history` will be empty. This is not a very "tight" design, but it's /// enforced at construction time. From the `query` field. pub(crate) legacy_query: Option, @@ -106,6 +110,7 @@ impl TryFrom for ValidPollWFTQResponse { attempt, previous_started_event_id, started_event_id, + started_time, query, queries, messages, @@ -130,6 +135,8 @@ impl TryFrom for ValidPollWFTQResponse { attempt: attempt as u32, previous_started_event_id, started_event_id, + started_time: started_time.and_then(|t| t.try_into().ok()), + received: Instant::now(), legacy_query: query, query_requests, messages, diff --git a/crates/sdk-core/src/replay/history_info.rs b/crates/sdk-core/src/replay/history_info.rs index e37e26dab2..2dc5a33ba2 100644 --- a/crates/sdk-core/src/replay/history_info.rs +++ b/crates/sdk-core/src/replay/history_info.rs @@ -177,6 +177,12 @@ impl HistoryInfo { }), previous_started_event_id: self.previous_started_event_id, started_event_id: self.workflow_task_started_event_id, + // The server sets this to the time of the task's started event. + started_time: self + .events + .iter() + .find(|e| e.event_id == self.workflow_task_started_event_id) + .and_then(|e| e.event_time), ..Default::default() } } diff --git a/crates/sdk-core/src/replay/mod.rs b/crates/sdk-core/src/replay/mod.rs index 6e4cadc506..11990d094a 100644 --- a/crates/sdk-core/src/replay/mod.rs +++ b/crates/sdk-core/src/replay/mod.rs @@ -79,6 +79,8 @@ where self.config.workflow_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(1)); self.config.task_types = WorkerTaskTypes::workflow_only(); self.config.skip_client_worker_set_check = true; + // Replay must not repeat the side effects of the hooks. + self.config.extension = None; let historator = Historator::new(self.history_stream); let post_activate = historator.get_post_activate_hook(); let shutdown_tok = historator.get_shutdown_setter(); diff --git a/crates/sdk-core/src/worker/extension.rs b/crates/sdk-core/src/worker/extension.rs new file mode 100644 index 0000000000..8d948449ae --- /dev/null +++ b/crates/sdk-core/src/worker/extension.rs @@ -0,0 +1,131 @@ +//! Hooks that let a plugin run code at fixed points of the worker's workflow task lifecycle. + +use std::{ + fmt::Debug, + sync::Arc, + time::{Instant, SystemTime}, +}; + +/// A component that lang passes to the worker through [crate::WorkerConfig::extension]. +/// +/// Core never inspects the extension beyond the hooks it asks for here. +pub trait WorkerExtension: Send + Sync + 'static { + /// A name that identifies the extension in logs. + fn name(&self) -> &str; + + /// The interceptor Core calls around each workflow task completion, if any. Core may ask + /// more than once per worker, so return the same interceptor each time. + fn workflow_task_interceptor(&self) -> Option> { + None + } +} + +/// Hooks around each `RespondWorkflowTaskCompleted` call that Core makes for a run. +/// +/// Core calls the hooks of one run in order, and the next workflow task of that run waits until +/// they return. Core waits at most five seconds on [WorkflowTaskInterceptor::after_complete], so +/// long work such as retries must run in the background. The replay worker never calls the hooks. +/// +/// A completion that lang sends waits only on its own run. Core sends the local activity heartbeat +/// completion from within the poll for the next activation, so while its hooks run, that poll +/// gives no activation to other runs. +#[async_trait::async_trait] +pub trait WorkflowTaskInterceptor: Send + Sync + 'static { + /// Called before every `RespondWorkflowTaskCompleted` for the run. This includes the + /// completions that have no lang activation, which are the local activity heartbeat + /// completion and the completion of output buffered before an eviction. + /// + /// An error, or a hook that has not returned by `task.deadline`, fails the workflow task + /// instead of completing it, and the run is evicted. + async fn before_complete( + &self, + task: &WorkflowTaskCompletionInfo, + ) -> Result<(), WorkflowTaskInterceptError>; + + /// Called after the server answered the completion, or after the call failed. + /// + /// `accepted` is false when the call failed, and also when the answer has a + /// `reset_history_event_id`, which means the server dropped a speculative task. + async fn after_complete(&self, task: &WorkflowTaskCompletionInfo, accepted: bool); + + /// Called when Core fails a workflow task of the run, whether or not it reports the failure. + fn on_task_failed(&self, _task: &WorkflowTaskFailedInfo) {} + + /// Called when Core removes the run from its cache. When an eviction sends buffered output, + /// this comes after [WorkflowTaskInterceptor::after_complete] for that completion. + fn on_evicted(&self, _run_id: &str) {} +} + +/// What Core knows about a workflow task when it is about to complete it. +#[derive(Clone, Debug, bon::Builder)] +#[builder(on(String, into))] +#[non_exhaustive] +pub struct WorkflowTaskCompletionInfo { + /// The namespace of the run. + pub namespace: String, + /// The workflow id of the run. + pub workflow_id: String, + /// The run id. + pub run_id: String, + /// The first run id of the run chain. + pub first_execution_run_id: String, + /// The workflow type. + pub workflow_type: String, + /// The id of the task's `WorkflowTaskStarted` event, from the poll response. The server writes + /// no events for transient attempts, so attempts after the second of one task share an id. + pub started_event_id: i64, + /// The task's attempt, from the poll response. + pub attempt: u32, + /// The time the server started the task, from the poll response. Unset if the server did not + /// send it. + pub started_time: Option, + /// Set when the completion has a command that closes the run. Continue-as-new does not + /// count. + pub closes_run: Option, + /// When the hook must have returned, so that the task does not time out on the server. Core + /// measures it from the local receipt of the task, with a margin. + pub deadline: Instant, +} + +/// The command with which a completion closes its run. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub enum RunCloseKind { + /// `CompleteWorkflowExecution` + Completed, + /// `FailWorkflowExecution`. A retry policy can still start a new run in the chain. + Failed, + /// `CancelWorkflowExecution` + Cancelled, +} + +/// What Core knows about a workflow task that it fails. +#[derive(Clone, Debug, bon::Builder)] +#[builder(on(String, into))] +#[non_exhaustive] +pub struct WorkflowTaskFailedInfo { + /// The run id. + pub run_id: String, + /// The attempt of the failed task. + pub attempt: u32, + /// True when Core sends `RespondWorkflowTaskFailed`. Core reports only first attempts and + /// lets later ones time out. + pub reported: bool, +} + +/// The error a [WorkflowTaskInterceptor] returns to fail a workflow task. +#[derive(Debug, thiserror::Error)] +#[error("{message}")] +pub struct WorkflowTaskInterceptError { + /// The message that goes into the task failure. + pub message: String, +} + +impl WorkflowTaskInterceptError { + /// Creates an error with a message. + pub fn new(message: impl Into) -> Self { + Self { + message: message.into(), + } + } +} diff --git a/crates/sdk-core/src/worker/mod.rs b/crates/sdk-core/src/worker/mod.rs index d31addaf47..085b54ffaa 100644 --- a/crates/sdk-core/src/worker/mod.rs +++ b/crates/sdk-core/src/worker/mod.rs @@ -1,5 +1,6 @@ mod activities; pub(crate) mod client; +mod extension; pub(crate) mod heartbeat; mod nexus; mod slot_provider; @@ -10,6 +11,10 @@ mod workflow; /// the error limit, shared so every conversion site reports the same identifier. pub(crate) const PAYLOADS_TOO_LARGE_FAILURE_TYPE: &str = "PayloadsTooLarge"; +pub use extension::{ + RunCloseKind, WorkerExtension, WorkflowTaskCompletionInfo, WorkflowTaskFailedInfo, + WorkflowTaskInterceptError, WorkflowTaskInterceptor, +}; use temporalio_client::{Connection, PayloadErrorLimits, worker::NamespaceDescriptionSource}; use temporalio_common::{ protos::{ @@ -299,6 +304,10 @@ pub struct WorkerConfig { /// NOTE: Experimental #[builder(default = false)] pub disable_payload_error_limit: bool, + + /// A plugin component that Core calls at fixed points of the workflow task lifecycle. See + /// [WorkerExtension]. + pub extension: Option>, } impl WorkerConfig { diff --git a/crates/sdk-core/src/worker/workflow/driven_workflow.rs b/crates/sdk-core/src/worker/workflow/driven_workflow.rs index 9a911d1136..470c8d82fd 100644 --- a/crates/sdk-core/src/worker/workflow/driven_workflow.rs +++ b/crates/sdk-core/src/worker/workflow/driven_workflow.rs @@ -50,6 +50,7 @@ impl DrivenWorkflow { debug!(run_id = %attribs.original_execution_run_id, "Driven WF start"); let started_info = WorkflowStartedInfo { workflow_task_timeout: attribs.workflow_task_timeout.try_into_or_none(), + first_execution_run_id: attribs.first_execution_run_id.clone(), memo: attribs.memo.clone(), search_attrs: attribs.search_attributes.clone(), retry_policy: attribs.retry_policy.clone(), diff --git a/crates/sdk-core/src/worker/workflow/history_update.rs b/crates/sdk-core/src/worker/workflow/history_update.rs index 693a9943b3..003324a37c 100644 --- a/crates/sdk-core/src/worker/workflow/history_update.rs +++ b/crates/sdk-core/src/worker/workflow/history_update.rs @@ -156,6 +156,8 @@ impl HistoryPaginator { let prepared = PreparedWFT { task_token: wft.task_token, attempt: wft.attempt, + started_time: wft.started_time, + received: wft.received, execution: wft.workflow_execution, workflow_type: wft.workflow_type, legacy_query: wft.legacy_query, diff --git a/crates/sdk-core/src/worker/workflow/managed_run.rs b/crates/sdk-core/src/worker/workflow/managed_run.rs index b04163982a..b35089cfc5 100644 --- a/crates/sdk-core/src/worker/workflow/managed_run.rs +++ b/crates/sdk-core/src/worker/workflow/managed_run.rs @@ -4,7 +4,7 @@ use crate::{ internal_flags::CoreInternalFlags, protosext::{WorkflowActivationExt, protocol_messages::IncomingProtocolMessage}, worker::{ - LEGACY_QUERY_ID, LocalActRequest, WorkflowErrorType, + LEGACY_QUERY_ID, LocalActRequest, WorkflowErrorType, WorkflowTaskCompletionInfo, workflow::{ ActivationAction, ActivationCompleteOutcome, ActivationCompleteResult, ActivationOrAuto, BufferedTasks, DrivenWorkflow, EvictionRequestResult, @@ -12,10 +12,11 @@ use crate::{ LocalActivityRequestSink, LocalResolution, NextPageReq, OutstandingActivation, OutstandingTask, PermittedWFT, RequestEvictMsg, RunBasics, ServerCommandsWithWorkflowInfo, TaskStorageMetrics, WFCommand, WFCommandVariant, - WFMachinesError, WFT_HEARTBEAT_TIMEOUT_FRACTION, WFTReportStatus, WftFailureKind, - WorkflowTaskInfo, + WFMachinesError, WFT_HEARTBEAT_TIMEOUT_FRACTION, WFT_INTERCEPTOR_DEADLINE_FRACTION, + WFTReportStatus, WftFailureKind, WorkflowTaskInfo, history_update::HistoryPaginator, machines::{MachinesWFTResponseContent, WorkflowMachines}, + run_close_kind, }, }, }; @@ -40,6 +41,7 @@ use temporalio_common::protos::{ workflow_completion, }, temporal::api::{ + command::v1::Command as ProtoCommand, enums::v1::{VersioningBehavior, WorkflowTaskFailedCause}, failure::v1::Failure, }, @@ -47,6 +49,10 @@ use temporalio_common::protos::{ use tokio::sync::oneshot; use tracing::Span; +/// The server's default workflow task timeout. The start event always has the timeout, so this +/// only keeps the hook's deadline finite if it is missing. +const DEFAULT_WFT_TIMEOUT: Duration = Duration::from_secs(10); + type Result = std::result::Result; pub(super) type RunUpdateAct = Option; @@ -194,6 +200,9 @@ impl ManagedRun { let is_incremental = work.is_incremental(); let wft_info = WorkflowTaskInfo { attempt: work.attempt, + started_event_id: work.update.wft_started_id, + started_time: work.started_time, + received: work.received, task_token: work.task_token, wf_id: work.execution.workflow_id.clone(), }; @@ -1232,7 +1241,7 @@ impl ManagedRun { force_new_wft = true; } - let outcome = if should_respond || has_query_responses { + let mut outcome = if should_respond || has_query_responses { // Broken evictions cannot pass should_respond or contain query responses. Keep this // fallback to avoid sending buffered output if those invariants regress. let (commands, messages) = if self.am_broken && data.activation_was_eviction { @@ -1256,6 +1265,7 @@ impl ManagedRun { sdk_metadata: machines_wft_response.metadata_for_complete(), versioning_behavior: data.versioning_behavior, attempt, + task_info: None, }, metrics: self.metrics.clone(), }) @@ -1263,6 +1273,18 @@ impl ManagedRun { ActivationCompleteOutcome::DoNothing }; let replaying = machines_wft_response.replaying; + if let ActivationCompleteOutcome::ReportWFTSuccess(ServerCommandsWithWorkflowInfo { + action: + ActivationAction::WftComplete { + commands, + task_info, + .. + }, + .. + }) = &mut outcome + { + *task_info = self.completion_info(commands).map(Box::new); + } if matches!(outcome, ActivationCompleteOutcome::ReportWFTSuccess(_)) { self.wfm.wft_completion_reported(); } @@ -1272,6 +1294,41 @@ impl ManagedRun { } } + /// The facts a workflow task interceptor gets about the task being completed. `None` when the + /// worker has no extension, or the run has no task. + fn completion_info(&self, commands: &[ProtoCommand]) -> Option { + self.config + .extension + .as_ref()? + .workflow_task_interceptor()?; + let wft = self.wft.as_ref()?; + let machines = &self.wfm.machines; + let started_info = machines.get_started_info(); + let wft_timeout = started_info + .and_then(|i| i.workflow_task_timeout) + .unwrap_or(DEFAULT_WFT_TIMEOUT); + Some( + WorkflowTaskCompletionInfo::builder() + .namespace(self.config.namespace.clone()) + .workflow_id(machines.workflow_id.clone()) + .run_id(machines.run_id.clone()) + .first_execution_run_id( + started_info + .map(|i| i.first_execution_run_id.clone()) + .unwrap_or_default(), + ) + .workflow_type(machines.workflow_type.clone()) + .started_event_id(wft.info.started_event_id) + .attempt(wft.info.attempt) + .maybe_started_time(wft.info.started_time) + .maybe_closes_run(run_close_kind(commands)) + .deadline( + wft.info.received + wft_timeout.mul_f32(WFT_INTERCEPTOR_DEADLINE_FRACTION), + ) + .build(), + ) + } + /// Pump some local activity requests into the sink, applying any immediate results to the /// workflow machines. fn sink_la_requests( diff --git a/crates/sdk-core/src/worker/workflow/mod.rs b/crates/sdk-core/src/worker/workflow/mod.rs index 7b79d9e442..b833ff9bf5 100644 --- a/crates/sdk-core/src/worker/workflow/mod.rs +++ b/crates/sdk-core/src/worker/workflow/mod.rs @@ -27,7 +27,8 @@ use crate::{ worker::{ ActivitySlotKind, CompleteWfError, LocalActRequest, LocalActivityExecutionResult, LocalActivityResolution, NamespaceCapabilities, PollError, PostActivateHookData, - WorkflowSlotKind, + RunCloseKind, WorkflowSlotKind, WorkflowTaskCompletionInfo, WorkflowTaskFailedInfo, + WorkflowTaskInterceptor, activities::{ActivitiesFromWFTsHandle, LocalActivityManager}, client::{ LegacyQueryResult, REQUEST_TOO_LARGE_KEY, WorkerClient, WorkflowTaskCompletion, @@ -62,7 +63,7 @@ use std::{ result, sync::{Arc, atomic, atomic::AtomicBool}, thread, - time::{Duration, Instant}, + time::{Duration, Instant, SystemTime}, }; use temporalio_client::MESSAGE_TOO_LARGE_KEY; use temporalio_common::{ @@ -88,7 +89,7 @@ use temporalio_common::{ Memo, MeteringMetadata, RetryPolicy, SearchAttributes, WorkflowExecution, }, enums::v1::{VersioningBehavior, WorkflowTaskFailedCause}, - failure::v1::{ApplicationFailureInfo, failure::FailureInfo}, + failure::v1::{ApplicationFailureInfo, Failure as ApiFailure, failure::FailureInfo}, protocol::v1::Message as ProtocolMessage, query::v1::WorkflowQuery, sdk::v1::{EventGroupMarker, UserMetadata, WorkflowTaskCompletedMetadata}, @@ -117,6 +118,11 @@ pub const LEGACY_QUERY_ID: &str = "legacy_query"; /// What percentage of a WFT timeout we are willing to wait before sending a WFT heartbeat when /// necessary. const WFT_HEARTBEAT_TIMEOUT_FRACTION: f32 = 0.8; +/// What percentage of a WFT timeout may pass before `before_complete` must have returned. It is +/// above the heartbeat fraction, so a heartbeat completion still has time for the hook. +const WFT_INTERCEPTOR_DEADLINE_FRACTION: f32 = 0.9; +/// How long Core waits on `after_complete` before it goes on with the run. +const WFT_INTERCEPTOR_AFTER_COMPLETE_TIMEOUT: Duration = Duration::from_secs(5); type Result = result::Result; type BoxedActivationStream = BoxStream<'static, Result>; @@ -147,6 +153,7 @@ pub(crate) struct Workflows { default_versioning_behavior: Option, namespace_capabilities: Arc, shutdown_token: CancellationToken, + wft_interceptor: Option>, } pub(crate) struct WorkflowBasics { @@ -196,6 +203,11 @@ impl Workflows { let default_versioning_behavior = basics.default_versioning_behavior; let namespace_capabilities = basics.namespace_capabilities.clone(); let shutdown_token = basics.shutdown_token.clone(); + let wft_interceptor = basics + .worker_config + .extension + .as_ref() + .and_then(|e| e.workflow_task_interceptor()); let extracted_wft_stream = WFTExtractor::build( client.clone(), basics.worker_config.fetching_concurrency, @@ -293,6 +305,7 @@ impl Workflows { default_versioning_behavior, namespace_capabilities, shutdown_token, + wft_interceptor, } } @@ -382,9 +395,46 @@ impl Workflows { sdk_metadata, mut versioning_behavior, attempt, + task_info, }, metrics: run_metrics, } => { + let interceptor = self.wft_interceptor.as_deref().zip(task_info.as_deref()); + let before = match (interceptor, self.wft_interceptor.is_some()) { + (Some((interceptor, info)), _) => run_before_complete(interceptor, info).await, + (None, false) => Ok(()), + // Completing without the hook would skip it, so fail the task instead. + (None, true) => { + dbg_panic!("Workflow task completion is missing its interceptor info"); + Err("Workflow task completion is missing its interceptor info".to_string()) + } + }; + if let Err(message) = before { + warn!(run_id, error = %message, "Workflow task interceptor failed the task"); + // The machines already count this task as completed, so the run must replay. + self.request_eviction( + run_id, + "Workflow task interceptor failed the task", + EvictionReason::Fatal, + ); + return self + .handle_activation_failed( + run_id, + completion_time, + FailedActivationWFTReport::new( + task_token, + attempt, + WorkflowTaskFailedCause::WorkflowWorkerUnhandledFailure, + Failure { + failure: Some(ApiFailure::application_failure(message, false)), + ..Default::default() + }, + WftFailureKind::Task, + &run_metrics, + ), + ) + .await; + } let reserved_act_permits = self.reserve_activity_slots_for_outgoing_commands(commands.as_mut_slice()); debug!(commands=%commands.display(), query_responses=%query_responses.display(), @@ -431,6 +481,7 @@ impl Workflows { completion.sticky_attributes = sticky_attrs; let mut reset_last_started_to = None; + let mut accepted = false; self.handle_wft_reporting_errs(run_id, || async { match self .client @@ -443,6 +494,8 @@ impl Workflows { } if response.reset_history_event_id > 0 { reset_last_started_to = Some(response.reset_history_event_id); + } else { + accepted = true; } if let Some(wft) = response.workflow_task { *wft_from_complete = Some(validate_wft(wft)?); @@ -498,6 +551,9 @@ impl Workflows { Ok(()) }) .await; + if let Some((interceptor, info)) = interceptor { + run_after_complete(interceptor, info, accepted).await; + } WFTReportStatus::Reported { reset_last_started_to, completion_time, @@ -549,9 +605,11 @@ impl Workflows { WftFailureKind::Task if attempt > 1 => { debug!(run_id=%run_id, attempt, failure=?failure, "Not reporting workflow task failure on non-first attempt"); + self.notify_task_failed(run_id, attempt, false); return WFTReportStatus::DropWft { completion_time }; } WftFailureKind::Task => { + self.notify_task_failed(run_id, attempt, true); warn!(run_id=%run_id, failure=?failure, "Failing workflow task"); self.handle_wft_reporting_errs(run_id, || async { self.client @@ -567,6 +625,18 @@ impl Workflows { } } + fn notify_task_failed(&self, run_id: &str, attempt: u32, reported: bool) { + if let Some(interceptor) = self.wft_interceptor.as_deref() { + interceptor.on_task_failed( + &WorkflowTaskFailedInfo::builder() + .run_id(run_id) + .attempt(attempt) + .reported(reported) + .build(), + ); + } + } + async fn handle_activation_completed_result( &self, run_id: &str, @@ -1019,6 +1089,8 @@ struct WFTWithPaginator { struct PreparedWFT { task_token: TaskToken, attempt: u32, + started_time: Option, + received: Instant, execution: WorkflowExecution, workflow_type: String, legacy_query: Option, @@ -1091,6 +1163,10 @@ pub(crate) enum OutstandingActivation { struct WorkflowTaskInfo { task_token: TaskToken, attempt: u32, + /// From the poll response, like the attempt. + started_event_id: i64, + started_time: Option, + received: Instant, /// Exists to allow easy tagging of spans with workflow ids. Is duplicative of info inside the /// run machines themselves, but that can't be accessed easily. Would be nice to somehow have a /// shared repository, or refcounts, or whatever, for strings like these that get duped all @@ -1178,6 +1254,8 @@ pub(crate) enum ActivationAction { sdk_metadata: WorkflowTaskCompletedMetadata, versioning_behavior: VersioningBehavior, attempt: u32, + /// Set when the worker has a workflow task interceptor. + task_info: Option>, }, /// We should respond to a legacy query request RespondLegacyQuery { result: Box }, @@ -1201,6 +1279,53 @@ impl EvictionRequestResult { } } +/// Runs the `before_complete` hook within the task's deadline, returning the failure message. +async fn run_before_complete( + interceptor: &dyn WorkflowTaskInterceptor, + info: &WorkflowTaskCompletionInfo, +) -> Result<(), String> { + let deadline = tokio::time::Instant::from_std(info.deadline); + match tokio::time::timeout_at(deadline, interceptor.before_complete(info)).await { + Ok(Ok(())) => Ok(()), + Ok(Err(e)) => Err(e.message), + Err(_) => Err("Workflow task interceptor missed the task deadline".to_string()), + } +} + +/// Runs the `after_complete` hook. The run's next task waits on it, so a hung hook must not stall +/// the run. +async fn run_after_complete( + interceptor: &dyn WorkflowTaskInterceptor, + info: &WorkflowTaskCompletionInfo, + accepted: bool, +) { + let after = interceptor.after_complete(info, accepted); + if tokio::time::timeout(WFT_INTERCEPTOR_AFTER_COMPLETE_TIMEOUT, after) + .await + .is_err() + { + warn!(run_id = %info.run_id, "Workflow task interceptor timed out after the completion"); + } +} + +/// The way the commands close the run, if they do. +fn run_close_kind(commands: &[ProtoCommand]) -> Option { + commands + .iter() + .find_map(|cmd| match cmd.attributes.as_ref() { + Some(Attributes::CompleteWorkflowExecutionCommandAttributes(_)) => { + Some(RunCloseKind::Completed) + } + Some(Attributes::FailWorkflowExecutionCommandAttributes(_)) => { + Some(RunCloseKind::Failed) + } + Some(Attributes::CancelWorkflowExecutionCommandAttributes(_)) => { + Some(RunCloseKind::Cancelled) + } + _ => None, + }) +} + #[allow(clippy::type_complexity)] fn detect_terminal_command( commands: &[ProtoCommand], @@ -1696,6 +1821,7 @@ enum CommandID { #[derive(Debug, Clone)] struct WorkflowStartedInfo { workflow_task_timeout: Option, + first_execution_run_id: String, memo: Option, search_attrs: Option, retry_policy: Option, diff --git a/crates/sdk-core/src/worker/workflow/workflow_stream.rs b/crates/sdk-core/src/worker/workflow/workflow_stream.rs index d00173f64e..53e6583993 100644 --- a/crates/sdk-core/src/worker/workflow/workflow_stream.rs +++ b/crates/sdk-core/src/worker/workflow/workflow_stream.rs @@ -36,6 +36,7 @@ pub(super) struct WFStream { history_fetch_refcounter: Arc, shutdown_token: CancellationToken, ignore_evicts_on_shutdown: bool, + wft_interceptor: Option>, metrics: MetricsContext, } @@ -97,6 +98,11 @@ impl WFStream { ), shutdown_token: basics.shutdown_token, ignore_evicts_on_shutdown: basics.worker_config.ignore_evicts_on_shutdown, + wft_interceptor: basics + .worker_config + .extension + .as_ref() + .and_then(|e| e.workflow_task_interceptor()), metrics: basics.metrics, runs_needing_fetching: Default::default(), history_fetch_refcounter: Arc::new(HistfetchRC {}), @@ -405,6 +411,9 @@ impl WFStream { if should_evict { debug!(run_id=%run_id, "Evicting run"); self.runs.remove(run_id); + if let Some(interceptor) = &self.wft_interceptor { + interceptor.on_evicted(run_id); + } } let maybe_ready_wft = maybe_buffered .get_next_wft()