Skip to main content

temporalio_workflow/workflow_context/
nexus.rs

1use super::*;
2use crate::{
3    runtime::{SdkGuardedFuture, model::NexusStartResult},
4    workflow_interceptors::{StartNexusOperationInput, call_start_nexus_operation},
5};
6use futures_util::{FutureExt, future::Shared};
7use temporalio_common_wasm::protos::coresdk::nexus::NexusOperationResult;
8
9impl BaseWorkflowContext {
10    pub(crate) fn start_nexus_operation(
11        &self,
12        opts: NexusOperationOptions,
13    ) -> impl CancellableFuture<Output = NexusStartResult> {
14        let input = StartNexusOperationInput::new(opts);
15        let base_ctx = self.clone();
16        let next = WorkflowNext::new(move |input: StartNexusOperationInput| {
17            let mut opts = input.into_options();
18            let cancellation_token = opts
19                .cancellation_token
20                .take()
21                .unwrap_or_else(|| base_ctx.cancellation_token());
22            let seq = base_ctx.inner.seq_nums.borrow_mut().next_nexus_op_seq();
23            let (result_future, unblocker) =
24                CancellableWFCommandFut::new(CancellableID::NexusOp(seq), base_ctx.clone());
25            base_ctx
26                .inner
27                .runtime
28                .register_unblocker(PendingCommandId::NexusOpComplete(seq), unblocker);
29            base_ctx
30                .inner
31                .runtime
32                .host
33                .push_command(opts.into_command(seq));
34            let result_future = CancellableWorkflowOutboundFuture::new(
35                result_future,
36                base_ctx.cancellation_handle(CancellableID::NexusOp(seq)),
37            )
38            .with_cancellation_token(cancellation_token)
39            .shared();
40            let (cmd, unblocker) = CancellableWFCommandFut::new_with_dat(
41                CancellableID::NexusOp(seq),
42                NexusUnblockData {
43                    result_future: result_future.clone(),
44                    schedule_seq: seq,
45                    base_ctx: base_ctx.clone(),
46                },
47                base_ctx.clone(),
48            );
49            base_ctx
50                .inner
51                .runtime
52                .register_unblocker(PendingCommandId::NexusOpStart(seq), unblocker);
53            cancellable_outbound(cmd)
54        });
55        let interceptors = self.inner.workflow_interceptors.clone();
56        let future = call_start_nexus_operation(
57            interceptors,
58            WorkflowInterceptorContext::new(self.clone()),
59            input,
60            next,
61        );
62        self.prepare_cancellable_outbound_future(future)
63    }
64}
65
66impl<W> SyncWorkflowContext<W> {
67    /// Start a Nexus operation.
68    pub fn start_nexus_operation(
69        &self,
70        opts: NexusOperationOptions,
71    ) -> impl CancellableFuture<Output = NexusStartResult> {
72        self.base.start_nexus_operation(opts)
73    }
74}
75
76impl<W> WorkflowContext<W> {
77    /// Start a Nexus operation.
78    pub fn start_nexus_operation(
79        &self,
80        opts: NexusOperationOptions,
81    ) -> impl CancellableFuture<Output = NexusStartResult> {
82        self.sync.start_nexus_operation(opts)
83    }
84}
85
86impl WfCtxProtectedDat {
87    fn next_nexus_op_seq(&mut self) -> u32 {
88        let seq = self.next_nexus_op_sequence_number;
89        self.next_nexus_op_sequence_number += 1;
90        seq
91    }
92}
93
94#[derive(derive_more::Debug)]
95#[debug("StartedNexusOperation{{ operation_token: {operation_token:?} }}")]
96/// Handle to a started Nexus operation.
97pub struct StartedNexusOperation {
98    /// The operation token, if the operation started asynchronously
99    pub operation_token: Option<String>,
100    #[debug(skip)]
101    pub(crate) result_future: Shared<CancellableWorkflowOutboundFuture<NexusOperationResult>>,
102    pub(crate) schedule_seq: u32,
103    #[debug(skip)]
104    pub(crate) base_ctx: BaseWorkflowContext,
105}
106
107pub(crate) struct NexusUnblockData {
108    pub(crate) result_future: Shared<CancellableWorkflowOutboundFuture<NexusOperationResult>>,
109    pub(crate) schedule_seq: u32,
110    pub(crate) base_ctx: BaseWorkflowContext,
111}
112
113impl StartedNexusOperation {
114    /// Wait for the operation result.
115    pub async fn result(&self) -> NexusOperationResult {
116        // The result future is a `Shared`; poll it inside an `SdkWakeGuard` (via
117        // `SdkGuardedFuture`) so its internal waker machinery isn't mistaken for a non-SDK wake on
118        // replay (which would fail the workflow task with TMPRL1100).
119        SdkGuardedFuture(self.result_future.clone()).await
120    }
121
122    /// Request cancellation of the operation.
123    pub fn cancel(&self) {
124        self.base_ctx
125            .cancel(CancellableID::NexusOp(self.schedule_seq));
126    }
127}