Address todos

Author: Sushisource

core-api/src/worker.rs

@@ -173,16 +173,7 @@ impl WorkerConfigBuilder {
         if self.max_concurrent_at_polls == Some(0) {
             return Err("`max_concurrent_at_polls` must be at least 1".to_owned());
         }
-        // TODO: Move these checks into config-for-default implementation
-        // if self.max_cached_workflows > Some(0)
-        //     && self.max_outstanding_workflow_tasks > self.max_cached_workflows
-        // {
-        //     return Err(
-        //         "Maximum concurrent workflow tasks cannot exceed the maximum number of cached \
-        //          workflows"
-        //             .to_owned(),
-        //     );
-        // }
+
         if let Some(Some(ref x)) = self.max_worker_activities_per_second {
             if !x.is_normal() || x.is_sign_negative() {
                 return Err(
@@ -195,6 +186,9 @@ impl WorkerConfigBuilder {
             .max_concurrent_wft_polls
             .unwrap_or(MAX_CONCURRENT_WFT_POLLS_DEFAULT);
 
+        // It wouldn't make any sense to have more outstanding polls than workflows we can possibly
+        // cache. If we allow this at low values it's possible for sticky pollers to reserve all
+        // available slots, crowding out the normal queue and gumming things up.
         if let Some(max_cache) = self.max_cached_workflows {
             if max_cache > 0 && max_wft_polls > max_cache {
                 return Err(
@@ -203,32 +197,6 @@ impl WorkerConfigBuilder {
             }
         }
 
-        // if matches!(self.max_concurrent_wft_polls, Some(1))
-        //     && self.max_cached_workflows > Some(0)
-        //     && self
-        //         .max_outstanding_workflow_tasks
-        //         .unwrap_or(MAX_OUTSTANDING_WFT_DEFAULT)
-        //         <= 1
-        // {
-        //     return Err(
-        //         "`max_outstanding_workflow_tasks` must be at at least 2 when \
-        //          `max_cached_workflows` is nonzero"
-        //             .to_owned(),
-        //     );
-        // }
-        // if self
-        //     .max_concurrent_wft_polls
-        //     .unwrap_or(MAX_CONCURRENT_WFT_POLLS_DEFAULT)
-        //     > self
-        //         .max_outstanding_workflow_tasks
-        //         .unwrap_or(MAX_OUTSTANDING_WFT_DEFAULT)
-        // {
-        //     return Err(
-        //         "`max_concurrent_wft_polls` cannot exceed `max_outstanding_workflow_tasks`"
-        //             .to_owned(),
-        //     );
-        // }
-
         if self.use_worker_versioning.unwrap_or_default()
             && self
                 .worker_build_id
@@ -245,6 +213,12 @@ impl WorkerConfigBuilder {
     }
 }
 
+/// Implementing this trait allows users to customize how many tasks of certain kinds the worker
+/// will perform concurrently.
+///
+/// Note that, for implementations on workflow tasks ([WorkflowSlotKind]), workers that have the
+/// workflow cache enabled should be willing to hand out _at least_ two slots, to avoid the worker
+/// becoming stuck only polling on the worker's sticky queue.
 #[async_trait::async_trait]
 pub trait SlotSupplier {
     type SlotKind: SlotKind;
@@ -281,10 +255,26 @@ pub trait SlotReservationContext: Send + Sync {
     fn num_issued_slots(&self) -> usize;
 }
 
-// TODO: Make this a struct with an optional field
-pub enum SlotSupplierPermit {
-    Data(Box<dyn Any + Send + Sync>),
-    NoData,
+#[derive(Default)]
+pub struct SlotSupplierPermit {
+    user_data: Option<Box<dyn Any + Send + Sync>>,
+}
+impl SlotSupplierPermit {
+    pub fn with_user_data<T: Any + Send + Sync>(user_data: T) -> Self {
+        Self {
+            user_data: Some(Box::new(user_data)),
+        }
+    }
+    /// Attempts to downcast the inner data, if any, into the provided type and returns it.
+    /// Returns none if there is no data or the data is not of the appropriate type.
+    pub fn user_data<T: Any + Send + Sync>(&self) -> Option<&T> {
+        self.user_data.as_ref().and_then(|b| b.downcast_ref())
+    }
+    /// Attempts to downcast the inner data, if any, into the provided type and returns it mutably.
+    /// Returns none if there is no data or the data is not of the appropriate type.
+    pub fn user_data_mut<T: Any + Send + Sync>(&mut self) -> Option<&mut T> {
+        self.user_data.as_mut().and_then(|b| b.downcast_mut())
+    }
 }
 
 pub enum SlotReleaseReason {
@@ -336,30 +326,3 @@ impl SlotKind for LocalActivitySlotKind {
         "local_activity"
     }
 }
-
-pub struct WorkflowSlotsInfo {
-    // TODO: Use wf slot info
-    pub used_slots: Vec<()>,
-    /// Current size of the workflow cache.
-    pub num_cached_workflows: usize,
-    /// The limit on the size of the cache, if any. This is important for users to know as discussed below in the section
-    /// on workflow cache management.
-    pub max_cache_size: Option<usize>,
-    // ... Possibly also metric information
-}
-
-pub trait WorkflowCacheSizer {
-    /// Return true if it is acceptable to cache a new workflow. Information about already-in-use
-    /// slots, and just-received task is provided. Will not be called for an already-cached workflow
-    /// who is receiving a new task.
-    ///
-    /// Because the number of available slots must be <= the number of workflows cached, if this
-    /// returns false when there are no idle workflows in the cache (IE: All other outstanding slots
-    /// are in use), we will buffer the task and wait for another to complete so we can evict it and
-    /// make room for the new one.
-    fn can_allow_workflow(
-        &self,
-        slots_info: &WorkflowSlotsInfo,
-        new_task: &WorkflowSlotInfo,
-    ) -> bool;
-}

core/src/abstractions.rs

@@ -104,8 +104,8 @@ where
         let supp_c_c = self.supplier.clone();
         let mets = self.metrics_ctx.clone();
         let metric_rec =
-            // When being called from the drop impl, the semaphore permit isn't actually dropped yet,
-            // so account for that. TODO: that should move somehow into fixed impl only
+            // When being called from the drop impl, the permit isn't actually dropped yet, so
+            // account for that with the `add_one` parameter.
             move |add_one: bool| {
                 let extra = usize::from(add_one);
                 let unused = uc_c.load(Ordering::Acquire);

core/src/core_tests/workflow_tasks.rs

@@ -2893,10 +2893,10 @@ async fn slot_provider_cant_hand_out_more_permits_than_cache_size() {
     impl SlotSupplier for EndlessSupplier {
         type SlotKind = WorkflowSlotKind;
         async fn reserve_slot(&self, _: &dyn SlotReservationContext) -> SlotSupplierPermit {
-            SlotSupplierPermit::NoData
+            SlotSupplierPermit::default()
         }
         fn try_reserve_slot(&self, _: &dyn SlotReservationContext) -> Option<SlotSupplierPermit> {
-            Some(SlotSupplierPermit::NoData)
+            Some(SlotSupplierPermit::default())
         }
         fn mark_slot_used(&self, _: <Self::SlotKind as SlotKind>::Info<'_>) {}
         fn release_slot(&self, _: SlotReleaseReason) {}

core/src/worker/slot_supplier.rs

@@ -6,7 +6,6 @@ use temporal_sdk_core_api::worker::{
 use tokio::sync::Semaphore;
 
 mod resource_based;
-mod tracking;
 
 pub use resource_based::{RealSysInfo, ResourceBasedSlots};
 use temporal_sdk_core_api::telemetry::metrics::TemporalMeter;
@@ -33,13 +32,18 @@ where
     type SlotKind = SK;
 
     async fn reserve_slot(&self, _: &dyn SlotReservationContext) -> SlotSupplierPermit {
-        let perm = self.sem.clone().acquire_owned().await.expect("todo");
-        SlotSupplierPermit::Data(Box::new(perm))
+        let perm = self
+            .sem
+            .clone()
+            .acquire_owned()
+            .await
+            .expect("inner semaphore is never closed");
+        SlotSupplierPermit::with_user_data(perm)
     }
 
     fn try_reserve_slot(&self, _: &dyn SlotReservationContext) -> Option<SlotSupplierPermit> {
         let perm = self.sem.clone().try_acquire_owned();
-        perm.ok().map(|p| SlotSupplierPermit::Data(Box::new(p)))
+        perm.ok().map(SlotSupplierPermit::with_user_data)
     }
 
     fn mark_slot_used(&self, _info: SK::Info<'_>) {}

core/src/worker/slot_supplier/resource_based.rs

@@ -13,7 +13,6 @@ use temporal_sdk_core_api::{
     telemetry::metrics::{CoreMeter, GaugeF64, MetricAttributes, TemporalMeter},
     worker::{
         SlotKind, SlotReleaseReason, SlotReservationContext, SlotSupplier, SlotSupplierPermit,
-        WorkflowCacheSizer, WorkflowSlotInfo, WorkflowSlotsInfo,
     },
 };
 use tokio::sync::watch;
@@ -212,7 +211,7 @@ where
 
     fn issue_slot(&self) -> SlotSupplierPermit {
         let _ = self.last_slot_issued_tx.send(Instant::now());
-        SlotSupplierPermit::NoData
+        SlotSupplierPermit::default()
     }
 
     fn time_since_last_issued(&self) -> Duration {
@@ -241,15 +240,6 @@ where
     }
 }
 
-impl<MI> WorkflowCacheSizer for ResourceBasedSlots<MI>
-where
-    MI: SystemResourceInfo + Sync + Send,
-{
-    fn can_allow_workflow(&self, _: &WorkflowSlotsInfo, _: &WorkflowSlotInfo) -> bool {
-        self.can_reserve()
-    }
-}
-
 impl<MI: SystemResourceInfo + Sync + Send> ResourceBasedSlots<MI> {
     /// Create a [ResourceBasedSlotsForType] for this instance which is willing to hand out
     /// `minimum` slots with no checks at all and `max` slots ever. Otherwise the underlying

core/src/worker/slot_supplier/tracking.rs

@@ -333,38 +333,6 @@ impl CoreWfStarter {
         &self.task_queue_name
     }
 
-    pub fn max_cached_workflows(&mut self, num: usize) -> &mut Self {
-        self.worker_config.max_cached_workflows(num);
-        self
-    }
-
-    // TODO: Eliminate these dumb pass-through ones
-    pub fn max_wft(&mut self, max: usize) -> &mut Self {
-        self.worker_config.max_outstanding_workflow_tasks(max);
-        self
-    }
-
-    pub fn max_at(&mut self, max: usize) -> &mut Self {
-        self.worker_config.max_outstanding_activities(max);
-        self
-    }
-
-    pub fn max_local_at(&mut self, max: usize) -> &mut Self {
-        self.worker_config.max_outstanding_local_activities(max);
-        self
-    }
-
-    pub fn max_at_polls(&mut self, max: usize) -> &mut Self {
-        self.worker_config.max_concurrent_at_polls(max);
-
-        self
-    }
-
-    pub fn no_remote_activities(&mut self) -> &mut Self {
-        self.worker_config.no_remote_activities(true);
-        self
-    }
-
     async fn get_or_init(&mut self) -> &InitializedWorker {
         self.initted_worker
             .get_or_init(|| async {

test-utils/src/lib.rs

@@ -84,20 +84,15 @@ async fn activity_load() {
 }
 
 #[tokio::test]
-async fn chunky_activities() {
+async fn chunky_activities_resource_based() {
     const WORKFLOWS: usize = 100;
 
-    let mut starter = CoreWfStarter::new("chunky_activities");
+    let mut starter = CoreWfStarter::new("chunky_activities_resource_based");
     starter
         .worker_config
         .max_concurrent_wft_polls(10_usize)
         .max_concurrent_at_polls(10_usize);
-    // starter
-    //     .worker_config
-    //     .max_outstanding_activities(25)
-    //     .max_outstanding_workflow_tasks(25);
-    // TODO: Fix /1 or /100 thing
-    let resource_slots = Arc::new(ResourceBasedSlots::new(0.7, 0.9));
+    let resource_slots = Arc::new(ResourceBasedSlots::new(0.7, 0.7));
     starter
         .worker_config
         .workflow_task_slot_supplier(resource_slots.as_kind(
@@ -332,7 +327,7 @@ pub async fn many_parallel_timers_longhist(ctx: WfContext) -> WorkflowResult<()>
 async fn can_paginate_long_history() {
     let wf_name = "can_paginate_long_history";
     let mut starter = CoreWfStarter::new(wf_name);
-    starter.no_remote_activities();
+    starter.worker_config.no_remote_activities(true);
     // Do not use sticky queues so we are forced to paginate once history gets long
     starter.max_cached_workflows(0);
 

tests/heavy_tests.rs

@@ -368,7 +368,7 @@ async fn query_of_closed_workflow_doesnt_tick_terminal_metric(
     let mut starter =
         CoreWfStarter::new_with_runtime("query_of_closed_workflow_doesnt_tick_terminal_metric", rt);
     // Disable cache to ensure replay happens completely
-    starter.max_cached_workflows(0);
+    starter.worker_config.max_cached_workflows(0_usize);
     let worker = starter.get_worker().await;
     let run_id = starter.start_wf().await;
     let task = worker.poll_workflow_activation().await.unwrap();