@temporalio/core-bridge 0.20.2 → 0.21.0

package/index.d.ts

@@ -162,6 +162,7 @@ export declare function newReplayWorker(
   callback: WorkerCallback
 ): void;
 export declare function workerShutdown(worker: Worker, callback: VoidCallback): void;
+export declare function clientClose(client: Client): void;
 export declare function runtimeShutdown(runtime: Runtime, callback: VoidCallback): void;
 export declare function pollLogs(runtime: Runtime, callback: LogsCallback): void;
 export declare function workerPollWorkflowActivation(

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@temporalio/core-bridge",
-  "version": "0.20.2",
+  "version": "0.21.0",
   "description": "Temporal.io SDK Core<>Node bridge",
   "main": "index.js",
   "types": "index.d.ts",
@@ -20,7 +20,7 @@
   "license": "MIT",
   "dependencies": {
     "@opentelemetry/api": "^1.0.3",
-    "@temporalio/internal-non-workflow-common": "^0.20.0",
+    "@temporalio/internal-non-workflow-common": "^0.21.0",
     "arg": "^5.0.1",
     "cargo-cp-artifact": "^0.1.4",
     "which": "^2.0.2"
@@ -43,5 +43,5 @@
   "publishConfig": {
     "access": "public"
   },
-  "gitHead": "ce4a3b017cdc327fb9f2d3812e3278304f6514b4"
+  "gitHead": "eb5901f47e16f8c8fe36c1154d5176c5f3205efc"
 }

package/sdk-core/core/src/core_tests/queries.rs

@@ -21,7 +21,8 @@ use temporal_sdk_core_protos::{
         history::v1::History,
         query::v1::WorkflowQuery,
         workflowservice::v1::{
-            RespondQueryTaskCompletedResponse, RespondWorkflowTaskCompletedResponse,
+            GetWorkflowExecutionHistoryResponse, RespondQueryTaskCompletedResponse,
+            RespondWorkflowTaskCompletedResponse,
         },
     },
 };
@@ -381,3 +382,108 @@ async fn legacy_query_after_complete(#[values(false, true)] full_history: bool)
 
     core.shutdown().await;
 }
+
+#[tokio::test]
+async fn query_cache_miss_causes_page_fetch_dont_reply_wft_too_early() {
+    let wfid = "fake_wf_id";
+    let query_resp = "response";
+    let t = canned_histories::single_timer("1");
+    let full_hist = t.get_full_history_info().unwrap();
+    let tasks = VecDeque::from(vec![{
+        // Create a partial task
+        let mut pr = hist_to_poll_resp(
+            &t,
+            wfid.to_owned(),
+            ResponseType::OneTask(2),
+            TEST_Q.to_string(),
+        );
+        pr.queries = HashMap::new();
+        pr.queries.insert(
+            "the-query".to_string(),
+            WorkflowQuery {
+                query_type: "query-type".to_string(),
+                query_args: Some(b"hi".into()),
+                header: None,
+            },
+        );
+        pr
+    }]);
+    let mut mock_client = mock_workflow_client();
+    mock_client
+        .expect_get_workflow_execution_history()
+        .returning(move |_, _, _| {
+            Ok(GetWorkflowExecutionHistoryResponse {
+                history: Some(full_hist.clone().into()),
+                ..Default::default()
+            })
+        });
+    mock_client
+        .expect_complete_workflow_task()
+        .times(1)
+        .returning(|resp| {
+            // Verify both the complete command and the query response are sent
+            assert_eq!(resp.commands.len(), 1);
+            assert_eq!(resp.query_responses.len(), 1);
+
+            Ok(RespondWorkflowTaskCompletedResponse::default())
+        });
+
+    let mut mock = MocksHolder::from_client_with_responses(mock_client, tasks, vec![]);
+    mock.worker_cfg(|wc| wc.max_cached_workflows = 10);
+    let core = mock_worker(mock);
+    let task = core.poll_workflow_activation().await.unwrap();
+    // The first task should *only* start the workflow. It should *not* have a query in it, which
+    // was the bug. Query should only appear after we have caught up on replay.
+    assert_matches!(
+        task.jobs.as_slice(),
+        [WorkflowActivationJob {
+            variant: Some(workflow_activation_job::Variant::StartWorkflow(_)),
+        }]
+    );
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
+        task.run_id,
+        start_timer_cmd(1, Duration::from_secs(1)),
+    ))
+    .await
+    .unwrap();
+
+    let task = core.poll_workflow_activation().await.unwrap();
+    assert_matches!(
+        task.jobs.as_slice(),
+        [WorkflowActivationJob {
+            variant: Some(workflow_activation_job::Variant::FireTimer(_)),
+        }]
+    );
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
+        task.run_id,
+        CompleteWorkflowExecution { result: None }.into(),
+    ))
+    .await
+    .unwrap();
+
+    // Now the query shall arrive
+    let task = core.poll_workflow_activation().await.unwrap();
+    assert_matches!(
+        task.jobs[0],
+        WorkflowActivationJob {
+            variant: Some(workflow_activation_job::Variant::QueryWorkflow(_)),
+        }
+    );
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
+        task.run_id,
+        QueryResult {
+            query_id: "the-query".to_string(),
+            variant: Some(
+                QuerySuccess {
+                    response: Some(query_resp.into()),
+                }
+                .into(),
+            ),
+        }
+        .into(),
+    ))
+    .await
+    .unwrap();
+
+    core.shutdown().await;
+}

package/sdk-core/core/src/workflow/workflow_tasks/mod.rs

@@ -57,7 +57,7 @@ pub struct WorkflowTaskManager {
     pending_activations: PendingActivations,
     /// Holds activations which are purely query activations needed to respond to legacy queries.
     /// Activations may only be added here for runs which do not have other pending activations.
-    pending_legacy_queries: SegQueue<WorkflowActivation>,
+    pending_queries: SegQueue<WorkflowActivation>,
     /// Holds poll wft responses from the server that need to be applied
     ready_buffered_wft: SegQueue<ValidPollWFTQResponse>,
     /// Used to wake blocked workflow task polling
@@ -74,9 +74,8 @@ pub struct WorkflowTaskManager {
 #[derive(Clone, Debug)]
 pub(crate) struct OutstandingTask {
     pub info: WorkflowTaskInfo,
-    /// If set the outstanding task has query from the old `query` field which must be fulfilled
-    /// upon finishing replay
-    pub legacy_query: Option<QueryWorkflow>,
+    /// Set if the outstanding task has quer(ies) which must be fulfilled upon finishing replay
+    pub pending_queries: Vec<QueryWorkflow>,
     start_time: Instant,
 }
 
@@ -179,7 +178,7 @@ impl WorkflowTaskManager {
         Self {
             workflow_machines: WorkflowConcurrencyManager::new(),
             pending_activations: Default::default(),
-            pending_legacy_queries: Default::default(),
+            pending_queries: Default::default(),
             ready_buffered_wft: Default::default(),
             pending_activations_notifier,
             cache_manager: Mutex::new(WorkflowCacheManager::new(eviction_policy, metrics.clone())),
@@ -188,8 +187,8 @@ impl WorkflowTaskManager {
     }
 
     pub(crate) fn next_pending_activation(&self) -> Option<WorkflowActivation> {
-        // Dispatch pending legacy queries first
-        if let leg_q @ Some(_) = self.pending_legacy_queries.pop() {
+        // Dispatch pending queries first
+        if let leg_q @ Some(_) = self.pending_queries.pop() {
             return leg_q;
         }
         // It is important that we do not issue pending activations for any workflows which already
@@ -333,33 +332,45 @@ impl WorkflowTaskManager {
             .take()
             .map(|q| query_to_job(LEGACY_QUERY_ID.to_string(), q));
 
-        let (info, mut next_activation) =
+        let (info, mut next_activation, mut pending_queries) =
             match self.instantiate_or_update_workflow(work, client).await {
-                Ok((info, next_activation)) => (info, next_activation),
+                Ok(res) => res,
                 Err(e) => {
                     return NewWfTaskOutcome::Evict(e);
                 }
             };
 
+        if !pending_queries.is_empty() && legacy_query.is_some() {
+            error!(
+                "Server issued both normal and legacy queries. This should not happen. Please \
+                 file a bug report."
+            );
+            return NewWfTaskOutcome::Evict(WorkflowUpdateError {
+                source: WFMachinesError::Fatal(
+                    "Server issued both normal and legacy query".to_string(),
+                ),
+                run_id: next_activation.run_id,
+            });
+        }
+
         // Immediately dispatch query activation if no other jobs
-        let legacy_query = if next_activation.jobs.is_empty() {
-            if let Some(lq) = legacy_query {
+        if let Some(lq) = legacy_query {
+            if next_activation.jobs.is_empty() {
                 debug!("Dispatching legacy query {}", &lq);
                 next_activation
                     .jobs
                     .push(workflow_activation_job::Variant::QueryWorkflow(lq).into());
+            } else {
+                pending_queries.push(lq);
             }
-            None
-        } else {
-            legacy_query
-        };
+        }
 
         self.workflow_machines
             .insert_wft(
                 &next_activation.run_id,
                 OutstandingTask {
                     info,
-                    legacy_query,
+                    pending_queries,
                     start_time: task_start_time,
                 },
             )
@@ -401,11 +412,11 @@ impl WorkflowTaskManager {
             return Ok(None);
         }
 
-        let (task_token, is_leg_query_task, start_time) =
+        let (task_token, has_pending_query, start_time) =
             if let Some(entry) = self.workflow_machines.get_task(run_id) {
                 (
                     entry.info.task_token.clone(),
-                    entry.legacy_query.is_some(),
+                    !entry.pending_queries.is_empty(),
                     entry.start_time,
                 )
             } else {
@@ -506,7 +517,7 @@ impl WorkflowTaskManager {
             let must_heartbeat = self
                 .wait_for_local_acts_or_heartbeat(run_id, wft_heartbeat_deadline)
                 .await;
-            let is_query_playback = is_leg_query_task && query_responses.is_empty();
+            let is_query_playback = has_pending_query && query_responses.is_empty();
 
             // We only actually want to send commands back to the server if there are no more
             // pending activations and we are caught up on replay. We don't want to complete a wft
@@ -592,7 +603,8 @@ impl WorkflowTaskManager {
         &self,
         poll_wf_resp: ValidPollWFTQResponse,
         client: Arc<WorkerClientBag>,
-    ) -> Result<(WorkflowTaskInfo, WorkflowActivation), WorkflowUpdateError> {
+    ) -> Result<(WorkflowTaskInfo, WorkflowActivation, Vec<QueryWorkflow>), WorkflowUpdateError>
+    {
         let run_id = poll_wf_resp.workflow_execution.run_id.clone();
 
         let wft_info = WorkflowTaskInfo {
@@ -607,10 +619,12 @@ impl WorkflowTaskManager {
             .map(|ev| ev.event_id > 1)
             .unwrap_or_default();
 
+        let mut did_miss_cache = false;
         let page_token = if !self.workflow_machines.exists(&run_id) && poll_resp_is_incremental {
             debug!(run_id=?run_id, "Workflow task has partial history, but workflow is not in \
                    cache. Will fetch history");
             self.metrics.sticky_cache_miss();
+            did_miss_cache = true;
             NextPageToken::FetchFromStart
         } else {
             poll_wf_resp.next_page_token.into()
@@ -639,16 +653,26 @@ impl WorkflowTaskManager {
             .await
         {
             Ok(mut activation) => {
-                // If there are in-poll queries, insert jobs for those queries into the activation
+                // If there are in-poll queries, insert jobs for those queries into the activation,
+                // but only if we hit the cache. If we didn't, those queries will need to be dealt
+                // with once replay is over
+                let mut pending_queries = vec![];
                 if !poll_wf_resp.query_requests.is_empty() {
-                    let query_jobs = poll_wf_resp
-                        .query_requests
-                        .into_iter()
-                        .map(|q| workflow_activation_job::Variant::QueryWorkflow(q).into());
-                    activation.jobs.extend(query_jobs);
+                    if !did_miss_cache {
+                        let query_jobs = poll_wf_resp
+                            .query_requests
+                            .into_iter()
+                            .map(|q| workflow_activation_job::Variant::QueryWorkflow(q).into());
+                        activation.jobs.extend(query_jobs);
+                    } else {
+                        poll_wf_resp
+                            .query_requests
+                            .into_iter()
+                            .for_each(|q| pending_queries.push(q));
+                    }
                 }
 
-                Ok((wft_info, activation))
+                Ok((wft_info, activation, pending_queries))
             }
             Err(source) => Err(WorkflowUpdateError { source, run_id }),
         }
@@ -675,16 +699,18 @@ impl WorkflowTaskManager {
         // removed from the outstanding tasks map
         let retme = if !self.pending_activations.has_pending(run_id) {
             if !just_evicted {
-                // Check if there was a legacy query which must be fulfilled, and if there is create
-                // a new pending activation for it.
+                // Check if there was a pending query which must be fulfilled, and if there is
+                // create a new pending activation for it.
                 if let Some(ref mut ot) = &mut *self
                     .workflow_machines
                     .get_task_mut(run_id)
                     .expect("Machine must exist")
                 {
-                    if let Some(query) = ot.legacy_query.take() {
-                        let na = create_query_activation(run_id.to_string(), [query]);
-                        self.pending_legacy_queries.push(na);
+                    if !ot.pending_queries.is_empty() {
+                        for query in ot.pending_queries.drain(..) {
+                            let na = create_query_activation(run_id.to_string(), [query]);
+                            self.pending_queries.push(na);
+                        }
                         self.pending_activations_notifier.notify_waiters();
                         return false;
                     }

package/src/errors.rs

@@ -10,6 +10,8 @@ pub static SHUTDOWN_ERROR: OnceCell<Root<JsFunction>> = OnceCell::new();
 pub static NO_WORKER_ERROR: OnceCell<Root<JsFunction>> = OnceCell::new();
 /// Something unexpected happened, considered fatal
 pub static UNEXPECTED_ERROR: OnceCell<Root<JsFunction>> = OnceCell::new();
+/// Used in different parts of the project to signal that something unexpected has happened
+pub static ILLEGAL_STATE_ERROR: OnceCell<Root<JsFunction>> = OnceCell::new();
 
 static ALREADY_REGISTERED_ERRORS: OnceCell<bool> = OnceCell::new();
 
@@ -70,9 +72,9 @@ pub fn register_errors(mut cx: FunctionContext) -> JsResult<JsUndefined> {
     let res = ALREADY_REGISTERED_ERRORS.set(true);
     if res.is_err() {
         // Don't do anything if errors are already registered
-        return Ok(cx.undefined())
+        return Ok(cx.undefined());
     }
-    
+
     let mapping = cx.argument::<JsObject>(0)?;
     let shutdown_error = mapping
         .get(&mut cx, "ShutdownError")?
@@ -90,11 +92,16 @@ pub fn register_errors(mut cx: FunctionContext) -> JsResult<JsUndefined> {
         .get(&mut cx, "UnexpectedError")?
         .downcast_or_throw::<JsFunction, FunctionContext>(&mut cx)?
         .root(&mut cx);
+    let illegal_state_error = mapping
+        .get(&mut cx, "IllegalStateError")?
+        .downcast_or_throw::<JsFunction, FunctionContext>(&mut cx)?
+        .root(&mut cx);
 
     TRANSPORT_ERROR.get_or_try_init(|| Ok(transport_error))?;
     SHUTDOWN_ERROR.get_or_try_init(|| Ok(shutdown_error))?;
     NO_WORKER_ERROR.get_or_try_init(|| Ok(no_worker_error))?;
     UNEXPECTED_ERROR.get_or_try_init(|| Ok(unexpected_error))?;
+    ILLEGAL_STATE_ERROR.get_or_try_init(|| Ok(illegal_state_error))?;
 
     Ok(cx.undefined())
 }

package/src/lib.rs

@@ -8,6 +8,7 @@ use once_cell::sync::OnceCell;
 use opentelemetry::trace::{FutureExt, SpanContext, TraceContextExt};
 use prost::Message;
 use std::{
+    cell::RefCell,
     fmt::Display,
     future::Future,
     sync::Arc,
@@ -135,7 +136,7 @@ struct Client {
     core_client: Arc<RawClient>,
 }
 
-type BoxedClient = JsBox<Client>;
+type BoxedClient = JsBox<RefCell<Option<Client>>>;
 impl Finalize for Client {}
 
 /// Worker struct, hold a reference for the channel sender responsible for sending requests from
@@ -291,10 +292,10 @@ fn start_bridge_loop(event_queue: Arc<EventQueue>, receiver: &mut UnboundedRecei
                         }
                         Ok(client) => {
                             send_result(event_queue.clone(), callback, |cx| {
-                                Ok(cx.boxed(Client {
+                                Ok(cx.boxed(RefCell::new(Some(Client {
                                     runtime,
                                     core_client: Arc::new(client),
-                                }))
+                                }))))
                             });
                         }
                     }
@@ -590,15 +591,23 @@ fn worker_new(mut cx: FunctionContext) -> JsResult<JsUndefined> {
     let callback = cx.argument::<JsFunction>(2)?;
 
     let config = worker_options.as_worker_config(&mut cx)?;
-
-    let request = Request::InitWorker {
-        client: client.core_client.clone(),
-        runtime: client.runtime.clone(),
-        config,
-        callback: callback.root(&mut cx),
-    };
-    if let Err(err) = client.runtime.sender.send(request) {
-        callback_with_unexpected_error(&mut cx, callback, err)?;
+    match &*client.borrow() {
+        None => {
+            callback_with_error(&mut cx, callback, move |cx| {
+                UNEXPECTED_ERROR.from_string(cx, "Tried to use closed Client".to_string())
+            })?;
+        }
+        Some(client) => {
+            let request = Request::InitWorker {
+                client: client.core_client.clone(),
+                runtime: client.runtime.clone(),
+                config,
+                callback: callback.root(&mut cx),
+            };
+            if let Err(err) = client.runtime.sender.send(request) {
+                callback_with_unexpected_error(&mut cx, callback, err)?;
+            };
+        }
     };
 
     Ok(cx.undefined())
@@ -783,13 +792,26 @@ fn worker_record_activity_heartbeat(mut cx: FunctionContext) -> JsResult<JsUndef
 fn worker_shutdown(mut cx: FunctionContext) -> JsResult<JsUndefined> {
     let worker = cx.argument::<BoxedWorker>(0)?;
     let callback = cx.argument::<JsFunction>(1)?;
-    match worker.runtime.sender.send(Request::ShutdownWorker {
+    if let Err(err) = worker.runtime.sender.send(Request::ShutdownWorker {
         worker: worker.core_worker.clone(),
         callback: callback.root(&mut cx),
     }) {
-        Err(err) => cx.throw_error(format!("{}", err)),
-        _ => Ok(cx.undefined()),
-    }
+        UNEXPECTED_ERROR
+            .from_error(&mut cx, err)
+            .and_then(|err| cx.throw(err))?;
+    };
+    Ok(cx.undefined())
+}
+
+/// Drop a reference to a Client, once all references are dropped, the Client will be closed.
+fn client_close(mut cx: FunctionContext) -> JsResult<JsUndefined> {
+    let client = cx.argument::<BoxedClient>(0)?;
+    if client.replace(None).is_none() {
+        ILLEGAL_STATE_ERROR
+            .from_error(&mut cx, "Client already closed")
+            .and_then(|err| cx.throw(err))?;
+    };
+    Ok(cx.undefined())
 }
 
 /// Convert Rust SystemTime into a JS array with 2 numbers (seconds, nanos)
@@ -824,6 +846,7 @@ fn main(mut cx: ModuleContext) -> NeonResult<()> {
     cx.export_function("newWorker", worker_new)?;
     cx.export_function("newReplayWorker", replay_worker_new)?;
     cx.export_function("workerShutdown", worker_shutdown)?;
+    cx.export_function("clientClose", client_close)?;
     cx.export_function("runtimeShutdown", runtime_shutdown)?;
     cx.export_function("pollLogs", poll_logs)?;
     cx.export_function(