metal: prevent command buffer exhaustion deadlocks

wgpu-hal/src/metal/adapter.rs

@@ -7,6 +7,7 @@ use parking_lot::Mutex;
 use wgt::{AstcBlock, AstcChannel};
 
 use alloc::sync::Arc;
+use core::sync::atomic;
 
 use super::TimestampQuerySupport;
 
@@ -27,7 +28,7 @@ use super::TimestampQuerySupport;
 /// <https://bugzilla.mozilla.org/show_bug.cgi?id=1971452>.
 ///
 /// [new command buffer]: https://developer.apple.com/documentation/metal/mtlcommandqueue/makecommandbuffer()?language=objc
-const MAX_COMMAND_BUFFERS: u64 = 4096;
+pub(super) const MAX_COMMAND_BUFFERS: u64 = 4096;
 
 unsafe impl Send for super::Adapter {}
 unsafe impl Sync for super::Adapter {}
@@ -87,6 +88,7 @@ impl crate::Adapter for super::Adapter {
             },
             queue: super::Queue {
                 raw: Arc::new(Mutex::new(queue)),
+                command_buffer_created_not_submitted: Arc::new(atomic::AtomicU64::new(0)),
                 timestamp_period,
             },
         })

wgpu-hal/src/metal/command.rs

@@ -1,10 +1,10 @@
-use super::{conv, AsNative, TimestampQuerySupport};
+use super::{adapter, conv, AsNative, TimestampQuerySupport};
 use crate::CommandEncoder as _;
 use alloc::{
     borrow::{Cow, ToOwned as _},
     vec::Vec,
 };
-use core::ops::Range;
+use core::{ops::Range, sync::atomic};
 use metal::{
     MTLIndexType, MTLLoadAction, MTLPrimitiveType, MTLScissorRect, MTLSize, MTLStoreAction,
     MTLViewport, MTLVisibilityResultMode, NSRange,
@@ -318,6 +318,25 @@ impl crate::CommandEncoder for super::CommandEncoder {
     unsafe fn begin_encoding(&mut self, label: crate::Label) -> Result<(), crate::DeviceError> {
         let queue = &self.raw_queue.lock();
         let retain_references = self.shared.settings.retain_command_buffer_references;
+
+        // Guard against exhausting Metal's command buffer budget. Use the hard
+        // limit (`MAX_COMMAND_BUFFERS`) so we fail before Metal can hang inside
+        // `new_command_buffer`.
+        let previous = self
+            .command_buffer_created_not_submitted
+            .fetch_add(1, atomic::Ordering::AcqRel);
+        if previous >= adapter::MAX_COMMAND_BUFFERS {
+            let current = previous + 1;
+            log::warn!(
+                "metal: refusing to create new command buffer; {current} outstanding command \
+                 buffers exceeds the limit of {}. Treating this as device lost. \
+                 Ensure command encoders are submitted or dropped rather than kept alive \
+                 to avoid exhausting Metal's command buffer budget.",
+                adapter::MAX_COMMAND_BUFFERS
+            );
+            return Err(crate::DeviceError::Lost);
+        }
+
         let raw = objc::rc::autoreleasepool(move || {
             let cmd_buf_ref = if retain_references {
                 queue.new_command_buffer()
@@ -345,7 +364,14 @@ impl crate::CommandEncoder for super::CommandEncoder {
         if let Some(encoder) = self.state.compute.take() {
             encoder.end_encoding();
         }
+        let had_command_buffer = self.raw_cmd_buf.is_some();
+        // Clear the Option first so the underlying `metal::CommandBuffer` is
+        // dropped before we update the counter.
         self.raw_cmd_buf = None;
+        if had_command_buffer {
+            self.command_buffer_created_not_submitted
+                .fetch_sub(1, atomic::Ordering::AcqRel);
+        }
     }
 
     unsafe fn end_encoding(&mut self) -> Result<super::CommandBuffer, crate::DeviceError> {

wgpu-hal/src/metal/device.rs

@@ -639,6 +639,9 @@ impl crate::Device for super::Device {
         Ok(super::CommandEncoder {
             shared: Arc::clone(&self.shared),
             raw_queue: Arc::clone(&desc.queue.raw),
+            command_buffer_created_not_submitted: Arc::clone(
+                &desc.queue.command_buffer_created_not_submitted,
+            ),
             raw_cmd_buf: None,
             state: super::CommandState::default(),
             temp: super::Temp::default(),

wgpu-hal/src/metal/mod.rs

@@ -362,6 +362,10 @@ pub struct Adapter {
 
 pub struct Queue {
     raw: Arc<Mutex<metal::CommandQueue>>,
+    // Tracks command buffers created via `CommandEncoder::begin_encoding` that
+    // have not yet been submitted or discarded. Used to proactively fail
+    // before hitting Metal's `maxCommandBufferCount`.
+    command_buffer_created_not_submitted: Arc<atomic::AtomicU64>,
     timestamp_period: f32,
 }
 
@@ -372,6 +376,7 @@ impl Queue {
     pub unsafe fn queue_from_raw(raw: metal::CommandQueue, timestamp_period: f32) -> Self {
         Self {
             raw: Arc::new(Mutex::new(raw)),
+            command_buffer_created_not_submitted: Arc::new(atomic::AtomicU64::new(0)),
             timestamp_period,
         }
     }
@@ -469,6 +474,13 @@ impl crate::Queue for Queue {
 
             for cmd_buffer in command_buffers {
                 cmd_buffer.raw.commit();
+                // One command buffer per `end_encoding` call moves from the
+                // "created but not yet submitted" bucket into the submitted
+                // set, so update the counter.
+                let previous = self
+                    .command_buffer_created_not_submitted
+                    .fetch_sub(1, atomic::Ordering::AcqRel);
+                debug_assert!(previous > 0);
             }
 
             if let Some(raw) = extra_command_buffer {
@@ -1012,6 +1024,9 @@ struct CommandState {
 pub struct CommandEncoder {
     shared: Arc<AdapterShared>,
     raw_queue: Arc<Mutex<metal::CommandQueue>>,
+    // Tracks how many command buffers have been created via this queue but not
+    // yet submitted or discarded. Shared across all encoders for the same queue.
+    command_buffer_created_not_submitted: Arc<atomic::AtomicU64>,
     raw_cmd_buf: Option<metal::CommandBuffer>,
     state: CommandState,
     temp: Temp,