@temporalio/core-bridge 0.19.2 → 0.20.2

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

@@ -10,6 +10,7 @@ pub extern crate assert_matches;
 #[macro_use]
 extern crate tracing;
 
+mod abstractions;
 mod log_export;
 mod pending_activations;
 mod pollers;
@@ -29,50 +30,30 @@ mod test_help;
 pub(crate) use temporal_sdk_core_api::errors;
 
 pub use pollers::{
-    ClientTlsConfig, RetryConfig, RetryGateway, ServerGateway, ServerGatewayApis,
-    ServerGatewayOptions, ServerGatewayOptionsBuilder, TlsConfig,
+    Client, ClientOptions, ClientOptionsBuilder, ClientTlsConfig, RetryClient, RetryConfig,
+    TlsConfig, WorkflowClientTrait,
+};
+pub use telemetry::{
+    fetch_global_buffered_logs, telemetry_init, TelemetryOptions, TelemetryOptionsBuilder,
 };
-pub use telemetry::{TelemetryOptions, TelemetryOptionsBuilder};
 pub use temporal_sdk_core_api as api;
 pub use temporal_sdk_core_protos as protos;
 pub use temporal_sdk_core_protos::TaskToken;
 pub use url::Url;
-pub use worker::{WorkerConfig, WorkerConfigBuilder};
+pub use worker::{Worker, WorkerConfig, WorkerConfigBuilder};
 
 use crate::{
-    telemetry::{
-        fetch_global_buffered_logs,
-        metrics::{MetricsContext, METRIC_METER},
-        telemetry_init,
-    },
-    worker::{Worker, WorkerDispatcher},
-};
-use std::{
-    ops::Deref,
-    sync::{
-        atomic::{AtomicBool, Ordering},
-        Arc,
-    },
+    replay::mock_client_from_history,
+    telemetry::metrics::{MetricsContext, METRIC_METER},
+    worker::client::WorkerClientBag,
 };
+use std::sync::Arc;
+use temporal_client::AnyClient;
 use temporal_sdk_core_api::{
-    errors::{
-        CompleteActivityError, CompleteWfError, CoreInitError, PollActivityError, PollWfError,
-        WorkerRegistrationError,
-    },
-    Core, CoreLog,
-};
-use temporal_sdk_core_protos::{
-    coresdk::{
-        activity_task::ActivityTask,
-        workflow_activation::{remove_from_cache::EvictionReason, WorkflowActivation},
-        workflow_completion::WorkflowActivationCompletion,
-        ActivityHeartbeat, ActivityTaskCompletion,
-    },
-    temporal::api::history::v1::History,
+    errors::{CompleteActivityError, PollActivityError, PollWfError},
+    CoreLog, Worker as WorkerTrait,
 };
-
-#[cfg(test)]
-use crate::test_help::MockWorker;
+use temporal_sdk_core_protos::{coresdk::ActivityHeartbeat, temporal::api::history::v1::History};
 
 lazy_static::lazy_static! {
     /// A process-wide unique string, which will be different on every startup
@@ -81,294 +62,68 @@ lazy_static::lazy_static! {
     };
 }
 
-/// Holds various configuration information required to call [init]
-#[derive(Debug, Clone, derive_builder::Builder)]
-#[builder(setter(into))]
-#[non_exhaustive]
-pub struct CoreInitOptions {
-    /// Options for the connection to the temporal server
-    pub gateway_opts: ServerGatewayOptions,
-    /// Options for telemetry (traces and metrics)
-    #[builder(default)]
-    pub telemetry_opts: TelemetryOptions,
-}
-
-/// Initializes an instance of the core sdk and establishes a connection to the temporal server.
-/// Expects that a tokio runtime exists.
-pub async fn init(opts: CoreInitOptions) -> Result<CoreSDK, CoreInitError> {
-    telemetry_init(&opts.telemetry_opts).map_err(CoreInitError::TelemetryInitError)?;
-    // Initialize server client
-    let server_gateway = opts.gateway_opts.connect(Some(&METRIC_METER)).await?;
-
-    Ok(CoreSDK::new(server_gateway, opts))
-}
-
-/// Initialize core using a provided gateway instance, which is typically a mock
-pub fn init_mock_gateway<SG: ServerGatewayApis + Send + Sync + 'static>(
-    opts: CoreInitOptions,
-    server_gateway: SG,
-) -> Result<CoreSDK, CoreInitError> {
-    telemetry_init(&opts.telemetry_opts).map_err(CoreInitError::TelemetryInitError)?;
-    Ok(CoreSDK::new(server_gateway, opts))
-}
-
-/// Implements the [Core] trait
-pub struct CoreSDK {
-    /// Options provided at initialization time
-    init_options: CoreInitOptions,
-    /// A client for interacting with the Temporal service.
-    server_gateway: Arc<dyn ServerGatewayApis + Send + Sync>,
-    /// Controls access to workers
-    workers: WorkerDispatcher,
-    /// Has shutdown been called and all workers drained of tasks?
-    whole_core_shutdown: AtomicBool,
-    /// Top-level metrics context
-    metrics: MetricsContext,
-}
-
-#[async_trait::async_trait]
-impl Core for CoreSDK {
-    fn register_worker(&self, config: WorkerConfig) -> Result<(), WorkerRegistrationError> {
-        info!(
-            task_queue = config.task_queue.as_str(),
-            "Registering worker"
-        );
-        let sticky_q = self.get_sticky_q_name_for_worker(&config);
-        self.workers.new_worker(
-            config,
-            sticky_q,
-            self.server_gateway.clone(),
-            self.metrics.clone(),
-        )
-    }
-
-    #[instrument(level = "debug", skip(self), fields(run_id))]
-    async fn poll_workflow_activation(
-        &self,
-        task_queue: &str,
-    ) -> Result<WorkflowActivation, PollWfError> {
-        let worker = self.worker(task_queue)?;
-        worker.next_workflow_activation().await
-    }
-
-    #[instrument(level = "debug", skip(self))]
-    async fn poll_activity_task(
-        &self,
-        task_queue: &str,
-    ) -> Result<ActivityTask, PollActivityError> {
-        loop {
-            if self.whole_core_shutdown.load(Ordering::Relaxed) {
-                return Err(PollActivityError::ShutDown);
-            }
-            let worker = self.worker(task_queue)?;
-            match worker.activity_poll().await.transpose() {
-                Some(r) => break r,
-                None => {
-                    tokio::task::yield_now().await;
-                    continue;
-                }
-            }
-        }
-    }
-
-    #[instrument(level = "debug", skip(self, completion),
-      fields(completion=%&completion, run_id=%completion.run_id))]
-    async fn complete_workflow_activation(
-        &self,
-        completion: WorkflowActivationCompletion,
-    ) -> Result<(), CompleteWfError> {
-        let worker = self.worker(&completion.task_queue)?;
-        worker.complete_workflow_activation(completion).await
-    }
-
-    #[instrument(level = "debug", skip(self))]
-    async fn complete_activity_task(
-        &self,
-        completion: ActivityTaskCompletion,
-    ) -> Result<(), CompleteActivityError> {
-        let task_token = TaskToken(completion.task_token);
-        let status = if let Some(s) = completion.result.and_then(|r| r.status) {
-            s
-        } else {
-            return Err(CompleteActivityError::MalformedActivityCompletion {
-                reason: "Activity completion had empty result/status field".to_owned(),
-                completion: None,
-            });
-        };
-
-        let worker = self.worker(&completion.task_queue)?;
-        worker.complete_activity(task_token, status).await
-    }
-
-    fn record_activity_heartbeat(&self, details: ActivityHeartbeat) {
-        if let Ok(w) = self.worker(&details.task_queue) {
-            w.record_heartbeat(details);
-        }
-    }
-
-    fn request_workflow_eviction(&self, task_queue: &str, run_id: &str) {
-        if let Ok(w) = self.worker(task_queue) {
-            w.request_wf_eviction(
-                run_id,
-                "Eviction explicitly requested by lang",
-                EvictionReason::LangRequested,
-            );
-        }
-    }
-
-    fn server_gateway(&self) -> Arc<dyn ServerGatewayApis + Send + Sync> {
-        self.server_gateway.clone()
-    }
-
-    async fn shutdown(&self) {
-        self.workers.shutdown_all().await;
-        self.whole_core_shutdown.store(true, Ordering::Relaxed);
-    }
-
-    async fn shutdown_worker(&self, task_queue: &str) {
-        self.workers.shutdown_one(task_queue).await;
-    }
-
-    fn fetch_buffered_logs(&self) -> Vec<CoreLog> {
-        fetch_global_buffered_logs()
-    }
-}
-
-impl CoreSDK {
-    pub(crate) fn new<SG: ServerGatewayApis + Send + Sync + 'static>(
-        server_gateway: SG,
-        init_options: CoreInitOptions,
-    ) -> Self {
-        let server_gateway = Arc::new(server_gateway);
-        let workers = WorkerDispatcher::default();
-        Self {
-            workers,
-            init_options,
-            whole_core_shutdown: AtomicBool::new(false),
-            metrics: MetricsContext::top_level(server_gateway.get_options().namespace.clone()),
-            server_gateway,
-        }
-    }
-
-    /// Register a worker for replaying a specific history. The worker should use a unique task
-    /// queue name. It will auto-shutdown as soon as the history has finished being replayed. The
-    /// provided gateway should be a mock, and this should only be used for workflow testing
-    /// purposes.
-    pub fn register_replay_worker(
-        &self,
-        config: WorkerConfig,
-        gateway: Arc<dyn ServerGatewayApis + Send + Sync>,
-        history: &History,
-    ) -> Result<(), anyhow::Error> {
-        info!(
-            task_queue = config.task_queue.as_str(),
-            "Registering replay worker"
-        );
-        // Could possibly just use mocked pollers here, but they'd need to be un-test-moded
-        let run_id = history.extract_run_id_from_start()?.to_string();
-        let last_event = history.last_event_id();
-        let tq = config.task_queue.clone();
-        let mut worker = Worker::new(config, None, gateway, self.metrics.clone());
-        worker.set_post_activate_hook(move |worker| {
-            if worker
-                .wft_manager
-                .most_recently_processed_event(&run_id)
-                .unwrap_or_default()
-                >= last_event
-            {
-                worker.initiate_shutdown();
-            }
-        });
-
-        self.workers.set_worker_for_task_queue(tq, worker)?;
-        Ok(())
-    }
-
-    /// Allow construction of workers with mocked poll responses during testing
-    #[cfg(test)]
-    pub(crate) fn reg_worker_sync(&self, worker: MockWorker) {
-        let sticky_q = self.get_sticky_q_name_for_worker(&worker.config);
-        let tq = worker.config.task_queue.clone();
-        let worker = Worker::new_with_pollers(
-            worker.config,
-            sticky_q,
-            self.server_gateway.clone(),
-            worker.wf_poller,
-            worker.act_poller,
-            self.metrics.clone(),
-        );
-        self.workers.set_worker_for_task_queue(tq, worker).unwrap();
-    }
-
-    #[cfg(test)]
-    pub(crate) fn outstanding_wfts(&self, tq: &str) -> usize {
-        self.worker(tq).unwrap().outstanding_workflow_tasks()
-    }
-    #[cfg(test)]
-    pub(crate) fn available_wft_permits(&self, tq: &str) -> usize {
-        self.worker(tq).unwrap().available_wft_permits()
-    }
-
-    fn get_sticky_q_name_for_worker(&self, config: &WorkerConfig) -> Option<String> {
-        if config.max_cached_workflows > 0 {
-            Some(format!(
-                "{}-{}-{}",
-                &self.init_options.gateway_opts.identity, &config.task_queue, *PROCCESS_UNIQ_ID
-            ))
-        } else {
-            None
+/// Initialize a worker bound to a task queue
+pub fn init_worker<CT>(worker_config: WorkerConfig, client: CT) -> Worker
+where
+    CT: Into<AnyClient>,
+{
+    let as_enum = client.into();
+    // TODO: Assert namespaces match
+    let client = match as_enum {
+        AnyClient::HighLevel(ac) => ac,
+        AnyClient::LowLevel(ll) => {
+            let client = Client::new(*ll, worker_config.namespace.clone());
+            let retry_client = RetryClient::new(client, RetryConfig::default());
+            Arc::new(retry_client)
         }
-    }
-
-    fn worker(&self, tq: &str) -> Result<impl Deref<Target = Worker>, WorkerLookupErr> {
-        let worker = self.workers.get(tq);
-        if worker.is_err() && self.whole_core_shutdown.load(Ordering::Relaxed) {
-            return Err(WorkerLookupErr::Shutdown(tq.to_owned()));
-        }
-        worker
-    }
-}
-
-#[derive(Debug)]
-enum WorkerLookupErr {
-    Shutdown(String),
-    NoWorker(String),
-}
-
-impl From<WorkerLookupErr> for PollWfError {
-    fn from(e: WorkerLookupErr) -> Self {
-        match e {
-            WorkerLookupErr::Shutdown(_) => Self::ShutDown,
-            WorkerLookupErr::NoWorker(s) => Self::NoWorkerForQueue(s),
-        }
-    }
-}
-
-impl From<WorkerLookupErr> for PollActivityError {
-    fn from(e: WorkerLookupErr) -> Self {
-        match e {
-            WorkerLookupErr::Shutdown(_) => Self::ShutDown,
-            WorkerLookupErr::NoWorker(s) => Self::NoWorkerForQueue(s),
-        }
-    }
+    };
+    let c_opts = client.get_options().clone();
+    if client.namespace() != worker_config.namespace {
+        panic!("Passed in client is not bound to the same namespace as the worker");
+    }
+    let client_bag = Arc::new(WorkerClientBag::new(
+        Box::new(client),
+        worker_config.namespace.clone(),
+    ));
+    let sticky_q = sticky_q_name_for_worker(&c_opts.identity, &worker_config);
+    let metrics = MetricsContext::top_level(worker_config.namespace.clone())
+        .with_task_q(worker_config.task_queue.clone());
+    Worker::new(worker_config, sticky_q, client_bag, metrics)
 }
 
-impl From<WorkerLookupErr> for CompleteWfError {
-    fn from(e: WorkerLookupErr) -> Self {
-        match e {
-            WorkerLookupErr::Shutdown(s) | WorkerLookupErr::NoWorker(s) => {
-                Self::NoWorkerForQueue(s)
-            }
-        }
-    }
+/// Create a worker for replaying a specific history. It will auto-shutdown as soon as the history
+/// has finished being replayed. The provided client should be a mock, and this should only be used
+/// for workflow testing purposes.
+pub fn init_replay_worker(
+    mut config: WorkerConfig,
+    history: &History,
+) -> Result<Worker, anyhow::Error> {
+    info!(
+        task_queue = config.task_queue.as_str(),
+        "Registering replay worker"
+    );
+    config.max_cached_workflows = 1;
+    config.max_concurrent_wft_polls = 1;
+    config.no_remote_activities = true;
+    // Could possibly just use mocked pollers here?
+    let client = mock_client_from_history(history, &config.task_queue);
+    let run_id = history.extract_run_id_from_start()?.to_string();
+    let last_event = history.last_event_id();
+    let mut worker = Worker::new(config, None, Arc::new(client), MetricsContext::default());
+    worker.set_shutdown_on_run_reaches_event(run_id, last_event);
+    Ok(worker)
 }
 
-impl From<WorkerLookupErr> for CompleteActivityError {
-    fn from(e: WorkerLookupErr) -> Self {
-        match e {
-            WorkerLookupErr::Shutdown(s) | WorkerLookupErr::NoWorker(s) => {
-                Self::NoWorkerForQueue(s)
-            }
-        }
+pub(crate) fn sticky_q_name_for_worker(
+    process_identity: &str,
+    config: &WorkerConfig,
+) -> Option<String> {
+    if config.max_cached_workflows > 0 {
+        Some(format!(
+            "{}-{}-{}",
+            &process_identity, &config.task_queue, *PROCCESS_UNIQ_ID
+        ))
+    } else {
+        None
     }
 }

package/sdk-core/core/src/pollers/mod.rs

@@ -1,14 +1,12 @@
 mod poll_buffer;
 
-pub use poll_buffer::{
-    new_activity_task_buffer, new_workflow_task_buffer, PollActivityTaskBuffer,
-    PollWorkflowTaskBuffer, WorkflowTaskPoller,
+pub(crate) use poll_buffer::{
+    new_activity_task_buffer, new_workflow_task_buffer, WorkflowTaskPoller,
 };
 pub use temporal_client::{
-    ClientTlsConfig, RetryConfig, RetryGateway, ServerGateway, ServerGatewayApis,
-    ServerGatewayOptions, ServerGatewayOptionsBuilder, TlsConfig,
+    Client, ClientOptions, ClientOptionsBuilder, ClientTlsConfig, RetryClient, RetryConfig,
+    TlsConfig, WorkflowClientTrait,
 };
-
 use temporal_sdk_core_protos::temporal::api::workflowservice::v1::{
     PollActivityTaskQueueResponse, PollWorkflowTaskQueueResponse,
 };

package/sdk-core/core/src/pollers/poll_buffer.rs

@@ -1,6 +1,6 @@
 use crate::{
     pollers::{self, Poller},
-    ServerGatewayApis,
+    worker::client::WorkerClientBag,
 };
 use futures::{prelude::stream::FuturesUnordered, StreamExt};
 use std::{
@@ -197,8 +197,8 @@ impl Poller<PollWorkflowTaskQueueResponse> for WorkflowTaskPoller {
 }
 
 pub type PollWorkflowTaskBuffer = LongPollBuffer<PollWorkflowTaskQueueResponse>;
-pub fn new_workflow_task_buffer(
-    sg: Arc<impl ServerGatewayApis + Send + Sync + 'static + ?Sized>,
+pub(crate) fn new_workflow_task_buffer(
+    client: Arc<WorkerClientBag>,
     task_queue: String,
     is_sticky: bool,
     concurrent_pollers: usize,
@@ -206,9 +206,9 @@ pub fn new_workflow_task_buffer(
 ) -> PollWorkflowTaskBuffer {
     LongPollBuffer::new(
         move || {
-            let sg = sg.clone();
+            let client = client.clone();
             let task_queue = task_queue.clone();
-            async move { sg.poll_workflow_task(task_queue, is_sticky).await }
+            async move { client.poll_workflow_task(task_queue, is_sticky).await }
         },
         concurrent_pollers,
         buffer_size,
@@ -216,17 +216,18 @@ pub fn new_workflow_task_buffer(
 }
 
 pub type PollActivityTaskBuffer = LongPollBuffer<PollActivityTaskQueueResponse>;
-pub fn new_activity_task_buffer(
-    sg: Arc<impl ServerGatewayApis + Send + Sync + 'static + ?Sized>,
+pub(crate) fn new_activity_task_buffer(
+    client: Arc<WorkerClientBag>,
     task_queue: String,
     concurrent_pollers: usize,
     buffer_size: usize,
+    max_tps: Option<f64>,
 ) -> PollActivityTaskBuffer {
     LongPollBuffer::new(
         move || {
-            let sg = sg.clone();
+            let client = client.clone();
             let task_queue = task_queue.clone();
-            async move { sg.poll_activity_task(task_queue).await }
+            async move { client.poll_activity_task(task_queue, max_tps).await }
         },
         concurrent_pollers,
         buffer_size,
@@ -236,15 +237,15 @@ pub fn new_activity_task_buffer(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::worker::client::mocks::mock_manual_workflow_client;
     use futures::FutureExt;
     use std::time::Duration;
-    use temporal_client::MockManualGateway;
     use tokio::{select, sync::mpsc::channel};
 
     #[tokio::test]
     async fn only_polls_once_with_1_poller() {
-        let mut mock_gateway = MockManualGateway::new();
-        mock_gateway
+        let mut mock_client = mock_manual_workflow_client();
+        mock_client
             .expect_poll_workflow_task()
             .times(2)
             .returning(move |_, _| {
@@ -254,9 +255,14 @@ mod tests {
                 }
                 .boxed()
             });
-        let mock_gateway = Arc::new(mock_gateway);
 
-        let pb = new_workflow_task_buffer(mock_gateway, "someq".to_string(), false, 1, 1);
+        let pb = new_workflow_task_buffer(
+            Arc::new(mock_client.into()),
+            "someq".to_string(),
+            false,
+            1,
+            1,
+        );
 
         // Poll a bunch of times, "interrupting" it each time, we should only actually have polled
         // once since the poll takes a while

package/sdk-core/core/src/protosext/mod.rs

@@ -19,12 +19,13 @@ use temporal_sdk_core_protos::{
         },
         external_data::LocalActivityMarkerData,
         workflow_activation::{
-            workflow_activation_job, QueryWorkflow, WorkflowActivation, WorkflowActivationJob,
+            query_to_job, workflow_activation_job, QueryWorkflow, WorkflowActivation,
+            WorkflowActivationJob,
         },
         workflow_commands::{
             query_result, ActivityCancellationType, QueryResult, ScheduleLocalActivity,
         },
-        workflow_completion, FromPayloadsExt,
+        workflow_completion,
     },
     temporal::api::{
         common::v1::{Payload, WorkflowExecution},
@@ -82,11 +83,7 @@ impl TryFrom<PollWorkflowTaskQueueResponse> for ValidPollWFTQResponse {
             } => {
                 let query_requests = queries
                     .into_iter()
-                    .map(|(id, q)| QueryWorkflow {
-                        query_id: id,
-                        query_type: q.query_type,
-                        arguments: Vec::from_payloads(q.query_args),
-                    })
+                    .map(|(id, q)| query_to_job(id, q))
                     .collect();
 
                 Ok(Self {
@@ -267,7 +264,7 @@ pub struct ValidScheduleLA {
     pub activity_type: String,
     pub attempt: u32,
     pub original_schedule_time: Option<SystemTime>,
-    pub header_fields: HashMap<String, SDKPayload>,
+    pub headers: HashMap<String, SDKPayload>,
     pub arguments: Vec<SDKPayload>,
     pub schedule_to_start_timeout: Option<Duration>,
     pub close_timeouts: LACloseTimeouts,
@@ -364,7 +361,7 @@ impl ValidScheduleLA {
         };
         let retry_policy = v
             .retry_policy
-            .ok_or(anyhow!("Retry policy must be defined!"))?;
+            .ok_or_else(|| anyhow!("Retry policy must be defined!"))?;
         let local_retry_threshold = v
             .local_retry_threshold
             .clone()
@@ -378,7 +375,7 @@ impl ValidScheduleLA {
             activity_type: v.activity_type,
             attempt: v.attempt,
             original_schedule_time,
-            header_fields: v.header_fields,
+            headers: v.headers,
             arguments: v.arguments,
             schedule_to_start_timeout,
             close_timeouts,