Do not lose wakeups for block device workers (#973)

The QueueMinder's `notify_workers` is the sole indicator that a worker
thread is asleep or not, so `take_notifications` transfers the
responsibility to handle these bits to the caller, which
`flush_notifications` did not fully handle.

Return remaining bits indicating idle workers to QueueMinder, so when
there's enough work to justify waking them, we actually do.

In the process, fix a less exciting bug where attempting to spuriously
wake already-woken worker threads would result in waking fewer workers
than desired.

And, along the way, add a few new probes that were useful in debugging.

Author: iximeow

lib/propolis/src/block/attachment.rs

@@ -74,27 +74,57 @@ impl QueueSlot {
             let _ = self.notify_count.fetch_add(count.get(), Ordering::Release);
         }
     }
-    fn take_notifications(&self) -> Option<(NonZeroUsize, Bitmap)> {
-        let state = self.state.lock().unwrap();
-        let minder = state.minder.as_ref()?;
-
-        if self.notify_count.load(Ordering::Relaxed) == 0 {
-            return None;
-        }
-        let wake_wids = minder.take_notifications()?;
-
-        let pending = self.notify_count.swap(0, Ordering::Acquire);
-        let count =
-            NonZeroUsize::new(pending).expect("notify count is non-zero");
-        Some((count, wake_wids))
-    }
     fn flush_notifications(&self) {
         let guard = self.workers.lock().unwrap();
         let Some(workers) = guard.as_ref() else {
             return;
         };
-        if let Some((count, wake_wids)) = self.take_notifications() {
-            workers.wake(wake_wids, Some(count), Some(self.queue_id));
+
+        let state = self.state.lock().unwrap();
+        let Some(minder) = state.minder.as_ref() else {
+            // The queue isn't associated with anything yet, so there are no
+            // interested workers to wake.
+            return;
+        };
+
+        let pending = self.notify_count.swap(0, Ordering::AcqRel);
+
+        let Some(pending) = NonZeroUsize::new(pending) else {
+            // We have not been asked to wake any workers since the last
+            // `flush_notifications`. This is relatively unlikely but
+            // legitimate, such as if this `QueueSlot` paused and resumed (as
+            // for migrations) repeatedly.
+            return;
+        };
+
+        // Take the full set of workers that may be idle and interested in this
+        // queue. At this point we are responsible for either waking workers
+        // here, or returning the idle-and-interested bit to `minder`.
+        let Some(wake_wids) = minder.take_notifications() else {
+            // `notify_count` was non-zero, but between checking the notify
+            // count and getting idle workers, we started pausing devices.
+            // Bummer. Request notification of as many workers as we were
+            // going to, and let a future `flush_notifications()` take care of
+            // it.
+            self.request_notify(Some(pending));
+            return;
+        };
+        drop(state);
+
+        let remaining_wids =
+            workers.wake(wake_wids, pending, Some(self.queue_id));
+
+        if !remaining_wids.is_empty() {
+            let state = self.state.lock().unwrap();
+            let Some(minder) = state.minder.as_ref() else {
+                // The queue no longer has a minder. This is unfortunate, but it
+                // is at least OK to discard `remaining_wids` here: if this
+                // queue is reassociated later, updating the queue collection's
+                // associations will wake all queues.
+                return;
+            };
+
+            minder.add_notifications(remaining_wids);
         }
     }
 }
@@ -975,17 +1005,33 @@ impl WorkerCollection {
     fn wake(
         &self,
         wake_wids: Bitmap,
-        limit: Option<NonZeroUsize>,
+        limit: NonZeroUsize,
         qid_hint: Option<QueueId>,
-    ) {
-        let _num_woke = wake_wids
-            .iter()
-            .take(limit.unwrap_or(MAX_WORKERS).get())
-            .filter_map(|wid| {
-                let slot = self.workers.get(wid)?;
-                slot.wake(None, qid_hint).then_some(())
-            })
-            .count();
+    ) -> Bitmap {
+        probes::block_worker_collection_wake!(|| (wake_wids.0, limit.get()));
+
+        let mut num_woken = 0;
+        let mut idle_wids = wake_wids.iter();
+
+        for wid in &mut idle_wids {
+            let Some(slot) = self.workers.get(wid) else {
+                continue;
+            };
+
+            if slot.wake(None, qid_hint) {
+                num_woken += 1;
+            }
+
+            if num_woken == limit.get() {
+                break;
+            }
+        }
+
+        let remainder = idle_wids.remainder();
+
+        probes::block_worker_collection_woken!(|| (remainder.0, num_woken));
+
+        remainder
     }
     fn update_queue_associations(&self, queues_associated: Versioned<Bitmap>) {
         let mut state = self.state.lock().unwrap();
@@ -1314,6 +1360,9 @@ impl Bitmap {
         assert!(idx < Self::TOP_BIT);
         self.0 &= !(1u64 << idx);
     }
+    pub fn set_all(&mut self, other: Bitmap) {
+        self.0 |= other.0;
+    }
     pub fn lowest_set(&self) -> Option<usize> {
         if self.0.count_ones() == 0 {
             None
@@ -1352,6 +1401,11 @@ impl Iterator for BitIter {
         Some(idx)
     }
 }
+impl BitIter {
+    fn remainder(self) -> Bitmap {
+        self.0
+    }
+}
 pub struct LoopIter {
     cur: Bitmap,
     orig: Bitmap,

lib/propolis/src/block/minder.rs

@@ -194,7 +194,7 @@ impl QueueMinder {
     /// that more requests are available.
     pub fn next_req(&self, wid: WorkerId) -> Option<DeviceRequest> {
         let mut state = self.state.lock().unwrap();
-        if state.paused || !state.notify_workers.is_empty() {
+        if state.paused {
             state.notify_workers.set(wid);
             return None;
         }
@@ -315,9 +315,14 @@ impl QueueMinder {
         }
     }
 
-    /// Get a bitmap of the workers which should be notified that this queue may
-    /// now have requests available.
-    pub(crate) fn take_notifications(&self) -> Option<Bitmap> {
+    /// Take the bitmap of the workers which should be notified that this queue
+    /// may now have requests available.
+    ///
+    /// Bits in this map correspond to workers that either should be
+    /// [`WorkerSlot::wake`]'d or returned to this `QueueMinder` via
+    /// [`add_notifications`]. Failure to do so will result in idle workers
+    /// never being woken for future work.
+    pub(in crate::block) fn take_notifications(&self) -> Option<Bitmap> {
         let mut state = self.state.lock().unwrap();
         if state.paused {
             state.notify_workers = Bitmap::ALL;
@@ -327,6 +332,18 @@ impl QueueMinder {
         }
     }
 
+    /// Add a set of workers to be notified when this queue may have requests
+    /// available.
+    ///
+    /// This should only be called with the remaining parts of a bitmap obtained
+    /// from an ealier [`take_notifications`]. Using other bit patterns may
+    /// result in wakeups to out-of-range worker IDs and subsequent panic.
+    pub(in crate::block) fn add_notifications(&self, worker_ids: Bitmap) {
+        let mut state = self.state.lock().unwrap();
+
+        state.notify_workers.set_all(worker_ids);
+    }
+
     /// Associate a [MetricConsumer] with this queue.
     ///
     /// It will be notified about each IO completion as they occur.

lib/propolis/src/block/mod.rs

@@ -86,6 +86,9 @@ mod probes {
     fn block_wake(dev_id: u32, worker_id: u64) {}
     fn block_notify(devq_id: u64) {}
     fn block_strategy(dev_id: u32, strat: String, generation: u64) {}
+
+    fn block_worker_collection_wake(wake_wids: u64, limit: usize) {}
+    fn block_worker_collection_woken(remaining_wids: u64, num_woken: usize) {}
 }
 
 /// Type of operations which may be issued to a virtual block device.

lib/propolis/src/hw/nvme/mod.rs

@@ -37,6 +37,7 @@ mod probes {
     fn nvme_doorbell_admin_cq(val: u16) {}
     fn nvme_doorbell_admin_sq(val: u16) {}
     fn nvme_admin_cmd(opcode: u8, prp1: u64, prp2: u64) {}
+    fn nvme_block_notify(sqid: u16, block_qid: u16, occupied_hint: u16) {}
 }
 
 /// The max number of MSI-X interrupts we support
@@ -1186,8 +1187,14 @@ impl PciNvme {
                     Some((sqid, sq.num_occupied()))
                 })
             {
+                let block_qid = queue::sqid_to_block_qid(sqid);
+                probes::nvme_block_notify!(|| (
+                    sqid,
+                    u16::from(block_qid),
+                    num_occupied
+                ));
                 self.block_attach.notify(
-                    queue::sqid_to_block_qid(sqid),
+                    block_qid,
                     NonZeroUsize::new(num_occupied as usize),
                 );
             }
@@ -1204,10 +1211,14 @@ impl PciNvme {
             drop(guard);
 
             assert_ne!(qid, queue::ADMIN_QUEUE_ID);
-            self.block_attach.notify(
-                queue::sqid_to_block_qid(qid),
-                NonZeroUsize::new(num_occupied as usize),
-            );
+            let block_qid = queue::sqid_to_block_qid(qid);
+            probes::nvme_block_notify!(|| (
+                qid,
+                u16::from(block_qid),
+                num_occupied
+            ));
+            self.block_attach
+                .notify(block_qid, NonZeroUsize::new(num_occupied as usize));
         };
         Ok(())
     }