temporalio 0.0.2 → 0.1.0

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

@@ -1,23 +1,21 @@
 use crate::telemetry::default_resource;
 use hyper::{
     header::CONTENT_TYPE,
+    server::conn::AddrIncoming,
     service::{make_service_fn, service_fn},
     Body, Method, Request, Response, Server,
 };
-use opentelemetry::{
-    metrics::MetricsError,
-    sdk::{
-        export::metrics::{aggregation::TemporalitySelector, AggregatorSelector},
-        metrics::{controllers, processors},
-    },
+use opentelemetry::sdk::{
+    export::metrics::{aggregation::TemporalitySelector, AggregatorSelector},
+    metrics::{controllers, processors},
 };
 use opentelemetry_prometheus::{ExporterBuilder, PrometheusExporter};
 use prometheus::{Encoder, TextEncoder};
-use std::{convert::Infallible, net::SocketAddr, sync::Arc};
+use std::{convert::Infallible, net::SocketAddr, sync::Arc, time::Duration};
 
 /// Exposes prometheus metrics for scraping
 pub(super) struct PromServer {
-    addr: SocketAddr,
+    bound_addr: AddrIncoming,
     pub exporter: Arc<PrometheusExporter>,
 }
 
@@ -26,19 +24,22 @@ impl PromServer {
         addr: SocketAddr,
         aggregation: impl AggregatorSelector + Send + Sync + 'static,
         temporality: impl TemporalitySelector + Send + Sync + 'static,
-    ) -> Result<Self, MetricsError> {
+    ) -> Result<Self, anyhow::Error> {
         let controller =
             controllers::basic(processors::factory(aggregation, temporality).with_memory(true))
+                // Because Prom is pull-based, make this always refresh
+                .with_collect_period(Duration::from_secs(0))
                 .with_resource(default_resource())
                 .build();
         let exporter = ExporterBuilder::new(controller).try_init()?;
+        let bound_addr = AddrIncoming::bind(&addr)?;
         Ok(Self {
             exporter: Arc::new(exporter),
-            addr,
+            bound_addr,
         })
     }
 
-    pub async fn run(&self) -> hyper::Result<()> {
+    pub async fn run(self) -> hyper::Result<()> {
         // Spin up hyper server to serve metrics for scraping. We use hyper since we already depend
         // on it via Tonic.
         let expclone = self.exporter.clone();
@@ -46,9 +47,13 @@ impl PromServer {
             let expclone = expclone.clone();
             async move { Ok::<_, Infallible>(service_fn(move |req| metrics_req(req, expclone.clone()))) }
         });
-        let server = Server::bind(&self.addr).serve(svc);
+        let server = Server::builder(self.bound_addr).serve(svc);
         server.await
     }
+
+    pub fn bound_addr(&self) -> SocketAddr {
+        self.bound_addr.local_addr()
+    }
 }
 
 /// Serves prometheus metrics in the expected format for scraping

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

@@ -14,6 +14,7 @@ use crate::{
     },
     TaskToken, Worker, WorkerConfig, WorkerConfigBuilder,
 };
+use async_trait::async_trait;
 use bimap::BiMap;
 use futures::{future::BoxFuture, stream, stream::BoxStream, FutureExt, Stream, StreamExt};
 use mockall::TimesRange;
@@ -29,6 +30,7 @@ use std::{
     task::{Context, Poll},
     time::Duration,
 };
+use temporal_sdk_core_api::errors::{PollActivityError, PollWfError};
 use temporal_sdk_core_api::Worker as WorkerTrait;
 use temporal_sdk_core_protos::{
     coresdk::{
@@ -52,7 +54,6 @@ use tokio_stream::wrappers::UnboundedReceiverStream;
 use tokio_util::sync::CancellationToken;
 
 pub const TEST_Q: &str = "q";
-pub static NO_MORE_WORK_ERROR_MSG: &str = "No more work to do";
 
 pub fn test_worker_cfg() -> WorkerConfigBuilder {
     let mut wcb = WorkerConfigBuilder::default();
@@ -148,6 +149,7 @@ pub(crate) fn mock_worker(mocks: MocksHolder) -> Worker {
         mocks.inputs.wft_stream,
         act_poller,
         MetricsContext::no_op(),
+        None,
         CancellationToken::new(),
     )
 }
@@ -301,7 +303,7 @@ where
             }
             .boxed()
         } else {
-            async { Some(Err(tonic::Status::cancelled(NO_MORE_WORK_ERROR_MSG))) }.boxed()
+            async { None }.boxed()
         }
     });
     Box::new(mock_poller) as BoxedPoller<T>
@@ -375,6 +377,7 @@ pub(crate) struct MockPollCfg {
     pub expect_fail_wft_matcher:
         Box<dyn Fn(&TaskToken, &WorkflowTaskFailedCause, &Option<Failure>) -> bool + Send>,
     pub completion_asserts: Option<Box<dyn Fn(&WorkflowTaskCompletion) + Send>>,
+    pub num_expected_completions: Option<TimesRange>,
     /// If being used with the Rust SDK, this is set true. It ensures pollers will not error out
     /// early with no work, since we cannot know the exact number of times polling will happen.
     /// Instead, they will just block forever.
@@ -396,6 +399,7 @@ impl MockPollCfg {
             mock_client: mock_workflow_client(),
             expect_fail_wft_matcher: Box::new(|_, _, _| true),
             completion_asserts: None,
+            num_expected_completions: None,
             using_rust_sdk: false,
             make_poll_stream_interminable: false,
         }
@@ -418,6 +422,7 @@ impl MockPollCfg {
             mock_client,
             expect_fail_wft_matcher: Box::new(|_, _, _| true),
             completion_asserts: None,
+            num_expected_completions: None,
             using_rust_sdk: false,
             make_poll_stream_interminable: false,
         }
@@ -587,16 +592,20 @@ pub(crate) fn build_mock_pollers(mut cfg: MockPollCfg) -> MocksHolder {
     );
 
     let outstanding = outstanding_wf_task_tokens.clone();
-    cfg.mock_client
-        .expect_complete_workflow_task()
-        .returning(move |comp| {
-            if let Some(ass) = cfg.completion_asserts.as_ref() {
-                // tee hee
-                ass(&comp)
-            }
-            outstanding.release_token(&comp.task_token);
-            Ok(RespondWorkflowTaskCompletedResponse::default())
-        });
+    let expect_completes = cfg.mock_client.expect_complete_workflow_task();
+    if let Some(range) = cfg.num_expected_completions {
+        expect_completes.times(range);
+    } else if cfg.completion_asserts.is_some() {
+        expect_completes.times(1..);
+    }
+    expect_completes.returning(move |comp| {
+        if let Some(ass) = cfg.completion_asserts.as_ref() {
+            // tee hee
+            ass(&comp)
+        }
+        outstanding.release_token(&comp.task_token);
+        Ok(RespondWorkflowTaskCompletedResponse::default())
+    });
     let outstanding = outstanding_wf_task_tokens.clone();
     cfg.mock_client
         .expect_fail_workflow_task()
@@ -842,6 +851,7 @@ pub(crate) fn gen_assert_and_fail(asserter: &dyn Fn(&WorkflowActivation)) -> Ass
                 message: "Intentional test failure".to_string(),
                 ..Default::default()
             }),
+            ..Default::default()
         }
         .into(),
     )
@@ -887,3 +897,43 @@ macro_rules! prost_dur {
             .expect("test duration fits")
     };
 }
+
+#[async_trait]
+pub(crate) trait WorkerExt {
+    /// Initiate shutdown, drain the pollers, and wait for shutdown to complete.
+    async fn drain_pollers_and_shutdown(self);
+    /// Initiate shutdown, drain the *activity* poller, and wait for shutdown to complete.
+    /// Takes a ref because of that one test that needs it.
+    async fn drain_activity_poller_and_shutdown(&self);
+}
+
+#[async_trait]
+impl WorkerExt for Worker {
+    async fn drain_pollers_and_shutdown(self) {
+        self.initiate_shutdown();
+        tokio::join!(
+            async {
+                assert_matches!(
+                    self.poll_activity_task().await.unwrap_err(),
+                    PollActivityError::ShutDown
+                );
+            },
+            async {
+                assert_matches!(
+                    self.poll_workflow_activation().await.unwrap_err(),
+                    PollWfError::ShutDown
+                );
+            }
+        );
+        self.finalize_shutdown().await;
+    }
+
+    async fn drain_activity_poller_and_shutdown(&self) {
+        self.initiate_shutdown();
+        assert_matches!(
+            self.poll_activity_task().await.unwrap_err(),
+            PollActivityError::ShutDown
+        );
+        self.shutdown().await;
+    }
+}

data/bridge/sdk-core/core/src/worker/activities/activity_heartbeat_manager.rs

@@ -26,9 +26,6 @@ use tokio_util::sync::CancellationToken;
 /// Used to supply new heartbeat events to the activity heartbeat manager, or to send a shutdown
 /// request.
 pub(crate) struct ActivityHeartbeatManager {
-    /// Cancellations that have been received when heartbeating are queued here and can be consumed
-    /// by [fetch_cancellations]
-    incoming_cancels: Mutex<UnboundedReceiver<PendingActivityCancel>>,
     shutdown_token: CancellationToken,
     /// Used during `shutdown` to await until all inflight requests are sent.
     join_handle: Mutex<Option<JoinHandle<()>>>,
@@ -74,15 +71,118 @@ pub enum ActivityHeartbeatError {
     /// to heartbeat.
     #[error("Unable to parse activity heartbeat timeout.")]
     InvalidHeartbeatTimeout,
-    /// Core is shutting down and thus new heartbeats are not accepted
-    #[error("New heartbeat requests are not accepted while shutting down")]
-    ShuttingDown,
 }
 
 /// Manages activity heartbeating for a worker. Allows sending new heartbeats or requesting and
 /// awaiting for the shutdown. When shutdown is requested, signal gets sent to all processors, which
 /// allows them to complete gracefully.
 impl ActivityHeartbeatManager {
+    /// Creates a new instance of an activity heartbeat manager and returns a handle to the user,
+    /// which allows to send new heartbeats and initiate the shutdown.
+    /// Returns the manager and a channel that buffers cancellation notifications to be sent to Lang.
+    pub(super) fn new(
+        client: Arc<dyn WorkerClient>,
+    ) -> (Self, UnboundedReceiver<PendingActivityCancel>) {
+        let (heartbeat_stream_state, heartbeat_tx_source, shutdown_token) =
+            HeartbeatStreamState::new();
+        let (cancels_tx, cancels_rx) = unbounded_channel();
+        let heartbeat_tx = heartbeat_tx_source.clone();
+
+        let join_handle = tokio::spawn(
+            // The stream of incoming heartbeats uses unfold to carry state across each item in the
+            // stream. The closure checks if, for any given activity, we should heartbeat or not
+            // depending on its delay and when we last issued a heartbeat for it.
+            futures::stream::unfold(heartbeat_stream_state, move |mut hb_states| {
+                async move {
+                    let hb = tokio::select! {
+                        biased;
+
+                        _ = hb_states.cancellation_token.cancelled() => {
+                            return None
+                        }
+                        hb = hb_states.incoming_hbs.recv() => match hb {
+                            None => return None,
+                            Some(hb) => hb,
+                        }
+                    };
+
+                    Some((
+                        match hb {
+                            HeartbeatAction::SendHeartbeat(hb) => hb_states.record(hb),
+                            HeartbeatAction::CompleteReport(tt) => hb_states.handle_report_completed(tt),
+                            HeartbeatAction::CompleteThrottle(tt) => hb_states.handle_throttle_completed(tt),
+                            HeartbeatAction::Evict{ token, on_complete } => hb_states.evict(token, on_complete),
+                        },
+                        hb_states,
+                    ))
+                }
+            })
+                // Filters out `None`s
+                .filter_map(|opt| async { opt })
+                .for_each_concurrent(None, move |action| {
+                    let heartbeat_tx = heartbeat_tx_source.clone();
+                    let sg = client.clone();
+                    let cancels_tx = cancels_tx.clone();
+                    async move {
+                        match action {
+                            HeartbeatExecutorAction::Sleep(tt, duration, cancellation_token) => {
+                                tokio::select! {
+                                _ = cancellation_token.cancelled() => (),
+                                _ = tokio::time::sleep(duration) => {
+                                    let _ = heartbeat_tx.send(HeartbeatAction::CompleteThrottle(tt));
+                                },
+                            };
+                            }
+                            HeartbeatExecutorAction::Report { task_token: tt, details } => {
+                                match sg
+                                    .record_activity_heartbeat(tt.clone(), details.into_payloads())
+                                    .await
+                                {
+                                    Ok(RecordActivityTaskHeartbeatResponse { cancel_requested }) => {
+                                        if cancel_requested {
+                                            cancels_tx
+                                                .send(PendingActivityCancel::new(
+                                                    tt.clone(),
+                                                    ActivityCancelReason::Cancelled,
+                                                ))
+                                                .expect(
+                                                    "Receive half of heartbeat cancels not blocked",
+                                                );
+                                        }
+                                    }
+                                    // Send cancels for any activity that learns its workflow already
+                                    // finished (which is one thing not found implies - other reasons
+                                    // would seem equally valid).
+                                    Err(s) if s.code() == tonic::Code::NotFound => {
+                                        debug!(task_token = %tt,
+                                           "Activity not found when recording heartbeat");
+                                        cancels_tx
+                                            .send(PendingActivityCancel::new(
+                                                tt.clone(),
+                                                ActivityCancelReason::NotFound,
+                                            ))
+                                            .expect("Receive half of heartbeat cancels not blocked");
+                                    }
+                                    Err(e) => {
+                                        warn!("Error when recording heartbeat: {:?}", e);
+                                    }
+                                };
+                                let _ = heartbeat_tx.send(HeartbeatAction::CompleteReport(tt));
+                            }
+                        }
+                    }
+                }),
+        );
+
+        (
+            Self {
+                join_handle: Mutex::new(Some(join_handle)),
+                shutdown_token,
+                heartbeat_tx,
+            },
+            cancels_rx,
+        )
+    }
     /// Records a new heartbeat, the first call will result in an immediate call to the server,
     /// while rapid successive calls would accumulate for up to `delay` and then latest heartbeat
     /// details will be sent to the server.
@@ -95,9 +195,6 @@ impl ActivityHeartbeatManager {
         hb: ActivityHeartbeat,
         throttle_interval: Duration,
     ) -> Result<(), ActivityHeartbeatError> {
-        if self.shutdown_token.is_cancelled() {
-            return Err(ActivityHeartbeatError::ShuttingDown);
-        }
         self.heartbeat_tx
             .send(HeartbeatAction::SendHeartbeat(ValidActivityHeartbeat {
                 task_token: TaskToken(hb.task_token),
@@ -121,13 +218,6 @@ impl ActivityHeartbeatManager {
         completed.notified().await;
     }
 
-    /// Returns a future that resolves any time there is a new activity cancel that must be
-    /// dispatched to lang
-    pub(super) async fn next_pending_cancel(&self) -> Option<PendingActivityCancel> {
-        self.incoming_cancels.lock().await.recv().await
-    }
-
-    // TODO: Can own self now!
     /// Initiates shutdown procedure by stopping lifecycle loop and awaiting for all in-flight
     /// heartbeat requests to be flushed to the server.
     pub(super) async fn shutdown(&self) {
@@ -301,110 +391,6 @@ impl HeartbeatStreamState {
     }
 }
 
-impl ActivityHeartbeatManager {
-    /// Creates a new instance of an activity heartbeat manager and returns a handle to the user,
-    /// which allows to send new heartbeats and initiate the shutdown.
-    pub fn new(client: Arc<dyn WorkerClient>) -> Self {
-        let (heartbeat_stream_state, heartbeat_tx_source, shutdown_token) =
-            HeartbeatStreamState::new();
-        let (cancels_tx, cancels_rx) = unbounded_channel();
-        let heartbeat_tx = heartbeat_tx_source.clone();
-
-        let join_handle = tokio::spawn(
-            // The stream of incoming heartbeats uses unfold to carry state across each item in the
-            // stream. The closure checks if, for any given activity, we should heartbeat or not
-            // depending on its delay and when we last issued a heartbeat for it.
-            futures::stream::unfold(heartbeat_stream_state, move |mut hb_states| {
-                async move {
-                    let hb = tokio::select! {
-                        biased;
-
-                        _ = hb_states.cancellation_token.cancelled() => {
-                            return None
-                        }
-                        hb = hb_states.incoming_hbs.recv() => match hb {
-                            None => return None,
-                            Some(hb) => hb,
-                        }
-                    };
-
-                    Some((
-                        match hb {
-                            HeartbeatAction::SendHeartbeat(hb) => hb_states.record(hb),
-                            HeartbeatAction::CompleteReport(tt) => hb_states.handle_report_completed(tt),
-                            HeartbeatAction::CompleteThrottle(tt) => hb_states.handle_throttle_completed(tt),
-                            HeartbeatAction::Evict{ token, on_complete } => hb_states.evict(token, on_complete),
-                        },
-                        hb_states,
-                    ))
-                }
-            })
-            // Filters out `None`s
-            .filter_map(|opt| async { opt })
-            .for_each_concurrent(None, move |action| {
-                let heartbeat_tx = heartbeat_tx_source.clone();
-                let sg = client.clone();
-                let cancels_tx = cancels_tx.clone();
-                async move {
-                    match action {
-                        HeartbeatExecutorAction::Sleep(tt, duration, cancellation_token) => {
-                            tokio::select! {
-                                _ = cancellation_token.cancelled() => (),
-                                _ = tokio::time::sleep(duration) => {
-                                    let _ = heartbeat_tx.send(HeartbeatAction::CompleteThrottle(tt));
-                                },
-                            };
-                        }
-                        HeartbeatExecutorAction::Report { task_token: tt, details } => {
-                            match sg
-                                .record_activity_heartbeat(tt.clone(), details.into_payloads())
-                                .await
-                            {
-                                Ok(RecordActivityTaskHeartbeatResponse { cancel_requested }) => {
-                                    if cancel_requested {
-                                        cancels_tx
-                                            .send(PendingActivityCancel::new(
-                                                tt.clone(),
-                                                ActivityCancelReason::Cancelled,
-                                            ))
-                                            .expect(
-                                                "Receive half of heartbeat cancels not blocked",
-                                            );
-                                    }
-                                }
-                                // Send cancels for any activity that learns its workflow already
-                                // finished (which is one thing not found implies - other reasons
-                                // would seem equally valid).
-                                Err(s) if s.code() == tonic::Code::NotFound => {
-                                    debug!(task_token = %tt,
-                                           "Activity not found when recording heartbeat");
-                                    cancels_tx
-                                        .send(PendingActivityCancel::new(
-                                            tt.clone(),
-                                            ActivityCancelReason::NotFound,
-                                        ))
-                                        .expect("Receive half of heartbeat cancels not blocked");
-                                }
-                                Err(e) => {
-                                    warn!("Error when recording heartbeat: {:?}", e);
-                                }
-                            };
-                            let _ = heartbeat_tx.send(HeartbeatAction::CompleteReport(tt));
-                        }
-                    }
-                }
-            }),
-        );
-
-        Self {
-            incoming_cancels: Mutex::new(cancels_rx),
-            join_handle: Mutex::new(Some(join_handle)),
-            shutdown_token,
-            heartbeat_tx,
-        }
-    }
-}
-
 #[cfg(test)]
 mod test {
     use super::*;
@@ -425,7 +411,7 @@ mod test {
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(2);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client));
+        let (hm, _) = ActivityHeartbeatManager::new(Arc::new(mock_client));
         let fake_task_token = vec![1, 2, 3];
         // Send 2 heartbeat requests for 20ms apart.
         // The first heartbeat should be sent right away, and
@@ -446,14 +432,13 @@ mod test {
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(3);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client));
+        let (hm, _) = ActivityHeartbeatManager::new(Arc::new(mock_client));
         let fake_task_token = vec![1, 2, 3];
         // Heartbeats always get sent if recorded less frequently than the throttle interval
         for i in 0_u8..3 {
             record_heartbeat(&hm, fake_task_token.clone(), i, Duration::from_millis(10));
             sleep(Duration::from_millis(20)).await;
         }
-        // sleep again to let heartbeats be flushed
         hm.shutdown().await;
     }
 
@@ -466,7 +451,7 @@ mod test {
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(1);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client));
+        let (hm, _) = ActivityHeartbeatManager::new(Arc::new(mock_client));
         let fake_task_token = vec![1, 2, 3];
         // Send a whole bunch of heartbeats very fast. We should still only send one total.
         for i in 0_u8..50 {
@@ -485,7 +470,7 @@ mod test {
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(2);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client));
+        let (hm, _) = ActivityHeartbeatManager::new(Arc::new(mock_client));
         let fake_task_token = vec![1, 2, 3];
         record_heartbeat(&hm, fake_task_token.clone(), 0, Duration::from_millis(100));
         sleep(Duration::from_millis(500)).await;
@@ -502,7 +487,7 @@ mod test {
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(2);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client));
+        let (hm, _) = ActivityHeartbeatManager::new(Arc::new(mock_client));
         let fake_task_token = vec![1, 2, 3];
         record_heartbeat(&hm, fake_task_token.clone(), 0, Duration::from_millis(100));
         // Let it propagate
@@ -522,42 +507,13 @@ mod test {
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(1);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client));
+        let (hm, _) = ActivityHeartbeatManager::new(Arc::new(mock_client));
         let fake_task_token = vec![1, 2, 3];
         record_heartbeat(&hm, fake_task_token.clone(), 0, Duration::from_millis(100));
         hm.evict(fake_task_token.clone().into()).await;
         hm.shutdown().await;
     }
 
-    /// Recording new heartbeats after shutdown is not allowed, and will result in error.
-    #[tokio::test]
-    async fn record_after_shutdown() {
-        let mut mock_client = mock_workflow_client();
-        mock_client
-            .expect_record_activity_heartbeat()
-            .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
-            .times(0);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client));
-        hm.shutdown().await;
-        match hm.record(
-            ActivityHeartbeat {
-                task_token: vec![1, 2, 3],
-                details: vec![Payload {
-                    // payload doesn't matter in this case, as it shouldn't get sent anyways.
-                    ..Default::default()
-                }],
-            },
-            Duration::from_millis(1000),
-        ) {
-            Ok(_) => {
-                unreachable!("heartbeat should not be recorded after the shutdown");
-            }
-            Err(e) => {
-                matches!(e, ActivityHeartbeatError::ShuttingDown);
-            }
-        }
-    }
-
     fn record_heartbeat(
         hm: &ActivityHeartbeatManager,
         task_token: Vec<u8>,

data/bridge/sdk-core/core/src/worker/activities/activity_task_poller_stream.rs

@@ -0,0 +1,89 @@
+use crate::abstractions::MeteredSemaphore;
+use crate::worker::activities::PermittedTqResp;
+use crate::{pollers::BoxedActPoller, MetricsContext};
+use futures::{stream, Stream};
+use governor::clock::DefaultClock;
+use governor::middleware::NoOpMiddleware;
+use governor::state::{InMemoryState, NotKeyed};
+use governor::RateLimiter;
+use std::sync::Arc;
+use temporal_sdk_core_protos::temporal::api::workflowservice::v1::PollActivityTaskQueueResponse;
+use tokio::select;
+use tokio_util::sync::CancellationToken;
+
+struct StreamState {
+    poller: BoxedActPoller,
+    semaphore: Arc<MeteredSemaphore>,
+    rate_limiter: Option<RateLimiter<NotKeyed, InMemoryState, DefaultClock, NoOpMiddleware>>,
+    metrics: MetricsContext,
+    shutdown_token: CancellationToken,
+    poller_was_shutdown: bool,
+}
+
+pub(crate) fn new_activity_task_poller(
+    poller: BoxedActPoller,
+    semaphore: Arc<MeteredSemaphore>,
+    rate_limiter: Option<RateLimiter<NotKeyed, InMemoryState, DefaultClock, NoOpMiddleware>>,
+    metrics: MetricsContext,
+    shutdown_token: CancellationToken,
+) -> impl Stream<Item = Result<PermittedTqResp, tonic::Status>> {
+    let state = StreamState {
+        poller,
+        semaphore,
+        rate_limiter,
+        metrics,
+        shutdown_token,
+        poller_was_shutdown: false,
+    };
+    stream::unfold(state, |mut state| async move {
+        loop {
+            let poll = async {
+                let permit = state
+                    .semaphore
+                    .acquire_owned()
+                    .await
+                    .expect("outstanding activity semaphore not closed");
+                if !state.poller_was_shutdown {
+                    if let Some(ref rl) = state.rate_limiter {
+                        rl.until_ready().await;
+                    }
+                }
+                loop {
+                    return match state.poller.poll().await {
+                        Some(Ok(resp)) => {
+                            if resp == PollActivityTaskQueueResponse::default() {
+                                // We get the default proto in the event that the long poll times out.
+                                debug!("Poll activity task timeout");
+                                state.metrics.act_poll_timeout();
+                                continue;
+                            }
+                            Some(Ok(PermittedTqResp { permit, resp }))
+                        }
+                        Some(Err(e)) => {
+                            warn!(error=?e, "Error while polling for activity tasks");
+                            Some(Err(e))
+                        }
+                        // If poller returns None, it's dead, thus we also return None to terminate this
+                        // stream.
+                        None => None,
+                    };
+                }
+            };
+            if state.poller_was_shutdown {
+                return poll.await.map(|res| (res, state));
+            }
+            select! {
+                biased;
+
+                _ = state.shutdown_token.cancelled() => {
+                    state.poller.notify_shutdown();
+                    state.poller_was_shutdown = true;
+                    continue;
+                }
+                res = poll => {
+                    return res.map(|res| (res, state));
+                }
+            }
+        }
+    })
+}