diff --git a/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto b/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto index 026844f0c1..a29005bfd7 100644 --- a/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto +++ b/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto @@ -31,6 +31,18 @@ service StreamService { // Waits until the stream notifications this process sent are out, for at most 10 seconds. A // notification goes in the background, so a process calls it before it stops. rpc FlushNotifications(FlushNotificationsRequest) returns (FlushNotificationsResponse); + // Attaches a Workflow's reader of a stream to the stream's notifier, so the notifier wakes the + // Workflow with a `StreamWake` on its next Workflow Task. A reader calls it before its first + // read. The process attaches each subscription once and answers a repeat from memory, and the + // server takes a repeat from another process as the same attach. + rpc Subscribe(SubscribeRequest) returns (SubscribeResponse); + // Detaches a Workflow's reader that stops before the stream ends. + rpc Unsubscribe(UnsubscribeRequest) returns (UnsubscribeResponse); + // Waits on the stream's notifier, for a client reader whose read armed the stream. Returns when + // the notifier's counter is above `after_counter`, when the stream closed, or at the server's + // wait limit (`streamnotifier.pollMaxWait`). The reader then reads again with `arm`. After the + // wait limit, that read finds records whose notification was lost, or arms again. + rpc WaitForNotification(WaitForNotificationRequest) returns (WaitForNotificationResponse); // What the store can do beyond the log every store gives, and its limits. rpc DescribeStore(DescribeStoreRequest) returns (DescribeStoreResponse); } @@ -256,6 +268,61 @@ message FlushNotificationsRequest {} message FlushNotificationsResponse {} +message SubscribeRequest { + // The stream to read. The reader's Workflow is in the stream's namespace. + StreamAddress stream = 1; + // The reader's Workflow and run. + string workflow_id = 2; + string run_id = 3; + // Names the reader in its run. Lang derives it from the reader's call number, so replay gives + // the same one. At most 256 bytes. + string subscription_id = 4; + // Counts the times the reader subscribed again after a wake with `resubscribe`. Each value is a + // new attach, since the server answers a repeat of a dropped attach as done. Lang counts it from + // the wakes, so replay gives the same one. + uint32 generation = 5; +} + +message SubscribeResponse { + // The first run of the stream's run chain, which keys the notifier the reader attached to. + string chain_first_run_id = 1; + // Whether this call sent the attach. False when this process attached the reader before. + bool attached = 2; +} + +message UnsubscribeRequest { + // The stream, the reader's run, the subscription and its generation, as the last Subscribe + // named them. + StreamAddress stream = 1; + string run_id = 2; + string subscription_id = 3; + uint32 generation = 4; +} + +message UnsubscribeResponse {} + +message WaitForNotificationRequest { + // The stream the reader reads. + StreamAddress stream = 1; + // The `wake_after` of the read that armed the stream. + int64 after_counter = 2; +} + +message WaitForNotificationResponse { + // The highest counter the notifier has seen. + int64 counter = 1; + // The position that came with `counter`. + string position = 2; + // The wait reached its limit with no counter above `after_counter` and no close. + bool timed_out = 3; + // The stream closed. The reader reads to its end. + bool closed = 4; + // The result the stream closed with, when it gave one. + temporal.api.common.v1.Payload close_result = 5; + // The failure the stream closed with, when it closed with one. + temporal.api.failure.v1.Failure close_failure = 6; +} + message DeleteOwnerRequest { string namespace = 1; StreamOwnerKind owner_kind = 2; diff --git a/crates/protos/protos/local/temporal/sdk/core/workflow_activation/workflow_activation.proto b/crates/protos/protos/local/temporal/sdk/core/workflow_activation/workflow_activation.proto index 7287ab1c1a..b0114c6ff7 100644 --- a/crates/protos/protos/local/temporal/sdk/core/workflow_activation/workflow_activation.proto +++ b/crates/protos/protos/local/temporal/sdk/core/workflow_activation/workflow_activation.proto @@ -140,6 +140,9 @@ message WorkflowActivationJob { ResolveNexusOperationStart resolve_nexus_operation_start = 15; // A nexus operation resolved. ResolveNexusOperation resolve_nexus_operation = 16; + // A stream that one of the Workflow's readers subscribed to moved or closed. Numbers 17 + // and 18 are left to the external stream replay job and to Nexus operation progress. + ResolveStreamWake resolve_stream_wake = 19; // Remove the workflow identified by the [WorkflowActivation] containing this job from the // cache after performing the activation. It is guaranteed that this will be the only job // in the activation if present. @@ -384,6 +387,34 @@ message ResolveNexusOperation { nexus.NexusOperationResult result = 2; } +// A wake for one of the Workflow's stream readers, as the server folded it onto this Workflow +// Task's scheduled event. +// +// Core issues at most one per subscription per Workflow Task, in the task's first activation. It +// comes from History, so replay issues the same job at the same point. A wake says that the +// stream moved, not what it holds, so the reader reads the stream itself. +message ResolveStreamWake { + // The subscription id the reader subscribed with. + string subscription_id = 1; + // The position that came with `counter`, in the producer's terms. + string position = 2; + // The highest counter the stream's notifier has seen. The reader reads when it is above the + // `wake_after` of the read that armed the stream. + int64 counter = 3; + // The metadata that came with `counter`. + map metadata = 4; + // The notifier asks the reader to read again whatever the counter, after a quiet period. + bool recheck = 5; + // The stream closed. The reader reads to its end, then sees the outcome. + bool closed = 6; + // The result the stream closed with, when it gave one. + temporal.api.common.v1.Payload close_result = 7; + // The failure the stream closed with, when it closed with one. + temporal.api.failure.v1.Failure close_failure = 8; + // The notifier dropped the subscription. The reader subscribes again. + bool resubscribe = 9; +} + message RemoveFromCache { string message = 1; diff --git a/crates/protos/protos/local/temporal/sdk/core/workflow_completion/workflow_completion.proto b/crates/protos/protos/local/temporal/sdk/core/workflow_completion/workflow_completion.proto index afe7d64427..f70c206e5a 100644 --- a/crates/protos/protos/local/temporal/sdk/core/workflow_completion/workflow_completion.proto +++ b/crates/protos/protos/local/temporal/sdk/core/workflow_completion/workflow_completion.proto @@ -33,6 +33,12 @@ message Success { repeated uint32 used_internal_flags = 6; // The versioning behavior this workflow is currently using temporal.api.enums.v1.VersioningBehavior versioning_behavior = 7; + // Asks Core to complete the workflow task now if local activities are still running, as it + // does at a local activity heartbeat. The completion carries the commands and markers so far + // and asks the server for a new task, and the local activities run on. Lang sets it from a + // count or a clock, so Core ignores it while replaying. Without running local activities, + // the task completes as it would anyway. + bool heartbeat_task = 8; } // Failure to activate or execute a workflow diff --git a/crates/protos/src/protos/mod.rs b/crates/protos/src/protos/mod.rs index bc07da0350..59716ec86c 100644 --- a/crates/protos/src/protos/mod.rs +++ b/crates/protos/src/protos/mod.rs @@ -81,6 +81,7 @@ pub mod coresdk { commands: v, used_internal_flags: vec![], versioning_behavior: VersioningBehavior::Unspecified.into(), + heartbeat_task: false, } } } @@ -1254,6 +1255,9 @@ pub mod coresdk { workflow_activation_job::Variant::ResolveNexusOperation(_) => { write!(f, "ResolveNexusOperation") } + workflow_activation_job::Variant::ResolveStreamWake(w) => { + write!(f, "ResolveStreamWake({}, {})", w.subscription_id, w.counter) + } } } } diff --git a/crates/sdk-core/CHANGELOG.md b/crates/sdk-core/CHANGELOG.md index 95a142a712..a18cd506dc 100644 --- a/crates/sdk-core/CHANGELOG.md +++ b/crates/sdk-core/CHANGELOG.md @@ -66,6 +66,26 @@ relevant information. finds the owner ended closes only the store. The notifier's owner check closes the notifier. * Experimental: `StreamStore` has `close_topic_with` and `close_stream`, which carry the close outcome to a wrapper that tells the notifier. +* Experimental: the new `Subscribe` and `Unsubscribe` stream calls attach a Workflow's reader to + a stream's notifier as a `workflow_target`, and detach it. The attach's request id is + `jetty//`, so a rerun attaches as the same reader. After a wake with + `resubscribe`, lang raises `SubscribeRequest.generation`, which makes a new attach. A process + attaches each subscription once and answers a repeat from memory. The attach names the chain + that holds the Workflow id now, and on an empty stream it arms the stream, so the first append + wakes the Workflow. +* Experimental: the new `WaitForNotification` stream call waits on a stream's notifier with the + server's `PollStreamNotifier` long poll, for a client reader whose read armed the stream. It + returns when the notifier's counter passes the read's `wake_after`, when the stream closed, or + with `timed_out` at the server's wait limit. The reader then reads again with `arm`. +* Experimental: a stream wake for one of a Workflow's readers reaches lang as a + `ResolveStreamWake` job, at most one per subscription per workflow task, folded from the + `stream_wakes` on the task's scheduled events. Replay delivers the same jobs in the same + activations. A task scheduled with wakes ends a workflow task heartbeat chain, so it keeps an + activation of its own on replay. +* Experimental: `Success.heartbeat_task` on a workflow activation completion asks Core to + complete the workflow task now while local activities run, as at a local activity heartbeat. + The task carries the commands and markers so far and asks for a new task. Core ignores the + flag while replaying, since lang sets it from a count or a clock. * Experimental: the stream service's `RetireProducer` call drops one producer attempt's dedupe state on a stream, which the Redis store keeps as one field of the topic's meta hash until retention. A producer that writes once per short-lived attempt calls it after the attempt's diff --git a/crates/sdk-core/src/core_tests/mod.rs b/crates/sdk-core/src/core_tests/mod.rs index f78ef7bcf8..76a9de9f81 100644 --- a/crates/sdk-core/src/core_tests/mod.rs +++ b/crates/sdk-core/src/core_tests/mod.rs @@ -4,6 +4,8 @@ mod event_groups; mod external_streams; mod queries; mod replay_flag; +mod stream_wakes; +mod task_heartbeat; mod updates; mod workers; mod workflow_cancels; diff --git a/crates/sdk-core/src/core_tests/stream_wakes.rs b/crates/sdk-core/src/core_tests/stream_wakes.rs new file mode 100644 index 0000000000..6d0ec452aa --- /dev/null +++ b/crates/sdk-core/src/core_tests/stream_wakes.rs @@ -0,0 +1,214 @@ +//! Stream wakes: the server folds them onto a Workflow Task's scheduled event, and Core hands them +//! to lang as one job per subscription per task, live and on replay alike. + +use crate::{ + replay::TestHistoryBuilder, + test_help::{MockPollCfg, ResponseType, build_mock_pollers, mock_worker}, + worker::client::mocks::mock_worker_client, +}; +use temporalio_common::{ + protos::{ + coresdk::{ + workflow_activation::{ResolveStreamWake, WorkflowActivation, workflow_activation_job}, + workflow_completion::WorkflowActivationCompletion, + }, + temporal::api::{ + common::v1::Payload, + enums::v1::{EventType, WorkflowTaskFailedCause}, + failure::v1::Failure, + stream::v1::StreamWake, + }, + }, + worker::WorkerTaskTypes, +}; + +fn wake(subscription_id: &str, counter: i64) -> StreamWake { + StreamWake { + subscription_id: subscription_id.to_string(), + position: format!("position-{counter}"), + counter, + metadata: [("records".to_string(), counter.to_string())].into(), + ..Default::default() + } +} + +/// The job lang should get for `wake`. +fn job_for(wake: &StreamWake) -> workflow_activation_job::Variant { + workflow_activation_job::Variant::ResolveStreamWake(ResolveStreamWake { + subscription_id: wake.subscription_id.clone(), + position: wake.position.clone(), + counter: wake.counter, + metadata: wake.metadata.clone(), + recheck: wake.recheck, + closed: wake.closed, + close_result: wake.close_result.clone(), + close_failure: wake.close_failure.clone(), + resubscribe: wake.resubscribe, + }) +} + +fn wake_jobs(activation: &WorkflowActivation) -> Vec { + activation + .jobs + .iter() + .filter_map(|job| match &job.variant { + v @ Some(workflow_activation_job::Variant::ResolveStreamWake(_)) => v.clone(), + _ => None, + }) + .collect() +} + +fn job_names(activation: &WorkflowActivation) -> Vec { + activation + .jobs + .iter() + .map(|job| match &job.variant { + Some(workflow_activation_job::Variant::ResolveStreamWake(w)) => { + format!("Wake({}, {})", w.subscription_id, w.counter) + } + Some(workflow_activation_job::Variant::SignalWorkflow(s)) => { + format!("Signal({})", s.signal_name) + } + Some(other) => other.to_string(), + None => "none".to_string(), + }) + .collect() +} + +fn worker(history: TestHistoryBuilder, batches: Vec) -> crate::Worker { + let mut mock = build_mock_pollers(MockPollCfg::from_resp_batches( + "fake_wf_id", + history, + batches, + mock_worker_client(), + )); + mock.worker_cfg(|w| { + w.task_types = WorkerTaskTypes::workflow_only(); + w.max_cached_workflows = 1; + }); + mock_worker(mock) +} + +async fn complete_empty(worker: &crate::Worker, run_id: &str) { + worker + .complete_workflow_activation(WorkflowActivationCompletion::empty(run_id.to_string())) + .await + .unwrap(); +} + +/// A first task with nothing for lang to do, then two tasks scheduled only because wakes arrived. +/// The first and the second complete with no commands, which is the shape of a heartbeat chain, so +/// replay must still give each wake task an activation of its own. +fn wakes_after_a_quiet_task() -> (TestHistoryBuilder, [StreamWake; 2]) { + let mut t = TestHistoryBuilder::default(); + t.add_by_type(EventType::WorkflowExecutionStarted); + t.add_full_wf_task(); + let first = wake("1", 41); + t.add_workflow_task_scheduled_with_stream_wakes(vec![first.clone()]); + t.add_workflow_task_started(); + t.add_workflow_task_completed(); + let second = wake("1", 42); + t.add_workflow_task_scheduled_with_stream_wakes(vec![second.clone()]); + t.add_workflow_task_started(); + (t, [first, second]) +} + +async fn run_wakes_after_a_quiet_task(batches: Vec, replaying: bool) { + let (history, [first, second]) = wakes_after_a_quiet_task(); + let worker = worker(history, batches); + + let init = worker.poll_workflow_activation().await.unwrap(); + let run_id = init.run_id.clone(); + assert!(wake_jobs(&init).is_empty(), "{:?}", job_names(&init)); + complete_empty(&worker, &run_id).await; + + let first_wake = worker.poll_workflow_activation().await.unwrap(); + assert_eq!( + wake_jobs(&first_wake), + vec![job_for(&first)], + "the wake-only task gets an activation of its own: {:?}", + job_names(&first_wake) + ); + assert_eq!(first_wake.is_replaying, replaying); + complete_empty(&worker, &run_id).await; + + let second_wake = worker.poll_workflow_activation().await.unwrap(); + assert_eq!(wake_jobs(&second_wake), vec![job_for(&second)]); + assert!(!second_wake.is_replaying); +} + +#[tokio::test] +async fn each_wake_task_gets_its_own_activation_live() { + run_wakes_after_a_quiet_task( + vec![ + ResponseType::ToTaskNum(1), + ResponseType::ToTaskNum(2), + ResponseType::ToTaskNum(3), + ], + false, + ) + .await; +} + +#[tokio::test] +async fn each_wake_task_gets_its_own_activation_on_replay() { + run_wakes_after_a_quiet_task(vec![ResponseType::AllHistory], true).await; +} + +#[tokio::test] +async fn a_failed_task_and_its_retry_fold_their_wakes_per_subscription() { + let mut t = TestHistoryBuilder::default(); + t.add_by_type(EventType::WorkflowExecutionStarted); + t.add_full_wf_task(); + t.add_workflow_task_scheduled_with_stream_wakes(vec![wake("a", 5), wake("b", 1)]); + t.add_workflow_task_started(); + t.add_workflow_task_failed_with_failure( + WorkflowTaskFailedCause::WorkflowWorkerUnhandledFailure, + Failure::default(), + ); + let closed = StreamWake { + closed: true, + close_result: Some(Payload { + data: b"summary".to_vec(), + ..Default::default() + }), + ..wake("a", 7) + }; + t.add_workflow_task_scheduled_with_stream_wakes(vec![closed.clone()]); + t.add_workflow_task_started(); + let worker = worker(t, vec![ResponseType::AllHistory]); + + let init = worker.poll_workflow_activation().await.unwrap(); + complete_empty(&worker, &init.run_id).await; + let woken = worker.poll_workflow_activation().await.unwrap(); + assert_eq!( + wake_jobs(&woken), + vec![job_for(&closed), job_for(&wake("b", 1))], + "{:?}", + job_names(&woken) + ); +} + +#[tokio::test] +async fn a_wake_and_a_signal_of_one_task_share_its_activation() { + let mut t = TestHistoryBuilder::default(); + t.add_by_type(EventType::WorkflowExecutionStarted); + t.add_full_wf_task(); + t.add_we_signaled("go", vec![]); + t.add_workflow_task_scheduled_with_stream_wakes(vec![StreamWake { + recheck: true, + ..wake("1", 0) + }]); + t.add_workflow_task_started(); + let worker = worker( + t, + vec![ResponseType::ToTaskNum(1), ResponseType::ToTaskNum(2)], + ); + + let init = worker.poll_workflow_activation().await.unwrap(); + complete_empty(&worker, &init.run_id).await; + let woken = worker.poll_workflow_activation().await.unwrap(); + let mut names = job_names(&woken); + names.sort(); + assert_eq!(names, ["Signal(go)", "Wake(1, 0)"]); +} diff --git a/crates/sdk-core/src/core_tests/task_heartbeat.rs b/crates/sdk-core/src/core_tests/task_heartbeat.rs new file mode 100644 index 0000000000..33c1950ebf --- /dev/null +++ b/crates/sdk-core/src/core_tests/task_heartbeat.rs @@ -0,0 +1,486 @@ +//! Lang's early workflow task heartbeat: `Success.heartbeat_task` completes the task while local +//! activities run, and replay must not depend on it. + +use crate::{ + Worker, + replay::TestHistoryBuilder, + test_help::{ + MockPollCfg, ResponseType, WorkerTestHelpers, build_mock_pollers, + drain_pollers_and_shutdown, mock_worker, schedule_local_activity_cmd, start_timer_cmd, + }, + worker::client::{WorkflowTaskCompletion, mocks::mock_worker_client}, +}; +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; +use temporalio_common::protos::{ + coresdk::{ + ActivityTaskCompletion, + activity_result::ActivityExecutionResult, + activity_task::activity_task, + workflow_activation::{WorkflowActivation, WorkflowActivationJob, workflow_activation_job}, + workflow_commands::{ + ActivityCancellationType, CompleteWorkflowExecution, workflow_command, + }, + workflow_completion::{WorkflowActivationCompletion, workflow_activation_completion}, + }, + temporal::api::{ + command::v1::command::Attributes, + enums::v1::{CommandType, EventType}, + history::v1::{MarkerRecordedEventAttributes, history_event}, + stream::v1::StreamWake, + workflowservice::v1::RespondWorkflowTaskCompletedResponse, + }, +}; + +// Long enough that Core's own heartbeat, at 80% of it, never fires during a test. +const WFT_TIMEOUT: Duration = Duration::from_secs(60); + +fn first_commands() -> Vec { + vec![ + start_timer_cmd(1, Duration::from_secs(3600)), + schedule_local_activity_cmd( + 1, + "1", + ActivityCancellationType::TryCancel, + Duration::from_secs(60), + ), + ] +} + +fn completion( + run_id: &str, + commands: Vec, + heartbeat_task: bool, +) -> WorkflowActivationCompletion { + let mut completion = WorkflowActivationCompletion::from_cmds(run_id, commands); + let Some(workflow_activation_completion::Status::Successful(success)) = + completion.status.as_mut() + else { + unreachable!("from_cmds builds a successful completion"); + }; + success.heartbeat_task = heartbeat_task; + completion +} + +fn is_init(activation: &WorkflowActivation) -> bool { + matches!( + activation.jobs.as_slice(), + [WorkflowActivationJob { + variant: Some(workflow_activation_job::Variant::InitializeWorkflow(_)), + }] + ) +} + +fn resolves_la_1(activation: &WorkflowActivation) -> bool { + matches!( + activation.jobs.as_slice(), + [WorkflowActivationJob { + variant: Some(workflow_activation_job::Variant::ResolveActivity(resolve)), + }] if resolve.seq == 1 && resolve.is_local + ) +} + +fn command_types(completion: &WorkflowTaskCompletion) -> Vec { + completion + .commands + .iter() + .map(|command| command.command_type()) + .collect() +} + +/// A worker over `history` that records every workflow task completion it reports. +fn recording_worker( + history: TestHistoryBuilder, + responses: impl IntoIterator, + expected_completions: usize, +) -> (Worker, Arc>>) { + let reported = Arc::new(Mutex::new(vec![])); + let mut cfg = + MockPollCfg::from_resp_batches("fake_wf_id", history, responses, mock_worker_client()); + cfg.num_expected_fails = 0; + cfg.num_expected_completions = Some(expected_completions.into()); + let recorder = reported.clone(); + cfg.completion_mock_fn = Some(Box::new(move |completion| { + recorder.lock().unwrap().push(completion.clone()); + Ok(RespondWorkflowTaskCompletedResponse::default()) + })); + let mut mock = build_mock_pollers(cfg); + mock.worker_cfg(|wc| wc.max_cached_workflows = 1); + (mock_worker(mock), reported) +} + +/// What a cut produces on the server: the first task completes with the timer and asks for a new +/// task while the local activity runs, and the next task records the marker. +fn cut_history_up_to_second_task() -> TestHistoryBuilder { + let mut t = TestHistoryBuilder::default(); + t.add_wfe_started_with_wft_timeout(WFT_TIMEOUT); + t.add_full_wf_task(); + t.add_timer_started("1".to_string()); + t.add_workflow_task_scheduled_and_started(); + t +} + +/// The history the server writes for `completions`, one workflow task each. +fn history_of(completions: &[WorkflowTaskCompletion]) -> TestHistoryBuilder { + history_with_wakes(completions, &[]) +} + +/// Like [history_of], with `wakes[i]` on the scheduled event of task `i`. +fn history_with_wakes( + completions: &[WorkflowTaskCompletion], + wakes: &[Vec], +) -> TestHistoryBuilder { + let mut t = TestHistoryBuilder::default(); + t.add_wfe_started_with_wft_timeout(WFT_TIMEOUT); + for (task, completion) in completions.iter().enumerate() { + match wakes.get(task) { + Some(wakes) if !wakes.is_empty() => { + t.add_workflow_task_scheduled_with_stream_wakes(wakes.clone()); + t.add_workflow_task_started(); + } + _ => t.add_workflow_task_scheduled_and_started(), + } + t.add_workflow_task_completed(); + let completed_id = t.current_event_id(); + if let Some(history_event::Attributes::WorkflowTaskCompletedEventAttributes(attrs)) = + t.last_event().and_then(|event| event.attributes.as_mut()) + { + attrs.sdk_metadata = Some(completion.sdk_metadata.clone()); + } + for command in &completion.commands { + match command.attributes.clone() { + Some(Attributes::StartTimerCommandAttributes(timer)) => { + t.add_timer_started(timer.timer_id); + } + Some(Attributes::RecordMarkerCommandAttributes(marker)) => { + t.add(MarkerRecordedEventAttributes { + marker_name: marker.marker_name, + details: marker.details, + workflow_task_completed_event_id: completed_id, + header: marker.header, + failure: marker.failure, + }); + } + Some(Attributes::CompleteWorkflowExecutionCommandAttributes(_)) => { + t.add_workflow_execution_completed(); + } + other => panic!("no history event for {other:?}"), + } + } + } + t +} + +/// Runs the cut live and answers the two completions it reported. +async fn run_cut_live() -> Vec { + let (core, reported) = recording_worker( + cut_history_up_to_second_task(), + [ResponseType::ToTaskNum(1), ResponseType::ToTaskNum(2)], + 2, + ); + let activation = core.poll_workflow_activation().await.unwrap(); + assert!(is_init(&activation), "{activation:?}"); + core.complete_workflow_activation(completion(&activation.run_id, first_commands(), true)) + .await + .unwrap(); + let task = core.poll_activity_task().await.unwrap(); + let Some(activity_task::Variant::Start(start)) = &task.variant else { + panic!("expected the local activity, got {task:?}"); + }; + assert!(start.is_local); + { + let reported = reported.lock().unwrap(); + let [cut] = reported.as_slice() else { + panic!("the cut reports one completion while the local activity runs: {reported:?}"); + }; + assert!(cut.force_create_new_workflow_task); + assert_eq!(command_types(cut), [CommandType::StartTimer]); + } + core.complete_activity_task(ActivityTaskCompletion { + task_token: task.task_token, + result: Some(ActivityExecutionResult::ok(b"done".to_vec().into())), + }) + .await + .unwrap(); + let activation = core.poll_workflow_activation().await.unwrap(); + assert!(resolves_la_1(&activation), "{activation:?}"); + core.complete_execution(&activation.run_id).await; + drain_pollers_and_shutdown(&core).await; + let reported = reported.lock().unwrap().clone(); + let [_, last] = reported.as_slice() else { + panic!("expected two completions, got {reported:?}"); + }; + assert!(!last.force_create_new_workflow_task); + assert_eq!( + command_types(last), + [ + CommandType::RecordMarker, + CommandType::CompleteWorkflowExecution + ] + ); + reported +} + +#[tokio::test] +async fn a_heartbeat_task_completes_the_task_while_the_local_activity_runs() { + run_cut_live().await; +} + +#[tokio::test] +async fn a_heartbeat_task_without_local_activities_completes_the_task_as_usual() { + let mut t = TestHistoryBuilder::default(); + t.add_wfe_started_with_wft_timeout(WFT_TIMEOUT); + t.add_workflow_task_scheduled_and_started(); + let (core, reported) = recording_worker(t, [ResponseType::ToTaskNum(1)], 1); + let activation = core.poll_workflow_activation().await.unwrap(); + core.complete_workflow_activation(completion( + &activation.run_id, + vec![start_timer_cmd(1, Duration::from_secs(3600))], + true, + )) + .await + .unwrap(); + drain_pollers_and_shutdown(&core).await; + let reported = reported.lock().unwrap(); + let [only] = reported.as_slice() else { + panic!("expected one completion, got {reported:?}"); + }; + assert!( + !only.force_create_new_workflow_task, + "with nothing running, there is nothing to heartbeat for" + ); + assert_eq!(command_types(only), [CommandType::StartTimer]); +} + +/// Replays `history` to its end with lang deciding the cut by `heartbeat_task`, and checks that +/// lang sees the activations of the live run and that nothing goes to the server. +async fn replay(history: TestHistoryBuilder, heartbeat_task: bool) { + let (core, reported) = recording_worker(history, [ResponseType::AllHistory], 0); + let activation = core.poll_workflow_activation().await.unwrap(); + assert!(is_init(&activation), "{activation:?}"); + core.complete_workflow_activation(completion( + &activation.run_id, + first_commands(), + heartbeat_task, + )) + .await + .unwrap(); + let activation = tokio::time::timeout(Duration::from_secs(10), core.poll_workflow_activation()) + .await + .expect("replay resolves the local activity from its marker") + .unwrap(); + assert!(resolves_la_1(&activation), "{activation:?}"); + core.complete_execution(&activation.run_id).await; + drain_pollers_and_shutdown(&core).await; + assert!(reported.lock().unwrap().is_empty()); +} + +#[rstest::rstest] +#[case::with_the_cut(true)] +#[case::without_the_cut(false)] +#[tokio::test] +async fn a_cut_history_replays_whatever_lang_decides(#[case] heartbeat_task: bool) { + let live = run_cut_live().await; + let history = history_of(&live); + let tasks = history + .get_full_history_info() + .unwrap() + .events() + .iter() + .filter(|event| event.event_type() == EventType::WorkflowTaskCompleted) + .count(); + assert_eq!(tasks, 2, "the cut adds a workflow task"); + replay(history, heartbeat_task).await; +} + +#[tokio::test] +async fn an_uncut_history_replays_when_lang_asks_for_a_cut() { + // The local activity resolved within the first task, so its marker sits in that task. A + // replay whose clock says to cut must not ask for the task that the live run never made. + let mut t = TestHistoryBuilder::default(); + t.add_wfe_started_with_wft_timeout(WFT_TIMEOUT); + t.add_full_wf_task(); + t.add_timer_started("1".to_string()); + t.add_local_activity_result_marker(1, "1", b"done".to_vec().into()); + t.add_workflow_execution_completed(); + replay(t, true).await; +} + +#[tokio::test] +async fn a_cut_task_runs_its_local_activity_after_a_cache_miss() { + // A new Worker gets the second task with the whole history. It replays the first task, where + // the cut is ignored, and then runs the local activity that the cut left without a marker. + let (core, reported) = recording_worker( + cut_history_up_to_second_task(), + [ResponseType::AllHistory], + 1, + ); + let activation = core.poll_workflow_activation().await.unwrap(); + assert!(is_init(&activation), "{activation:?}"); + core.complete_workflow_activation(completion(&activation.run_id, first_commands(), true)) + .await + .unwrap(); + let task = core.poll_activity_task().await.unwrap(); + assert!(matches!( + &task.variant, + Some(activity_task::Variant::Start(start)) if start.is_local + )); + core.complete_activity_task(ActivityTaskCompletion { + task_token: task.task_token, + result: Some(ActivityExecutionResult::ok(b"done".to_vec().into())), + }) + .await + .unwrap(); + let activation = core.poll_workflow_activation().await.unwrap(); + assert!(resolves_la_1(&activation), "{activation:?}"); + core.complete_workflow_activation(completion( + &activation.run_id, + vec![CompleteWorkflowExecution { result: None }.into()], + true, + )) + .await + .unwrap(); + drain_pollers_and_shutdown(&core).await; + let reported = reported.lock().unwrap(); + let [only] = reported.as_slice() else { + panic!("expected one completion, got {reported:?}"); + }; + assert!(!only.force_create_new_workflow_task); + assert_eq!( + command_types(only), + [ + CommandType::RecordMarker, + CommandType::CompleteWorkflowExecution + ] + ); +} + +#[tokio::test] +async fn a_cut_history_has_the_events_the_server_writes() { + let live = run_cut_live().await; + let events: Vec = history_of(&live) + .get_full_history_info() + .unwrap() + .events() + .iter() + .map(|event| event.event_type()) + .collect(); + assert_eq!( + events, + [ + EventType::WorkflowExecutionStarted, + EventType::WorkflowTaskScheduled, + EventType::WorkflowTaskStarted, + EventType::WorkflowTaskCompleted, + EventType::TimerStarted, + EventType::WorkflowTaskScheduled, + EventType::WorkflowTaskStarted, + EventType::WorkflowTaskCompleted, + EventType::MarkerRecorded, + EventType::WorkflowExecutionCompleted, + ] + ); +} + +fn wake() -> StreamWake { + StreamWake { + subscription_id: "1".to_string(), + position: "position-42".to_string(), + counter: 42, + ..Default::default() + } +} + +fn wakes_reader_1(activation: &WorkflowActivation) -> bool { + matches!( + activation.jobs.as_slice(), + [WorkflowActivationJob { + variant: Some(workflow_activation_job::Variant::ResolveStreamWake(wake)), + }] if wake.subscription_id == "1" && wake.counter == 42 + ) +} + +/// A wake that lands while the cut task is open comes on the task the cut asked for, with the +/// local activity still running, and the activity resolves in a later activation of that task. +async fn run_cut_then_wake_live() -> Vec { + let mut t = TestHistoryBuilder::default(); + t.add_wfe_started_with_wft_timeout(WFT_TIMEOUT); + t.add_full_wf_task(); + t.add_timer_started("1".to_string()); + t.add_workflow_task_scheduled_with_stream_wakes(vec![wake()]); + t.add_workflow_task_started(); + let (core, reported) = recording_worker( + t, + [ResponseType::ToTaskNum(1), ResponseType::ToTaskNum(2)], + 2, + ); + let activation = core.poll_workflow_activation().await.unwrap(); + core.complete_workflow_activation(completion(&activation.run_id, first_commands(), true)) + .await + .unwrap(); + let task = core.poll_activity_task().await.unwrap(); + let woken = core.poll_workflow_activation().await.unwrap(); + assert!(wakes_reader_1(&woken), "{woken:?}"); + core.complete_workflow_activation(completion(&woken.run_id, vec![], false)) + .await + .unwrap(); + core.complete_activity_task(ActivityTaskCompletion { + task_token: task.task_token, + result: Some(ActivityExecutionResult::ok(b"done".to_vec().into())), + }) + .await + .unwrap(); + let activation = core.poll_workflow_activation().await.unwrap(); + assert!(resolves_la_1(&activation), "{activation:?}"); + core.complete_execution(&activation.run_id).await; + drain_pollers_and_shutdown(&core).await; + let reported = reported.lock().unwrap().clone(); + let [cut, last] = reported.as_slice() else { + panic!("expected two completions, got {reported:?}"); + }; + assert!(cut.force_create_new_workflow_task); + assert!(!last.force_create_new_workflow_task); + reported +} + +#[tokio::test] +async fn a_wake_on_the_task_after_a_cut_reaches_lang_while_the_local_activity_runs() { + run_cut_then_wake_live().await; +} + +#[rstest::rstest] +#[case::with_the_cut(true)] +#[case::without_the_cut(false)] +#[tokio::test] +async fn a_cut_then_a_wake_replays_in_the_live_order(#[case] heartbeat_task: bool) { + let live = run_cut_then_wake_live().await; + let history = history_with_wakes(&live, &[vec![], vec![wake()]]); + let (core, reported) = recording_worker(history, [ResponseType::AllHistory], 0); + let activation = core.poll_workflow_activation().await.unwrap(); + assert!(is_init(&activation), "{activation:?}"); + core.complete_workflow_activation(completion( + &activation.run_id, + first_commands(), + heartbeat_task, + )) + .await + .unwrap(); + let woken = core.poll_workflow_activation().await.unwrap(); + assert!( + wakes_reader_1(&woken), + "the wake comes first, as live: {woken:?}" + ); + core.complete_workflow_activation(completion(&woken.run_id, vec![], false)) + .await + .unwrap(); + let activation = tokio::time::timeout(Duration::from_secs(10), core.poll_workflow_activation()) + .await + .expect("replay resolves the local activity from its marker") + .unwrap(); + assert!(resolves_la_1(&activation), "{activation:?}"); + core.complete_execution(&activation.run_id).await; + drain_pollers_and_shutdown(&core).await; + assert!(reported.lock().unwrap().is_empty()); +} diff --git a/crates/sdk-core/src/replay/history_builder.rs b/crates/sdk-core/src/replay/history_builder.rs index a14c92b3c5..355c0b6ad6 100644 --- a/crates/sdk-core/src/replay/history_builder.rs +++ b/crates/sdk-core/src/replay/history_builder.rs @@ -111,6 +111,17 @@ impl TestHistoryBuilder { self.workflow_task_scheduled_event_id = self.add_by_type(EventType::WorkflowTaskScheduled); } + /// Add a workflow task scheduled event carrying folded stream wakes. + pub fn add_workflow_task_scheduled_with_stream_wakes( + &mut self, + stream_wakes: Vec, + ) { + self.workflow_task_scheduled_event_id = self.add(WorkflowTaskScheduledEventAttributes { + stream_wakes, + ..Default::default() + }); + } + /// Add a workflow task started event. pub fn add_workflow_task_started(&mut self) { self.final_workflow_task_started_event_id = self.add(WorkflowTaskStartedEventAttributes { diff --git a/crates/sdk-core/src/streams.rs b/crates/sdk-core/src/streams.rs index b4c8ea3980..936a8d1927 100644 --- a/crates/sdk-core/src/streams.rs +++ b/crates/sdk-core/src/streams.rs @@ -8,20 +8,23 @@ use temporalio_common::protos::google::rpc::Status as RpcStatus; use temporalio_common::protos::temporal::api::{ common::v1::WorkflowExecution, enums::v1::StreamOwnerKind, - stream::v1::StreamReference, + stream::v1::{StreamReference, WorkflowWakeTarget}, workflowservice::v1::{ - DescribeWorkflowExecutionRequest, DescribeWorkflowExecutionResponse, NotifyStreamRequest, + AttachStreamCallbackRequest, DescribeWorkflowExecutionRequest, + DescribeWorkflowExecutionResponse, DetachStreamCallbackRequest, NotifyStreamRequest, + PollStreamNotifierRequest, PollStreamNotifierResponse, }, }; use temporalio_streams::{ - CloseOutcome, NotifierClient, NotifyError, OwnerClient, OwnerDescription, OwnerError, - StreamNotification, + CloseOutcome, NotifierClient, NotifierPoll, NotifyError, OwnerClient, OwnerDescription, + OwnerError, StreamNotification, WorkflowAttachment, proto::ChainId, }; pub use temporalio_streams::{StreamError, StreamResult, StreamService, StreamStore, proto}; use tonic::IntoRequest; /// Connects to the store `config` names. The service asks `connection`'s server about streams' -/// owners and sends the stream notifications the config turns on to it. +/// owners, sends the stream notifications the config turns on to it, and attaches Workflow +/// readers there. pub async fn connect_stream_service( config: proto::StreamStoreConfig, connection: Connection, @@ -51,19 +54,107 @@ impl NotifierClient for ConnectionNotifierClient { .map(|_| ()) .map_err(notify_error) } + + async fn attach_workflow(&self, attachment: WorkflowAttachment) -> Result<(), NotifyError> { + self.connection + .workflow_service() + .attach_stream_callback(attach_request(attachment).into_request()) + .await + .map(|_| ()) + .map_err(notify_error) + } + + async fn detach( + &self, + chain: ChainId, + topic: String, + request_id: String, + ) -> Result<(), NotifyError> { + self.connection + .workflow_service() + .detach_stream_callback( + DetachStreamCallbackRequest { + namespace: chain.namespace.clone(), + stream_ref: Some(chain_reference(&chain, topic)), + request_id, + ..Default::default() + } + .into_request(), + ) + .await + .map(|_| ()) + .map_err(notify_error) + } + + async fn poll( + &self, + chain: ChainId, + topic: String, + after_counter: i64, + ) -> Result, NotifyError> { + let answer = self + .connection + .workflow_service() + .poll_stream_notifier( + PollStreamNotifierRequest { + namespace: chain.namespace.clone(), + stream_ref: Some(chain_reference(&chain, topic)), + after_counter, + identity: self.connection.identity().to_string(), + } + .into_request(), + ) + .await; + match answer { + Ok(response) => Ok(Some(poll_answer(response.into_inner()))), + // The server answers at its own wait limit, before the client's deadline. A deadline + // that passes anyway asks for the same re-read. + Err(status) if status.code() == tonic::Code::DeadlineExceeded => Ok(None), + Err(status) => Err(notify_error(status)), + } + } +} + +fn poll_answer(response: PollStreamNotifierResponse) -> NotifierPoll { + NotifierPoll { + counter: response.counter, + position: response.position, + closed: response.closed, + close_result: response.close_result, + close_failure: response.close_failure, + } +} + +/// The notifier of `chain`'s stream on `topic`. The run chain's first run keys the notifier, so one +/// notifier serves the stream across Continue-as-New and a new chain on the same Workflow id gets +/// its own. +fn chain_reference(chain: &ChainId, topic: String) -> StreamReference { + StreamReference { + owner_kind: StreamOwnerKind::Workflow as i32, + workflow_id: chain.workflow_id.clone(), + run_id: chain.first_run_id.clone(), + topic, + } +} + +fn attach_request(attachment: WorkflowAttachment) -> AttachStreamCallbackRequest { + AttachStreamCallbackRequest { + namespace: attachment.chain.namespace.clone(), + stream_ref: Some(chain_reference(&attachment.chain, attachment.topic)), + request_id: attachment.request_id, + workflow_target: Some(WorkflowWakeTarget { + workflow_id: attachment.workflow_id, + run_id: attachment.run_id, + subscription_id: attachment.subscription_id, + }), + ..Default::default() + } } fn notify_request(notification: StreamNotification) -> NotifyStreamRequest { NotifyStreamRequest { + stream_ref: Some(chain_reference(¬ification.chain, notification.topic)), namespace: notification.chain.namespace, - // The run chain's first run keys the notifier, so one notifier serves the stream across - // Continue-as-New and a new chain on the same Workflow id gets its own. - stream_ref: Some(StreamReference { - owner_kind: StreamOwnerKind::Workflow as i32, - workflow_id: notification.chain.workflow_id, - run_id: notification.chain.first_run_id, - topic: notification.topic, - }), position: notification.position, counter: notification.counter, close: notification.close.is_some(), @@ -295,6 +386,42 @@ mod tests { assert!(request.close_result.is_none()); } + #[test] + fn an_attach_names_the_chains_notifier_and_wakes_the_reader() { + let request = attach_request(WorkflowAttachment { + chain: proto::ChainId { + namespace: "ns".to_string(), + workflow_id: "owner".to_string(), + first_run_id: "chain-1".to_string(), + }, + topic: "tokens".to_string(), + request_id: "jetty/reader-run/1".to_string(), + workflow_id: "reader".to_string(), + run_id: "reader-run".to_string(), + subscription_id: "1".to_string(), + }); + assert_eq!(request.namespace, "ns"); + assert_eq!( + request.stream_ref, + Some(StreamReference { + owner_kind: StreamOwnerKind::Workflow as i32, + workflow_id: "owner".to_string(), + run_id: "chain-1".to_string(), + topic: "tokens".to_string(), + }) + ); + assert_eq!(request.request_id, "jetty/reader-run/1"); + assert_eq!( + request.workflow_target, + Some(WorkflowWakeTarget { + workflow_id: "reader".to_string(), + run_id: "reader-run".to_string(), + subscription_id: "1".to_string(), + }) + ); + assert!(request.callback.is_none(), "exactly one target"); + } + #[test] fn only_answers_a_retry_cant_change_refuse_a_notification() { assert!(matches!( diff --git a/crates/sdk-core/src/worker/workflow/history_update.rs b/crates/sdk-core/src/worker/workflow/history_update.rs index 8565899465..c60fdf39c1 100644 --- a/crates/sdk-core/src/worker/workflow/history_update.rs +++ b/crates/sdk-core/src/worker/workflow/history_update.rs @@ -741,10 +741,14 @@ fn find_end_index_of_next_wft_seq( if let Some(next_next_event) = events.get(ix + 2) { if !saw_command && next_next_event.event_type() == EventType::WorkflowTaskScheduled + && !scheduled_with_stream_wakes(next_next_event) { // If we've never seen an interesting event and the next two events are // a completion followed immediately again by scheduled, then this is a - // WFT heartbeat and also doesn't conclude the sequence. + // WFT heartbeat and also doesn't conclude the sequence. A task + // scheduled with stream wakes had an activation of its own live, so + // folding it into this one on replay would hand lang the wakes too + // early. continue; } else { // If we see an update accepted command after WFT completed, we want to @@ -801,6 +805,13 @@ fn find_end_index_of_next_wft_seq( NextWFTSeqEndIndex::Incomplete(last_index) } +fn scheduled_with_stream_wakes(e: &HistoryEvent) -> bool { + matches!( + &e.attributes, + Some(Attributes::WorkflowTaskScheduledEventAttributes(a)) if !a.stream_wakes.is_empty() + ) +} + #[cfg(test)] mod tests { use super::*; @@ -941,6 +952,43 @@ mod tests { assert_eq!(seq.len(), 2); } + #[test] + fn a_task_scheduled_with_stream_wakes_ends_a_heartbeat_chain() { + let history = |with_wakes: bool| { + let mut t = TestHistoryBuilder::default(); + t.add_by_type(EventType::WorkflowExecutionStarted); + t.add_full_wf_task(); + t.add_full_wf_task(); // wft started 6 + if with_wakes { + t.add_workflow_task_scheduled_with_stream_wakes(vec![Default::default()]); + t.add_workflow_task_started(); + } else { + t.add_workflow_task_scheduled_and_started(); + } + t.add_workflow_task_completed(); // wft started 9 + t.add_we_signaled("whee", vec![]); + t.add_full_wf_task(); + t.add_workflow_execution_completed(); + t.as_history_update() + }; + + let mut heartbeat = history(false); + next_check_peek(&mut heartbeat, 0); + let seq = next_check_peek(&mut heartbeat, 6); + assert_eq!(seq.last().unwrap().event_id, 13); + + let mut woken = history(true); + next_check_peek(&mut woken, 0); + let seq = next_check_peek(&mut woken, 6); + assert_eq!( + seq.last().unwrap().event_id, + 9, + "the wake task must end the sequence at its own started event" + ); + let seq = next_check_peek(&mut woken, 9); + assert_eq!(seq.last().unwrap().event_id, 13); + } + #[test] fn heartbeat_marker_end() { let mut t = TestHistoryBuilder::default(); diff --git a/crates/sdk-core/src/worker/workflow/machines/workflow_machines.rs b/crates/sdk-core/src/worker/workflow/machines/workflow_machines.rs index 9286d70486..b21508219b 100644 --- a/crates/sdk-core/src/worker/workflow/machines/workflow_machines.rs +++ b/crates/sdk-core/src/worker/workflow/machines/workflow_machines.rs @@ -59,7 +59,8 @@ use temporalio_common::{ coresdk::{ common::{NamespacedWorkflowExecution, VersioningIntent}, workflow_activation::{ - self, NotifyHasPatch, UpdateRandomSeed, WorkflowActivation, workflow_activation_job, + self, NotifyHasPatch, ResolveStreamWake, UpdateRandomSeed, WorkflowActivation, + workflow_activation_job, }, workflow_commands::ContinueAsNewWorkflowExecution, }, @@ -70,6 +71,7 @@ use temporalio_common::{ history::v1::{HistoryEvent, history_event}, protocol::v1::{Message as ProtocolMessage, message::SequencingId}, sdk::v1::WorkflowTaskCompletedMetadata, + stream::v1::StreamWake, }, }, worker::WorkerDeploymentVersion, @@ -120,6 +122,10 @@ pub(crate) struct WorkflowMachines { last_history_from_server: HistoryUpdate, /// Protocol messages that have yet to be processed for the current WFT. protocol_msgs: Vec, + /// Stream wakes read from the scheduled events of the task being applied, one per + /// subscription. A failed task's scheduled event and its retry's are applied together, and + /// the server clears what it put on the first, so both count. + pending_stream_wakes: Vec, /// EventId of the last handled WorkflowTaskStarted event current_started_event_id: i64, /// The WorkflowTaskStarted event id of the last task History shows completed, 0 before one @@ -320,6 +326,7 @@ impl WorkflowMachines { Self { last_history_from_server: basics.history, protocol_msgs: vec![], + pending_stream_wakes: vec![], workflow_id: basics.workflow_id, workflow_type: basics.workflow_type, run_id: basics.run_id, @@ -562,6 +569,27 @@ impl WorkflowMachines { } } + /// Hands lang one job per subscription for the wakes folded from this task's scheduled + /// events. The wakes come from History, so replay hands lang the same jobs at the same point. + fn deliver_stream_wakes(&mut self) { + for wake in std::mem::take(&mut self.pending_stream_wakes) { + self.drive_me.send_job( + workflow_activation_job::Variant::ResolveStreamWake(ResolveStreamWake { + subscription_id: wake.subscription_id, + position: wake.position, + counter: wake.counter, + metadata: wake.metadata, + recheck: wake.recheck, + closed: wake.closed, + close_result: wake.close_result, + close_failure: wake.close_failure, + resubscribe: wake.resubscribe, + }) + .into(), + ); + } + } + pub(crate) fn has_pending_jobs(&self) -> bool { !self.drive_me.peek_pending_jobs().is_empty() } @@ -847,6 +875,7 @@ impl WorkflowMachines { } self.last_processed_event = eid; } + self.deliver_stream_wakes(); // Needed to delay mutation of self until after we've iterated over peeked events. #[allow(clippy::large_enum_variant)] @@ -1129,6 +1158,14 @@ impl WorkflowMachines { } } Ok(EventType::WorkflowTaskScheduled) => { + if let Some(history_event::Attributes::WorkflowTaskScheduledEventAttributes( + ref attrs, + )) = event_dat.event.attributes + { + for wake in &attrs.stream_wakes { + fold_stream_wake(&mut self.pending_stream_wakes, wake); + } + } let wf_task_sm = WorkflowTaskMachine::new(self.next_started_event_id); let key = self.all_machines.insert(wf_task_sm.into()); self.submachine_handle_event(key, event_dat)?; @@ -1962,6 +1999,96 @@ enum CommandIdKind { NeverResolves, } +/// Keeps one wake per subscription, in order of first appearance. The highest counter keeps its +/// position and metadata. A recheck, a close or a resubscribe on any of the folded wakes stays, +/// since each one asks the reader to act whatever the counter. +fn fold_stream_wake(folded: &mut Vec, wake: &StreamWake) { + let Some(held) = folded + .iter_mut() + .find(|held| held.subscription_id == wake.subscription_id) + else { + folded.push(wake.clone()); + return; + }; + if wake.counter > held.counter { + held.counter = wake.counter; + held.position = wake.position.clone(); + held.metadata = wake.metadata.clone(); + } + held.recheck |= wake.recheck; + held.resubscribe |= wake.resubscribe; + if wake.closed && !held.closed { + held.closed = true; + held.close_result = wake.close_result.clone(); + held.close_failure = wake.close_failure.clone(); + } +} + +#[cfg(test)] +mod stream_wake_tests { + use super::fold_stream_wake; + use temporalio_common::protos::temporal::api::{common::v1::Payload, stream::v1::StreamWake}; + + fn wake(subscription_id: &str, counter: i64, position: &str) -> StreamWake { + StreamWake { + subscription_id: subscription_id.to_string(), + position: position.to_string(), + counter, + ..Default::default() + } + } + + #[test] + fn the_highest_counter_per_subscription_is_kept_in_first_seen_order() { + let mut folded = vec![]; + for w in [ + wake("b", 2, "a"), + wake("a", 1, "b"), + wake("b", 4, "c"), + wake("b", 3, "d"), + wake("a", 1, "e"), + ] { + fold_stream_wake(&mut folded, &w); + } + assert_eq!(folded, vec![wake("b", 4, "c"), wake("a", 1, "b")]); + } + + #[test] + fn a_recheck_a_close_or_a_resubscribe_survives_a_higher_counter() { + let result = Payload { + data: b"done".to_vec(), + ..Default::default() + }; + let mut folded = vec![]; + fold_stream_wake( + &mut folded, + &StreamWake { + recheck: true, + closed: true, + close_result: Some(result.clone()), + ..wake("a", 5, "closed") + }, + ); + fold_stream_wake( + &mut folded, + &StreamWake { + resubscribe: true, + ..wake("a", 7, "later") + }, + ); + assert_eq!( + folded, + vec![StreamWake { + recheck: true, + closed: true, + close_result: Some(result), + resubscribe: true, + ..wake("a", 7, "later") + }] + ); + } +} + #[cfg(test)] mod stream_attempt_tests { use super::*; diff --git a/crates/sdk-core/src/worker/workflow/managed_run.rs b/crates/sdk-core/src/worker/workflow/managed_run.rs index caa8d46bb4..33f91b6d8a 100644 --- a/crates/sdk-core/src/worker/workflow/managed_run.rs +++ b/crates/sdk-core/src/worker/workflow/managed_run.rs @@ -441,6 +441,7 @@ impl ManagedRun { versioning_behavior: VersioningBehavior, resp_chan: Option>, is_forced_failure: bool, + heartbeat_task: bool, ) -> Result> { let activation_was_only_eviction = self.activation_is_eviction(); if activation_was_only_eviction && !commands.is_empty() { @@ -520,6 +521,7 @@ impl ManagedRun { resp_chan, is_forced_failure, versioning_behavior, + heartbeat_task, }; // Verify we can actually apply the next workflow task, which will happen as part of @@ -689,6 +691,7 @@ impl ManagedRun { VersioningBehavior::Unspecified, // Doesn't matter since we're failing wf resp_chan, true, + false, ) .unwrap_or_else(|e| { dbg_panic!("Got next page request when auto-failing workflow: {e:?}"); @@ -764,7 +767,10 @@ impl ManagedRun { let completing_heartbeat_autocomplete = matches!(self.activation, Some(OutstandingActivation::Autocomplete)) && self.waiting_on_la.is_some(); + // Lang decides a task heartbeat from a count or a clock, so replay must not depend on it. + let lang_heartbeat = completion.heartbeat_task && !self.wfm.machines.replaying; let completing_la_heartbeat = completing_heartbeat_autocomplete + || lang_heartbeat || self .waiting_on_la .as_ref() @@ -1844,6 +1850,8 @@ struct RunActivationCompletion { /// unblock. Must always be `Some` when initialized. resp_chan: Option>, versioning_behavior: VersioningBehavior, + /// Lang asked to complete the workflow task now while local activities run. + heartbeat_task: bool, } #[derive(Debug, derive_more::From)] enum ActOrFulfill { diff --git a/crates/sdk-core/src/worker/workflow/mod.rs b/crates/sdk-core/src/worker/workflow/mod.rs index f68227d88a..460e25dc43 100644 --- a/crates/sdk-core/src/worker/workflow/mod.rs +++ b/crates/sdk-core/src/worker/workflow/mod.rs @@ -1566,6 +1566,7 @@ fn validate_completion( commands, used_flags: success.used_internal_flags, versioning_behavior: success.versioning_behavior.try_into().unwrap_or_default(), + heartbeat_task: success.heartbeat_task, }) } Some(workflow_activation_completion::Status::Failed(failure)) => { @@ -1590,6 +1591,7 @@ enum ValidatedCompletion { commands: Vec, used_flags: Vec, versioning_behavior: VersioningBehavior, + heartbeat_task: bool, }, Fail { run_id: String, diff --git a/crates/sdk-core/src/worker/workflow/workflow_stream.rs b/crates/sdk-core/src/worker/workflow/workflow_stream.rs index 7c45c2fd19..34bec33b6e 100644 --- a/crates/sdk-core/src/worker/workflow/workflow_stream.rs +++ b/crates/sdk-core/src/worker/workflow/workflow_stream.rs @@ -306,6 +306,7 @@ impl WFStream { commands, used_flags, versioning_behavior, + heartbeat_task, .. } => match rh.successful_completion( commands, @@ -313,6 +314,7 @@ impl WFStream { versioning_behavior, complete.response_tx, false, + heartbeat_task, ) { Ok(acts) => acts, Err(npr) => { diff --git a/crates/sdk/src/workflow_future.rs b/crates/sdk/src/workflow_future.rs index f18604ce0b..9a9e89acf4 100644 --- a/crates/sdk/src/workflow_future.rs +++ b/crates/sdk/src/workflow_future.rs @@ -238,6 +238,7 @@ impl WorkflowFuture { commands: activation_cmds, used_internal_flags: vec![], versioning_behavior: VersioningBehavior::Unspecified.into(), + heartbeat_task: false, }, )), ..Default::default() @@ -336,6 +337,11 @@ impl WorkflowFuture { .context("Nexus operation must have result")?; push_polled_context!(ActivationJobContext::Passive); } + // This SDK has no stream reader yet. A wake carries no records, only a hint to + // read the stream, so dropping it loses nothing. + Variant::ResolveStreamWake(_) => { + push_polled_context!(ActivationJobContext::Passive); + } Variant::RemoveFromCache(_) => { unreachable!("Cache removal should happen higher up"); } diff --git a/crates/streams/src/handle.rs b/crates/streams/src/handle.rs index 565e4baab7..2e1f4287a9 100644 --- a/crates/streams/src/handle.rs +++ b/crates/streams/src/handle.rs @@ -88,31 +88,35 @@ struct ProducerKey { } #[derive(Debug, Clone)] -struct Checked { - first_run_id: String, +pub(crate) struct Checked { + pub(crate) first_run_id: String, at: Instant, } /// Remembers when each key was last checked, forgetting the oldest when full. #[derive(Debug)] -struct Remembered { +pub(crate) struct Remembered { entries: Mutex>, limit: usize, } impl Remembered { - fn new(limit: usize) -> Self { + pub(crate) fn new(limit: usize) -> Self { Self { entries: Mutex::default(), limit: limit.max(1), } } - fn get(&self, key: &K) -> Option { + pub(crate) fn get(&self, key: &K) -> Option { self.entries.lock().unwrap().get(key).cloned() } - fn put(&self, key: K, first_run_id: &str, fresh_for: Duration) { + pub(crate) fn forget(&self, key: &K) { + self.entries.lock().unwrap().remove(key); + } + + pub(crate) fn put(&self, key: K, first_run_id: &str, fresh_for: Duration) { let mut entries = self.entries.lock().unwrap(); if entries.len() >= self.limit && !entries.contains_key(&key) { entries.retain(|_, checked| checked.at.elapsed() < fresh_for); @@ -133,10 +137,6 @@ impl Remembered { }, ); } - - fn forget(&self, key: &K) { - self.entries.lock().unwrap().remove(key); - } } /// The producer attempt a call writes as, and the topic it writes, so a close at its owner check @@ -653,8 +653,23 @@ impl Streams { )) } + /// The run chain a checked stream address names now. A pinned run's chain comes from the + /// owner cache, since it never changes. The latest run's chain comes from a fresh describe, + /// for a caller that must not land on a chain that already ended. + pub(crate) async fn chain_now(&self, stream: &StreamAddress) -> StreamResult { + let key = owner_key(stream); + if !key.run_id.is_empty() { + return self.chain(stream).await; + } + let described = self.describe(&key, "").await.map_err(not_found)?; + let first_run_id = first_run_id(&key, &described)?; + self.chains + .put(key, &first_run_id, self.options.owner_recheck); + Ok(chain_of(stream, &first_run_id)) + } + /// The chain a stream address names now. - async fn chain(&self, stream: &StreamAddress) -> StreamResult { + pub(crate) async fn chain(&self, stream: &StreamAddress) -> StreamResult { let key = owner_key(stream); // A pinned run's chain never changes. The latest run's chain changes when a new chain // takes the Workflow id. @@ -704,7 +719,7 @@ impl Streams { } /// The address of a stream this release keeps, checked. -fn address(stream: Option<&StreamAddress>) -> StreamResult<&StreamAddress> { +pub(crate) fn address(stream: Option<&StreamAddress>) -> StreamResult<&StreamAddress> { let stream = stream.ok_or_else(|| StreamError::invalid_argument("a stream call needs its stream"))?; if stream.owner_kind != StreamOwnerKind::Workflow as i32 { diff --git a/crates/streams/src/lib.rs b/crates/streams/src/lib.rs index 4162f88f2a..b077edebc8 100644 --- a/crates/streams/src/lib.rs +++ b/crates/streams/src/lib.rs @@ -21,6 +21,7 @@ mod record; mod redis; mod service; mod store; +mod subscribe; #[cfg(feature = "redis")] pub use self::redis::{RedisStore, RedisStoreOptions}; @@ -31,8 +32,8 @@ pub use handle::{OWNER_RECHECK, OWNER_START_GRACE, Streams, StreamsOptions}; pub use memory::MemoryStore; pub use notify::{ Arm, CLOSE_RETRY_CAP, CLOSE_RETRY_FIRST, CLOSE_RETRY_LIMIT, DEFAULT_MAX_NOTIFIERS, FLUSH_LIMIT, - NOTIFIER_REFRESH, Notifier, NotifierClient, NotifyError, NotifyingStore, StreamNotification, - progress_counter, + NOTIFIER_REFRESH, Notifier, NotifierClient, NotifierPoll, NotifyError, NotifyingStore, + StreamNotification, WorkflowAttachment, progress_counter, }; pub use owner::{OwnerClient, OwnerDescription, OwnerError}; pub use reader::{ diff --git a/crates/streams/src/notify.rs b/crates/streams/src/notify.rs index 9c6ec21b2a..edd15f21b3 100644 --- a/crates/streams/src/notify.rs +++ b/crates/streams/src/notify.rs @@ -86,11 +86,80 @@ pub enum NotifyError { Unsupported(String), } -/// Sends notifications to the stream notifier on the server. +/// A Workflow's reader, to attach to a stream's notifier as a wake target. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WorkflowAttachment { + /// The run chain whose stream it reads. Its first run keys the notifier. + pub chain: ChainId, + /// The stream's topic. + pub topic: String, + /// The attach is idempotent by it, and a detach names it. + pub request_id: String, + /// The reader's Workflow, in the chain's namespace. + pub workflow_id: String, + /// The reader's run. + pub run_id: String, + /// Names the reader in its run. + pub subscription_id: String, +} + +/// What a stream's notifier answered a client's wait with. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct NotifierPoll { + /// The highest counter the notifier has seen. + pub counter: i64, + /// The position that came with `counter`. + pub position: String, + /// Whether the stream closed. + pub closed: bool, + /// The result the stream closed with, when it gave one. + pub close_result: Option, + /// The failure the stream closed with, when it closed with one. + pub close_failure: Option, +} + +/// Sends notifications to the stream notifier on the server, attaches and detaches the +/// Workflows it wakes, and waits on it for clients. #[async_trait::async_trait] pub trait NotifierClient: Send + Sync { /// Sends one notification. async fn notify(&self, notification: StreamNotification) -> Result<(), NotifyError>; + + /// Attaches a Workflow's reader to a stream's notifier. Idempotent by its request id. + async fn attach_workflow(&self, attachment: WorkflowAttachment) -> Result<(), NotifyError> { + let _ = attachment; + Err(NotifyError::Unsupported( + "this client can't attach readers".to_string(), + )) + } + + /// Detaches what the attach with `request_id` attached to `chain`'s stream on `topic`. + async fn detach( + &self, + chain: ChainId, + topic: String, + request_id: String, + ) -> Result<(), NotifyError> { + let _ = (chain, topic, request_id); + Err(NotifyError::Unsupported( + "this client can't detach readers".to_string(), + )) + } + + /// Waits until the notifier of `chain`'s stream on `topic` has a counter above + /// `after_counter`, or the stream closed. Answers `None` when the wait reached its limit with + /// no answer from the server. + async fn poll( + &self, + chain: ChainId, + topic: String, + after_counter: i64, + ) -> Result, NotifyError> { + let _ = (chain, topic, after_counter); + Err(NotifyError::Unsupported( + "this client can't wait on a notifier".to_string(), + )) + } } /// The notifier counter for a store position. diff --git a/crates/streams/src/service.rs b/crates/streams/src/service.rs index a09821f2ee..c99da7c409 100644 --- a/crates/streams/src/service.rs +++ b/crates/streams/src/service.rs @@ -7,8 +7,9 @@ use crate::{ AppendRequest, CloseRequest, DeleteOwnerRequest, DescribeStoreRequest, DescribeStoreResponse, FlushNotificationsRequest, FlushNotificationsResponse, LatestRequest, ReadRequest, RetireProducerRequest, StreamFailure, StreamStoreConfig, - stream_store_config, + SubscribeRequest, UnsubscribeRequest, WaitForNotificationRequest, stream_store_config, }, + subscribe::Subscriptions, }; use prost::Message; use std::{sync::Arc, time::Duration}; @@ -28,6 +29,7 @@ const MAX_READ_WAIT: Duration = Duration::from_secs(60); pub struct StreamService { store: Arc, streams: Streams, + subscriptions: Subscriptions, } impl std::fmt::Debug for StreamService { @@ -49,7 +51,8 @@ impl StreamService { } /// Like [StreamService::connect], closing each stream's notifier through `notifier` after - /// the store closes the stream, and sending progress too when the config turns it on. + /// the store closes the stream, sending progress too when the config turns it on, attaching + /// and detaching Workflow readers through it, and waiting on it for client readers. pub async fn connect_with_notifier( config: StreamStoreConfig, owner: Arc, @@ -75,7 +78,7 @@ impl StreamService { "stream notifications need a client to send them with", )); } - let store: Arc = match notifier { + let store: Arc = match notifier.clone() { Some(client) => { let max_notifiers = match config.max_notifiers { 0 => DEFAULT_MAX_NOTIFIERS, @@ -95,7 +98,11 @@ impl StreamService { StreamError::refused("the stream store's owner start grace must not be negative") })?; } - Ok(Self::new(store, owner, options)) + let service = Self::new(store, owner, options); + Ok(match notifier { + Some(client) => service.with_subscriptions(client), + None => service, + }) } /// Serves calls on `store`. @@ -105,7 +112,17 @@ impl StreamService { options: StreamsOptions, ) -> Self { let streams = Streams::new(store.clone(), owner, options); - Self { store, streams } + Self { + store, + streams, + subscriptions: Subscriptions::default(), + } + } + + /// Attaches and detaches Workflow readers, and waits for client readers, through `client`. + pub fn with_subscriptions(mut self, client: Arc) -> Self { + self.subscriptions = Subscriptions::new(Some(client)); + self } /// The store, for the Workers that append Workflow output to it. @@ -155,6 +172,28 @@ impl StreamService { .await; Ok(FlushNotificationsResponse {}.encode_to_vec()) } + "Subscribe" => self + .subscriptions + .subscribe( + &self.streams, + self.store.as_ref(), + decode::(rpc, request)?, + ) + .await + .map(|r| r.encode_to_vec()), + "Unsubscribe" => self + .subscriptions + .unsubscribe(&self.streams, decode::(rpc, request)?) + .await + .map(|r| r.encode_to_vec()), + "WaitForNotification" => self + .subscriptions + .wait( + &self.streams, + decode::(rpc, request)?, + ) + .await + .map(|r| r.encode_to_vec()), "DescribeStore" => { decode::(rpc, request)?; Ok(DescribeStoreResponse { diff --git a/crates/streams/src/subscribe.rs b/crates/streams/src/subscribe.rs new file mode 100644 index 0000000000..5232efead3 --- /dev/null +++ b/crates/streams/src/subscribe.rs @@ -0,0 +1,743 @@ +//! Readers that sleep on a stream's notifier. +//! +//! A Workflow's reader attaches to the stream's notifier before its first read, so the notifier +//! wakes the Workflow on its next Workflow Task when the stream moves. The attach names the reader +//! by its run and its subscription id, which replay gives back, so a rerun attaches as the same +//! reader and the server keeps one callback. A client reader has no task to wake, so it waits on +//! the notifier with a long poll. + +use crate::{ + NotifierClient, NotifyError, OWNER_RECHECK, StreamError, StreamResult, StreamStore, Streams, + StreamsOptions, WorkflowAttachment, + handle::{Remembered, address}, + proto::{ + StreamAddress, SubscribeRequest, SubscribeResponse, UnsubscribeRequest, + UnsubscribeResponse, WaitForNotificationRequest, WaitForNotificationResponse, + }, +}; +use std::sync::Arc; + +/// The longest subscription id the server takes. +const MAX_SUBSCRIPTION_ID_BYTES: usize = 256; + +/// The subscriptions this process attached, so a reader's later reads send no attach. A reader +/// that ends or a run that closes never detaches, so the oldest are forgotten. A reader this +/// process forgot attaches again, and the server dedupes that attach. +pub(crate) struct Subscriptions { + client: Option>, + attached: Remembered, +} + +impl Default for Subscriptions { + fn default() -> Self { + Self::new(None) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct SubscriptionKey { + namespace: String, + owner_workflow_id: String, + owner_run_id: String, + topic: String, + request_id: String, +} + +impl SubscriptionKey { + fn new(stream: &StreamAddress, request_id: &str) -> Self { + Self { + namespace: stream.namespace.clone(), + owner_workflow_id: stream.workflow_id.clone(), + owner_run_id: stream.run_id.clone(), + topic: stream.topic.clone(), + request_id: request_id.to_string(), + } + } +} + +/// The request id of a reader's attach. The server dedupes attaches by it, so it depends only on +/// what replay gives back. A later generation is a new attach, after the notifier dropped the +/// reader and the server would answer the old id as already done. +fn request_id(run_id: &str, subscription_id: &str, generation: u32) -> String { + match generation { + 0 => format!("jetty/{run_id}/{subscription_id}"), + generation => format!("jetty/{run_id}/{subscription_id}/{generation}"), + } +} + +fn check_subscription(run_id: &str, subscription_id: &str) -> StreamResult<()> { + if run_id.is_empty() || subscription_id.is_empty() { + return Err(StreamError::invalid_argument( + "a subscription needs the reader's run id and its subscription id", + )); + } + if subscription_id.len() > MAX_SUBSCRIPTION_ID_BYTES { + return Err(StreamError::invalid_argument(format!( + "a subscription id is at most {MAX_SUBSCRIPTION_ID_BYTES} bytes, not {}", + subscription_id.len() + ))); + } + Ok(()) +} + +fn subscribe_error(error: NotifyError) -> StreamError { + match error { + NotifyError::Refused(message) => StreamError::refused(message), + NotifyError::Unsupported(message) => StreamError::unsupported(message), + NotifyError::Failed(message) => StreamError::storage(message), + } +} + +impl Subscriptions { + pub(crate) fn new(client: Option>) -> Self { + Self { + client, + attached: Remembered::new(StreamsOptions::default().remembered), + } + } + + fn client(&self) -> StreamResult<&Arc> { + self.client.as_ref().ok_or_else(|| { + StreamError::unsupported("a Workflow reader needs a client connected to the server") + }) + } + + /// Attaches the reader once per process. On an empty stream the attach also arms the stream, + /// in the same store step that finds it empty, so the first append wakes the Workflow. + pub(crate) async fn subscribe( + &self, + streams: &Streams, + store: &dyn StreamStore, + request: SubscribeRequest, + ) -> StreamResult { + let stream = address(request.stream.as_ref())?; + if request.workflow_id.is_empty() { + return Err(StreamError::invalid_argument( + "a subscription needs the reader's Workflow id", + )); + } + check_subscription(&request.run_id, &request.subscription_id)?; + let request_id = request_id( + &request.run_id, + &request.subscription_id, + request.generation, + ); + let key = SubscriptionKey::new(stream, &request_id); + if let Some(attached) = self.attached.get(&key) { + return Ok(SubscribeResponse { + chain_first_run_id: attached.first_run_id, + attached: false, + }); + } + let client = self.client()?; + // Fresh for the latest run, since the Workflow id may have moved on to a new chain since a + // read warmed the cache, and an attach to the ended chain would never wake the reader. + let chain = streams.chain_now(stream).await?; + client + .attach_workflow(WorkflowAttachment { + chain: chain.clone(), + topic: stream.topic.clone(), + request_id, + workflow_id: request.workflow_id, + run_id: request.run_id, + subscription_id: request.subscription_id, + }) + .await + .map_err(subscribe_error)?; + // The reader's first read arms an empty stream too, so a failed arm costs nothing. + if let Err(error) = store.newest_or_arm(&chain, &stream.topic).await { + tracing::warn!( + workflow_id = %chain.workflow_id, + topic = %stream.topic, + %error, + "Could not arm the stream a Workflow reader subscribed to" + ); + } + self.attached.put(key, &chain.first_run_id, OWNER_RECHECK); + Ok(SubscribeResponse { + chain_first_run_id: chain.first_run_id, + attached: true, + }) + } + + /// Detaches the reader from the chain it attached to, or from the chain the stream names now + /// when this process never attached it. + pub(crate) async fn unsubscribe( + &self, + streams: &Streams, + request: UnsubscribeRequest, + ) -> StreamResult { + let stream = address(request.stream.as_ref())?; + check_subscription(&request.run_id, &request.subscription_id)?; + let request_id = request_id( + &request.run_id, + &request.subscription_id, + request.generation, + ); + let client = self.client()?; + let key = SubscriptionKey::new(stream, &request_id); + let chain = match self.attached.get(&key) { + Some(attached) => crate::proto::ChainId { + namespace: stream.namespace.clone(), + workflow_id: stream.workflow_id.clone(), + first_run_id: attached.first_run_id, + }, + None => streams.chain_now(stream).await?, + }; + client + .detach(chain, stream.topic.clone(), request_id) + .await + .map_err(subscribe_error)?; + self.attached.forget(&key); + Ok(UnsubscribeResponse {}) + } + + /// Waits on the notifier of the chain the stream names. A wait that reached its limit answers + /// `timed_out`, and the reader's next read checks the store itself, which covers a + /// notification that a producer never sent. + pub(crate) async fn wait( + &self, + streams: &Streams, + request: WaitForNotificationRequest, + ) -> StreamResult { + let stream = address(request.stream.as_ref())?; + let client = self.client()?; + let chain = streams.chain(stream).await?; + let polled = client + .poll(chain, stream.topic.clone(), request.after_counter) + .await + .map_err(subscribe_error)?; + let Some(polled) = polled else { + return Ok(WaitForNotificationResponse { + timed_out: true, + ..Default::default() + }); + }; + Ok(WaitForNotificationResponse { + timed_out: !polled.closed && polled.counter <= request.after_counter, + counter: polled.counter, + position: polled.position, + closed: polled.closed, + close_result: polled.close_result, + close_failure: polled.close_failure, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + MemoryStore, NotifierPoll, OwnerClient, OwnerDescription, OwnerError, StreamNotification, + StreamService, StreamsOptions, + proto::{ChainId, StoreAppendRequest, StreamFailureKind, StreamOwnerKind}, + }; + use prost::Message; + use std::sync::Mutex; + use temporalio_common::protos::temporal::api::enums::v1::WorkflowExecutionStatus; + + /// The owner's chain is `first_run_id`, which a test moves to a new chain. + struct Owner(Mutex); + + #[async_trait::async_trait] + impl OwnerClient for Owner { + async fn describe( + &self, + _: &str, + _: &str, + _: &str, + ) -> Result { + let first_run_id = self.0.lock().unwrap().clone(); + Ok(OwnerDescription { + run_id: first_run_id.clone(), + first_run_id, + status: WorkflowExecutionStatus::Running, + }) + } + } + + /// Records attaches and detaches, and answers each with `answer`. + #[derive(Default)] + struct Server { + attached: Mutex>, + detached: Mutex>, + answer: Mutex>, + polled: Mutex>, + poll_answer: Mutex, NotifyError>>>, + } + + #[async_trait::async_trait] + impl NotifierClient for Server { + async fn notify(&self, _: StreamNotification) -> Result<(), NotifyError> { + Ok(()) + } + + async fn attach_workflow(&self, attachment: WorkflowAttachment) -> Result<(), NotifyError> { + if let Some(error) = self.answer.lock().unwrap().clone() { + return Err(error); + } + self.attached.lock().unwrap().push(attachment); + Ok(()) + } + + async fn detach( + &self, + chain: ChainId, + topic: String, + request_id: String, + ) -> Result<(), NotifyError> { + self.detached + .lock() + .unwrap() + .push((chain, topic, request_id)); + Ok(()) + } + + async fn poll( + &self, + chain: ChainId, + topic: String, + after_counter: i64, + ) -> Result, NotifyError> { + self.polled + .lock() + .unwrap() + .push((chain, topic, after_counter)); + self.poll_answer + .lock() + .unwrap() + .clone() + .expect("the test sets the poll's answer") + } + } + + struct Setup { + service: StreamService, + store: Arc, + owner: Arc, + server: Arc, + } + + fn setup() -> Setup { + let store = Arc::new(MemoryStore::new()); + let owner = Arc::new(Owner(Mutex::new("chain-1".to_string()))); + let server = Arc::new(Server::default()); + let service = StreamService::new(store.clone(), owner.clone(), StreamsOptions::default()) + .with_subscriptions(server.clone()); + Setup { + service, + store, + owner, + server, + } + } + + fn stream() -> StreamAddress { + StreamAddress { + namespace: "ns".to_string(), + owner_kind: StreamOwnerKind::Workflow as i32, + workflow_id: "owner".to_string(), + run_id: String::new(), + topic: "tokens".to_string(), + } + } + + fn subscribe_request(subscription_id: &str) -> SubscribeRequest { + SubscribeRequest { + stream: Some(stream()), + workflow_id: "reader".to_string(), + run_id: "reader-run".to_string(), + subscription_id: subscription_id.to_string(), + generation: 0, + } + } + + async fn subscribe( + service: &StreamService, + request: SubscribeRequest, + ) -> Result { + service + .call("Subscribe", &request.encode_to_vec()) + .await + .map(|answer| SubscribeResponse::decode(answer.as_slice()).unwrap()) + } + + fn chain(first_run_id: &str) -> ChainId { + ChainId { + namespace: "ns".to_string(), + workflow_id: "owner".to_string(), + first_run_id: first_run_id.to_string(), + } + } + + /// The arm the next append on the stream reports. + async fn arm_of_next_append(store: &MemoryStore, first_run_id: &str) -> Option { + store + .append(StoreAppendRequest { + chain: Some(chain(first_run_id)), + topic: "tokens".to_string(), + producer_id: "p".to_string(), + attempt: 1, + sequence: 1, + digest: vec![1], + records: vec![b"record".to_vec()], + }) + .await + .unwrap() + .arm + } + + #[tokio::test] + async fn a_full_process_forgets_its_subscriptions_and_attaches_again() { + let setup = setup(); + let subscriptions = Subscriptions { + client: Some(setup.server.clone()), + attached: Remembered::new(2), + }; + let streams = Streams::new( + setup.store.clone(), + setup.owner.clone(), + StreamsOptions::default(), + ); + let subscribe = |id: &str| { + subscriptions.subscribe(&streams, setup.store.as_ref(), subscribe_request(id)) + }; + for id in ["1", "2", "3"] { + assert!(subscribe(id).await.unwrap().attached); + } + assert!(!subscribe("3").await.unwrap().attached); + assert!( + subscribe("1").await.unwrap().attached, + "the oldest was forgotten" + ); + assert_eq!(setup.server.attached.lock().unwrap().len(), 4); + } + + #[tokio::test] + async fn a_subscription_attaches_the_workflow_target_once_per_process() { + let setup = setup(); + let first = subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + assert!(first.attached); + assert_eq!(first.chain_first_run_id, "chain-1"); + let again = subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + assert!(!again.attached, "a later read sends no attach"); + assert_eq!(again.chain_first_run_id, "chain-1"); + let attached = setup.server.attached.lock().unwrap().clone(); + assert_eq!( + attached, + [WorkflowAttachment { + chain: chain("chain-1"), + topic: "tokens".to_string(), + request_id: "jetty/reader-run/1".to_string(), + workflow_id: "reader".to_string(), + run_id: "reader-run".to_string(), + subscription_id: "1".to_string(), + }] + ); + } + + #[tokio::test] + async fn a_new_generation_attaches_again_after_the_notifier_dropped_the_reader() { + let setup = setup(); + subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + let again = subscribe( + &setup.service, + SubscribeRequest { + generation: 1, + ..subscribe_request("1") + }, + ) + .await + .unwrap(); + assert!(again.attached, "a resubscribe sends a new attach"); + let attached = setup.server.attached.lock().unwrap().clone(); + let ids: Vec<&str> = attached.iter().map(|a| a.request_id.as_str()).collect(); + assert_eq!(ids, ["jetty/reader-run/1", "jetty/reader-run/1/1"]); + assert!( + attached.iter().all(|a| a.subscription_id == "1"), + "the wakes still name the same reader" + ); + } + + #[tokio::test] + async fn readers_of_one_run_attach_apart() { + let setup = setup(); + subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + subscribe(&setup.service, subscribe_request("2")) + .await + .unwrap(); + let ids: Vec = setup + .server + .attached + .lock() + .unwrap() + .iter() + .map(|a| a.request_id.clone()) + .collect(); + assert_eq!(ids, ["jetty/reader-run/1", "jetty/reader-run/2"]); + } + + #[tokio::test] + async fn a_subscription_to_an_empty_stream_arms_it_for_the_first_append() { + let setup = setup(); + subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + assert!( + arm_of_next_append(&setup.store, "chain-1") + .await + .is_some_and(|arm| arm > 0), + "the first append wakes the Workflow" + ); + } + + #[tokio::test] + async fn a_subscription_to_a_stream_with_records_does_not_arm() { + let setup = setup(); + arm_of_next_append(&setup.store, "chain-1").await; + subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + let landed = setup + .store + .append(StoreAppendRequest { + chain: Some(chain("chain-1")), + topic: "tokens".to_string(), + producer_id: "other".to_string(), + attempt: 1, + sequence: 1, + digest: vec![1], + records: vec![b"record".to_vec()], + }) + .await + .unwrap(); + assert_eq!(landed.arm, Some(0), "the reader reads the records first"); + } + + #[tokio::test] + async fn a_subscription_attaches_to_the_chain_that_holds_the_workflow_id_now() { + let setup = setup(); + // A read warms the owner cache with the first chain. + setup + .service + .call( + "Latest", + &crate::proto::LatestRequest { + stream: Some(stream()), + } + .encode_to_vec(), + ) + .await + .unwrap(); + *setup.owner.0.lock().unwrap() = "chain-2".to_string(); + let answer = subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + assert_eq!(answer.chain_first_run_id, "chain-2"); + } + + #[tokio::test] + async fn a_refused_attach_fails_the_subscription_and_is_not_remembered() { + let setup = setup(); + *setup.server.answer.lock().unwrap() = Some(NotifyError::Refused( + "the execution has 32 stream wake sources".to_string(), + )); + let error = subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap_err(); + assert_eq!(error.kind, StreamFailureKind::Refused as i32); + *setup.server.answer.lock().unwrap() = None; + let answer = subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + assert!(answer.attached, "the next read tries the attach again"); + } + + #[tokio::test] + async fn an_attach_that_may_land_later_fails_as_storage() { + let setup = setup(); + *setup.server.answer.lock().unwrap() = Some(NotifyError::Failed("down".to_string())); + let error = subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap_err(); + assert_eq!(error.kind, StreamFailureKind::Storage as i32); + } + + #[tokio::test] + async fn a_subscription_needs_its_reader() { + let setup = setup(); + for request in [ + SubscribeRequest { + workflow_id: String::new(), + ..subscribe_request("1") + }, + SubscribeRequest { + run_id: String::new(), + ..subscribe_request("1") + }, + subscribe_request(""), + subscribe_request(&"x".repeat(MAX_SUBSCRIPTION_ID_BYTES + 1)), + SubscribeRequest { + stream: None, + ..subscribe_request("1") + }, + ] { + let error = subscribe(&setup.service, request).await.unwrap_err(); + assert_eq!(error.kind, StreamFailureKind::InvalidArgument as i32); + } + assert!(setup.server.attached.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn a_subscription_without_a_client_is_unsupported() { + let service = StreamService::new( + Arc::new(MemoryStore::new()), + Arc::new(Owner(Mutex::new("chain-1".to_string()))), + StreamsOptions::default(), + ); + let error = subscribe(&service, subscribe_request("1")) + .await + .unwrap_err(); + assert_eq!(error.kind, StreamFailureKind::Unsupported as i32); + } + + #[tokio::test] + async fn an_unsubscribe_detaches_from_the_chain_it_attached_to_and_forgets_it() { + let setup = setup(); + subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + *setup.owner.0.lock().unwrap() = "chain-2".to_string(); + let unsubscribe = UnsubscribeRequest { + stream: Some(stream()), + run_id: "reader-run".to_string(), + subscription_id: "1".to_string(), + ..Default::default() + }; + setup + .service + .call("Unsubscribe", &unsubscribe.encode_to_vec()) + .await + .unwrap(); + assert_eq!( + setup.server.detached.lock().unwrap().clone(), + [( + chain("chain-1"), + "tokens".to_string(), + "jetty/reader-run/1".to_string() + )] + ); + let answer = subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + assert!(answer.attached, "a new subscription attaches again"); + assert_eq!(answer.chain_first_run_id, "chain-2"); + } + + async fn wait( + setup: &Setup, + after_counter: i64, + answer: Result, NotifyError>, + ) -> Result { + *setup.server.poll_answer.lock().unwrap() = Some(answer); + setup + .service + .call( + "WaitForNotification", + &WaitForNotificationRequest { + stream: Some(stream()), + after_counter, + } + .encode_to_vec(), + ) + .await + .map(|answer| WaitForNotificationResponse::decode(answer.as_slice()).unwrap()) + } + + #[tokio::test] + async fn a_wait_asks_the_notifier_of_the_streams_chain_for_a_counter_above_the_arm() { + let setup = setup(); + let answer = wait( + &setup, + 41, + Ok(Some(NotifierPoll { + counter: 42, + position: "cursor-42".to_string(), + ..Default::default() + })), + ) + .await + .unwrap(); + assert_eq!( + setup.server.polled.lock().unwrap().clone(), + [(chain("chain-1"), "tokens".to_string(), 41)] + ); + assert!(!answer.timed_out && !answer.closed); + assert_eq!( + (answer.counter, answer.position.as_str()), + (42, "cursor-42") + ); + } + + #[tokio::test] + async fn a_wait_with_no_new_counter_times_out_for_a_re_read() { + let setup = setup(); + let unchanged = wait( + &setup, + 41, + Ok(Some(NotifierPoll { + counter: 41, + ..Default::default() + })), + ) + .await + .unwrap(); + assert!(unchanged.timed_out); + let deadline = wait(&setup, 41, Ok(None)).await.unwrap(); + assert!(deadline.timed_out && !deadline.closed); + } + + #[tokio::test] + async fn a_wait_on_a_closed_stream_carries_the_outcome() { + let setup = setup(); + let failure = temporalio_common::protos::temporal::api::failure::v1::Failure { + message: "the agent gave up".to_string(), + ..Default::default() + }; + let answer = wait( + &setup, + 41, + Ok(Some(NotifierPoll { + counter: 41, + closed: true, + close_failure: Some(failure.clone()), + ..Default::default() + })), + ) + .await + .unwrap(); + assert!(answer.closed && !answer.timed_out); + assert_eq!(answer.close_failure, Some(failure)); + } + + #[tokio::test] + async fn a_wait_on_a_server_without_the_poll_is_unsupported() { + let setup = setup(); + let error = wait( + &setup, + 0, + Err(NotifyError::Unsupported("unimplemented".to_string())), + ) + .await + .unwrap_err(); + assert_eq!(error.kind, StreamFailureKind::Unsupported as i32); + } +} diff --git a/crates/workflow/src/runtime/instance.rs b/crates/workflow/src/runtime/instance.rs index 7c6b3a0332..3baaf46741 100644 --- a/crates/workflow/src/runtime/instance.rs +++ b/crates/workflow/src/runtime/instance.rs @@ -1119,6 +1119,9 @@ where ActivationJobResult::None } Some(ActivationVariant::RemoveFromCache(_)) => ActivationJobResult::None, + // This runtime has no stream reader yet. A wake carries no records, only a hint + // to read the stream, so dropping it loses nothing. + Some(ActivationVariant::ResolveStreamWake(_)) => ActivationJobResult::None, None => { return Err(Box::new(Failure { message: "Activation job missing variant".to_string(),