temporalio 0.0.1 → 0.1.0

data/bridge/sdk-core/core/src/replay/mod.rs

@@ -2,70 +2,214 @@
 //! to replay canned histories. It should be used by Lang SDKs to provide replay capabilities to
 //! users during testing.
 
-use crate::worker::client::{mocks::mock_manual_workflow_client, WorkerClient};
-use futures::FutureExt;
+use crate::{
+    worker::client::{mocks::mock_manual_workflow_client, WorkerClient},
+    Worker,
+};
+use futures::{FutureExt, Stream, StreamExt};
+use once_cell::sync::OnceCell;
+use parking_lot::Mutex;
 use std::{
-    sync::{
-        atomic::{AtomicBool, Ordering},
-        Arc,
-    },
-    time::Duration,
+    collections::HashMap,
+    pin::Pin,
+    sync::Arc,
+    task::{Context, Poll},
 };
-use temporal_sdk_core_protos::temporal::api::{
-    common::v1::WorkflowExecution,
-    history::v1::History,
-    workflowservice::v1::{
-        RespondWorkflowTaskCompletedResponse, RespondWorkflowTaskFailedResponse,
+use temporal_sdk_core_protos::{
+    coresdk::workflow_activation::remove_from_cache::EvictionReason,
+    temporal::api::{
+        common::v1::WorkflowExecution,
+        history::v1::History,
+        workflowservice::v1::{
+            RespondWorkflowTaskCompletedResponse, RespondWorkflowTaskFailedResponse,
+        },
     },
 };
 pub use temporal_sdk_core_protos::{
     default_wes_attribs, HistoryInfo, TestHistoryBuilder, DEFAULT_WORKFLOW_TYPE,
 };
+use tokio::sync::{mpsc, mpsc::UnboundedSender, Mutex as TokioMutex};
+use tokio_stream::wrappers::UnboundedReceiverStream;
+use tokio_util::sync::CancellationToken;
+
+/// A history which will be used during replay verification. Since histories do not include the
+/// workflow id, it must be manually attached.
+#[derive(Debug, Clone, derive_more::Constructor)]
+pub struct HistoryForReplay {
+    hist: History,
+    workflow_id: String,
+}
+
+/// Allows lang to feed histories into the replayer one at a time. Simply drop the feeder to signal
+/// to the worker that you're done and it should initiate shutdown.
+pub struct HistoryFeeder {
+    tx: mpsc::Sender<HistoryForReplay>,
+}
+/// The stream half of a [HistoryFeeder]
+pub struct HistoryFeederStream {
+    rcvr: mpsc::Receiver<HistoryForReplay>,
+}
+
+impl HistoryFeeder {
+    /// Make a new history feeder, which will store at most `buffer_size` histories before `feed`
+    /// blocks.
+    ///
+    /// Returns a feeder which will be used to feed in histories, and a stream you can pass to
+    /// one of the replay worker init functions.
+    pub fn new(buffer_size: usize) -> (Self, HistoryFeederStream) {
+        let (tx, rcvr) = mpsc::channel(buffer_size);
+        (Self { tx }, HistoryFeederStream { rcvr })
+    }
+    /// Feed a new history into the replayer, blocking if there is not room to accept another
+    /// history.
+    pub async fn feed(&self, history: HistoryForReplay) -> anyhow::Result<()> {
+        self.tx.send(history).await?;
+        Ok(())
+    }
+}
+
+impl Stream for HistoryFeederStream {
+    type Item = HistoryForReplay;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        self.rcvr.poll_recv(cx)
+    }
+}
 
-/// Create a mock client which can be used by a replay worker to serve up canned history.
-/// It will return the entire history in one workflow task, after that it will return default
-/// responses (with a 10s wait). If a workflow task failure is sent to the mock, it will send
-/// the complete response again.
-pub(crate) fn mock_client_from_history(
-    history: &History,
-    task_queue: impl Into<String>,
-) -> impl WorkerClient {
+/// Create a mock client which can be used by a replay worker to serve up canned histories. It will
+/// return the entire history in one workflow task. If a workflow task failure is sent to the mock,
+/// it will send the complete response again.
+///
+/// Once it runs out of histories to return, it will serve up default responses after a 10s delay
+pub(crate) fn mock_client_from_histories(historator: Historator) -> impl WorkerClient {
     let mut mg = mock_manual_workflow_client();
 
-    let hist_info = HistoryInfo::new_from_history(history, None).unwrap();
-    let wf = WorkflowExecution {
-        workflow_id: "fake_wf_id".to_string(),
-        run_id: hist_info.orig_run_id().to_string(),
-    };
+    let hist_allow_tx = historator.replay_done_tx.clone();
+    let historator = Arc::new(TokioMutex::new(historator));
 
-    let did_send = Arc::new(AtomicBool::new(false));
-    let did_send_clone = did_send.clone();
-    let tq = task_queue.into();
     mg.expect_poll_workflow_task().returning(move |_, _| {
-        let hist_info = hist_info.clone();
-        let wf = wf.clone();
-        let did_send_clone = did_send_clone.clone();
-        let tq = tq.clone();
+        let historator = historator.clone();
         async move {
-            if !did_send_clone.swap(true, Ordering::AcqRel) {
-                let mut resp = hist_info.as_poll_wft_response(tq);
-                resp.workflow_execution = Some(wf.clone());
+            let mut hlock = historator.lock().await;
+            // Always wait for permission before dispatching the next task
+            let _ = hlock.allow_stream.next().await;
+
+            if let Some(history) = hlock.next().await {
+                let hist_info = HistoryInfo::new_from_history(&history.hist, None).unwrap();
+                let mut resp = hist_info.as_poll_wft_response();
+                resp.workflow_execution = Some(WorkflowExecution {
+                    workflow_id: history.workflow_id,
+                    run_id: hist_info.orig_run_id().to_string(),
+                });
                 Ok(resp)
             } else {
-                tokio::time::sleep(Duration::from_secs(10)).await;
+                if let Some(wc) = hlock.worker_closer.get() {
+                    wc.cancel();
+                }
                 Ok(Default::default())
             }
         }
         .boxed()
     });
 
-    mg.expect_complete_workflow_task()
-        .returning(|_| async move { Ok(RespondWorkflowTaskCompletedResponse::default()) }.boxed());
+    mg.expect_complete_workflow_task().returning(move |_| {
+        async move { Ok(RespondWorkflowTaskCompletedResponse::default()) }.boxed()
+    });
     mg.expect_fail_workflow_task().returning(move |_, _, _| {
-        // We'll need to re-send the history if WFT fails
-        did_send.store(false, Ordering::Release);
-        async move { Ok(RespondWorkflowTaskFailedResponse {}) }.boxed()
+        hist_allow_tx.send("Failed".to_string()).unwrap();
+        async move { Ok(RespondWorkflowTaskFailedResponse::default()) }.boxed()
     });
 
     mg
 }
+
+pub(crate) struct Historator {
+    iter: Pin<Box<dyn Stream<Item = HistoryForReplay> + Send>>,
+    allow_stream: UnboundedReceiverStream<String>,
+    worker_closer: Arc<OnceCell<CancellationToken>>,
+    dat: Arc<Mutex<HistoratorDat>>,
+    replay_done_tx: UnboundedSender<String>,
+}
+impl Historator {
+    pub(crate) fn new(histories: impl Stream<Item = HistoryForReplay> + Send + 'static) -> Self {
+        let dat = Arc::new(Mutex::new(HistoratorDat::default()));
+        let (replay_done_tx, replay_done_rx) = mpsc::unbounded_channel();
+        // Need to allow the first history item
+        replay_done_tx.send("fake".to_string()).unwrap();
+        Self {
+            iter: Box::pin(histories.fuse()),
+            allow_stream: UnboundedReceiverStream::new(replay_done_rx),
+            worker_closer: Arc::new(OnceCell::new()),
+            dat,
+            replay_done_tx,
+        }
+    }
+
+    /// Returns a callback that can be used as the post-activation hook for a worker to indicate
+    /// we're ready to replay the next history, or whatever else.
+    pub(crate) fn get_post_activate_hook(&self) -> impl Fn(&Worker, &str, usize) + Send + Sync {
+        let dat = self.dat.clone();
+        let done_tx = self.replay_done_tx.clone();
+        move |worker, activated_run_id, last_processed_event| {
+            // We can't hold the lock while evaluating the hook, or we'd deadlock.
+            let last_event_in_hist = dat
+                .lock()
+                .run_id_to_last_event_num
+                .get(activated_run_id)
+                .cloned();
+            if let Some(le) = last_event_in_hist {
+                if last_processed_event >= le {
+                    worker.request_wf_eviction(
+                        activated_run_id,
+                        "Always evict workflows after replay",
+                        EvictionReason::LangRequested,
+                    );
+                    done_tx.send(activated_run_id.to_string()).unwrap();
+                }
+            }
+        }
+    }
+
+    pub(crate) fn get_shutdown_setter(&self) -> impl FnOnce(CancellationToken) + 'static {
+        let wc = self.worker_closer.clone();
+        move |ct| {
+            wc.set(ct).expect("Shutdown token must only be set once");
+        }
+    }
+}
+
+impl Stream for Historator {
+    type Item = HistoryForReplay;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        match self.iter.poll_next_unpin(cx) {
+            Poll::Ready(Some(history)) => {
+                let run_id = history
+                    .hist
+                    .extract_run_id_from_start()
+                    .expect(
+                        "Histories provided for replay must contain run ids in their workflow \
+                                 execution started events",
+                    )
+                    .to_string();
+                let last_event = history.hist.last_event_id();
+                self.dat
+                    .lock()
+                    .run_id_to_last_event_num
+                    .insert(run_id, last_event as usize);
+                Poll::Ready(Some(history))
+            }
+            Poll::Ready(None) => {
+                self.dat.lock().all_dispatched = true;
+                Poll::Ready(None)
+            }
+            o => o,
+        }
+    }
+}
+
+#[derive(Default)]
+struct HistoratorDat {
+    run_id_to_last_event_num: HashMap<String, usize>,
+    all_dispatched: bool,
+}

data/bridge/sdk-core/core/src/telemetry/log_export.rs

@@ -0,0 +1,190 @@
+use parking_lot::Mutex;
+use ringbuf::{Consumer, HeapRb, Producer};
+use std::{collections::HashMap, sync::Arc, time::SystemTime};
+use temporal_sdk_core_api::telemetry::CoreLog;
+use tracing_subscriber::Layer;
+
+const RB_SIZE: usize = 2048;
+
+pub(super) type CoreLogsOut = Consumer<CoreLog, Arc<HeapRb<CoreLog>>>;
+
+pub(super) struct CoreLogExportLayer {
+    logs_in: Mutex<Producer<CoreLog, Arc<HeapRb<CoreLog>>>>,
+}
+
+#[derive(Debug)]
+struct CoreLogFieldStorage(HashMap<String, serde_json::Value>);
+
+impl CoreLogExportLayer {
+    pub(super) fn new() -> (Self, CoreLogsOut) {
+        let (lin, lout) = HeapRb::new(RB_SIZE).split();
+        (
+            Self {
+                logs_in: Mutex::new(lin),
+            },
+            lout,
+        )
+    }
+}
+
+impl<S> Layer<S> for CoreLogExportLayer
+where
+    S: tracing::Subscriber,
+    S: for<'lookup> tracing_subscriber::registry::LookupSpan<'lookup>,
+{
+    fn on_new_span(
+        &self,
+        attrs: &tracing::span::Attributes<'_>,
+        id: &tracing::span::Id,
+        ctx: tracing_subscriber::layer::Context<'_, S>,
+    ) {
+        let span = ctx.span(id).unwrap();
+        let mut fields = HashMap::new();
+        let mut visitor = JsonVisitor(&mut fields);
+        attrs.record(&mut visitor);
+        let storage = CoreLogFieldStorage(fields);
+        let mut extensions = span.extensions_mut();
+        extensions.insert(storage);
+    }
+
+    fn on_record(
+        &self,
+        id: &tracing::span::Id,
+        values: &tracing::span::Record<'_>,
+        ctx: tracing_subscriber::layer::Context<'_, S>,
+    ) {
+        let span = ctx.span(id).unwrap();
+
+        let mut extensions_mut = span.extensions_mut();
+        let custom_field_storage: &mut CoreLogFieldStorage =
+            extensions_mut.get_mut::<CoreLogFieldStorage>().unwrap();
+        let json_data = &mut custom_field_storage.0;
+
+        let mut visitor = JsonVisitor(json_data);
+        values.record(&mut visitor);
+    }
+
+    fn on_event(&self, event: &tracing::Event<'_>, ctx: tracing_subscriber::layer::Context<'_, S>) {
+        let mut fields = HashMap::new();
+        let mut visitor = JsonVisitor(&mut fields);
+        event.record(&mut visitor);
+
+        let mut spans = vec![];
+        if let Some(scope) = ctx.event_scope(event) {
+            for span in scope.from_root() {
+                let extensions = span.extensions();
+                let storage = extensions.get::<CoreLogFieldStorage>().unwrap();
+                let field_data = &storage.0;
+                for (k, v) in field_data {
+                    fields.insert(k.to_string(), v.clone());
+                }
+                spans.push(span.name().to_string());
+            }
+        }
+
+        // "message" is the magic default field keyname for the string passed to the event
+        let message = fields.remove("message").unwrap_or_default();
+        let log = CoreLog {
+            target: event.metadata().target().to_string(),
+            // This weird as_str dance prevents adding extra quotes
+            message: message.as_str().unwrap_or_default().to_string(),
+            timestamp: SystemTime::now(),
+            level: *event.metadata().level(),
+            fields,
+            span_contexts: spans,
+        };
+        let _ = self.logs_in.lock().push(log);
+    }
+}
+
+struct JsonVisitor<'a>(&'a mut HashMap<String, serde_json::Value>);
+
+impl<'a> tracing::field::Visit for JsonVisitor<'a> {
+    fn record_f64(&mut self, field: &tracing::field::Field, value: f64) {
+        self.0
+            .insert(field.name().to_string(), serde_json::json!(value));
+    }
+
+    fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
+        self.0
+            .insert(field.name().to_string(), serde_json::json!(value));
+    }
+
+    fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
+        self.0
+            .insert(field.name().to_string(), serde_json::json!(value));
+    }
+
+    fn record_bool(&mut self, field: &tracing::field::Field, value: bool) {
+        self.0
+            .insert(field.name().to_string(), serde_json::json!(value));
+    }
+
+    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
+        self.0
+            .insert(field.name().to_string(), serde_json::json!(value));
+    }
+
+    fn record_error(
+        &mut self,
+        field: &tracing::field::Field,
+        value: &(dyn std::error::Error + 'static),
+    ) {
+        self.0.insert(
+            field.name().to_string(),
+            serde_json::json!(value.to_string()),
+        );
+    }
+
+    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
+        self.0.insert(
+            field.name().to_string(),
+            serde_json::json!(format!("{value:?}")),
+        );
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{telemetry::construct_filter_string, telemetry_init};
+    use temporal_sdk_core_api::telemetry::{CoreTelemetry, Logger, TelemetryOptionsBuilder};
+    use tracing::Level;
+
+    #[instrument(fields(bros = "brohemian"))]
+    fn instrumented(thing: &str) {
+        warn!("warn");
+        info!(foo = "bar", "info");
+        debug!("debug");
+    }
+
+    #[tokio::test]
+    async fn test_forwarding_output() {
+        let opts = TelemetryOptionsBuilder::default()
+            .logging(Logger::Forward {
+                filter: construct_filter_string(Level::INFO, Level::WARN),
+            })
+            .build()
+            .unwrap();
+        let instance = telemetry_init(opts).unwrap();
+        let _g = tracing::subscriber::set_default(instance.trace_subscriber.clone());
+
+        let top_span = span!(Level::INFO, "yayspan", huh = "wat");
+        let _guard = top_span.enter();
+        info!("Whata?");
+        instrumented("hi");
+        info!("Donezo");
+
+        let logs = instance.fetch_buffered_logs();
+        // Verify debug log was not forwarded
+        assert!(!logs.iter().any(|l| l.message == "debug"));
+        assert_eq!(logs.len(), 4);
+        // Ensure fields are attached to events properly
+        let info_msg = &logs[2];
+        assert_eq!(info_msg.message, "info");
+        assert_eq!(info_msg.fields.len(), 4);
+        assert_eq!(info_msg.fields.get("huh"), Some(&"wat".into()));
+        assert_eq!(info_msg.fields.get("foo"), Some(&"bar".into()));
+        assert_eq!(info_msg.fields.get("bros"), Some(&"brohemian".into()));
+        assert_eq!(info_msg.fields.get("thing"), Some(&"hi".into()));
+    }
+}