@temporalio/core-bridge 1.13.2 → 1.14.0

package/sdk-core/client/src/worker_registry/mod.rs

@@ -1,282 +0,0 @@
-//! This module enables the tracking of workers that are associated with a client instance.
-//! This is needed to implement Eager Workflow Start, a latency optimization in which the client,
-//!  after reserving a slot, directly forwards a WFT to a local worker.
-
-use parking_lot::RwLock;
-use slotmap::SlotMap;
-use std::collections::{HashMap, hash_map::Entry::Vacant};
-
-use temporal_sdk_core_protos::temporal::api::workflowservice::v1::PollWorkflowTaskQueueResponse;
-
-slotmap::new_key_type! {
-    /// Registration key for a worker
-    pub struct WorkerKey;
-}
-
-/// This trait is implemented by an object associated with a worker, which provides WFT processing slots.
-#[cfg_attr(test, mockall::automock)]
-pub trait SlotProvider: std::fmt::Debug {
-    /// The namespace for the WFTs that it can process.
-    fn namespace(&self) -> &str;
-    /// The task queue this provider listens to.
-    fn task_queue(&self) -> &str;
-    /// Try to reserve a slot on this worker.
-    fn try_reserve_wft_slot(&self) -> Option<Box<dyn Slot + Send>>;
-    /// Get the worker deployment options for this worker, if using deployment-based versioning.
-    fn deployment_options(&self) -> Option<temporal_sdk_core_api::worker::WorkerDeploymentOptions>;
-}
-
-/// This trait represents a slot reserved for processing a WFT by a worker.
-#[cfg_attr(test, mockall::automock)]
-pub trait Slot {
-    /// Consumes this slot by dispatching a WFT to its worker. This can only be called once.
-    fn schedule_wft(
-        self: Box<Self>,
-        task: PollWorkflowTaskQueueResponse,
-    ) -> Result<(), anyhow::Error>;
-}
-
-/// Result of reserving a workflow task slot, including deployment options if applicable.
-pub(crate) struct SlotReservation {
-    /// The reserved slot for processing the workflow task
-    pub slot: Box<dyn Slot + Send>,
-    /// Worker deployment options, if the worker is using deployment-based versioning
-    pub deployment_options: Option<temporal_sdk_core_api::worker::WorkerDeploymentOptions>,
-}
-
-#[derive(PartialEq, Eq, Hash, Debug, Clone)]
-struct SlotKey {
-    namespace: String,
-    task_queue: String,
-}
-
-impl SlotKey {
-    fn new(namespace: String, task_queue: String) -> SlotKey {
-        SlotKey {
-            namespace,
-            task_queue,
-        }
-    }
-}
-
-/// This is an inner class for [SlotManager] needed to hide the mutex.
-#[derive(Default, Debug)]
-struct SlotManagerImpl {
-    /// Maps keys, i.e., namespace#task_queue, to provider.
-    providers: HashMap<SlotKey, Box<dyn SlotProvider + Send + Sync>>,
-    /// Maps ids to keys in `providers`.
-    index: SlotMap<WorkerKey, SlotKey>,
-}
-
-impl SlotManagerImpl {
-    /// Factory method.
-    fn new() -> Self {
-        Self {
-            index: Default::default(),
-            providers: Default::default(),
-        }
-    }
-
-    fn try_reserve_wft_slot(
-        &self,
-        namespace: String,
-        task_queue: String,
-    ) -> Option<SlotReservation> {
-        let key = SlotKey::new(namespace, task_queue);
-        if let Some(p) = self.providers.get(&key)
-            && let Some(slot) = p.try_reserve_wft_slot()
-        {
-            let deployment_options = p.deployment_options();
-            return Some(SlotReservation {
-                slot,
-                deployment_options,
-            });
-        }
-        None
-    }
-
-    fn register(&mut self, provider: Box<dyn SlotProvider + Send + Sync>) -> Option<WorkerKey> {
-        let key = SlotKey::new(
-            provider.namespace().to_string(),
-            provider.task_queue().to_string(),
-        );
-        if let Vacant(p) = self.providers.entry(key.clone()) {
-            p.insert(provider);
-            Some(self.index.insert(key))
-        } else {
-            warn!("Ignoring registration for worker: {key:?}.");
-            None
-        }
-    }
-
-    fn unregister(&mut self, id: WorkerKey) -> Option<Box<dyn SlotProvider + Send + Sync>> {
-        if let Some(key) = self.index.remove(id) {
-            self.providers.remove(&key)
-        } else {
-            None
-        }
-    }
-
-    #[cfg(test)]
-    fn num_providers(&self) -> (usize, usize) {
-        (self.index.len(), self.providers.len())
-    }
-}
-
-/// Enables local workers to make themselves visible to a shared client instance.
-/// There can only be one worker registered per namespace+queue_name+client, others will get ignored.
-/// It also provides a convenient method to find compatible slots within the collection.
-#[derive(Default, Debug)]
-pub struct SlotManager {
-    manager: RwLock<SlotManagerImpl>,
-}
-
-impl SlotManager {
-    /// Factory method.
-    pub fn new() -> Self {
-        Self {
-            manager: RwLock::new(SlotManagerImpl::new()),
-        }
-    }
-
-    /// Try to reserve a compatible processing slot in any of the registered workers.
-    /// Returns the slot and the worker's deployment options (if using deployment-based versioning).
-    pub(crate) fn try_reserve_wft_slot(
-        &self,
-        namespace: String,
-        task_queue: String,
-    ) -> Option<SlotReservation> {
-        self.manager
-            .read()
-            .try_reserve_wft_slot(namespace, task_queue)
-    }
-
-    /// Register a local worker that can provide WFT processing slots.
-    pub fn register(&self, provider: Box<dyn SlotProvider + Send + Sync>) -> Option<WorkerKey> {
-        self.manager.write().register(provider)
-    }
-
-    /// Unregister a provider, typically when its worker starts shutdown.
-    pub fn unregister(&self, id: WorkerKey) -> Option<Box<dyn SlotProvider + Send + Sync>> {
-        self.manager.write().unregister(id)
-    }
-
-    #[cfg(test)]
-    /// Returns (num_providers, num_buckets), where a bucket key is namespace+task_queue.
-    /// There is only one provider per bucket so `num_providers` should be equal to `num_buckets`.
-    pub fn num_providers(&self) -> (usize, usize) {
-        self.manager.read().num_providers()
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    fn new_mock_slot(with_error: bool) -> Box<MockSlot> {
-        let mut mock_slot = MockSlot::new();
-        if with_error {
-            mock_slot
-                .expect_schedule_wft()
-                .returning(|_| Err(anyhow::anyhow!("Changed my mind")));
-        } else {
-            mock_slot.expect_schedule_wft().returning(|_| Ok(()));
-        }
-        Box::new(mock_slot)
-    }
-
-    fn new_mock_provider(
-        namespace: String,
-        task_queue: String,
-        with_error: bool,
-        no_slots: bool,
-    ) -> MockSlotProvider {
-        let mut mock_provider = MockSlotProvider::new();
-        mock_provider
-            .expect_try_reserve_wft_slot()
-            .returning(move || {
-                if no_slots {
-                    None
-                } else {
-                    Some(new_mock_slot(with_error))
-                }
-            });
-        mock_provider.expect_namespace().return_const(namespace);
-        mock_provider.expect_task_queue().return_const(task_queue);
-        mock_provider.expect_deployment_options().return_const(None);
-        mock_provider
-    }
-
-    #[test]
-    fn registry_respects_registration_order() {
-        let mock_provider1 =
-            new_mock_provider("foo".to_string(), "bar_q".to_string(), false, false);
-        let mock_provider2 = new_mock_provider("foo".to_string(), "bar_q".to_string(), false, true);
-
-        let manager = SlotManager::new();
-        let some_slots = manager.register(Box::new(mock_provider1));
-        let no_slots = manager.register(Box::new(mock_provider2));
-        assert!(no_slots.is_none());
-
-        let mut found = 0;
-        for _ in 0..10 {
-            if manager
-                .try_reserve_wft_slot("foo".to_string(), "bar_q".to_string())
-                .is_some()
-            {
-                found += 1;
-            }
-        }
-        assert_eq!(found, 10);
-        assert_eq!((1, 1), manager.num_providers());
-
-        manager.unregister(some_slots.unwrap());
-        assert_eq!((0, 0), manager.num_providers());
-
-        let mock_provider1 =
-            new_mock_provider("foo".to_string(), "bar_q".to_string(), false, false);
-        let mock_provider2 = new_mock_provider("foo".to_string(), "bar_q".to_string(), false, true);
-
-        let no_slots = manager.register(Box::new(mock_provider2));
-        let some_slots = manager.register(Box::new(mock_provider1));
-        assert!(some_slots.is_none());
-
-        let mut not_found = 0;
-        for _ in 0..10 {
-            if manager
-                .try_reserve_wft_slot("foo".to_string(), "bar_q".to_string())
-                .is_none()
-            {
-                not_found += 1;
-            }
-        }
-        assert_eq!(not_found, 10);
-        assert_eq!((1, 1), manager.num_providers());
-        manager.unregister(no_slots.unwrap());
-        assert_eq!((0, 0), manager.num_providers());
-    }
-
-    #[test]
-    fn registry_keeps_one_provider_per_namespace() {
-        let manager = SlotManager::new();
-        let mut worker_keys = vec![];
-        for i in 0..10 {
-            let namespace = format!("myId{}", i % 3);
-            let mock_provider = new_mock_provider(namespace, "bar_q".to_string(), false, false);
-            worker_keys.push(manager.register(Box::new(mock_provider)));
-        }
-        assert_eq!((3, 3), manager.num_providers());
-
-        let count = worker_keys
-            .iter()
-            .filter(|key| key.is_some())
-            .fold(0, |count, key| {
-                manager.unregister(key.unwrap());
-                // Should be idempotent
-                manager.unregister(key.unwrap());
-                count + 1
-            });
-        assert_eq!(3, count);
-        assert_eq!((0, 0), manager.num_providers());
-    }
-}

package/sdk-core/core/src/worker/heartbeat.rs

@@ -1,230 +0,0 @@
-use crate::{WorkerClient, abstractions::dbg_panic};
-use gethostname::gethostname;
-use parking_lot::Mutex;
-use prost_types::Duration as PbDuration;
-use std::{
-    sync::{Arc, OnceLock},
-    time::{Duration, SystemTime},
-};
-use temporal_sdk_core_api::worker::WorkerConfig;
-use temporal_sdk_core_protos::temporal::api::worker::v1::{WorkerHeartbeat, WorkerHostInfo};
-use tokio::{sync::Notify, task::JoinHandle, time::MissedTickBehavior};
-use uuid::Uuid;
-
-pub(crate) type HeartbeatFn = Box<dyn Fn() -> Option<WorkerHeartbeat> + Send + Sync>;
-
-pub(crate) struct WorkerHeartbeatManager {
-    heartbeat_handle: JoinHandle<()>,
-}
-
-impl WorkerHeartbeatManager {
-    pub(crate) fn new(
-        config: WorkerConfig,
-        identity: String,
-        heartbeat_fn: Arc<OnceLock<HeartbeatFn>>,
-        client: Arc<dyn WorkerClient>,
-    ) -> Self {
-        let sdk_name_and_ver = client.sdk_name_and_version();
-        let reset_notify = Arc::new(Notify::new());
-        let data = Arc::new(Mutex::new(WorkerHeartbeatData::new(
-            config,
-            identity,
-            sdk_name_and_ver,
-            reset_notify.clone(),
-        )));
-        let data_clone = data.clone();
-
-        let heartbeat_handle = tokio::spawn(async move {
-            let mut ticker = tokio::time::interval(data_clone.lock().heartbeat_interval);
-            ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
-            loop {
-                tokio::select! {
-                    _ = ticker.tick() => {
-                        let heartbeat = if let Some(heartbeat) = data_clone.lock().capture_heartbeat_if_needed() {
-                            heartbeat
-                        } else {
-                            continue
-                        };
-                        if let Err(e) = client.clone().record_worker_heartbeat(heartbeat).await {
-                            if matches!(
-                                e.code(),
-                                tonic::Code::Unimplemented
-                            ) {
-                                return;
-                            }
-                            warn!(error=?e, "Network error while sending worker heartbeat");
-                        }
-                    }
-                    _ = reset_notify.notified() => {
-                        ticker.reset();
-                    }
-                }
-            }
-        });
-
-        let data_clone = data.clone();
-        if heartbeat_fn
-            .set(Box::new(move || {
-                data_clone.lock().capture_heartbeat_if_needed()
-            }))
-            .is_err()
-        {
-            dbg_panic!(
-                "Failed to set heartbeat_fn, heartbeat_fn should only be set once, when a singular WorkerHeartbeatInfo is created"
-            );
-        }
-
-        Self { heartbeat_handle }
-    }
-
-    pub(crate) fn shutdown(&self) {
-        self.heartbeat_handle.abort()
-    }
-}
-
-#[derive(Debug, Clone)]
-struct WorkerHeartbeatData {
-    worker_instance_key: String,
-    worker_identity: String,
-    host_info: WorkerHostInfo,
-    // Time of the last heartbeat. This is used to both for heartbeat_time and last_heartbeat_time
-    heartbeat_time: Option<SystemTime>,
-    task_queue: String,
-    /// SDK name
-    sdk_name: String,
-    /// SDK version
-    sdk_version: String,
-    /// Worker start time
-    start_time: SystemTime,
-    heartbeat_interval: Duration,
-    reset_notify: Arc<Notify>,
-}
-
-impl WorkerHeartbeatData {
-    fn new(
-        worker_config: WorkerConfig,
-        worker_identity: String,
-        sdk_name_and_ver: (String, String),
-        reset_notify: Arc<Notify>,
-    ) -> Self {
-        Self {
-            worker_identity,
-            host_info: WorkerHostInfo {
-                host_name: gethostname().to_string_lossy().to_string(),
-                process_id: std::process::id().to_string(),
-                ..Default::default()
-            },
-            sdk_name: sdk_name_and_ver.0,
-            sdk_version: sdk_name_and_ver.1,
-            task_queue: worker_config.task_queue.clone(),
-            start_time: SystemTime::now(),
-            heartbeat_time: None,
-            worker_instance_key: Uuid::new_v4().to_string(),
-            heartbeat_interval: worker_config
-                .heartbeat_interval
-                .expect("WorkerHeartbeatData is only called when heartbeat_interval is Some"),
-            reset_notify,
-        }
-    }
-
-    fn capture_heartbeat_if_needed(&mut self) -> Option<WorkerHeartbeat> {
-        let now = SystemTime::now();
-        let elapsed_since_last_heartbeat = if let Some(heartbeat_time) = self.heartbeat_time {
-            let dur = now.duration_since(heartbeat_time).unwrap_or(Duration::ZERO);
-
-            // Only send poll data if it's nearly been a full interval since this data has been sent
-            // In this case, "nearly" is 90% of the interval
-            if dur.as_secs_f64() < 0.9 * self.heartbeat_interval.as_secs_f64() {
-                return None;
-            }
-            Some(PbDuration {
-                seconds: dur.as_secs() as i64,
-                nanos: dur.subsec_nanos() as i32,
-            })
-        } else {
-            None
-        };
-
-        self.heartbeat_time = Some(now);
-
-        self.reset_notify.notify_one();
-
-        Some(WorkerHeartbeat {
-            worker_instance_key: self.worker_instance_key.clone(),
-            worker_identity: self.worker_identity.clone(),
-            host_info: Some(self.host_info.clone()),
-            task_queue: self.task_queue.clone(),
-            sdk_name: self.sdk_name.clone(),
-            sdk_version: self.sdk_version.clone(),
-            status: 0,
-            start_time: Some(self.start_time.into()),
-            heartbeat_time: Some(SystemTime::now().into()),
-            elapsed_since_last_heartbeat,
-            ..Default::default()
-        })
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::{
-        test_help::{WorkerExt, test_worker_cfg},
-        worker,
-        worker::client::mocks::mock_worker_client,
-    };
-    use std::{sync::Arc, time::Duration};
-    use temporal_sdk_core_api::worker::PollerBehavior;
-    use temporal_sdk_core_protos::temporal::api::workflowservice::v1::RecordWorkerHeartbeatResponse;
-
-    #[tokio::test]
-    async fn worker_heartbeat() {
-        let mut mock = mock_worker_client();
-        mock.expect_record_worker_heartbeat()
-            .times(2)
-            .returning(move |heartbeat| {
-                let host_info = heartbeat.host_info.clone().unwrap();
-                assert_eq!("test-identity", heartbeat.worker_identity);
-                assert!(!heartbeat.worker_instance_key.is_empty());
-                assert_eq!(
-                    host_info.host_name,
-                    gethostname::gethostname().to_string_lossy().to_string()
-                );
-                assert_eq!(host_info.process_id, std::process::id().to_string());
-                assert_eq!(heartbeat.sdk_name, "test-core");
-                assert_eq!(heartbeat.sdk_version, "0.0.0");
-                assert_eq!(heartbeat.task_queue, "q");
-                assert!(heartbeat.heartbeat_time.is_some());
-                assert!(heartbeat.start_time.is_some());
-
-                Ok(RecordWorkerHeartbeatResponse {})
-            });
-
-        let config = test_worker_cfg()
-            .activity_task_poller_behavior(PollerBehavior::SimpleMaximum(1_usize))
-            .max_outstanding_activities(1_usize)
-            .heartbeat_interval(Duration::from_millis(200))
-            .build()
-            .unwrap();
-
-        let heartbeat_fn = Arc::new(OnceLock::new());
-        let client = Arc::new(mock);
-        let worker = worker::Worker::new(config, None, client, None, Some(heartbeat_fn.clone()));
-        heartbeat_fn.get().unwrap()();
-
-        // heartbeat timer fires once
-        advance_time(Duration::from_millis(300)).await;
-        // it hasn't been >90% of the interval since the last heartbeat, so no data should be returned here
-        assert_eq!(None, heartbeat_fn.get().unwrap()());
-        // heartbeat timer fires once
-        advance_time(Duration::from_millis(300)).await;
-
-        worker.drain_activity_poller_and_shutdown().await;
-    }
-
-    async fn advance_time(dur: Duration) {
-        tokio::time::pause();
-        tokio::time::advance(dur).await;
-        tokio::time::resume();
-    }
-}