Support CPU count

Currently one queue item is one CPU count.  This change ensures that the cpu_count in the properties is used.

Author: chrisstaite-menlo

src/action_info.rs

@@ -22,7 +22,7 @@ use tokio::sync::mpsc::Receiver;
 pub(crate) type PropertyName = String;
 pub(crate) type PropertyValue = String;
 pub(crate) type PropertySet = BTreeMap<PropertyName, PropertyValue>;
-pub(crate) type OperationCount = usize;
+pub(crate) type OperationCount = u64;
 type OperationId = String;
 
 /// The name of the field in the Redis hash that stores the data.
@@ -144,7 +144,8 @@ fn monitor_changes(
     rx
 }
 
-type QueuedOperations = HashSet<OperationId>;
+type CpuCount = u64; // This is scaled by 1000.
+type QueuedOperations = HashMap<OperationId, CpuCount>;
 type ActiveOperations = HashMap<PropertySet, QueuedOperations>;
 
 #[derive(Deserialize)]
@@ -193,6 +194,7 @@ struct Operation {
     operation_id: OperationId,
     state: OperationState,
     properties: PropertySet,
+    cpu_count: CpuCount,
 }
 
 impl RedisOperationId {
@@ -254,6 +256,16 @@ impl FromValue for RedisCursorData {
     }
 }
 
+fn get_cpu_count(properties: &HashMap<PropertyName, PropertyValue>) -> CpuCount {
+    properties
+        .get("cpu_count")
+        .and_then(|count| {
+            // Round up to the nearest 1000.
+            count.parse::<u64>().ok().map(|v| v.div_ceil(1000))
+        })
+        .unwrap_or(1)
+}
+
 async fn list_operations(
     redis_client: &Pool,
     group_by: &HashSet<PropertyName>,
@@ -306,11 +318,12 @@ async fn list_operations(
                         continue;
                     }
                 };
+                let cpu_count = get_cpu_count(&redis_operation.action_info.platform_properties);
                 let properties =
                     filter_properties(group_by, redis_operation.action_info.platform_properties);
 
                 let queued_operations = active_operations.entry(properties).or_default();
-                queued_operations.insert(redis_operation.operation_id.into_string());
+                queued_operations.insert(redis_operation.operation_id.into_string(), cpu_count);
             }
         }
         if data.cursor == 0 {
@@ -359,15 +372,21 @@ async fn get_operation(
         .await
         .ok()??;
     let redis_operation: RedisOperation = serde_json::from_slice(&data).ok()?;
+    let cpu_count = get_cpu_count(&redis_operation.action_info.platform_properties);
     let properties = filter_properties(group_by, redis_operation.action_info.platform_properties);
     let state = redis_operation.state.stage.into();
     Some(Operation {
         operation_id: redis_operation.operation_id.into_string(),
         state,
         properties,
+        cpu_count,
     })
 }
 
+fn count_queue(operations: &QueuedOperations) -> u64 {
+    operations.values().sum()
+}
+
 fn operation_manager(
     redis_client: Pool,
     mut operation_channel: Receiver<OperationId>,
@@ -388,8 +407,9 @@ fn operation_manager(
                     // If there were no notifications, then just re-populate.
                     let new_active_operations = list_operations(&redis_client, &group_by, &index_name).await;
                     for (properties, queued) in &new_active_operations {
-                        tracing::info!(queue=queued.len(), properties=?properties, "Refreshed queue");
-                        if active_operations.get(properties).is_none_or(|previously_queued: &QueuedOperations| previously_queued.len() != queued.len()) && tx.send((properties.clone(), queued.len())).await.is_err() {
+                        let queue_length = count_queue(queued);
+                        tracing::info!(queue=queue_length, properties=?properties, "Refreshed queue");
+                        if active_operations.get(properties).is_none_or(|previously_queued: &QueuedOperations| count_queue(previously_queued) != queue_length) && tx.send((properties.clone(), queue_length)).await.is_err() {
                             return;
                         }
                     }
@@ -412,8 +432,9 @@ fn operation_manager(
                         // Re-populate with all operations.
                         let new_active_operations = list_operations(&redis_client, &group_by, &index_name).await;
                         for (properties, queued) in &new_active_operations {
-                            tracing::info!(queue=queued.len(), properties=?properties, "Refreshed queue");
-                            if active_operations.get(properties).is_none_or(|previously_queued: &QueuedOperations| previously_queued.len() != queued.len()) && tx.send((properties.clone(), queued.len())).await.is_err() {
+                            let queue_length = count_queue(queued);
+                            tracing::info!(queue=queue_length, properties=?properties, "Refreshed queue");
+                            if active_operations.get(properties).is_none_or(|previously_queued: &QueuedOperations| count_queue(previously_queued) != queue_length) && tx.send((properties.clone(), queue_length)).await.is_err() {
                                 return;
                             }
                         }
@@ -434,19 +455,19 @@ fn operation_manager(
                     };
                     let mut entry = match active_operations.entry(operation.properties) {
                         std::collections::hash_map::Entry::Occupied(entry) => entry,
-                        std::collections::hash_map::Entry::Vacant(entry) => entry.insert_entry(HashSet::new()),
+                        std::collections::hash_map::Entry::Vacant(entry) => entry.insert_entry(QueuedOperations::new()),
                     };
                     let queued_operations = entry.get_mut();
-                    let original_size = queued_operations.len();
+                    let original_size = count_queue(queued_operations);
                     match operation.state {
                         OperationState::Queued => {
-                            queued_operations.insert(operation.operation_id);
+                            queued_operations.insert(operation.operation_id, operation.cpu_count);
                         }
                         _ => {
                             queued_operations.remove(&operation.operation_id);
                         }
                     }
-                    let new_size = queued_operations.len();
+                    let new_size = count_queue(queued_operations);
                     if new_size != original_size && tx.send((entry.key().clone(), new_size)).await.is_err() {
                         return;
                     }

src/main.rs

@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 use std::env;
 use std::fs::File;
-use std::num::NonZeroUsize;
+use std::num::NonZeroU64;
 use std::path::Path;
 use std::pin::Pin;
 use std::time::{Duration, Instant};
@@ -29,16 +29,16 @@ struct Config {
     namespace: Option<String>,
     /// The maximum number of instances per container image, this does not
     /// include the base worker.
-    max_pods: NonZeroUsize,
+    max_pods: NonZeroU64,
     /// The number of tasks in the queue per allocated CPU.
-    queue_per_cpu: NonZeroUsize,
+    queue_per_cpu: NonZeroU64,
     /// The number of CPUs to allocate on the base worker, this is a worker
     /// that remains for two days after a property set is used.
-    base_worker_cpu: usize,
+    base_worker_cpu: action_info::OperationCount,
     /// The number of GB of RAM to allocate per CPU allocated to a worker.
-    memory_to_cpu: NonZeroUsize,
+    memory_to_cpu: NonZeroU64,
     /// The number of CPUs to allocate to a standard worker.
-    worker_cpu: NonZeroUsize,
+    worker_cpu: NonZeroU64,
 }
 
 const DOCKER_IMAGE_PREFIX: &str = "docker://";
@@ -60,7 +60,7 @@ struct WorkerInfo {
 
 #[derive(Debug)]
 struct PropertyWorkers {
-    queue_size: usize,
+    queue_size: action_info::OperationCount,
     workers: HashMap<WorkerName, WorkerInfo>,
     base_worker: Option<WorkerName>,
     last_queue_update: Instant,
@@ -71,7 +71,7 @@ impl PropertyWorkers {
         Self::new_with_workers(0, HashMap::new())
     }
 
-    fn new_with_workers(queue_size: usize, workers: HashMap<WorkerName, WorkerInfo>) -> Self {
+    fn new_with_workers(queue_size: action_info::OperationCount, workers: HashMap<WorkerName, WorkerInfo>) -> Self {
         Self {
             queue_size,
             workers,
@@ -104,7 +104,7 @@ impl PropertyWorkers {
             .map(|(name, info)| (name, info.start_time))
     }
 
-    fn update_queue(&mut self, queue_size: usize) {
+    fn update_queue(&mut self, queue_size: action_info::OperationCount) {
         if self.queue_size != 0 || queue_size != 0 {
             self.last_queue_update = Instant::now();
         }
@@ -119,8 +119,8 @@ impl PropertyWorkers {
         self.base_worker.is_none() && self.base_worker_required()
     }
 
-    fn workers_count(&self) -> usize {
-        self.workers.len()
+    fn workers_count(&self) -> action_info::OperationCount {
+        self.workers.len() as u64
     }
 
     fn set_base_worker(&mut self, name: WorkerName) -> bool {
@@ -150,19 +150,19 @@ impl PropertyWorkers {
         }
     }
 
-    fn scale_down_to(&mut self, workers: usize) -> Vec<(WorkerName, WorkerInfo)> {
-        if self.workers.len() <= workers {
+    fn scale_down_to(&mut self, workers: action_info::OperationCount) -> Vec<(WorkerName, WorkerInfo)> {
+        if self.workers_count() <= workers {
             return Vec::new();
         }
-        let to_remove = self.workers.len() - workers;
+        let to_remove = self.workers_count() - workers;
         let mut all_workers: Vec<_> = self.workers.iter().collect();
         all_workers.sort_by(|a, b| b.1.start_time.cmp(&a.1.start_time));
         let keep_names: Vec<_> = all_workers
             .iter()
-            .skip(to_remove)
+            .skip(to_remove as usize)
             .filter_map(|(name, info)| info.death_time.is_some().then_some((*name).clone()))
             .collect();
-        all_workers.truncate(to_remove);
+        all_workers.truncate(to_remove as usize);
         let now: Instant = Instant::now();
         let (to_add_death_time, workers_to_remove): (Vec<_>, Vec<_>) = all_workers
             .into_iter()
@@ -194,22 +194,22 @@ impl PropertyWorkers {
         removed_workers
     }
 
-    fn get_queue_size(&self) -> usize {
+    fn get_queue_size(&self) -> action_info::OperationCount {
         self.queue_size
     }
 }
 
 type WorkerMap = HashMap<action_info::PropertySet, PropertyWorkers>;
 
 struct ScaleInfo {
-    max_pods: usize,
-    queue_per_cpu: usize,
-    base_worker_cpu: usize,
-    worker_cpu: usize,
+    max_pods: action_info::OperationCount,
+    queue_per_cpu: action_info::OperationCount,
+    base_worker_cpu: action_info::OperationCount,
+    worker_cpu: action_info::OperationCount,
 }
 
 impl ScaleInfo {
-    fn required_workers(&self, queue_size: usize) -> usize {
+    fn required_workers(&self, queue_size: action_info::OperationCount) -> action_info::OperationCount {
         let base_queue = self.queue_per_cpu * self.base_worker_cpu;
         let required_workers = if queue_size > base_queue {
             let queue_per_pod = self.queue_per_cpu * self.worker_cpu;
@@ -265,7 +265,7 @@ async fn get_worker_pods(
         .collect()
 }
 
-async fn enumerate_existing_workers(pods: &Api<Pod>, base_cpu: usize) -> WorkerMap {
+async fn enumerate_existing_workers(pods: &Api<Pod>, base_cpu: action_info::OperationCount) -> WorkerMap {
     let pod_list = match pods.list(&ListParams::default()).await {
         Ok(list) => list,
         Err(e) => {
@@ -287,7 +287,7 @@ async fn enumerate_existing_workers(pods: &Api<Pod>, base_cpu: usize) -> WorkerM
             .as_ref()
             .and_then(|resources| resources.requests.as_ref())
             .and_then(|requests| requests.get("cpu"))
-            .and_then(|cpu| cpu.0.parse::<usize>().ok())
+            .and_then(|cpu| cpu.0.parse::<action_info::OperationCount>().ok())
             .is_some_and(|request_cpu| request_cpu == base_cpu);
         Some((image_name, name, is_base))
     });
@@ -321,9 +321,9 @@ struct WorkerManager {
     workers: WorkerMap,
     operations_change: Receiver<(action_info::PropertySet, action_info::OperationCount)>,
     watcher: Pin<Box<dyn Stream<Item = kube::runtime::watcher::Result<Event<Pod>>>>>,
-    memory_to_cpu: usize,
-    worker_cpu: usize,
-    base_worker_cpu: usize,
+    memory_to_cpu: action_info::OperationCount,
+    worker_cpu: action_info::OperationCount,
+    base_worker_cpu: action_info::OperationCount,
 }
 
 impl WorkerManager {
@@ -409,7 +409,7 @@ impl WorkerManager {
         None
     }
 
-    fn configure_pod(pod: &Pod, cpu: usize, memory: usize) -> Pod {
+    fn configure_pod(pod: &Pod, cpu: action_info::OperationCount, memory: action_info::OperationCount) -> Pod {
         let mut this_pod = pod.clone();
         // Create a unique name for this pod.
         this_pod.metadata.name = Some(format!(
@@ -453,7 +453,7 @@ impl WorkerManager {
         this_pod
     }
 
-    async fn maybe_scale_up(&mut self, properties: action_info::PropertySet, queue_size: usize) {
+    async fn maybe_scale_up(&mut self, properties: action_info::PropertySet, queue_size: action_info::OperationCount) {
         let mut workers_entry = match self.workers.entry(properties.clone()) {
             std::collections::hash_map::Entry::Occupied(mut occupied_entry) => {
                 occupied_entry.get_mut().update_queue(queue_size);