[scheduler] Fix race condition in multi-scheduler environments. (AcademySoftwareFoundation#2191)

When multiple instances or async tasks are updating the same host, there
was a race condition that would cause over booking errors instead of
gracefully treating the condition that caused the race.

This PR is intended to reduce the number database errors like the
following:

```
WARN cue_scheduler::pipeline::dispatcher::actor: Wasn't able to dispatch all frames: FailedToUpdateResources( × (57b07e8d-6813-4fc5-8f71-132e968f4e29) Failed to update host resources 
├─▶ error returned from database: unable to allocate additional memory 
╰─▶ unable to allocate additional memory
)
```

### AI Disclosure

Claude Code Opus was used to brainstorm the root cause of this problem
and also to review the changes before pushing the PR.

Author: DiegoTavares

.gitignore

@@ -46,3 +46,7 @@ CLAUDE.md
 docs/nav_order_index.txt
 rust/.cargo/config.toml
 sandbox/pgadmin-data/*
+.claude/settings.json
+plans/*
+.entire/.gitignore
+.entire/settings.json

rust/crates/scheduler/src/dao/host_dao.rs

@@ -17,6 +17,7 @@ use chrono::{DateTime, Utc};
 use miette::{Context, IntoDiagnostic, Result};
 use opencue_proto::host::ThreadMode;
 use sqlx::{Pool, Postgres, Transaction};
+use thiserror::Error;
 use tracing::trace;
 use uuid::Uuid;
 
@@ -27,6 +28,45 @@ use crate::{
     pgpool::connection_pool,
 };
 
+#[derive(Debug, Error)]
+pub enum HostDaoError {
+    #[error("Host resources exhausted (likely booked by another scheduler)")]
+    HostResourcesExhausted,
+
+    #[error("Resource limit exceeded: {message}")]
+    ResourceLimitExceeded { message: String },
+
+    #[error("{context}: {source}")]
+    DbFailure {
+        context: &'static str,
+        source: sqlx::Error,
+    },
+}
+
+/// Known PostgreSQL trigger messages that indicate expected resource limit enforcement.
+const RESOURCE_LIMIT_MESSAGES: &[&str] = &[
+    "job has exceeded max cores",
+    "job has exceeded max GPU units",
+    "subscription has exceeded burst size",
+    "unable to allocate additional core units",
+    "unable to allocate additional memory",
+    "unable to allocate additional GPU units",
+    "unable to allocate additional GPU memory",
+];
+
+fn check_resource_limit_error(err: sqlx::Error, context: &'static str) -> HostDaoError {
+    let msg = err.to_string();
+    match RESOURCE_LIMIT_MESSAGES.iter().find(|m| msg.contains(*m)) {
+        Some(m) => HostDaoError::ResourceLimitExceeded {
+            message: m.to_string(),
+        },
+        None => HostDaoError::DbFailure {
+            context,
+            source: err,
+        },
+    }
+}
+
 /// Data Access Object for host operations in the job dispatch system.
 ///
 /// Manages database operations related to render hosts, including:
@@ -114,16 +154,19 @@ impl From<HostModel> for Host {
 //  - memory and core fields on table host are only updated when booking procs (update_host_resources)
 //  - the table host_stat contains memory fields that are updated by cuebot on HostReportHandler
 //
-// In summary, use host_stat for most up to date memory stats
+// We use LEAST(h.int_mem_idle, hs.int_mem_free) to get the most conservative memory estimate.
+// h.int_mem_idle reflects bookings (decremented on dispatch), while hs.int_mem_free reflects
+// actual OS-reported free memory. Using the minimum avoids overestimating available memory
+// when another scheduler (e.g. Cuebot) has booked resources that haven't been consumed yet.
 static QUERY_HOST_BY_SHOW_FACILITY_AND_TAG: &str = r#"
 SELECT DISTINCT
     h.pk_host,
     h.str_name,
     hs.str_os,
     h.int_cores_idle,
-    hs.int_mem_free,
+    LEAST(h.int_mem_idle, hs.int_mem_free) as int_mem_free,
     h.int_gpus_idle,
-    hs.int_gpu_mem_free,
+    LEAST(h.int_gpu_mem_idle, hs.int_gpu_mem_free) as int_gpu_mem_free,
     h.int_cores,
     hs.int_mem_total,
     h.int_thread_mode,
@@ -149,6 +192,10 @@ SET int_cores_idle = int_cores_idle - $1,
     int_gpus_idle = int_gpus_idle - $3,
     int_gpu_mem_idle = int_gpu_mem_idle - $4
 WHERE pk_host = $5
+    AND int_cores_idle >= $1
+    AND int_mem_idle >= $2
+    AND int_gpus_idle >= $3
+    AND int_gpu_mem_idle >= $4
 RETURNING int_cores_idle, int_mem_idle, int_gpus_idle, int_gpu_mem_idle, NOW()
 "#;
 
@@ -324,24 +371,22 @@ impl HostDao {
         transaction: &mut Transaction<'_, Postgres>,
         host_id: &Uuid,
         virtual_proc: &VirtualProc,
-        dispatch_id: Uuid,
-    ) -> Result<UpdatedHostResources> {
-        let (cores_idle, mem_idle, gpus_idle, gpu_mem_idle, last_updated): (
-            i64,
-            i64,
-            i64,
-            i64,
-            DateTime<Utc>,
-        ) = sqlx::query_as(UPDATE_HOST_RESOURCES)
-            .bind(virtual_proc.cores_reserved.value())
-            .bind((virtual_proc.memory_reserved.as_u64() / KB) as i64)
-            .bind(virtual_proc.gpus_reserved as i32)
-            .bind(virtual_proc.gpu_memory_reserved.as_u64() as i64)
-            .bind(host_id.to_string())
-            .fetch_one(&mut **transaction)
-            .await
-            .into_diagnostic()
-            .wrap_err(format!("({dispatch_id}) Failed to update host resources"))?;
+    ) -> Result<UpdatedHostResources, HostDaoError> {
+        let row: Option<(i64, i64, i64, i64, DateTime<Utc>)> =
+            sqlx::query_as(UPDATE_HOST_RESOURCES)
+                .bind(virtual_proc.cores_reserved.value())
+                .bind((virtual_proc.memory_reserved.as_u64() / KB) as i64)
+                .bind(virtual_proc.gpus_reserved as i32)
+                .bind(virtual_proc.gpu_memory_reserved.as_u64() as i64)
+                .bind(host_id.to_string())
+                .fetch_optional(&mut **transaction)
+                .await
+                .map_err(|err| {
+                    check_resource_limit_error(err, "Failed to update host resources")
+                })?;
+
+        let (cores_idle, mem_idle, gpus_idle, gpu_mem_idle, last_updated) =
+            row.ok_or(HostDaoError::HostResourcesExhausted)?;
 
         if CONFIG.host_cache.update_stat_on_book {
             sqlx::query(UPDATE_HOST_STAT)
@@ -350,8 +395,7 @@ impl HostDao {
                 .bind(host_id.to_string())
                 .execute(&mut **transaction)
                 .await
-                .into_diagnostic()
-                .wrap_err("Failed to update host stat")?;
+                .map_err(|err| check_resource_limit_error(err, "Failed to update host stat"))?;
         }
 
         sqlx::query(UPDATE_SUBSCRIPTION)
@@ -361,35 +405,33 @@ impl HostDao {
             .bind(virtual_proc.alloc_id.to_string())
             .execute(&mut **transaction)
             .await
-            .into_diagnostic()
-            .wrap_err("Failed to update subscription resources")?;
+            .map_err(|err| {
+                check_resource_limit_error(err, "Failed to update subscription resources")
+            })?;
 
         sqlx::query(UPDATE_LAYER_RESOURCE)
             .bind(virtual_proc.cores_reserved.value())
             .bind(virtual_proc.gpus_reserved as i32)
             .bind(virtual_proc.layer_id.to_string())
             .execute(&mut **transaction)
             .await
-            .into_diagnostic()
-            .wrap_err("Failed to update layer resources")?;
+            .map_err(|err| check_resource_limit_error(err, "Failed to update layer resources"))?;
 
         sqlx::query(UPDATE_JOB_RESOURCE)
             .bind(virtual_proc.cores_reserved.value())
             .bind(virtual_proc.gpus_reserved as i32)
             .bind(virtual_proc.job_id.to_string())
             .execute(&mut **transaction)
             .await
-            .into_diagnostic()
-            .wrap_err("Failed to update job resources")?;
+            .map_err(|err| check_resource_limit_error(err, "Failed to update job resources"))?;
 
         sqlx::query(UPDATE_FOLDER_RESOURCE)
             .bind(virtual_proc.cores_reserved.value())
             .bind(virtual_proc.gpus_reserved as i32)
             .bind(virtual_proc.job_id.to_string())
             .execute(&mut **transaction)
             .await
-            .into_diagnostic()
-            .wrap_err("Failed to update folder resources")?;
+            .map_err(|err| check_resource_limit_error(err, "Failed to update folder resources"))?;
 
         sqlx::query(UPDATE_POINT)
             .bind(virtual_proc.cores_reserved.value())
@@ -398,8 +440,7 @@ impl HostDao {
             .bind(virtual_proc.show_id.to_string())
             .execute(&mut **transaction)
             .await
-            .into_diagnostic()
-            .wrap_err("Failed to update point resources")?;
+            .map_err(|err| check_resource_limit_error(err, "Failed to update point resources"))?;
 
         Ok(UpdatedHostResources {
             cores_idle,

rust/crates/scheduler/src/dao/mod.rs

@@ -29,4 +29,5 @@ pub use proc_dao::ProcDao;
 
 pub use allocation_dao::{AllocationName, ShowId};
 pub use frame_dao::FrameDaoError;
+pub use host_dao::HostDaoError;
 pub use host_dao::UpdatedHostResources;

rust/crates/scheduler/src/pipeline/dispatcher/actor.rs

@@ -347,16 +347,40 @@ impl RqdDispatcherService {
                             if last_error.is_some() {
                                 break;
                             }
-                            warn!(
-                                "({dispatch_id}) Failed to start frame {} on Db. {}",
-                                frame_str, err
-                            );
+                            match &err {
+                                DispatchError::FailedToUpdateResources(_)
+                                | DispatchError::FailedToCreateProc { .. } => {
+                                    // Resource contention during DB updates is expected in
+                                    // multi-scheduler environments.
+                                    info!(
+                                        "({dispatch_id}) Failed to start frame {} on Db. {}",
+                                        frame_str, err
+                                    );
+                                }
+                                _ => {
+                                    warn!(
+                                        "({dispatch_id}) Failed to start frame {} on Db. {}",
+                                        frame_str, err
+                                    );
+                                }
+                            }
                             last_error = Some(err);
                             // IMPORTANT: Do NOT update last_host_version here since the transaction
                             // rolled back and we didn't actually consume any resources from the database.
                             // The next iteration should use the same host state as before.
                             continue;
                         }
+                        DispatchVirtualProcError::HostResourcesExhausted => {
+                            // Host resources were booked by another scheduler (e.g. Cuebot)
+                            // between cache refresh and dispatch. Break the loop and try another host.
+                            info!(
+                                "({dispatch_id}) Host resources exhausted for frame {} (likely booked by another scheduler)",
+                                frame_str
+                            );
+                            last_error = Some(DispatchError::HostResourcesExhausted(()));
+
+                            break;
+                        }
                         DispatchVirtualProcError::FrameCouldNotBeUpdated => {
                             // The entire transaction is probably compromised, stop working on this layer
                             info!(
@@ -366,6 +390,13 @@ impl RqdDispatcherService {
                             non_retrieable_frames.push(frame.id);
                             break;
                         }
+                        DispatchVirtualProcError::ResourceLimitExceeded(err) => {
+                            // Resource limit enforced by database trigger (e.g. job max cores,
+                            // subscription burst size). This is expected in normal operation.
+                            info!("({dispatch_id}) {frame_str} {err}");
+                            last_error = Some(err);
+                            break;
+                        }
                         DispatchVirtualProcError::RqdConnectionFailed { host, error } => {
                             // An error here means connection with this host is probably broken,
                             // there's no reason to attempt the next frame
@@ -403,7 +434,18 @@ impl RqdDispatcherService {
         }
 
         if let Some(error) = last_error {
-            warn!("Wasn't able to dispatch all frames: {:?}", error)
+            match &error {
+                DispatchError::ResourceLimitExceeded(_)
+                | DispatchError::AllocationOverBurst(_)
+                | DispatchError::HostResourcesExhausted(_)
+                | DispatchError::FailedToUpdateResources(_)
+                | DispatchError::FailedToCreateProc { .. } => {
+                    info!("Wasn't able to dispatch all frames: {:?}", error)
+                }
+                _ => {
+                    warn!("Wasn't able to dispatch all frames: {:?}", error)
+                }
+            }
         }
         Ok((last_host_version, layer))
     }
@@ -452,7 +494,7 @@ impl RqdDispatcherService {
 
         // Update database
         let updated_resources = self
-            .update_database_for_dispatch(transaction, &virtual_proc, host.id, dispatch_id)
+            .update_database_for_dispatch(transaction, &virtual_proc, host.id)
             .await?;
 
         // When running on dry_run_mode, just log the outcome
@@ -501,7 +543,6 @@ impl RqdDispatcherService {
     /// * `transaction` - Database transaction for atomic updates
     /// * `virtual_proc` - The virtual proc being dispatched
     /// * `host_id` - ID of the host receiving the dispatch
-    /// * `dispatch_id` - Unique identifier for this dispatch operation
     ///
     /// # Returns
     /// On success, returns the updated host resources from the database.
@@ -511,7 +552,6 @@ impl RqdDispatcherService {
         transaction: &mut Transaction<'_, Postgres>,
         virtual_proc: &VirtualProc,
         host_id: Uuid,
-        dispatch_id: Uuid,
     ) -> Result<UpdatedHostResources, DispatchVirtualProcError> {
         self.frame_dao
             .update_frame_started(transaction, virtual_proc)
@@ -538,12 +578,31 @@ impl RqdDispatcherService {
 
         let updated_resources = self
             .host_dao
-            .update_resources(transaction, &host_id, virtual_proc, dispatch_id)
+            .update_resources(transaction, &host_id, virtual_proc)
             .await
-            .map_err(|err| {
-                DispatchVirtualProcError::FailedToStartOnDb(DispatchError::FailedToUpdateResources(
-                    err,
-                ))
+            .map_err(|err| match err {
+                crate::dao::HostDaoError::HostResourcesExhausted => {
+                    DispatchVirtualProcError::HostResourcesExhausted
+                }
+                crate::dao::HostDaoError::ResourceLimitExceeded { message } => {
+                    DispatchVirtualProcError::ResourceLimitExceeded(
+                        DispatchError::ResourceLimitExceeded(message),
+                    )
+                }
+                crate::dao::HostDaoError::DbFailure { context, source } => {
+                    DispatchVirtualProcError::FailedToStartOnDb(
+                        DispatchError::FailedToUpdateResources(
+                            // It would be great to simply use
+                            // `miette::Report::new(source).wrap_err(context)` here, but
+                            // sqlx::Error doesn't implement Diagnostics, therefore needs to be
+                            // wrapped and unwrapped to add context.
+                            Result::<(), _>::Err(source)
+                                .into_diagnostic()
+                                .unwrap_err()
+                                .wrap_err(context),
+                        ),
+                    )
+                }
             })?;
 
         Ok(updated_resources)

rust/crates/scheduler/src/pipeline/dispatcher/error.rs

@@ -51,6 +51,12 @@ pub enum DispatchError {
 
     #[error("DispatchError: Failed to execute command on GRPC interface")]
     GrpcFailure(tonic::Status),
+
+    #[error("DispatchError: Host Resources exhausted before updating the database")]
+    HostResourcesExhausted(()),
+
+    #[error("DispatchError: Resource limit exceeded: {0}")]
+    ResourceLimitExceeded(String),
 }
 
 #[derive(Debug, Error, Diagnostic)]
@@ -61,9 +67,15 @@ pub enum DispatchVirtualProcError {
     #[error("Failed to start frame on database")]
     FailedToStartOnDb(DispatchError),
 
+    #[error("Host resources exhausted (likely booked by another scheduler)")]
+    HostResourcesExhausted,
+
     #[error("Failed to lock frame on database")]
     FrameCouldNotBeUpdated,
 
     #[error("Failed to connect to RQD on host {host}")]
     RqdConnectionFailed { host: String, error: Error },
+
+    #[error("Resource limit exceeded")]
+    ResourceLimitExceeded(DispatchError),
 }

rust/crates/scheduler/src/pipeline/matcher.rs

@@ -485,6 +485,18 @@ impl MatchingService {
                     ))
                 );
             }
+            DispatchError::HostResourcesExhausted(_) => {
+                info!(
+                    "Host resources exhausted before updating database when dispatching {} on {}.",
+                    layer_display, host
+                );
+            }
+            DispatchError::ResourceLimitExceeded(msg) => {
+                info!(
+                    "Resource limit exceeded when dispatching {} on {}: {}",
+                    layer_display, host, msg
+                );
+            }
         }
     }
 }