temporalio 0.0.1

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

@@ -0,0 +1,145 @@
+use crate::{
+    telemetry::metrics::workflow_type,
+    worker::workflow::{
+        managed_run::WorkflowManager, HistoryUpdate, LocalActivityRequestSink, ManagedRunHandle,
+        NewIncomingWFT, RunUpdateResponse,
+    },
+    MetricsContext,
+};
+use lru::LruCache;
+use std::{num::NonZeroUsize, time::Instant};
+use tokio::sync::mpsc::UnboundedSender;
+
+pub(super) struct RunCache {
+    max: usize,
+    namespace: String,
+    run_update_tx: UnboundedSender<RunUpdateResponse>,
+    /// Run id -> Data
+    runs: LruCache<String, ManagedRunHandle>,
+    local_activity_request_sink: LocalActivityRequestSink,
+
+    metrics: MetricsContext,
+}
+
+impl RunCache {
+    pub fn new(
+        max_cache_size: usize,
+        namespace: String,
+        run_update_tx: UnboundedSender<RunUpdateResponse>,
+        local_activity_request_sink: LocalActivityRequestSink,
+        metrics: MetricsContext,
+    ) -> Self {
+        // The cache needs room for at least one run, otherwise we couldn't do anything. In
+        // "0" size mode, the run is evicted once the workflow task is complete.
+        let lru_size = if max_cache_size > 0 {
+            max_cache_size
+        } else {
+            1
+        };
+        Self {
+            max: max_cache_size,
+            namespace,
+            run_update_tx,
+            runs: LruCache::new(
+                NonZeroUsize::new(lru_size).expect("LRU size is guaranteed positive"),
+            ),
+            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 {
+        let cur_num_cached_runs = self.runs.len();
+
+        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,
+            });
+            self.metrics.cache_size(cur_num_cached_runs as u64);
+            return run_handle;
+        }
+
+        // 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(),
+            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() {
+            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()
+    }
+    pub fn remove(&mut self, k: &str) -> Option<ManagedRunHandle> {
+        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> {
+        self.runs.get_mut(k)
+    }
+    pub fn get(&mut self, k: &str) -> Option<&ManagedRunHandle> {
+        self.runs.get(k)
+    }
+
+    /// Returns the current least-recently-used run. Returns `None` when cache empty.
+    pub fn current_lru_run(&self) -> Option<&str> {
+        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)> {
+        self.runs.iter().rev().map(|(k, v)| (k.as_str(), v))
+    }
+    pub fn peek(&self, k: &str) -> Option<&ManagedRunHandle> {
+        self.runs.peek(k)
+    }
+    pub fn has_run(&self, k: &str) -> bool {
+        self.runs.contains(k)
+    }
+    pub fn handles(&self) -> impl Iterator<Item = &ManagedRunHandle> {
+        self.runs.iter().map(|(_, v)| v)
+    }
+    pub fn is_full(&self) -> bool {
+        self.runs.cap().get() == self.runs.len()
+    }
+    pub fn len(&self) -> usize {
+        self.runs.len()
+    }
+    pub fn cache_capacity(&self) -> usize {
+        self.max
+    }
+}

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

@@ -0,0 +1,88 @@
+use crate::{pollers::BoxedWFPoller, protosext::ValidPollWFTQResponse, MetricsContext};
+use futures::{stream, Stream};
+use temporal_sdk_core_protos::temporal::api::workflowservice::v1::PollWorkflowTaskQueueResponse;
+
+pub(crate) fn new_wft_poller(
+    poller: BoxedWFPoller,
+    metrics: MetricsContext,
+) -> impl Stream<Item = Result<ValidPollWFTQResponse, tonic::Status>> {
+    stream::unfold((poller, metrics), |(poller, metrics)| async move {
+        loop {
+            return match poller.poll().await {
+                Some(Ok(wft)) => {
+                    if wft == PollWorkflowTaskQueueResponse::default() {
+                        // We get the default proto in the event that the long poll times out.
+                        debug!("Poll wft timeout");
+                        metrics.wf_tq_poll_empty();
+                        continue;
+                    }
+                    if let Some(dur) = wft.sched_to_start() {
+                        metrics.wf_task_sched_to_start_latency(dur);
+                    }
+                    let work = match validate_wft(wft) {
+                        Ok(w) => w,
+                        Err(e) => {
+                            error!(error=?e, "Server returned an unparseable workflow task");
+                            continue;
+                        }
+                    };
+                    metrics.wf_tq_poll_ok();
+                    Some((Ok(work), (poller, metrics)))
+                }
+                Some(Err(e)) => {
+                    warn!(error=?e, "Error while polling for workflow tasks");
+                    Some((Err(e), (poller, metrics)))
+                }
+                // If poller returns None, it's dead, thus we also return None to terminate this
+                // stream.
+                None => None,
+            };
+        }
+    })
+}
+
+pub(crate) fn validate_wft(
+    wft: PollWorkflowTaskQueueResponse,
+) -> Result<ValidPollWFTQResponse, tonic::Status> {
+    wft.try_into().map_err(|resp| {
+        tonic::Status::new(
+            tonic::Code::DataLoss,
+            format!(
+                "Server returned a poll WFT response we couldn't interpret: {:?}",
+                resp
+            ),
+        )
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::test_help::mock_poller;
+    use futures::{pin_mut, StreamExt};
+
+    #[tokio::test]
+    async fn poll_timeouts_do_not_produce_responses() {
+        let mut mock_poller = mock_poller();
+        mock_poller
+            .expect_poll()
+            .times(1)
+            .returning(|| Some(Ok(PollWorkflowTaskQueueResponse::default())));
+        mock_poller.expect_poll().times(1).returning(|| None);
+        let stream = new_wft_poller(Box::new(mock_poller), Default::default());
+        pin_mut!(stream);
+        assert_matches!(stream.next().await, None);
+    }
+
+    #[tokio::test]
+    async fn poll_errors_do_produce_responses() {
+        let mut mock_poller = mock_poller();
+        mock_poller
+            .expect_poll()
+            .times(1)
+            .returning(|| Some(Err(tonic::Status::internal("ahhh"))));
+        let stream = new_wft_poller(Box::new(mock_poller), Default::default());
+        pin_mut!(stream);
+        assert_matches!(stream.next().await, Some(Err(_)));
+    }
+}