formating

Author: yuandrew

core-c-bridge/src/runtime.rs

@@ -1,31 +1,29 @@
-use crate::ByteArray;
-use crate::ByteArrayRef;
-use crate::MetadataRef;
-use crate::metric::CustomMetricMeter;
-use crate::metric::CustomMetricMeterRef;
+use crate::{
+    ByteArray, ByteArrayRef, MetadataRef,
+    metric::{CustomMetricMeter, CustomMetricMeterRef},
+};
 
 use serde_json::json;
-use std::collections::HashMap;
-use std::fmt;
-use std::net::SocketAddr;
-use std::str::FromStr;
-use std::sync::atomic::AtomicBool;
-use std::sync::atomic::Ordering;
-use std::sync::{Arc, Mutex};
-use std::time::Duration;
-use std::time::UNIX_EPOCH;
-use temporal_sdk_core::CoreRuntime;
-use temporal_sdk_core::RuntimeOptions as CoreRuntimeOptions;
-use temporal_sdk_core::RuntimeOptionsBuilder as CoreRuntimeOptionsBuilder;
-use temporal_sdk_core::TokioRuntimeBuilder;
-use temporal_sdk_core::telemetry::{build_otlp_metric_exporter, start_prometheus_metric_exporter};
-use temporal_sdk_core_api::telemetry::HistogramBucketOverrides;
-use temporal_sdk_core_api::telemetry::MetricTemporality;
-use temporal_sdk_core_api::telemetry::metrics::CoreMeter;
-use temporal_sdk_core_api::telemetry::{CoreLog, CoreLogConsumer};
+use std::{
+    collections::HashMap,
+    fmt,
+    net::SocketAddr,
+    str::FromStr,
+    sync::{
+        Arc, Mutex,
+        atomic::{AtomicBool, Ordering},
+    },
+    time::{Duration, UNIX_EPOCH},
+};
+use temporal_sdk_core::{
+    CoreRuntime, RuntimeOptions as CoreRuntimeOptions,
+    RuntimeOptionsBuilder as CoreRuntimeOptionsBuilder, TokioRuntimeBuilder,
+    telemetry::{build_otlp_metric_exporter, start_prometheus_metric_exporter},
+};
 use temporal_sdk_core_api::telemetry::{
-    Logger, OtelCollectorOptionsBuilder, PrometheusExporterOptionsBuilder,
-    TelemetryOptions as CoreTelemetryOptions, TelemetryOptionsBuilder,
+    CoreLog, CoreLogConsumer, HistogramBucketOverrides, Logger, MetricTemporality,
+    OtelCollectorOptionsBuilder, PrometheusExporterOptionsBuilder,
+    TelemetryOptions as CoreTelemetryOptions, TelemetryOptionsBuilder, metrics::CoreMeter,
 };
 use tracing::Level;
 use url::Url;

core/src/lib.rs

@@ -30,7 +30,6 @@ mod core_tests;
 #[macro_use]
 pub mod test_help;
 
-use std::collections::HashMap;
 pub(crate) use temporal_sdk_core_api::errors;
 
 pub use pollers::{
@@ -119,7 +118,7 @@ where
         client_ident.clone(),
         worker_config.versioning_strategy.clone(),
     ));
-    
+
     let worker = Worker::new(
         worker_config.clone(),
         sticky_q,
@@ -370,10 +369,6 @@ impl CoreRuntime {
     pub fn telemetry_mut(&mut self) -> &mut TelemetryInstance {
         &mut self.telemetry
     }
-
-    fn task_queue_key(&self) -> Uuid {
-        self.process_key
-    }
 }
 
 impl Drop for CoreRuntime {

core/src/telemetry/metrics.rs

@@ -1,4 +1,4 @@
-use crate::{abstractions::dbg_panic, telemetry::TelemetryInstance};
+use crate::abstractions::dbg_panic;
 
 use std::{
     fmt::{Debug, Display},

core/src/worker/client.rs

@@ -2,14 +2,13 @@
 
 pub(crate) mod mocks;
 use crate::protosext::legacy_query_failure;
-use parking_lot::{RwLock};
+use parking_lot::RwLock;
 use std::{sync::Arc, time::Duration};
 use temporal_client::{
     Client, ClientWorkerSet, IsWorkerTaskLongPoll, Namespace, NamespacedClient, NoRetryOnMatching,
     RetryClient, WorkflowService,
 };
 use temporal_sdk_core_api::worker::WorkerVersioningStrategy;
-use temporal_sdk_core_protos::temporal::api::worker::v1::WorkerHeartbeat;
 use temporal_sdk_core_protos::{
     TaskToken,
     coresdk::{workflow_commands::QueryResult, workflow_completion},
@@ -29,6 +28,7 @@ use temporal_sdk_core_protos::{
         query::v1::WorkflowQueryResult,
         sdk::v1::WorkflowTaskCompletedMetadata,
         taskqueue::v1::{StickyExecutionAttributes, TaskQueue, TaskQueueMetadata},
+        worker::v1::WorkerHeartbeat,
         workflowservice::v1::{get_system_info_response::Capabilities, *},
     },
 };
@@ -126,7 +126,7 @@ impl WorkerClientBag {
 
 /// This trait contains everything workers need to interact with Temporal, and hence provides a
 /// minimal mocking surface. Delegates to [WorkflowClientTrait] so see that for details.
-#[cfg_attr(test, mockall::automock)]
+#[cfg_attr(any(feature = "test-utilities", test), mockall::automock)]
 #[async_trait::async_trait]
 pub trait WorkerClient: Sync + Send {
     /// Poll workflow tasks

core/src/worker/client/mocks.rs

@@ -35,6 +35,7 @@ pub fn mock_worker_client() -> MockWorkerClient {
         .returning(|| ("test-core".to_string(), "0.0.0".to_string()));
     r.expect_get_identity()
         .returning(|| "test-identity".to_string());
+    r.expect_get_process_key().returning(Uuid::new_v4);
     r
 }
 

core/src/worker/heartbeat.rs

@@ -3,9 +3,11 @@ use crate::worker::{TaskPollers, WorkerTelemetry};
 use parking_lot::Mutex;
 use prost_types::Duration as PbDuration;
 use std::collections::HashMap;
-use std::fmt;
-use std::sync::Arc;
-use std::time::{Duration, SystemTime};
+use std::{
+    fmt,
+    sync::Arc,
+    time::{Duration, SystemTime},
+};
 use temporal_client::SharedNamespaceWorkerTrait;
 use temporal_sdk_core_api::worker::{WorkerConfigBuilder, WorkerVersioningStrategy};
 use temporal_sdk_core_protos::temporal::api::worker::v1::WorkerHeartbeat;
@@ -29,12 +31,16 @@ impl SharedNamespaceWorker {
         heartbeat_interval: Duration,
         telemetry: Option<WorkerTelemetry>,
     ) -> Self {
-        println!("SharedNamespaceWorker::new() {:?}\n\tsdk_name_and_version key{:?}", client.get_identity(), client.sdk_name_and_version());
+        println!(
+            "SharedNamespaceWorker::new() {:?}\n\tsdk_name_and_version key{:?}",
+            client.get_identity(),
+            client.sdk_name_and_version()
+        );
         let config = WorkerConfigBuilder::default()
             .namespace(namespace.clone())
             .task_queue(format!(
                 "temporal-sys/worker-commands/{namespace}/{}",
-                client.get_process_key().to_string()
+                client.get_process_key()
             ))
             .no_remote_activities(true)
             .max_outstanding_nexus_tasks(5_usize)
@@ -68,7 +74,9 @@ impl SharedNamespaceWorker {
             loop {
                 // TODO: Race condition here, can technically shut down before anything is ever initialized
                 if heartbeat_map_clone.lock().is_empty() {
-                    println!("// TODO: Race condition here, can technically shut down before anything is ever initialized");
+                    println!(
+                        "// TODO: Race condition here, can technically shut down before anything is ever initialized"
+                    );
                     worker.shutdown().await;
                     return;
                 }
@@ -169,17 +177,20 @@ impl fmt::Debug for SharedNamespaceWorker {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::test_help::WorkerExt;
-    use crate::test_help::test_worker_cfg;
-    use crate::worker;
-    use crate::worker::client::mocks::mock_worker_client;
-    use std::sync::Arc;
-    use std::sync::atomic::AtomicUsize;
-    use std::sync::atomic::Ordering;
-    use std::time::Duration;
+    use crate::{
+        test_help::{WorkerExt, test_worker_cfg},
+        worker,
+        worker::client::mocks::mock_worker_client,
+    };
+    use std::{
+        sync::{
+            Arc,
+            atomic::{AtomicUsize, Ordering},
+        },
+        time::Duration,
+    };
     use temporal_sdk_core_api::worker::PollerBehavior;
     use temporal_sdk_core_protos::temporal::api::workflowservice::v1::RecordWorkerHeartbeatResponse;
-    use uuid::Uuid;
 
     #[tokio::test]
     async fn worker_heartbeat_basic() {
@@ -218,29 +229,13 @@ mod tests {
             .into();
 
         let client = Arc::new(mock);
-        let worker = worker::Worker::new(config, None, client.clone(), None, None);
-
-        let namespace = "test-namespace".to_string();
-        let process_key = Uuid::new_v4();
-        // let mut shared_namespace_worker = SharedNamespaceWorker::new(
-        //     client,
-        //     WorkerHeartbeatIdentity::new(
-        //         "test-endpoint".to_string(),
-        //         namespace.clone(),
-        //         process_key.to_string(),
-        //         "test-identity".to_string(),
-        //     ),
-        //     Duration::from_millis(100),
-        //     None,
-        // );
-        // TODO: translate to new way
-        // let worker_instance_key = worker.worker_instance_key().unwrap();
-        // shared_namespace_worker.register_callback(
-        //     worker_instance_key,
-        //     worker
-        //         .get_heartbeat_callback()
-        //         .expect("heartbeat callback should be set"),
-        // );
+        let worker = worker::Worker::new(
+            config,
+            None,
+            client.clone(),
+            None,
+            Some(Duration::from_millis(100)),
+        );
 
         tokio::time::sleep(Duration::from_millis(250)).await;
         worker.drain_activity_poller_and_shutdown().await;

core/src/worker/mod.rs

@@ -70,7 +70,6 @@ use temporal_sdk_core_api::{
     errors::{CompleteNexusError, WorkerValidationError},
     worker::PollerBehavior,
 };
-use temporal_sdk_core_protos::temporal::api::cloud::account::v1::Metrics;
 use temporal_sdk_core_protos::temporal::api::worker::v1::WorkerHostInfo;
 use temporal_sdk_core_protos::{
     TaskToken,
@@ -87,7 +86,7 @@ use temporal_sdk_core_protos::{
         taskqueue::v1::{StickyExecutionAttributes, TaskQueue},
     },
 };
-use tokio::sync::{OnceCell, mpsc::unbounded_channel, watch};
+use tokio::sync::{mpsc::unbounded_channel, watch};
 use tokio_stream::wrappers::UnboundedReceiverStream;
 use tokio_util::sync::CancellationToken;
 use tracing::Subscriber;
@@ -170,9 +169,6 @@ impl WorkerTelemetry {
 struct WorkerHeartbeat {
     /// Instance key used to identify this worker in worker heartbeating
     worker_instance_key: Uuid,
-    /// Used to remove this worker from the parent map used to track this worker for
-    /// worker heartbeat
-    shutdown_callback: OnceCell<Arc<dyn Fn() + Send + Sync>>,
     /// Heartbeat interval, defaults to 60s
     heartbeat_interval: Duration,
     /// Telemetry instance, needed to initialize [SharedNamespaceWorker] when replacing client
@@ -391,15 +387,11 @@ impl Worker {
         telem_instance: Option<&TelemetryInstance>,
         worker_heartbeat_interval: Option<Duration>,
     ) -> Self {
-        let worker_telemetry = if let Some(telem) = telem_instance {
-            Some(WorkerTelemetry {
-                metric_meter: telem.get_metric_meter(),
-                temporal_metric_meter: telem.get_temporal_metric_meter(),
-                trace_subscriber: telem.trace_subscriber(),
-            })
-        } else {
-            None
-        };
+        let worker_telemetry = telem_instance.map(|telem| WorkerTelemetry {
+            metric_meter: telem.get_metric_meter(),
+            temporal_metric_meter: telem.get_temporal_metric_meter(),
+            trace_subscriber: telem.trace_subscriber(),
+        });
 
         Worker::new_with_pollers_inner(
             config,
@@ -664,8 +656,7 @@ impl Worker {
 
             Some(WorkerHeartbeat {
                 worker_instance_key,
-                shutdown_callback: OnceCell::new(),
-                heartbeat_interval: heartbeat_interval,
+                heartbeat_interval,
                 telemetry: worker_telemetry.clone(),
             })
         } else {

core/src/worker/workflow/mod.rs

@@ -23,7 +23,7 @@ use crate::{
     internal_flags::InternalFlags,
     pollers::TrackedPermittedTqResp,
     protosext::{ValidPollWFTQResponse, protocol_messages::IncomingProtocolMessage},
-    telemetry::{TelemetryInstance, VecDisplayer, set_trace_subscriber_for_current_thread},
+    telemetry::{VecDisplayer, set_trace_subscriber_for_current_thread},
     worker::{
         LocalActRequest, LocalActivityExecutionResult, LocalActivityResolution,
         PostActivateHookData,