@temporalio/core-bridge 1.6.0 → 1.7.1

package/sdk-core/test-utils/src/lib.rs

@@ -34,6 +34,7 @@ use temporal_sdk_core::{
     replay::HistoryForReplay,
     ClientOptions, ClientOptionsBuilder, CoreRuntime, WorkerConfig, WorkerConfigBuilder,
 };
+use temporal_sdk_core_api::errors::{PollActivityError, PollWfError};
 use temporal_sdk_core_api::{
     telemetry::{
         Logger, MetricsExporter, OtelCollectorOptions, TelemetryOptions, TelemetryOptionsBuilder,
@@ -50,6 +51,7 @@ use temporal_sdk_core_protos::{
         workflow_completion::WorkflowActivationCompletion,
     },
     temporal::api::{common::v1::Payload, history::v1::History},
+    DEFAULT_ACTIVITY_TYPE,
 };
 use tokio::sync::{mpsc::unbounded_channel, OnceCell};
 use url::Url;
@@ -58,8 +60,8 @@ pub const NAMESPACE: &str = "default";
 pub const TEST_Q: &str = "q";
 /// The env var used to specify where the integ tests should point
 pub const INTEG_SERVER_TARGET_ENV_VAR: &str = "TEMPORAL_SERVICE_ADDRESS";
-/// This env var is set (to any value) if temporalite is in use
-pub const INTEG_TEMPORALITE_USED_ENV_VAR: &str = "INTEG_TEMPORALITE_ON";
+/// This env var is set (to any value) if temporal CLI dev server is in use
+pub const INTEG_TEMPORAL_DEV_SERVER_USED_ENV_VAR: &str = "INTEG_TEMPORAL_DEV_SERVER_ON";
 /// This env var is set (to any value) if the test server is in use
 pub const INTEG_TEST_SERVER_USED_ENV_VAR: &str = "INTEG_TEST_SERVER_ON";
 
@@ -301,6 +303,11 @@ impl CoreWfStarter {
         self
     }
 
+    pub fn no_remote_activities(&mut self) -> &mut Self {
+        self.worker_config.no_remote_activities = true;
+        self
+    }
+
     pub fn enable_wf_state_input_recording(&mut self) -> &mut Self {
         let (ser_tx, ser_rx) = unbounded_channel();
         let worker_cfg_clone = self.worker_config.clone();
@@ -569,7 +576,7 @@ pub fn get_integ_telem_options() -> TelemetryOptions {
 
 pub fn default_cached_download() -> EphemeralExe {
     EphemeralExe::CachedDownload {
-        version: EphemeralExeVersion::Default {
+        version: EphemeralExeVersion::SDKDefault {
             sdk_name: "sdk-rust".to_string(),
             sdk_version: "0.1.0".to_string(),
         },
@@ -588,7 +595,7 @@ pub fn schedule_activity_cmd(
     ScheduleActivity {
         seq,
         activity_id: activity_id.to_string(),
-        activity_type: "test_activity".to_string(),
+        activity_type: DEFAULT_ACTIVITY_TYPE.to_string(),
         task_queue: task_q.to_owned(),
         schedule_to_start_timeout: Some(activity_timeout.try_into().expect("duration fits")),
         start_to_close_timeout: Some(activity_timeout.try_into().expect("duration fits")),
@@ -609,7 +616,7 @@ pub fn schedule_local_activity_cmd(
     ScheduleLocalActivity {
         seq,
         activity_id: activity_id.to_string(),
-        activity_type: "test_activity".to_string(),
+        activity_type: DEFAULT_ACTIVITY_TYPE.to_string(),
         schedule_to_start_timeout: Some(activity_timeout.try_into().expect("duration fits")),
         start_to_close_timeout: Some(activity_timeout.try_into().expect("duration fits")),
         schedule_to_close_timeout: Some(activity_timeout.try_into().expect("duration fits")),
@@ -680,3 +687,23 @@ where
         .unwrap();
     }
 }
+
+/// Initiate shutdown, drain the pollers, and wait for shutdown to complete.
+pub async fn drain_pollers_and_shutdown(worker: &Arc<dyn CoreWorker>) {
+    worker.initiate_shutdown();
+    tokio::join!(
+        async {
+            assert!(matches!(
+                worker.poll_activity_task().await.unwrap_err(),
+                PollActivityError::ShutDown
+            ));
+        },
+        async {
+            assert!(matches!(
+                worker.poll_workflow_activation().await.unwrap_err(),
+                PollWfError::ShutDown,
+            ));
+        }
+    );
+    worker.shutdown().await;
+}

package/sdk-core/tests/heavy_tests.rs

@@ -1,11 +1,9 @@
-use assert_matches::assert_matches;
 use futures::{future::join_all, sink, stream::FuturesUnordered, StreamExt};
 use std::time::{Duration, Instant};
 use temporal_client::{WfClientExt, WorkflowClientTrait, WorkflowOptions};
 use temporal_sdk::{ActContext, ActivityOptions, WfContext, WorkflowResult};
 use temporal_sdk_core_protos::coresdk::{
-    activity_result::ActivityExecutionResult, activity_task::activity_task as act_task,
-    workflow_commands::ActivityCancellationType, ActivityTaskCompletion, AsJsonPayloadExt,
+    workflow_commands::ActivityCancellationType, AsJsonPayloadExt,
 };
 use temporal_sdk_core_test_utils::{workflows::la_problem_workflow, CoreWfStarter};
 
@@ -25,18 +23,16 @@ async fn activity_load() {
 
     let activity_id = "act-1";
     let activity_timeout = Duration::from_secs(8);
-    let payload_dat = b"hello".to_vec();
     let task_queue = starter.get_task_queue().to_owned();
 
-    let pd = payload_dat.clone();
     let wf_fn = move |ctx: WfContext| {
         let task_queue = task_queue.clone();
-        let payload_dat = pd.clone();
-
+        let payload = "yo".as_json_payload().unwrap();
         async move {
             let activity = ActivityOptions {
                 activity_id: Some(activity_id.to_string()),
                 activity_type: "test_activity".to_string(),
+                input: payload.clone(),
                 task_queue,
                 schedule_to_start_timeout: Some(activity_timeout),
                 start_to_close_timeout: Some(activity_timeout),
@@ -46,7 +42,7 @@ async fn activity_load() {
                 ..Default::default()
             };
             let res = ctx.activity(activity).await.unwrap_ok_payload();
-            assert_eq!(res.data, payload_dat);
+            assert_eq!(res.data, payload.data);
             Ok(().into())
         }
     };
@@ -54,6 +50,10 @@ async fn activity_load() {
     let starting = Instant::now();
     let wf_type = "activity_load";
     worker.register_wf(wf_type.to_owned(), wf_fn);
+    worker.register_activity(
+        "test_activity",
+        |_ctx: ActContext, echo: String| async move { Ok(echo) },
+    );
     join_all((0..CONCURRENCY).map(|i| {
         let worker = &worker;
         let wf_id = format!("activity_load_{i}");
@@ -73,42 +73,8 @@ async fn activity_load() {
     dbg!(starting.elapsed());
 
     let running = Instant::now();
-    let core = starter.get_worker().await;
 
-    // Poll for and complete all activities
-    let c2 = core.clone();
-    let all_acts = async move {
-        let mut act_complete_futs = vec![];
-        for _ in 0..CONCURRENCY {
-            let task = c2.poll_activity_task().await.unwrap();
-            assert_matches!(
-                task.variant,
-                Some(act_task::Variant::Start(ref start_activity)) => {
-                    assert_eq!(start_activity.activity_type, "test_activity")
-                }
-            );
-            let pd = payload_dat.clone();
-            let core = c2.clone();
-            act_complete_futs.push(tokio::spawn(async move {
-                core.complete_activity_task(ActivityTaskCompletion {
-                    task_token: task.task_token,
-                    result: Some(ActivityExecutionResult::ok(pd.into())),
-                })
-                .await
-                .unwrap()
-            }));
-        }
-        join_all(act_complete_futs)
-            .await
-            .into_iter()
-            .for_each(|h| h.unwrap());
-    };
-    tokio::join! {
-        async {
-            worker.run_until_done().await.unwrap();
-        },
-        all_acts
-    };
+    worker.run_until_done().await.unwrap();
     dbg!(running.elapsed());
 }
 
@@ -262,6 +228,7 @@ pub async fn many_parallel_timers_longhist(ctx: WfContext) -> WorkflowResult<()>
 async fn can_paginate_long_history() {
     let wf_name = "can_paginate_long_history";
     let mut starter = CoreWfStarter::new(wf_name);
+    starter.no_remote_activities();
     // Do not use sticky queues so we are forced to paginate once history gets long
     starter.max_cached_workflows(0);
 

package/sdk-core/tests/integ_tests/ephemeral_server_tests.rs

@@ -1,13 +1,35 @@
 use std::time::{SystemTime, UNIX_EPOCH};
 use temporal_client::{ClientOptionsBuilder, TestService, WorkflowService};
 use temporal_sdk_core::ephemeral_server::{
-    EphemeralExe, EphemeralExeVersion, EphemeralServer, TemporaliteConfigBuilder,
-    TestServerConfigBuilder,
+    EphemeralExe, EphemeralExeVersion, EphemeralServer, TemporalDevServerConfigBuilder,
+    TemporaliteConfigBuilder, TestServerConfigBuilder,
 };
 use temporal_sdk_core_protos::temporal::api::workflowservice::v1::DescribeNamespaceRequest;
 use temporal_sdk_core_test_utils::{default_cached_download, NAMESPACE};
 use url::Url;
 
+#[tokio::test]
+async fn temporal_cli_default() {
+    let config = TemporalDevServerConfigBuilder::default()
+        .exe(default_cached_download())
+        .build()
+        .unwrap();
+    let mut server = config.start_server().await.unwrap();
+    assert_ephemeral_server(&server).await;
+    server.shutdown().await.unwrap();
+}
+
+#[tokio::test]
+async fn temporal_cli_fixed() {
+    let config = TemporalDevServerConfigBuilder::default()
+        .exe(fixed_cached_download("v0.4.0"))
+        .build()
+        .unwrap();
+    let mut server = config.start_server().await.unwrap();
+    assert_ephemeral_server(&server).await;
+    server.shutdown().await.unwrap();
+}
+
 #[tokio::test]
 async fn temporalite_default() {
     let config = TemporaliteConfigBuilder::default()
@@ -22,7 +44,7 @@ async fn temporalite_default() {
 #[tokio::test]
 async fn temporalite_fixed() {
     let config = TemporaliteConfigBuilder::default()
-        .exe(fixed_cached_download("v0.1.1"))
+        .exe(fixed_cached_download("v0.2.0"))
         .build()
         .unwrap();
     let mut server = config.start_server().await.unwrap();

package/sdk-core/tests/integ_tests/heartbeat_tests.rs

@@ -17,9 +17,11 @@ use temporal_sdk_core_protos::{
         common::v1::{Payload, RetryPolicy},
         enums::v1::TimeoutType,
     },
+    DEFAULT_ACTIVITY_TYPE,
 };
 use temporal_sdk_core_test_utils::{
-    init_core_and_create_wf, schedule_activity_cmd, CoreWfStarter, WorkerTestHelpers,
+    drain_pollers_and_shutdown, init_core_and_create_wf, schedule_activity_cmd, CoreWfStarter,
+    WorkerTestHelpers,
 };
 use tokio::time::sleep;
 
@@ -49,7 +51,7 @@ async fn activity_heartbeat() {
     assert_matches!(
         task.variant,
         Some(activity_task::Variant::Start(start_activity)) => {
-            assert_eq!(start_activity.activity_type, "test_activity".to_string())
+            assert_eq!(start_activity.activity_type, DEFAULT_ACTIVITY_TYPE.to_string())
         }
     );
     // Heartbeat timeout is set to 1 second, this loop is going to send heartbeat every 100ms.
@@ -169,7 +171,7 @@ async fn many_act_fails_with_heartbeats() {
         },]
     );
     core.complete_execution(&task.run_id).await;
-    core.shutdown().await;
+    drain_pollers_and_shutdown(&core).await;
 }
 
 #[tokio::test]

package/sdk-core/tests/integ_tests/metrics_tests.rs

@@ -1,18 +1,34 @@
-use temporal_client::WorkflowService;
-use temporal_sdk_core::CoreRuntime;
-use temporal_sdk_core_api::telemetry::MetricsExporter;
-use temporal_sdk_core_protos::temporal::api::workflowservice::v1::ListNamespacesRequest;
-use temporal_sdk_core_test_utils::{get_integ_server_options, get_integ_telem_options};
+use std::{sync::Arc, time::Duration};
+use temporal_client::{WorkflowClientTrait, WorkflowOptions, WorkflowService};
+use temporal_sdk_core::{init_worker, CoreRuntime};
+use temporal_sdk_core_api::{telemetry::MetricsExporter, worker::WorkerConfigBuilder, Worker};
+use temporal_sdk_core_protos::{
+    coresdk::{
+        activity_result::ActivityExecutionResult,
+        workflow_commands::{ScheduleActivity, ScheduleLocalActivity},
+        workflow_completion::WorkflowActivationCompletion,
+        ActivityTaskCompletion,
+    },
+    temporal::api::{enums::v1::WorkflowIdReusePolicy, workflowservice::v1::ListNamespacesRequest},
+};
+use temporal_sdk_core_test_utils::{get_integ_server_options, get_integ_telem_options, NAMESPACE};
+use tokio::sync::Barrier;
+
+static ANY_PORT: &str = "127.0.0.1:0";
+
+async fn get_text(endpoint: String) -> String {
+    reqwest::get(endpoint).await.unwrap().text().await.unwrap()
+}
 
 #[tokio::test]
 async fn prometheus_metrics_exported() {
     let mut telemopts = get_integ_telem_options();
-    let addr = "127.0.0.1:10919";
-    telemopts.metrics = Some(MetricsExporter::Prometheus(addr.parse().unwrap()));
+    telemopts.metrics = Some(MetricsExporter::Prometheus(ANY_PORT.parse().unwrap()));
     let rt = CoreRuntime::new_assume_tokio(telemopts).unwrap();
+    let addr = rt.telemetry().prom_port().unwrap();
     let opts = get_integ_server_options();
     let mut raw_client = opts
-        .connect_no_namespace(rt.metric_meter(), None)
+        .connect_no_namespace(rt.metric_meter().as_deref(), None)
         .await
         .unwrap();
     assert!(raw_client.get_client().capabilities().is_some());
@@ -22,16 +38,202 @@ async fn prometheus_metrics_exported() {
         .await
         .unwrap();
 
-    let body = reqwest::get(format!("http://{addr}/metrics"))
-        .await
-        .unwrap()
-        .text()
-        .await
-        .unwrap();
+    let body = get_text(format!("http://{addr}/metrics")).await;
     assert!(body.contains(
-        "request_latency_count{operation=\"ListNamespaces\",service_name=\"temporal-core-sdk\"} 1"
+        "temporal_request_latency_count{operation=\"ListNamespaces\",service_name=\"temporal-core-sdk\"} 1"
     ));
     assert!(body.contains(
-        "request_latency_count{operation=\"GetSystemInfo\",service_name=\"temporal-core-sdk\"} 1"
+        "temporal_request_latency_count{operation=\"GetSystemInfo\",service_name=\"temporal-core-sdk\"} 1"
     ));
 }
+
+#[tokio::test]
+async fn one_slot_worker_reports_available_slot() {
+    let mut telemopts = get_integ_telem_options();
+    let tq = "one_slot_worker_tq";
+    telemopts.metrics = Some(MetricsExporter::Prometheus(ANY_PORT.parse().unwrap()));
+    let rt = CoreRuntime::new_assume_tokio(telemopts).unwrap();
+    let addr = rt.telemetry().prom_port().unwrap();
+
+    let worker_cfg = WorkerConfigBuilder::default()
+        .namespace(NAMESPACE)
+        .task_queue(tq)
+        .worker_build_id("test_build_id")
+        .max_cached_workflows(1_usize)
+        .max_outstanding_activities(1_usize)
+        .max_outstanding_local_activities(1_usize)
+        .max_outstanding_workflow_tasks(1_usize)
+        .build()
+        .unwrap();
+
+    let client = Arc::new(
+        get_integ_server_options()
+            .connect(worker_cfg.namespace.clone(), None, None)
+            .await
+            .expect("Must connect"),
+    );
+    let worker = init_worker(&rt, worker_cfg, client.clone()).expect("Worker inits cleanly");
+    let wf_task_barr = Barrier::new(2);
+    let act_task_barr = Barrier::new(2);
+
+    let wf_polling = async {
+        let task = worker.poll_workflow_activation().await.unwrap();
+        wf_task_barr.wait().await;
+        wf_task_barr.wait().await;
+        worker
+            .complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
+                task.run_id,
+                ScheduleActivity {
+                    seq: 1,
+                    activity_id: "1".to_string(),
+                    activity_type: "test_act".to_string(),
+                    task_queue: tq.to_string(),
+                    start_to_close_timeout: Some(prost_dur!(from_secs(30))),
+                    ..Default::default()
+                }
+                .into(),
+            ))
+            .await
+            .unwrap();
+        wf_task_barr.wait().await;
+
+        let task = worker.poll_workflow_activation().await.unwrap();
+        worker
+            .complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
+                task.run_id,
+                ScheduleLocalActivity {
+                    seq: 2,
+                    activity_id: "2".to_string(),
+                    activity_type: "test_act".to_string(),
+                    start_to_close_timeout: Some(prost_dur!(from_secs(30))),
+                    ..Default::default()
+                }
+                .into(),
+            ))
+            .await
+            .unwrap();
+    };
+
+    let act_polling = async {
+        let task = worker.poll_activity_task().await.unwrap();
+        act_task_barr.wait().await;
+        worker
+            .complete_activity_task(ActivityTaskCompletion {
+                task_token: task.task_token,
+                result: Some(ActivityExecutionResult::ok(vec![1].into())),
+            })
+            .await
+            .unwrap();
+        act_task_barr.wait().await;
+
+        let task = worker.poll_activity_task().await.unwrap();
+        act_task_barr.wait().await;
+        act_task_barr.wait().await;
+        worker
+            .complete_activity_task(ActivityTaskCompletion {
+                task_token: task.task_token,
+                result: Some(ActivityExecutionResult::ok(vec![1].into())),
+            })
+            .await
+            .unwrap();
+        act_task_barr.wait().await;
+    };
+
+    let testing = async {
+        // Wait just a beat for the poller to initiate
+        tokio::time::sleep(Duration::from_millis(50)).await;
+        let body = get_text(format!("http://{addr}/metrics")).await;
+        assert!(body.contains(&format!(
+            "temporal_worker_task_slots_available{{namespace=\"{NAMESPACE}\",\
+             service_name=\"temporal-core-sdk\",task_queue=\"one_slot_worker_tq\",\
+             worker_type=\"WorkflowWorker\"}} 1"
+        )));
+
+        // Start a workflow so that a task will get delivered
+        client
+            .start_workflow(
+                vec![],
+                tq.to_owned(),
+                "one_slot_metric_test".to_owned(),
+                "whatever".to_string(),
+                None,
+                WorkflowOptions {
+                    id_reuse_policy: WorkflowIdReusePolicy::TerminateIfRunning,
+                    execution_timeout: Some(Duration::from_secs(5)),
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+
+        wf_task_barr.wait().await;
+
+        // At this point the workflow task is outstanding, so there should be 0 slots, and
+        // the activities haven't started, so there should be 1 each.
+        let body = get_text(format!("http://{addr}/metrics")).await;
+        assert!(body.contains(&format!(
+            "temporal_worker_task_slots_available{{namespace=\"{NAMESPACE}\",\
+             service_name=\"temporal-core-sdk\",task_queue=\"one_slot_worker_tq\",\
+             worker_type=\"WorkflowWorker\"}} 0"
+        )));
+        assert!(body.contains(&format!(
+            "temporal_worker_task_slots_available{{namespace=\"{NAMESPACE}\",\
+             service_name=\"temporal-core-sdk\",task_queue=\"one_slot_worker_tq\",\
+             worker_type=\"ActivityWorker\"}} 1"
+        )));
+        assert!(body.contains(&format!(
+            "temporal_worker_task_slots_available{{namespace=\"{NAMESPACE}\",\
+             service_name=\"temporal-core-sdk\",task_queue=\"one_slot_worker_tq\",\
+             worker_type=\"LocalActivityWorker\"}} 1"
+        )));
+
+        // Now we allow the complete to proceed. Once it goes through, there should be 1 WFT slot
+        // open but 0 activity slots
+        wf_task_barr.wait().await;
+        wf_task_barr.wait().await;
+        // Sometimes the recording takes an extra bit. 🤷
+        tokio::time::sleep(Duration::from_millis(100)).await;
+        let body = get_text(format!("http://{addr}/metrics")).await;
+        assert!(body.contains(&format!(
+            "temporal_worker_task_slots_available{{namespace=\"{NAMESPACE}\",\
+             service_name=\"temporal-core-sdk\",task_queue=\"one_slot_worker_tq\",\
+             worker_type=\"WorkflowWorker\"}} 1"
+        )));
+        assert!(body.contains(&format!(
+            "temporal_worker_task_slots_available{{namespace=\"{NAMESPACE}\",\
+             service_name=\"temporal-core-sdk\",task_queue=\"one_slot_worker_tq\",\
+             worker_type=\"ActivityWorker\"}} 0"
+        )));
+
+        // Now complete the activity and watch it go up
+        act_task_barr.wait().await;
+        // Wait for completion to finish
+        act_task_barr.wait().await;
+        let body = get_text(format!("http://{addr}/metrics")).await;
+        assert!(body.contains(&format!(
+            "temporal_worker_task_slots_available{{namespace=\"{NAMESPACE}\",\
+             service_name=\"temporal-core-sdk\",task_queue=\"one_slot_worker_tq\",\
+             worker_type=\"ActivityWorker\"}} 1"
+        )));
+
+        // Proceed to local activity command
+        act_task_barr.wait().await;
+        // Ensure that, once we have the LA task, slots are 0
+        let body = get_text(format!("http://{addr}/metrics")).await;
+        assert!(body.contains(&format!(
+            "temporal_worker_task_slots_available{{namespace=\"{NAMESPACE}\",\
+             service_name=\"temporal-core-sdk\",task_queue=\"one_slot_worker_tq\",\
+             worker_type=\"LocalActivityWorker\"}} 0"
+        )));
+        // When completion is done, we have 1 again
+        act_task_barr.wait().await;
+        act_task_barr.wait().await;
+        let body = get_text(format!("http://{addr}/metrics")).await;
+        assert!(body.contains(&format!(
+            "temporal_worker_task_slots_available{{namespace=\"{NAMESPACE}\",\
+             service_name=\"temporal-core-sdk\",task_queue=\"one_slot_worker_tq\",\
+             worker_type=\"LocalActivityWorker\"}} 1"
+        )));
+    };
+    tokio::join!(wf_polling, act_polling, testing);
+}

package/sdk-core/tests/integ_tests/polling_tests.rs

@@ -40,15 +40,10 @@ async fn out_of_order_completion_doesnt_hang() {
     )
     .await
     .unwrap();
-    // Poll activity and verify that it's been scheduled with correct parameters, we don't expect to
-    // complete it in this test as activity is try-cancelled.
+    // Poll activity and verify that it's been scheduled, we don't expect to complete it in this
+    // test as activity is try-cancelled.
     let activity_task = core.poll_activity_task().await.unwrap();
-    assert_matches!(
-        activity_task.variant,
-        Some(act_task::Variant::Start(start_activity)) => {
-            assert_eq!(start_activity.activity_type, "test_activity".to_string())
-        }
-    );
+    assert_matches!(activity_task.variant, Some(act_task::Variant::Start(_)));
     // Poll workflow task and verify that activity has failed.
     let task = core.poll_workflow_activation().await.unwrap();
     assert_matches!(

package/sdk-core/tests/integ_tests/queries_tests.rs

@@ -10,7 +10,9 @@ use temporal_sdk_core_protos::{
     },
     temporal::api::{failure::v1::Failure, query::v1::WorkflowQuery},
 };
-use temporal_sdk_core_test_utils::{init_core_and_create_wf, WorkerTestHelpers};
+use temporal_sdk_core_test_utils::{
+    drain_pollers_and_shutdown, init_core_and_create_wf, WorkerTestHelpers,
+};
 
 #[tokio::test]
 async fn simple_query_legacy() {
@@ -205,7 +207,7 @@ async fn query_after_execution_complete(#[case] do_evict: bool) {
         query_futs.push(do_workflow(true).map(|_| ()).boxed());
     }
     while query_futs.next().await.is_some() {}
-    core.shutdown().await;
+    drain_pollers_and_shutdown(core).await;
 }
 
 #[tokio::test]

package/sdk-core/tests/integ_tests/visibility_tests.rs

@@ -8,7 +8,8 @@ use temporal_sdk_core_protos::coresdk::workflow_activation::{
     workflow_activation_job, WorkflowActivationJob,
 };
 use temporal_sdk_core_test_utils::{
-    get_integ_server_options, CoreWfStarter, WorkerTestHelpers, NAMESPACE,
+    drain_pollers_and_shutdown, get_integ_server_options, CoreWfStarter, WorkerTestHelpers,
+    NAMESPACE,
 };
 use tokio::time::sleep;
 
@@ -51,29 +52,39 @@ async fn client_list_open_closed_workflow_executions() {
 
     // Complete workflow
     core.complete_execution(&task.run_id).await;
-    core.shutdown().await;
+    drain_pollers_and_shutdown(&core).await;
 
-    // List above CLOSED workflow
-    let start_time_filter = StartTimeFilter {
-        earliest_time: Some(earliest).and_then(|t| t.try_into().ok()),
-        latest_time: Some(latest).and_then(|t| t.try_into().ok()),
-    };
-    let filter = ListClosedFilters::ExecutionFilter(WorkflowExecutionFilter {
-        workflow_id: wf_name.clone(),
-        run_id,
-    });
-    let closed_workflows = client
-        .list_closed_workflow_executions(
-            1,
-            Default::default(),
-            Some(start_time_filter),
-            Some(filter),
-        )
-        .await
-        .unwrap();
-    assert_eq!(closed_workflows.executions.len(), 1);
-    let workflow = closed_workflows.executions[0].clone();
-    assert_eq!(workflow.execution.as_ref().unwrap().workflow_id, wf_name);
+    // List above CLOSED workflow. Visibility doesn't always update immediately so we give this a
+    // few tries.
+    let mut passed = false;
+    for _ in 1..=5 {
+        let closed_workflows = client
+            .list_closed_workflow_executions(
+                1,
+                Default::default(),
+                Some(StartTimeFilter {
+                    earliest_time: Some(earliest).and_then(|t| t.try_into().ok()),
+                    latest_time: Some(latest).and_then(|t| t.try_into().ok()),
+                }),
+                Some(ListClosedFilters::ExecutionFilter(
+                    WorkflowExecutionFilter {
+                        workflow_id: wf_name.clone(),
+                        run_id: run_id.clone(),
+                    },
+                )),
+            )
+            .await
+            .unwrap();
+        if closed_workflows.executions.len() == 1 {
+            let workflow = &closed_workflows.executions[0];
+            if workflow.execution.as_ref().unwrap().workflow_id == wf_name {
+                passed = true;
+                break;
+            }
+        }
+        sleep(Duration::from_millis(100)).await;
+    }
+    assert!(passed);
 }
 
 #[tokio::test]