temporalio 0.0.2 → 0.1.0

data/bridge/sdk-core/core/src/worker/workflow/run_cache.rs

@@ -1,22 +1,22 @@
 use crate::{
     telemetry::metrics::workflow_type,
     worker::workflow::{
-        managed_run::WorkflowManager, HistoryUpdate, LocalActivityRequestSink, ManagedRunHandle,
-        NewIncomingWFT, RunUpdateResponse,
+        managed_run::{ManagedRun, RunUpdateAct},
+        HistoryUpdate, LocalActivityRequestSink, PermittedWFT, RunBasics,
     },
     MetricsContext,
 };
 use lru::LruCache;
-use std::{num::NonZeroUsize, time::Instant};
-use tokio::sync::mpsc::UnboundedSender;
+use std::{mem, num::NonZeroUsize, rc::Rc};
+use temporal_sdk_core_protos::temporal::api::workflowservice::v1::get_system_info_response;
 
 pub(super) struct RunCache {
     max: usize,
     namespace: String,
-    run_update_tx: UnboundedSender<RunUpdateResponse>,
+    server_capabilities: get_system_info_response::Capabilities,
     /// Run id -> Data
-    runs: LruCache<String, ManagedRunHandle>,
-    local_activity_request_sink: LocalActivityRequestSink,
+    runs: LruCache<String, ManagedRun>,
+    local_activity_request_sink: Rc<dyn LocalActivityRequestSink>,
 
     metrics: MetricsContext,
 }
@@ -25,8 +25,8 @@ impl RunCache {
     pub fn new(
         max_cache_size: usize,
         namespace: String,
-        run_update_tx: UnboundedSender<RunUpdateResponse>,
-        local_activity_request_sink: LocalActivityRequestSink,
+        server_capabilities: get_system_info_response::Capabilities,
+        local_activity_request_sink: impl LocalActivityRequestSink,
         metrics: MetricsContext,
     ) -> Self {
         // The cache needs room for at least one run, otherwise we couldn't do anything. In
@@ -39,80 +39,63 @@ impl RunCache {
         Self {
             max: max_cache_size,
             namespace,
-            run_update_tx,
+            server_capabilities,
             runs: LruCache::new(
                 NonZeroUsize::new(lru_size).expect("LRU size is guaranteed positive"),
             ),
-            local_activity_request_sink,
+            local_activity_request_sink: Rc::new(local_activity_request_sink),
             metrics,
         }
     }
 
-    pub fn instantiate_or_update(
-        &mut self,
-        run_id: &str,
-        workflow_id: &str,
-        wf_type: &str,
-        history_update: HistoryUpdate,
-        start_time: Instant,
-    ) -> &mut ManagedRunHandle {
+    pub fn instantiate_or_update(&mut self, mut pwft: PermittedWFT) -> RunUpdateAct {
         let cur_num_cached_runs = self.runs.len();
+        let run_id = &pwft.work.execution.run_id;
 
-        if self.runs.contains(run_id) {
-            // For some weird reason, maybe a NLL bug, there are double-mutable-borrow errors if I
-            // use get_mut above instead of in here (even though we always return from this branch).
-            // So, forced to do this.
-            let run_handle = self.runs.get_mut(run_id).unwrap();
-
-            run_handle.metrics.sticky_cache_hit();
-            run_handle.incoming_wft(NewIncomingWFT {
-                history_update: Some(history_update),
-                start_time,
-            });
+        if let Some(run_handle) = self.runs.get_mut(run_id) {
+            let rur = run_handle.incoming_wft(pwft);
             self.metrics.cache_size(cur_num_cached_runs as u64);
-            return run_handle;
+            return rur;
         }
 
         // Create a new workflow machines instance for this workflow, initialize it, and
         // track it.
         let metrics = self
             .metrics
-            .with_new_attrs([workflow_type(wf_type.to_string())]);
-        let wfm = WorkflowManager::new(
-            history_update,
-            self.namespace.clone(),
-            workflow_id.to_string(),
-            wf_type.to_string(),
-            run_id.to_string(),
-            metrics.clone(),
-        );
-        let mut mrh = ManagedRunHandle::new(
-            wfm,
-            self.run_update_tx.clone(),
+            .with_new_attrs([workflow_type(pwft.work.workflow_type.clone())]);
+        // Replace the update in the wft with a dummy one, since we must instantiate the machines
+        // with the update.
+        let history_update = mem::replace(&mut pwft.work.update, HistoryUpdate::dummy());
+        let mut mrh = ManagedRun::new(
+            RunBasics {
+                namespace: self.namespace.clone(),
+                workflow_id: pwft.work.execution.workflow_id.clone(),
+                workflow_type: pwft.work.workflow_type.clone(),
+                run_id: pwft.work.execution.run_id.clone(),
+                history: history_update,
+                metrics,
+                capabilities: &self.server_capabilities,
+            },
             self.local_activity_request_sink.clone(),
-            metrics,
         );
-        mrh.incoming_wft(NewIncomingWFT {
-            history_update: None,
-            start_time,
-        });
-        if self.runs.push(run_id.to_string(), mrh).is_some() {
+        let run_id = run_id.to_string();
+        let rur = mrh.incoming_wft(pwft);
+        if self.runs.push(run_id, mrh).is_some() {
             panic!("Overflowed run cache! Cache owner is expected to avoid this!");
         }
         self.metrics.cache_size(cur_num_cached_runs as u64 + 1);
-        // This is safe, we just inserted.
-        self.runs.get_mut(run_id).unwrap()
+        rur
     }
-    pub fn remove(&mut self, k: &str) -> Option<ManagedRunHandle> {
+    pub fn remove(&mut self, k: &str) -> Option<ManagedRun> {
         let r = self.runs.pop(k);
         self.metrics.cache_size(self.len() as u64);
         r
     }
 
-    pub fn get_mut(&mut self, k: &str) -> Option<&mut ManagedRunHandle> {
+    pub fn get_mut(&mut self, k: &str) -> Option<&mut ManagedRun> {
         self.runs.get_mut(k)
     }
-    pub fn get(&mut self, k: &str) -> Option<&ManagedRunHandle> {
+    pub fn get(&mut self, k: &str) -> Option<&ManagedRun> {
         self.runs.get(k)
     }
 
@@ -121,16 +104,16 @@ impl RunCache {
         self.runs.peek_lru().map(|(run_id, _)| run_id.as_str())
     }
     /// Returns an iterator yielding cached runs in LRU order
-    pub fn runs_lru_order(&self) -> impl Iterator<Item = (&str, &ManagedRunHandle)> {
+    pub fn runs_lru_order(&self) -> impl Iterator<Item = (&str, &ManagedRun)> {
         self.runs.iter().rev().map(|(k, v)| (k.as_str(), v))
     }
-    pub fn peek(&self, k: &str) -> Option<&ManagedRunHandle> {
+    pub fn peek(&self, k: &str) -> Option<&ManagedRun> {
         self.runs.peek(k)
     }
     pub fn has_run(&self, k: &str) -> bool {
         self.runs.contains(k)
     }
-    pub fn handles(&self) -> impl Iterator<Item = &ManagedRunHandle> {
+    pub fn handles(&self) -> impl Iterator<Item = &ManagedRun> {
         self.runs.iter().map(|(_, v)| v)
     }
     pub fn is_full(&self) -> bool {

data/bridge/sdk-core/core/src/worker/workflow/wft_extraction.rs

@@ -0,0 +1,125 @@
+use crate::{
+    abstractions::OwnedMeteredSemPermit,
+    protosext::ValidPollWFTQResponse,
+    worker::{
+        client::WorkerClient,
+        workflow::{
+            history_update::HistoryPaginator, CacheMissFetchReq, HistoryUpdate, NextPageReq,
+            PermittedWFT,
+        },
+    },
+};
+use futures::Stream;
+use futures_util::{stream, stream::PollNext, FutureExt, StreamExt};
+use std::{future, sync::Arc};
+use tracing::Span;
+
+/// Transforms incoming validated WFTs and history fetching requests into [PermittedWFT]s ready
+/// for application to workflow state
+pub(super) struct WFTExtractor {}
+
+pub(super) enum WFTExtractorOutput {
+    NewWFT(PermittedWFT),
+    FetchResult(PermittedWFT, Arc<HistfetchRC>),
+    NextPage {
+        paginator: HistoryPaginator,
+        update: HistoryUpdate,
+        span: Span,
+        rc: Arc<HistfetchRC>,
+    },
+    FailedFetch {
+        run_id: String,
+        err: tonic::Status,
+    },
+    PollerDead,
+}
+
+type WFTStreamIn = (
+    Result<ValidPollWFTQResponse, tonic::Status>,
+    OwnedMeteredSemPermit,
+);
+#[derive(derive_more::From, Debug)]
+pub(super) enum HistoryFetchReq {
+    Full(CacheMissFetchReq, Arc<HistfetchRC>),
+    NextPage(NextPageReq, Arc<HistfetchRC>),
+}
+/// Used inside of `Arc`s to ensure we don't shutdown while there are outstanding fetches.
+#[derive(Debug)]
+pub(super) struct HistfetchRC {}
+
+impl WFTExtractor {
+    pub(super) fn build(
+        client: Arc<dyn WorkerClient>,
+        max_fetch_concurrency: usize,
+        wft_stream: impl Stream<Item = WFTStreamIn> + Send + 'static,
+        fetch_stream: impl Stream<Item = HistoryFetchReq> + Send + 'static,
+    ) -> impl Stream<Item = Result<WFTExtractorOutput, tonic::Status>> + Send + 'static {
+        let fetch_client = client.clone();
+        let wft_stream = wft_stream
+            .map(move |(wft, permit)| {
+                let client = client.clone();
+                async move {
+                    match wft {
+                        Ok(wft) => {
+                            let run_id = wft.workflow_execution.run_id.clone();
+                            Ok(match HistoryPaginator::from_poll(wft, client).await {
+                                Ok((pag, prep)) => WFTExtractorOutput::NewWFT(PermittedWFT {
+                                    work: prep,
+                                    permit: permit.into_used(),
+                                    paginator: pag,
+                                }),
+                                Err(err) => WFTExtractorOutput::FailedFetch { run_id, err },
+                            })
+                        }
+                        Err(e) => Err(e),
+                    }
+                }
+                // This is... unattractive, but lets us avoid boxing all the futs in the stream
+                .left_future()
+                .left_future()
+            })
+            .chain(stream::iter([future::ready(Ok(
+                WFTExtractorOutput::PollerDead,
+            ))
+            .right_future()
+            .left_future()]));
+
+        stream::select_with_strategy(
+            wft_stream,
+            fetch_stream.map(move |fetchreq: HistoryFetchReq| {
+                let client = fetch_client.clone();
+                async move {
+                    Ok(match fetchreq {
+                        // It's OK to simply drop the refcounters in the event of fetch
+                        // failure. We'll just proceed with shutdown.
+                        HistoryFetchReq::Full(req, rc) => {
+                            let run_id = req.original_wft.work.execution.run_id.clone();
+                            match HistoryPaginator::from_fetchreq(req, client).await {
+                                Ok(r) => WFTExtractorOutput::FetchResult(r, rc),
+                                Err(err) => WFTExtractorOutput::FailedFetch { run_id, err },
+                            }
+                        }
+                        HistoryFetchReq::NextPage(mut req, rc) => {
+                            match req.paginator.extract_next_update().await {
+                                Ok(update) => WFTExtractorOutput::NextPage {
+                                    paginator: req.paginator,
+                                    update,
+                                    span: req.span,
+                                    rc,
+                                },
+                                Err(err) => WFTExtractorOutput::FailedFetch {
+                                    run_id: req.paginator.run_id,
+                                    err,
+                                },
+                            }
+                        }
+                    })
+                }
+                .right_future()
+            }),
+            // Priority always goes to the fetching stream
+            |_: &mut ()| PollNext::Right,
+        )
+        .buffer_unordered(max_fetch_concurrency)
+    }
+}

data/bridge/sdk-core/core/src/worker/workflow/wft_poller.rs

@@ -47,10 +47,7 @@ pub(crate) fn validate_wft(
     wft.try_into().map_err(|resp| {
         tonic::Status::new(
             tonic::Code::DataLoss,
-            format!(
-                "Server returned a poll WFT response we couldn't interpret: {:?}",
-                resp
-            ),
+            format!("Server returned a poll WFT response we couldn't interpret: {resp:?}"),
         )
     })
 }

data/bridge/sdk-core/core/src/worker/workflow/workflow_stream/saved_wf_inputs.rs

@@ -0,0 +1,117 @@
+use crate::{
+    telemetry::metrics::MetricsContext,
+    worker::{
+        client::mocks::DEFAULT_TEST_CAPABILITIES,
+        workflow::{
+            workflow_stream::{WFStream, WFStreamInput},
+            LAReqSink, LocalActivityRequestSink,
+        },
+        LocalActRequest, LocalActivityResolution,
+    },
+};
+use crossbeam::queue::SegQueue;
+use futures::Stream;
+use futures_util::StreamExt;
+use serde::{Deserialize, Serialize};
+use std::{future, sync::Arc};
+use temporal_sdk_core_api::worker::WorkerConfig;
+use tokio::sync::mpsc::UnboundedSender;
+use tokio_util::sync::CancellationToken;
+
+/// Replay everything that happened to internal workflow state. Useful for 100% deterministic
+/// reproduction of bugs.
+///
+/// Use `CoreWfStarter::enable_wf_state_input_recording` from the integration test utilities to
+/// activate saving the data to disk, and use the `wf_input_replay` example binary to replay.
+pub async fn replay_wf_state_inputs(mut config: WorkerConfig, inputs: impl Stream<Item = Vec<u8>>) {
+    use crate::worker::build_wf_basics;
+
+    let la_resp_q = Arc::new(SegQueue::new());
+    let la_resp_q_clone = la_resp_q.clone();
+    let inputs = inputs
+        .map(|bytes| {
+            rmp_serde::from_slice::<StoredWFStateInputDeSer>(&bytes)
+                .expect("Can decode wf stream input")
+        })
+        .filter_map(|si| {
+            future::ready(match si {
+                StoredWFStateInputDeSer::Stream(wfsi) => Some(wfsi),
+                StoredWFStateInputDeSer::ImmediateLASinkResolutions(lares) => {
+                    la_resp_q_clone.push(lares);
+                    None
+                }
+            })
+        });
+    let basics = build_wf_basics(
+        &mut config,
+        MetricsContext::no_op(),
+        CancellationToken::new(),
+        DEFAULT_TEST_CAPABILITIES.clone(),
+    );
+    let sink = ReadingFromFileLaReqSink {
+        resolutions: la_resp_q,
+    };
+    info!("Beginning workflow stream internal state replay");
+    let stream = WFStream::build_internal(inputs, basics, sink);
+    stream
+        .for_each(|o| async move { trace!("Stream output: {:?}", o) })
+        .await;
+}
+
+impl WFStream {
+    pub(super) fn prep_input(&mut self, action: &WFStreamInput) -> Option<PreppedInputWrite> {
+        // Remove the channel, we'll put it back, avoiding a clone
+        self.wf_state_inputs.take().map(|chan| PreppedInputWrite {
+            data: rmp_serde::to_vec(&StoredWFStateInputSer::Stream(action))
+                .expect("WF Inputs are serializable"),
+            chan,
+        })
+    }
+    pub(super) fn flush_write(&mut self, w: PreppedInputWrite) {
+        let _ = w.chan.send(w.data);
+        self.wf_state_inputs = Some(w.chan);
+    }
+}
+pub(super) struct PreppedInputWrite {
+    data: Vec<u8>,
+    chan: UnboundedSender<Vec<u8>>,
+}
+
+#[derive(Serialize)]
+enum StoredWFStateInputSer<'a> {
+    Stream(&'a WFStreamInput),
+    ImmediateLASinkResolutions(&'a Vec<LocalActivityResolution>),
+}
+
+#[derive(Deserialize)]
+enum StoredWFStateInputDeSer {
+    Stream(WFStreamInput),
+    ImmediateLASinkResolutions(Vec<LocalActivityResolution>),
+}
+
+struct ReadingFromFileLaReqSink {
+    resolutions: Arc<SegQueue<Vec<LocalActivityResolution>>>,
+}
+impl LocalActivityRequestSink for ReadingFromFileLaReqSink {
+    fn sink_reqs(&self, reqs: Vec<LocalActRequest>) -> Vec<LocalActivityResolution> {
+        if !reqs.is_empty() {
+            self.resolutions
+                .pop()
+                .expect("LA sink was called, but there's no stored immediate response")
+        } else {
+            vec![]
+        }
+    }
+}
+
+impl LAReqSink {
+    pub(crate) fn write_req(&self, res: &Vec<LocalActivityResolution>) {
+        if let Some(r) = self.recorder.as_ref() {
+            r.send(
+                rmp_serde::to_vec(&StoredWFStateInputSer::ImmediateLASinkResolutions(res))
+                    .expect("LA immediate resolutions are serializable"),
+            )
+            .expect("WF input serialization channel is available for immediate LA result storage");
+        }
+    }
+}

data/bridge/sdk-core/core/src/worker/workflow/workflow_stream/tonic_status_serde.rs

@@ -0,0 +1,24 @@
+use tonic::{Code, Status};
+
+#[derive(serde::Serialize, serde::Deserialize)]
+#[serde(remote = "Status")]
+pub(super) struct SerdeStatus {
+    #[serde(getter = "get_code")]
+    code: i32,
+    #[serde(getter = "get_owned_msg")]
+    message: String,
+}
+
+fn get_owned_msg(v: &Status) -> String {
+    v.message().to_string()
+}
+
+fn get_code(v: &Status) -> i32 {
+    v.code() as i32
+}
+
+impl From<SerdeStatus> for Status {
+    fn from(v: SerdeStatus) -> Self {
+        Status::new(Code::from(v.code), v.message)
+    }
+}