Hackathon wf dash

core-api/src/builtin_queries.rs

@@ -0,0 +1,145 @@
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use temporal_sdk_core_protos::{
+    coresdk::{
+        workflow_activation::{workflow_activation_job, QueryWorkflow, WorkflowActivationJob},
+        workflow_commands::{query_result, QueryResult, QuerySuccess},
+    },
+    ENHANCED_STACK_QUERY,
+};
+
+pub static FAKE_ENHANCED_STACK_QUERY_ID: &str = "__fake_enhanced_stack";
+// must start with the above
+pub static FAKE_ENHANCED_STACK_QUERY_ID_FINAL_LEGACY: &str = "__fake_enhanced_stack_final";
+
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+pub struct SDKInfo {
+    pub name: String,
+    pub version: String,
+}
+
+/// Represents a slice of a file starting at line_offset
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+#[serde(rename_all = "camelCase")]
+pub struct FileSlice {
+    /// slice of a file with `\n` (newline) line terminator.
+    pub content: String,
+    /// Only used possible to trim the file without breaking syntax highlighting.
+    pub line_offset: u64,
+}
+
+/// A pointer to a location in a file
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+#[serde(rename_all = "camelCase")]
+pub struct FileLocation {
+    /// Path to source file (absolute or relative).
+    /// When using a relative path, make sure all paths are relative to the same root.
+    pub file_path: Option<String>,
+    /// If possible, SDK should send this, required for displaying the code location.
+    pub line: Option<u64>,
+    /// If possible, SDK should send this.
+    pub column: Option<u64>,
+    /// Function name this line belongs to (if applicable).
+    /// Used for falling back to stack trace view.
+    pub function_name: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum InternalCommandType {
+    ScheduleActivity,
+    StartTimer,
+    StartChildWorkflow,
+}
+
+/// An internal (Lang<->Core) command identifier
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct InternalCommand {
+    pub r#type: InternalCommandType,
+    pub seq: u32,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+pub struct InternalStackTrace {
+    pub locations: Vec<FileLocation>,
+    pub commands: Vec<InternalCommand>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+pub struct InternalEnhancedStackTrace {
+    pub sdk: SDKInfo,
+    /// Mapping of file path to file contents. SDK may choose to send no, some or all sources.
+    /// Sources might be trimmed, and some time only the file(s) of the top element of the trace
+    /// will be sent.
+    pub sources: HashMap<String, Vec<FileSlice>>,
+    pub stacks: Vec<InternalStackTrace>,
+}
+
+impl TryFrom<&QueryResult> for InternalEnhancedStackTrace {
+    type Error = ();
+
+    fn try_from(qr: &QueryResult) -> Result<Self, Self::Error> {
+        if let Some(query_result::Variant::Succeeded(QuerySuccess {
+            response: Some(ref payload),
+        })) = qr.variant
+        {
+            if payload.is_json_payload() {
+                if let Ok(internal_trace) =
+                    serde_json::from_slice::<InternalEnhancedStackTrace>(payload.data.as_slice())
+                {
+                    return Ok(internal_trace);
+                }
+            }
+        }
+        Err(())
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+#[serde(rename_all = "camelCase")]
+pub struct StackTrace {
+    pub locations: Vec<FileLocation>,
+    pub correlating_event_ids: Vec<i64>,
+}
+
+// Used as the result for the enhanced stack trace query
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+#[serde(rename_all = "camelCase")]
+pub struct EnhancedStackTrace {
+    pub sdk: SDKInfo,
+    /// Mapping of file path to file contents. SDK may choose to send no, some or all sources.
+    /// Sources might be trimmed, and some time only the file(s) of the top element of the trace
+    /// will be sent.
+    pub sources: HashMap<String, Vec<FileSlice>>,
+    pub stacks: Vec<StackTrace>,
+}
+
+// Used as the result for the time travel stack trace query
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+#[serde(rename_all = "camelCase")]
+pub struct TimeTravelStackTrace {
+    pub sdk: SDKInfo,
+    /// Mapping of file path to file contents. SDK may choose to send no, some or all sources.
+    /// Sources might be trimmed, and some time only the file(s) of the top element of the trace
+    /// will be sent.
+    pub sources: HashMap<String, Vec<FileSlice>>,
+    /// Maps WFT started event ids to active stack traces upon completion of that WFT
+    pub stacks: HashMap<u32, Vec<StackTrace>>,
+}
+
+/// Generate an activation job with an enhanced stack trace query.
+/// * `final_legacy` - Set to true if this is the final query before aggregation will be completed,
+///   and we should respond to the *legacy query* with the aggregation result when ready.
+pub fn enhanced_stack_query_job(final_legacy: bool) -> WorkflowActivationJob {
+    let id = if final_legacy {
+        FAKE_ENHANCED_STACK_QUERY_ID_FINAL_LEGACY
+    } else {
+        FAKE_ENHANCED_STACK_QUERY_ID
+    };
+    workflow_activation_job::Variant::QueryWorkflow(QueryWorkflow {
+        query_id: id.to_string(),
+        query_type: ENHANCED_STACK_QUERY.to_string(),
+        arguments: vec![],
+        headers: Default::default(),
+    })
+    .into()
+}

core-api/src/lib.rs

@@ -1,3 +1,4 @@
+pub mod builtin_queries;
 pub mod errors;
 pub mod telemetry;
 pub mod worker;

core/src/abstractions.rs

@@ -268,6 +268,22 @@ macro_rules! dbg_panic {
 }
 pub(crate) use dbg_panic;
 
+/// Get rid of me once https://doc.rust-lang.org/std/vec/struct.Vec.html#method.drain_filter is
+/// stable
+pub(crate) fn lame_drain_filter<T>(vec: &mut Vec<T>, pred: impl Fn(&T) -> bool) -> Vec<T> {
+    let mut drained = vec![];
+    let mut i = 0;
+    while i < vec.len() {
+        if pred(&vec[i]) {
+            let val = vec.remove(i);
+            drained.push(val);
+        } else {
+            i += 1;
+        }
+    }
+    drained
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

core/src/test_help/mod.rs

@@ -794,6 +794,7 @@ pub(crate) async fn poll_and_reply_clears_outstanding_evicts<'a>(
                 // Eviction plus some work, we still want to issue the reply
                 WorkflowActivationCompletion {
                     run_id: res.run_id.clone(),
+                    alternate_cache_key: "".to_string(),
                     status: Some(reply.clone()),
                 }
             };

core/src/worker/mod.rs

@@ -549,11 +549,12 @@ impl Worker {
     /// Request a workflow eviction
     pub(crate) fn request_wf_eviction(
         &self,
-        run_id: &str,
+        cache_key: &str,
         message: impl Into<String>,
         reason: EvictionReason,
     ) {
-        self.workflows.request_eviction(run_id, message, reason);
+        self.workflows
+            .request_eviction(cache_key.parse().unwrap(), message, reason);
     }
 
     /// Sets a function to be called at the end of each activation completion
@@ -583,8 +584,9 @@ impl Worker {
         )
     }
 
-    fn notify_local_result(&self, run_id: &str, res: LocalResolution) {
-        self.workflows.notify_of_local_result(run_id, res);
+    fn notify_local_result(&self, cache_key: &str, res: LocalResolution) {
+        self.workflows
+            .notify_of_local_result(cache_key.parse().unwrap(), res);
     }
 }
 

core/src/worker/workflow/history_update.rs

@@ -2,7 +2,7 @@ use crate::{
     protosext::ValidPollWFTQResponse,
     worker::{
         client::WorkerClient,
-        workflow::{CacheMissFetchReq, PermittedWFT, PreparedWFT},
+        workflow::{run_cache::RunCacheKey, CacheMissFetchReq, PermittedWFT, PreparedWFT},
     },
 };
 use futures::{future::BoxFuture, FutureExt, Stream};
@@ -17,9 +17,12 @@ use std::{
     sync::Arc,
     task::{Context, Poll},
 };
-use temporal_sdk_core_protos::temporal::api::{
-    enums::v1::EventType,
-    history::v1::{history_event, History, HistoryEvent, WorkflowTaskCompletedEventAttributes},
+use temporal_sdk_core_protos::{
+    temporal::api::{
+        enums::v1::EventType,
+        history::v1::{history_event, History, HistoryEvent, WorkflowTaskCompletedEventAttributes},
+    },
+    TIME_TRAVEL_QUERY,
 };
 use tracing::Instrument;
 
@@ -73,15 +76,15 @@ pub enum NextWFT {
 }
 
 #[derive(derive_more::DebugCustom)]
-#[debug(fmt = "HistoryPaginator(run_id: {run_id})")]
+#[debug(fmt = "HistoryPaginator(cache_key: {cache_key})")]
 #[cfg_attr(
     feature = "save_wf_inputs",
     derive(serde::Serialize, serde::Deserialize),
     serde(default = "HistoryPaginator::fake_deserialized")
 )]
 pub struct HistoryPaginator {
     pub(crate) wf_id: String,
-    pub(crate) run_id: String,
+    pub(crate) cache_key: RunCacheKey,
     pub(crate) previous_wft_started_id: i64,
     pub(crate) wft_started_event_id: i64,
 
@@ -122,9 +125,27 @@ impl HistoryPaginator {
     /// Use a new poll response to create a new [WFTPaginator], returning it and the
     /// [PreparedWFT] extracted from it that can be fed into workflow state.
     pub(super) async fn from_poll(
-        wft: ValidPollWFTQResponse,
+        mut wft: ValidPollWFTQResponse,
         client: Arc<dyn WorkerClient>,
     ) -> Result<(Self, PreparedWFT), tonic::Status> {
+        let mut alternate_cache = 0;
+        // Intercept and special time travel query
+        let mut all_queries = wft
+            .query_requests
+            .iter()
+            .map(|q| q.query_type.as_str())
+            .chain(wft.legacy_query.iter().map(|q| q.query_type.as_str()));
+        if all_queries.any(|qt| qt == TIME_TRAVEL_QUERY) {
+            // If we see one of these, we need to fetch *all* history, attach the alternate cache
+            // key, and then break things into made up WFTs and issue enhanced stack trace queries
+            // for each one
+            info!("Time travel query from poll");
+
+            // This will force fetching all history.
+            wft.history.events = vec![];
+            alternate_cache = 1;
+        }
+
         let empty_hist = wft.history.events.is_empty();
         let npt = if empty_hist {
             NextPageToken::FetchFromStart
@@ -162,6 +183,7 @@ impl HistoryPaginator {
             legacy_query: wft.legacy_query,
             query_requests: wft.query_requests,
             update,
+            alternate_cache,
         };
         Ok((paginator, prepared))
     }
@@ -172,7 +194,7 @@ impl HistoryPaginator {
     ) -> Result<PermittedWFT, tonic::Status> {
         let mut paginator = Self {
             wf_id: req.original_wft.work.execution.workflow_id.clone(),
-            run_id: req.original_wft.work.execution.run_id.clone(),
+            cache_key: req.original_wft.work.cache_key(),
             previous_wft_started_id: req.original_wft.work.update.previous_wft_started_id,
             wft_started_event_id: req.original_wft.work.update.wft_started_id,
             client,
@@ -206,7 +228,9 @@ impl HistoryPaginator {
             client,
             event_queue,
             wf_id,
-            run_id,
+            // TODO: At least, with what's happening for this hack, the key never needs to start
+            //  with an alternate, but this is a bit ugly
+            cache_key: run_id.parse().unwrap(),
             next_page_token,
             final_events,
             previous_wft_started_id,
@@ -221,7 +245,10 @@ impl HistoryPaginator {
             client: Arc::new(mock_manual_workflow_client()),
             event_queue: Default::default(),
             wf_id: "".to_string(),
-            run_id: "".to_string(),
+            cache_key: RunCacheKey {
+                run_id: "".to_string(),
+                alternate_key: 0,
+            },
             next_page_token: NextPageToken::FetchFromStart,
             final_events: vec![],
             previous_wft_started_id: -2,
@@ -292,10 +319,14 @@ impl HistoryPaginator {
                 NextPageToken::FetchFromStart => vec![],
                 NextPageToken::Next(v) => v,
             };
-            debug!(run_id=%self.run_id, "Fetching new history page");
+            debug!(run_id=%self.cache_key, "Fetching new history page");
             let fetch_res = self
                 .client
-                .get_workflow_execution_history(self.wf_id.clone(), Some(self.run_id.clone()), npt)
+                .get_workflow_execution_history(
+                    self.wf_id.clone(),
+                    Some(self.cache_key.run_id.clone()),
+                    npt,
+                )
                 .instrument(span!(tracing::Level::TRACE, "fetch_history_in_paginator"))
                 .await?;