Improve can_consume check

jschwe · jschwe · commit 07b1bf561ed4 · 2023-02-21T15:54:22.000+01:00
diff --git a/bwosqueue/src/lib.rs b/bwosqueue/src/lib.rs
@@ -87,10 +87,17 @@ pub struct Stealer<E, const NUM_BLOCKS: usize, const ENTRIES_PER_BLOCK: usize> {
     queue: Pin<Arc<BwsQueue<E, NUM_BLOCKS, ENTRIES_PER_BLOCK>>>,
 }
 
-/// An iterator over elements of one Block
+/// An iterator over elements of one Block.
+///
+/// The iterator borrows all elements up to `committed` to allows batched
+/// operations on the elements. When the iterator is dropped the entries
+/// are marked as consumed in one atomic operation.
 pub struct BlockIter<'a, E, const ENTRIES_PER_BLOCK: usize> {
     buffer: &'a [UnsafeCell<MaybeUninit<E>>; ENTRIES_PER_BLOCK],
+    /// Index if the next to be consumed entry in the buffer.
     i: usize,
+    /// Number of committed entries in the buffer.
+    committed: usize,
 }
 
 /// An iterator over elements of one Block of a stealer
@@ -263,107 +270,61 @@ impl<E, const NUM_BLOCKS: usize, const ENTRIES_PER_BLOCK: usize>
     /// Try to dequeue the oldest element in the queue.
     #[inline(always)]
     pub fn dequeue(&mut self) -> Option<E> {
-        loop {
-            // SAFETY: `ccache` always points to a valid `Block` in the queue. We never create a mutable reference
-            // to a Block, so it is safe to construct a shared reference here.
-            let blk = unsafe { &**self.ccache };
 
-            // check if the block is fully consumed already
-            let consumed = blk.consumed.load(Relaxed);
-            let consumed_idx = consumed.raw_index();
+        let (blk, consumed_opt) = self.next_dequeable_entry();
+        *self.ccache = blk;
+        let consumed = consumed_opt?;
 
-            // Fastpath (Block is not fully consumed yet)
-            if let Some(entry_cell) = blk.entries.get(consumed_idx) {
-                // we know the block is not full, but most first check if there is an entry to
-                // dequeue.
-                let committed_idx = blk.committed.load(Relaxed).raw_index();
-                if consumed_idx == committed_idx {
-                    return None;
-                }
+        // To assign to `ccache` we need to end the borrow of the previous `blk` variable.
+        let blk = unsafe { &**self.ccache };
 
-                /* There is an entry to dequeue */
+        // We trust that the correct index is passed to us here.
+        let entry_cell = &blk.entries[consumed.raw_index()];
+        // SAFETY: We know there is an entry to dequeue, so we know the entry is a valid initialized `E`.
+        let item = unsafe { entry_cell.with(|entry| entry.read().assume_init()) };
+        // SAFETY: We already checked that `consumed_idx < ENTRIES_PER_BLOCK`.
+        let new_consumed = unsafe { consumed.index_add_unchecked(1) };
+        blk.consumed.store(new_consumed, Relaxed);
+        #[cfg(feature = "stats")]
+        self.queue.stats.increment_dequeued(1);
+        return Some(item);
 
-                // SAFETY: We know there is an entry to dequeue, so we know the entry is a valid initialized `E`.
-                let item = unsafe { entry_cell.with(|entry| entry.read().assume_init()) };
-                // SAFETY: We already checked that `consumed_idx < ENTRIES_PER_BLOCK`.
-                let new_consumed = unsafe { consumed.index_add_unchecked(1) };
-                blk.consumed.store(new_consumed, Relaxed);
-                #[cfg(feature = "stats")]
-                self.queue.stats.increment_dequeued(1);
-                return Some(item);
-            }
-
-            /* Slow-path */
-
-            /* Consumer head may never pass the Producer head and Consumer/Stealer tail */
-            let nblk = unsafe { &*blk.next() };
-            if self.try_advance_consumer_block(nblk, consumed).is_err() {
-                return None;
-            }
-            /* We advanced to the next block - loop around and try again */
-        }
     }
 
-    /// Tru to dequeue a whole block
+    /// Try to dequeue all remaining committed entries in the current block.
     pub fn dequeue_block(&mut self) -> Option<BlockIter<'_, E, ENTRIES_PER_BLOCK>> {
-        loop {
-            // SAFETY: `ccache` always points to a valid `Block` in the queue. We never create a mutable reference
-            // to a Block, so it is safe to construct a shared reference here.
-            let blk = unsafe { &**self.ccache };
-
-            // check if the block is fully consumed already
-            let consumed = blk.consumed.load(Relaxed);
-            let consumed_idx = consumed.raw_index();
-
-            if consumed_idx < ENTRIES_PER_BLOCK {
-                // for now just return none. We could also return  consumed_idx..committed_idx
-                if !(blk.committed.load(Relaxed).index().is_full()) {
-                    return None;
-                }
-
-                // We are claiming the tasks **before** reading them out of the buffer.
-                // This is safe because only the **current** thread is able to push new
-                // tasks.
-                //
-                // There isn't really any need for memory ordering... Relaxed would
-                // work. This is because all tasks are pushed into the queue from the
-                // current thread (or memory has been acquired if the local queue handle
-                // moved).
-                let new_consumed = consumed.set_full();
-                blk.consumed.store(new_consumed, Relaxed);
-                #[cfg(feature = "stats")]
-                self.queue
-                    .stats
-                    .increment_dequeued(new_consumed.raw_index() - consumed_idx);
-
-                // Pre-advance ccache for the next time
-                let nblk = unsafe { &*blk.next() };
-                // We don't care if this fails. The consumer can try again next time.
-                let _ = self.try_advance_consumer_block(nblk, new_consumed);
-
-                return Some(BlockIter {
-                    buffer: &blk.entries,
-                    i: consumed_idx,
-                });
-            }
-
-            /* Slow-path */
-
-            /* Consumer head may never pass the Producer head and Consumer/Stealer tail */
-            let nblk = unsafe { &*blk.next() };
-            if self.try_advance_consumer_block(nblk, consumed).is_err() {
-                return None;
-            }
-
-            /* We advanced to the next block - loop around and try again */
-        }
+        let (blk, consumed_opt) = self.next_dequeable_entry();
+        *self.ccache = blk;
+        let consumed = consumed_opt?;
+
+        // To assign to `ccache` we need to end the borrow of the previous `blk` variable.
+        let blk = unsafe { &**self.ccache };
+
+        let committed = blk.committed.load(Relaxed);
+
+        // We are claiming the tasks **before** reading them out of the buffer.
+        // This is safe because only the **current** thread is able to push new
+        // tasks.
+        //
+        // There isn't really any need for memory ordering... Relaxed would
+        // work. This is because all tasks are pushed into the queue from the
+        // current thread (or memory has been acquired if the local queue handle
+        // moved).
+        blk.consumed.store(committed, Relaxed);
+
+        return Some(BlockIter {
+            buffer: &blk.entries,
+            i: consumed.raw_index(),
+            committed: committed.raw_index()
+        });
     }
-    /// Advance consumer to the next block, unless the producer has not reached the block yet.
-    fn try_advance_consumer_block(
-        &mut self,
+
+     /// Advance consumer to the next block, unless the producer has not reached the block yet.
+     fn can_advance_consumer_block(
+        &self,
         next_block: &Block<E, ENTRIES_PER_BLOCK>,
         curr_consumed: IndexAndVersion<ENTRIES_PER_BLOCK>,
-    ) -> Result<(), ()> {
+    ) -> bool {
         let next_cons_vsn = curr_consumed
             .version()
             .wrapping_add(next_block.is_head() as usize);
@@ -373,7 +334,7 @@ impl<E, const NUM_BLOCKS: usize, const ENTRIES_PER_BLOCK: usize>
         // and we must wait.
         let next_reserved_vsn = next_block.reserved.load(Relaxed).version();
         if next_reserved_vsn != next_cons_vsn {
-            return Err(());
+            return false;
         }
 
         /* stop stealers */
@@ -394,10 +355,24 @@ impl<E, const NUM_BLOCKS: usize, const ENTRIES_PER_BLOCK: usize>
         /* advance the block and try again */
         // The consumer must skip already reserved entries.
         next_block.consumed.store(reserved_old, Relaxed);
-        *self.ccache = next_block;
-        Ok(())
+        true
     }
 
+    // /// Advance consumer to the next block, unless the producer has not reached the block yet.
+    // fn try_advance_consumer_block(
+    //     &mut self,
+    //     next_block: &Block<E, ENTRIES_PER_BLOCK>,
+    //     curr_consumed: IndexAndVersion<ENTRIES_PER_BLOCK>,
+    // ) -> Result<(), ()> {
+    //     if self.can_advance_consumer_block(next_block, curr_consumed) {
+    //         *self.ccache = next_block;
+    //         Ok(())
+    //     } else {
+    //         Err(())
+    //     }
+    // }
+
+    /// Todo: Ideally we would not have this function.
     pub fn has_stealers(&self) -> bool {
         let curr_spos = self.spos.load(Relaxed);
         // spos increments beyond NUM_BLOCKS to prevent ABA problems.
@@ -419,6 +394,8 @@ impl<E, const NUM_BLOCKS: usize, const ENTRIES_PER_BLOCK: usize>
     /// Check if there is a block available for stealing in the queue.
     ///
     /// Note that stealing may still fail for a number of reasons even if this function returned true
+    /// Todo: the overhead could be reduced, if we allow this function to return false in some
+    /// cases when the queue size is low.
     #[cfg(feature = "stats")]
     pub fn has_stealable_block(&self) -> bool {
         let n = self.queue.stats.curr_enqueued();
@@ -430,12 +407,57 @@ impl<E, const NUM_BLOCKS: usize, const ENTRIES_PER_BLOCK: usize>
         n > (committed_idx - consumed_idx)
     }
 
-    /// Check if there are items in the queue available for the consumer.
+    /// `true` if there is at least one entry that can be dequeued.
     ///
-    /// This function may sporadically provide a wrong result.
-    #[cfg(feature = "stats")]
+    /// This function is accurate and not racy, since stealers can not steal from the same block
+    /// as the consumer.
     pub fn can_consume(&self) -> bool {
-        self.queue.stats.curr_enqueued() > 0
+        self.next_dequeable_entry().1.is_some()
+    }
+
+    /// Find the next dequeable entry without advancing the consumer cache.
+    ///
+    /// To allow this function to be usable without mutable access to self, it is left to the caller
+    /// to update `self.ccache` with the returned block.
+    /// If a dequable value is found, the second tuple element will be some and contain the index
+    /// of the element in the block.
+    fn next_dequeable_entry(&self) -> (&Block<E, ENTRIES_PER_BLOCK>, Option<IndexAndVersion<ENTRIES_PER_BLOCK>>) {
+        // SAFETY: `ccache` always points to a valid `Block` in the queue. We never create a mutable reference
+        // to a Block, so it is safe to construct a shared reference here.
+        let current_blk_cache = unsafe { &**self.ccache };
+        let mut blk = current_blk_cache;
+        for _ in 0..NUM_BLOCKS {
+              // check if the block is fully consumed already
+              let consumed = blk.consumed.load(Relaxed);
+              let consumed_idx = consumed.raw_index();
+
+              // Fastpath (Block is not fully consumed yet)
+            if consumed_idx < ENTRIES_PER_BLOCK {
+                // we know the block is not full, but we must first check if there is an entry to
+                // dequeue.
+                let committed_idx = blk.committed.load(Relaxed).raw_index();
+                if consumed_idx == committed_idx {
+                    return (blk, None);
+                }
+
+                /* There is an entry to dequeue */
+                return (blk, Some(consumed));
+            }
+
+            /* Slow-path */
+
+            /* Consumer head may never pass the Producer head and Consumer/Stealer tail */
+            let nblk = unsafe { &*blk.next() };
+            if self.can_advance_consumer_block(nblk, consumed) {
+                blk = nblk;
+            } else {
+                return (blk, None);
+            }
+            /* We advanced to the next block - loop around and try again */
+        }
+        // Since there is no concurrent enqueuing and the buffer is bounded, we should reach
+        // one of the exit conditions in at most NUM_BLOCKS iterations.
+        unreachable!()
     }
 }
 
@@ -604,12 +626,16 @@ impl<'a, E, const ENTRIES_PER_BLOCK: usize> Iterator for BlockIter<'a, E, ENTRIE
     fn next(&mut self) -> Option<E> {
         let i = self.i;
         self.i += 1;
-        self.buffer.get(i).map(|entry_cell| {
-            entry_cell.with(|entry| {
-                // SAFETY: we claimed the entries
-                unsafe { entry.read().assume_init() }
+        if i < self.committed {
+            self.buffer.get(i).map(|entry_cell| {
+                entry_cell.with(|entry| {
+                    // SAFETY: we claimed the entries
+                    unsafe { entry.read().assume_init() }
+                })
             })
-        })
+        } else {
+            None
+        }
     }
 }
 
