Clean up temp fix / other minor cleanup

Sushisource · Sushisource · commit 038ed5b5efdc · 2026-04-15T10:26:10.000-07:00
diff --git a/crates/sdk-core/src/core_tests/workers.rs b/crates/sdk-core/src/core_tests/workers.rs
@@ -1322,42 +1322,3 @@ async fn graceful_shutdown_sends_shutdown_worker_rpc_during_initiate() {
 
     worker.finalize_shutdown().await;
 }
-
-/// The all_permits_tracker timeout must accommodate the 5s TEMP_FIX graceful poll timeout.
-/// When a poll holds a permit during the graceful shutdown window, all_permits_tracker must
-/// not fire before the permit is released. This test acquires a permit, holds it for 3s
-/// (simulating a poll blocked during graceful shutdown), and verifies shutdown() completes
-/// without dbg_panic!.
-#[tokio::test]
-async fn all_permits_tracker_timeout_accommodates_graceful_poll_delay() {
-    use crate::{abstractions::tests::fixed_size_permit_dealer, worker::WorkflowSlotKind};
-
-    let dealer = fixed_size_permit_dealer::<WorkflowSlotKind>(5);
-    let rcv = dealer.get_extant_count_rcv();
-
-    // Acquire a permit (simulating a LongPollBuffer poll task holding one)
-    let permit = dealer.acquire_owned().await;
-
-    // Start the same wait logic as shutdown()'s all_permits_tracker
-    // This must match the timeout in Worker::shutdown()'s all_permits_tracker select!
-    let wait_result = tokio::time::timeout(Duration::from_secs(6), async {
-        let mut rcv = rcv;
-        let _ = rcv.wait_for(|x| *x == 0).await;
-    });
-
-    // Release the permit after 3s (simulating the graceful poll timeout releasing it)
-    let release_task = tokio::spawn(async move {
-        tokio::time::sleep(Duration::from_secs(3)).await;
-        drop(permit);
-    });
-
-    let result = wait_result.await;
-    release_task.await.unwrap();
-
-    // With the 6s timeout, this should succeed (permit released at 3s < 6s).
-    // With the old 1s timeout, this would fail (permit released at 3s > 1s).
-    assert!(
-        result.is_ok(),
-        "all_permits_tracker should complete within 6s when permit is held for 3s"
-    );
-}
diff --git a/crates/sdk-core/src/pollers/poll_buffer.rs b/crates/sdk-core/src/pollers/poll_buffer.rs
@@ -367,8 +367,6 @@ where
                         };
                     let graceful_shutdown = graceful_shutdown.clone();
                     let poll_task = tokio::spawn(async move {
-                        let shutdown_clone = shutdown.clone();
-
                         let r = if graceful_shutdown.load(Ordering::Relaxed) {
                             pf(timeout_override).await
                         } else {
@@ -389,10 +387,11 @@ where
                         }
                         let (should_forward, backoff_duration) = report_handle.poll_result(&r);
                         if let Some(duration) = backoff_duration {
-                            // Apply backoff BEFORE dropping active_guard to prevent next poll from starting
+                            // Apply backoff BEFORE dropping active_guard to prevent next poll from
+                            // starting
                             tokio::select! {
                                 _ = tokio::time::sleep(duration) => return,
-                                _ = shutdown_clone.cancelled() => (),
+                                _ = shutdown.cancelled() => (),
                             };
                         }
                         drop(active_guard);
diff --git a/crates/sdk-core/src/worker/mod.rs b/crates/sdk-core/src/worker/mod.rs
@@ -978,8 +978,7 @@ impl Worker {
         // Wait for all permits to be released, but don't totally hang real-world shutdown.
         tokio::select! {
             _ = async { self.all_permits_tracker.lock().await.all_done().await } => {},
-            // TEMP: Shutdown can take 5s longer due to TEMP_FIX in poll_buffer.rs
-            _ = tokio::time::sleep(Duration::from_secs(6)) => {
+            _ = tokio::time::sleep(Duration::from_secs(1)) => {
                 dbg_panic!("Waiting for all slot permits to release took too long!");
             }
         }
@@ -1417,7 +1416,6 @@ impl Worker {
             .as_ref()
             .map(|hm| hm.heartbeat_callback.clone()());
         let handle = tokio::spawn(async move {
-            let start_time = std::time::Instant::now();
             match client
                 .shutdown_worker(sticky_name, task_queue, task_queue_types, heartbeat)
                 .await
@@ -1435,10 +1433,6 @@ impl Worker {
                 }
                 _ => {}
             }
-            warn!(
-                "shutdown_worker rpc completed, took {:?}",
-                start_time.elapsed()
-            )
         });
         *guard = Some(handle);
     }
diff --git a/crates/sdk-core/src/worker/workflow/wft_poller.rs b/crates/sdk-core/src/worker/workflow/wft_poller.rs
@@ -216,10 +216,6 @@ fn new_wft_poller(
                 return match poller.poll().await {
                     Some(Ok((wft, permit))) => {
                         if wft == PollWorkflowTaskQueueResponse::default() {
-                            if shutdown_token.is_cancelled() {
-                                poller.shutdown_box().await;
-                                return None;
-                            }
                             // We get the default proto in the event that the long poll times out.
                             debug!("Poll wft timeout");
                             metrics.wf_tq_poll_empty();
@@ -297,50 +293,24 @@ mod tests {
         assert_matches!(stream.next().await, None);
     }
 
-    /// empty responses after shutdown are retried indefinitely in
-    /// new_wft_poller, causing the WFT stream to spin. This is what caused the workflow_load
-    /// heavy test to hang on CI when the server has enableCancelWorkerPollsOnShutdown.
+    /// When the underlying poller returns None (indicating shutdown), the wrapping WFT stream
+    /// should also return None to terminate.
     #[tokio::test]
-    async fn empty_response_after_shutdown_terminates_wft_stream() {
-        use std::sync::atomic::{AtomicUsize, Ordering};
-
-        let poll_count = Arc::new(AtomicUsize::new(0));
-        let poll_count_clone = poll_count.clone();
-
+    async fn poller_returning_none_terminates_wft_stream() {
         let mut mock_poller = mock_poller();
-        mock_poller.expect_poll().returning(move || {
-            poll_count_clone.fetch_add(1, Ordering::SeqCst);
-            Some(Ok(PollWorkflowTaskQueueResponse::default()))
-        });
-        mock_poller.expect_shutdown().returning(|| ());
+        mock_poller.expect_poll().times(1).returning(|| None);
+        mock_poller.expect_shutdown().times(1).returning(|| ());
 
         let sem = Arc::new(fixed_size_permit_dealer::<WorkflowSlotKind>(10));
-        let shutdown_token = CancellationToken::new();
 
         let stream = new_wft_poller(
             Box::new(MockPermittedPollBuffer::new(sem, mock_poller)),
             MetricsContext::no_op(),
-            shutdown_token.clone(),
+            CancellationToken::new(),
         );
         pin_mut!(stream);
 
-        shutdown_token.cancel();
-
-        let result = tokio::time::timeout(std::time::Duration::from_secs(2), stream.next()).await;
-        assert!(
-            result.is_ok(),
-            "WFT stream should terminate promptly after shutdown, not spin on retries"
-        );
-        assert!(
-            result.unwrap().is_none(),
-            "WFT stream should return None on empty response after shutdown"
-        );
-
-        let total = poll_count.load(Ordering::SeqCst);
-        assert!(
-            total < 5,
-            "Expected WFT stream to terminate quickly, but poller was called {total} times"
-        );
+        assert_matches!(stream.next().await, None);
     }
 
     #[tokio::test]
diff --git a/crates/sdk-core/tests/integ_tests/worker_tests.rs b/crates/sdk-core/tests/integ_tests/worker_tests.rs
@@ -980,163 +980,6 @@ async fn shutdown_worker_not_retried() {
     assert_eq!(shutdown_call_count.load(Ordering::Relaxed), 1);
 }
 
-/// Reproduces the server-side race in temporalio/temporal#9545 where a poll completes naturally
-/// (e.g. long-poll timeout) right as shutdown begins, causing the poller to re-poll AFTER the
-/// ShutdownWorker RPC has already been processed. The server never flushes this second poll,
-/// so without the 5s TEMP_FIX timeout the worker would hang for 60s.
-///
-/// Sequence:
-///   1. Worker starts polling with graceful_poll_shutdown enabled
-///   2. initiate_shutdown() fires → shutdown token cancelled, ShutdownWorker RPC spawned
-///   3. First poll returns empty (natural timeout) — but this races with ShutdownWorker
-///   4. Poller re-polls (graceful mode: no select! against shutdown token)
-///   5. ShutdownWorker RPC completes — server flushes nothing (first poll already returned)
-///   6. Second poll hangs forever — server doesn't know about it
-///   7. TEMP_FIX: after 5s the graceful_interruptor cancels the hanging poll
-///
-/// This test verifies shutdown completes within 10s (not 60s), proving the temp fix works.
-/// Remove this test when temporalio/temporal#9545 is fully deployed.
-#[tokio::test]
-async fn graceful_shutdown_race_temp_fix_prevents_60s_hang() {
-    use prost::Message;
-    use std::sync::atomic::AtomicUsize;
-    use temporalio_common::protos::temporal::api::{
-        namespace::v1::{NamespaceInfo, namespace_info::Capabilities},
-        workflowservice::v1::DescribeNamespaceResponse,
-    };
-    use tokio::sync::Notify;
-
-    fn grpc_ok_empty() -> tonic::codegen::http::Response<tonic::body::Body> {
-        tonic::codegen::http::Response::builder()
-            .header("content-type", "application/grpc")
-            .header("grpc-status", "0")
-            .body(tonic::body::Body::empty())
-            .unwrap()
-    }
-    fn grpc_ok_proto(msg: &impl Message) -> tonic::codegen::http::Response<tonic::body::Body> {
-        let encoded = msg.encode_to_vec();
-        let mut buf = Vec::with_capacity(5 + encoded.len());
-        buf.push(0);
-        buf.extend_from_slice(&(encoded.len() as u32).to_be_bytes());
-        buf.extend_from_slice(&encoded);
-        tonic::codegen::http::Response::builder()
-            .header("content-type", "application/grpc")
-            .header("grpc-status", "0")
-            .body(tonic::body::Body::new(http_body_util::Full::new(
-                bytes::Bytes::from(buf),
-            )))
-            .unwrap()
-    }
-
-    // Track poll count to distinguish first poll (returns empty) from re-polls (hang forever)
-    let poll_count = Arc::new(AtomicUsize::new(0));
-    let poll_count_clone = poll_count.clone();
-    // Signal from initiate_shutdown (via shutdown_worker RPC) to release the first poll
-    let shutdown_signal = Arc::new(Notify::new());
-    let shutdown_signal_for_rpc = shutdown_signal.clone();
-    let shutdown_signal_for_poll = shutdown_signal.clone();
-
-    let fs = fake_server(move |req| {
-        let uri = req.uri().to_string();
-        let poll_count = poll_count_clone.clone();
-        let shutdown_signal_for_poll = shutdown_signal_for_poll.clone();
-        let shutdown_signal_for_rpc = shutdown_signal_for_rpc.clone();
-
-        if uri.contains("DescribeNamespace") {
-            let resp = DescribeNamespaceResponse {
-                namespace_info: Some(NamespaceInfo {
-                    capabilities: Some(Capabilities {
-                        worker_poll_complete_on_shutdown: true,
-                        ..Capabilities::default()
-                    }),
-                    ..NamespaceInfo::default()
-                }),
-                ..DescribeNamespaceResponse::default()
-            };
-            async move { grpc_ok_proto(&resp) }.boxed()
-        } else if uri.contains("Poll") {
-            async move {
-                let n = poll_count.fetch_add(1, Ordering::SeqCst);
-                if n == 0 {
-                    // First poll: wait for shutdown to start, then return empty.
-                    // This simulates the poll timing out naturally right as shutdown begins.
-                    shutdown_signal_for_poll.notified().await;
-                    grpc_ok_empty()
-                } else {
-                    // Re-poll after shutdown: hang forever. This is the race —
-                    // the server already processed ShutdownWorker and won't flush this poll.
-                    futures_util::future::pending().await
-                }
-            }
-            .boxed()
-        } else if uri.contains("ShutdownWorker") {
-            // ShutdownWorker arrives — signal the first poll to return (simulating the race
-            // where poll returns right as/after ShutdownWorker is processed).
-            async move {
-                shutdown_signal_for_rpc.notify_waiters();
-                grpc_ok_empty()
-            }
-            .boxed()
-        } else {
-            async { grpc_ok_empty() }.boxed()
-        }
-    })
-    .await;
-
-    let mut opts = get_integ_server_options();
-    opts.target = format!("http://localhost:{}", fs.addr.port())
-        .parse::<url::Url>()
-        .unwrap();
-    opts.set_skip_get_system_info(true);
-    let connection = Connection::connect(opts).await.unwrap();
-    let client_opts = temporalio_client::ClientOptions::new("ns").build();
-    let client = temporalio_client::Client::new(connection, client_opts).unwrap();
-
-    let wf_type = "graceful_shutdown_race";
-    let mut starter = CoreWfStarter::new_with_overrides(wf_type, None, Some(client));
-    let worker = starter.get_worker().await;
-
-    // Enable graceful poll shutdown via validate()
-    worker.validate().await.unwrap();
-
-    // Start polling BEFORE initiating shutdown so the poll is in-flight.
-    // poll_workflow_activation triggers LongPollBuffer to start, which spawns a poll task
-    // that hits the fake server and blocks on the first poll (waiting for shutdown_signal).
-    let poll_handle = tokio::spawn({
-        let w = worker.clone();
-        async move {
-            // This will block until shutdown causes PollError::ShutDown
-            let _ = w.poll_workflow_activation().await;
-        }
-    });
-    let act_handle = tokio::spawn({
-        let w = worker.clone();
-        async move {
-            let _ = w.poll_activity_task().await;
-        }
-    });
-
-    // Give polls time to reach the fake server before initiating shutdown
-    tokio::time::sleep(Duration::from_millis(500)).await;
-
-    // Shutdown should complete within 10s (5s TEMP_FIX + margin).
-    // Without the temp fix, the re-poll hangs for 60s.
-    let result = tokio::time::timeout(Duration::from_secs(10), async {
-        worker.shutdown().await;
-    })
-    .await;
-
-    assert!(
-        result.is_ok(),
-        "Shutdown should complete within 10s. If it hangs, the TEMP_FIX graceful poll \
-         timeout is not working and the server race (temporal#9545) caused a 60s hang."
-    );
-
-    let _ = poll_handle.await;
-    let _ = act_handle.await;
-    fs.shutdown().await;
-}
-
 #[test]
 fn test_default_build_id() {
     let o = WorkerOptions::new("task_queue").build();
@@ -1153,7 +996,7 @@ impl ShutdownTimerActivityLoopWf {
     #[run]
     async fn run(ctx: &mut WorkflowContext<Self>) -> WorkflowResult<()> {
         loop {
-            ctx.timer(Duration::from_millis(100)).await;
+            ctx.timer(Duration::from_millis(10)).await;
             ctx.start_activity(
                 StdActivities::no_op,
                 (),

Original file line number	Diff line number	Diff line change
`@@ -978,8 +978,7 @@ impl Worker {`
`978`	`978`	`// Wait for all permits to be released, but don't totally hang real-world shutdown.`
`979`	`979`	`tokio::select! {`
`980`	`980`	`_ = async { self.all_permits_tracker.lock().await.all_done().await } => {},`
`981`		`- // TEMP: Shutdown can take 5s longer due to TEMP_FIX in poll_buffer.rs`
`982`		`- _ = tokio::time::sleep(Duration::from_secs(6)) => {`
	`981`	`+ _ = tokio::time::sleep(Duration::from_secs(1)) => {`
`983`	`982`	`dbg_panic!("Waiting for all slot permits to release took too long!");`
`984`	`983`	`}`
`985`	`984`	`}`
`@@ -1417,7 +1416,6 @@ impl Worker {`
`1417`	`1416`	`.as_ref()`
`1418`	`1417`	`.map(\|hm\| hm.heartbeat_callback.clone()());`
`1419`	`1418`	`let handle = tokio::spawn(async move {`
`1420`		`- let start_time = std::time::Instant::now();`
`1421`	`1419`	`match client`
`1422`	`1420`	`.shutdown_worker(sticky_name, task_queue, task_queue_types, heartbeat)`
`1423`	`1421`	`.await`
`@@ -1435,10 +1433,6 @@ impl Worker {`
`1435`	`1433`	`}`
`1436`	`1434`	`_ => {}`
`1437`	`1435`	`}`
`1438`		`- warn!(`
`1439`		`- "shutdown_worker rpc completed, took {:?}",`
`1440`		`- start_time.elapsed()`
`1441`		`- )`
`1442`	`1436`	`});`
`1443`	`1437`	`*guard = Some(handle);`
`1444`	`1438`	`}`