sync: add mpsc::Receiver::recv_many (#6010)

benches/sync_mpsc.rs

@@ -73,6 +73,33 @@ fn contention_bounded(g: &mut BenchmarkGroup<WallTime>) {
     });
 }
 
+fn contention_bounded_recv_many(g: &mut BenchmarkGroup<WallTime>) {
+    let rt = rt();
+
+    g.bench_function("bounded_recv_many", |b| {
+        b.iter(|| {
+            rt.block_on(async move {
+                let (tx, mut rx) = mpsc::channel::<usize>(1_000_000);
+
+                for _ in 0..5 {
+                    let tx = tx.clone();
+                    tokio::spawn(async move {
+                        for i in 0..1000 {
+                            tx.send(i).await.unwrap();
+                        }
+                    });
+                }
+
+                let mut buffer = Vec::<usize>::with_capacity(5_000);
+                let mut total = 0;
+                while total < 1_000 * 5 {
+                    total += rx.recv_many(&mut buffer, 5_000).await;
+                }
+            })
+        })
+    });
+}
+
 fn contention_bounded_full(g: &mut BenchmarkGroup<WallTime>) {
     let rt = rt();
 
@@ -98,6 +125,33 @@ fn contention_bounded_full(g: &mut BenchmarkGroup<WallTime>) {
     });
 }
 
+fn contention_bounded_full_recv_many(g: &mut BenchmarkGroup<WallTime>) {
+    let rt = rt();
+
+    g.bench_function("bounded_full_recv_many", |b| {
+        b.iter(|| {
+            rt.block_on(async move {
+                let (tx, mut rx) = mpsc::channel::<usize>(100);
+
+                for _ in 0..5 {
+                    let tx = tx.clone();
+                    tokio::spawn(async move {
+                        for i in 0..1000 {
+                            tx.send(i).await.unwrap();
+                        }
+                    });
+                }
+
+                let mut buffer = Vec::<usize>::with_capacity(5_000);
+                let mut total = 0;
+                while total < 1_000 * 5 {
+                    total += rx.recv_many(&mut buffer, 5_000).await;
+                }
+            })
+        })
+    });
+}
+
 fn contention_unbounded(g: &mut BenchmarkGroup<WallTime>) {
     let rt = rt();
 
@@ -123,6 +177,33 @@ fn contention_unbounded(g: &mut BenchmarkGroup<WallTime>) {
     });
 }
 
+fn contention_unbounded_recv_many(g: &mut BenchmarkGroup<WallTime>) {
+    let rt = rt();
+
+    g.bench_function("unbounded_recv_many", |b| {
+        b.iter(|| {
+            rt.block_on(async move {
+                let (tx, mut rx) = mpsc::unbounded_channel::<usize>();
+
+                for _ in 0..5 {
+                    let tx = tx.clone();
+                    tokio::spawn(async move {
+                        for i in 0..1000 {
+                            tx.send(i).unwrap();
+                        }
+                    });
+                }
+
+                let mut buffer = Vec::<usize>::with_capacity(5_000);
+                let mut total = 0;
+                while total < 1_000 * 5 {
+                    total += rx.recv_many(&mut buffer, 5_000).await;
+                }
+            })
+        })
+    });
+}
+
 fn uncontented_bounded(g: &mut BenchmarkGroup<WallTime>) {
     let rt = rt();
 
@@ -143,6 +224,28 @@ fn uncontented_bounded(g: &mut BenchmarkGroup<WallTime>) {
     });
 }
 
+fn uncontented_bounded_recv_many(g: &mut BenchmarkGroup<WallTime>) {
+    let rt = rt();
+
+    g.bench_function("bounded_recv_many", |b| {
+        b.iter(|| {
+            rt.block_on(async move {
+                let (tx, mut rx) = mpsc::channel::<usize>(1_000_000);
+
+                for i in 0..5000 {
+                    tx.send(i).await.unwrap();
+                }
+
+                let mut buffer = Vec::<usize>::with_capacity(5_000);
+                let mut total = 0;
+                while total < 1_000 * 5 {
+                    total += rx.recv_many(&mut buffer, 5_000).await;
+                }
+            })
+        })
+    });
+}
+
 fn uncontented_unbounded(g: &mut BenchmarkGroup<WallTime>) {
     let rt = rt();
 
@@ -163,6 +266,28 @@ fn uncontented_unbounded(g: &mut BenchmarkGroup<WallTime>) {
     });
 }
 
+fn uncontented_unbounded_recv_many(g: &mut BenchmarkGroup<WallTime>) {
+    let rt = rt();
+
+    g.bench_function("unbounded_recv_many", |b| {
+        b.iter(|| {
+            rt.block_on(async move {
+                let (tx, mut rx) = mpsc::unbounded_channel::<usize>();
+
+                for i in 0..5000 {
+                    tx.send(i).unwrap();
+                }
+
+                let mut buffer = Vec::<usize>::with_capacity(5_000);
+                let mut total = 0;
+                while total < 1_000 * 5 {
+                    total += rx.recv_many(&mut buffer, 5_000).await;
+                }
+            })
+        })
+    });
+}
+
 fn bench_create_medium(c: &mut Criterion) {
     let mut group = c.benchmark_group("create_medium");
     create_medium::<1>(&mut group);
@@ -181,15 +306,20 @@ fn bench_send(c: &mut Criterion) {
 fn bench_contention(c: &mut Criterion) {
     let mut group = c.benchmark_group("contention");
     contention_bounded(&mut group);
+    contention_bounded_recv_many(&mut group);
     contention_bounded_full(&mut group);
+    contention_bounded_full_recv_many(&mut group);
     contention_unbounded(&mut group);
+    contention_unbounded_recv_many(&mut group);
     group.finish();
 }
 
 fn bench_uncontented(c: &mut Criterion) {
     let mut group = c.benchmark_group("uncontented");
     uncontented_bounded(&mut group);
+    uncontented_bounded_recv_many(&mut group);
     uncontented_unbounded(&mut group);
+    uncontented_unbounded_recv_many(&mut group);
     group.finish();
 }
 

tokio/src/sync/mpsc/bounded.rs

@@ -44,10 +44,10 @@ pub struct Sender<T> {
 /// async fn main() {
 ///     let (tx, _rx) = channel::<i32>(15);
 ///     let tx_weak = tx.downgrade();
-///   
+///
 ///     // Upgrading will succeed because `tx` still exists.
 ///     assert!(tx_weak.upgrade().is_some());
-///   
+///
 ///     // If we drop `tx`, then it will fail.
 ///     drop(tx);
 ///     assert!(tx_weak.clone().upgrade().is_none());
@@ -230,6 +230,82 @@ impl<T> Receiver<T> {
         poll_fn(|cx| self.chan.recv(cx)).await
     }
 
+    /// Receives the next values for this receiver and extends `buffer`.
+    ///
+    /// This method extends `buffer` by no more than a fixed number of values
+    /// as specified by `limit`. If `limit` is zero, the function immediately
+    /// returns `0`. The return value is the number of values added to `buffer`.
+    ///
+    /// For `limit > 0`, if there are no messages in the channel's queue, but
+    /// the channel has not yet been closed, this method will sleep until a
+    /// message is sent or the channel is closed. Note that if [`close`] is
+    /// called, but there are still outstanding [`Permits`] from before it was
+    /// closed, the channel is not considered closed by `recv_many` until the
+    /// permits are released.
+    ///
+    /// For non-zero values of `limit`, this method will never return `0` unless
+    /// the channel has been closed and there are no remaining messages in the
+    /// channel's queue. This indicates that no further values can ever be
+    /// received from this `Receiver`. The channel is closed when all senders
+    /// have been dropped, or when [`close`] is called.
+    ///
+    /// The capacity of `buffer` is increased as needed.
+    ///
+    /// # Cancel safety
+    ///
+    /// This method is cancel safe. If `recv_many` is used as the event in a
+    /// [`tokio::select!`](crate::select) statement and some other branch
+    /// completes first, it is guaranteed that no messages were received on this
+    /// channel.
+    ///
+    /// [`close`]: Self::close
+    /// [`Permits`]: struct@crate::sync::mpsc::Permit
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tokio::sync::mpsc;
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     let mut buffer: Vec<&str> = Vec::with_capacity(2);
+    ///     let limit = 2;
+    ///     let (tx, mut rx) = mpsc::channel(100);
+    ///     let tx2 = tx.clone();
+    ///     tx2.send("first").await.unwrap();
+    ///     tx2.send("second").await.unwrap();
+    ///     tx2.send("third").await.unwrap();
+    ///
+    ///     // Call `recv_many` to receive up to `limit` (2) values.
+    ///     assert_eq!(2, rx.recv_many(&mut buffer, limit).await);
+    ///     assert_eq!(vec!["first", "second"], buffer);
+    ///
+    ///     // If the buffer is full, the next call to `recv_many`
+    ///     // reserves additional capacity.
+    ///     assert_eq!(1, rx.recv_many(&mut buffer, 1).await);
+    ///
+    ///     tokio::spawn(async move {
+    ///         tx.send("fourth").await.unwrap();
+    ///     });
+    ///
+    ///     // 'tx' is dropped, but `recv_many`
+    ///     // is guaranteed not to return 0 as the channel
+    ///     // is not yet closed.
+    ///     assert_eq!(1, rx.recv_many(&mut buffer, 1).await);
+    ///     assert_eq!(vec!["first", "second", "third", "fourth"], buffer);
+    ///
+    ///     // Once the last sender is dropped, the channel is
+    ///     // closed and `recv_many` returns 0, capacity unchanged.
+    ///     drop(tx2);
+    ///     assert_eq!(0, rx.recv_many(&mut buffer, limit).await);
+    ///     assert_eq!(vec!["first", "second", "third", "fourth"], buffer);
+    /// }
+    /// ```
+    pub async fn recv_many(&mut self, buffer: &mut Vec<T>, limit: usize) -> usize {
+        use crate::future::poll_fn;
+        poll_fn(|cx| self.chan.recv_many(cx, buffer, limit)).await
+    }
+
     /// Tries to receive the next value for this receiver.
     ///
     /// This method returns the [`Empty`] error if the channel is currently
@@ -1072,7 +1148,7 @@ impl<T> Sender<T> {
     /// #[tokio::main]
     /// async fn main() {
     ///     let (tx, _rx) = mpsc::channel::<()>(5);
-    ///      
+    ///
     ///     // both max capacity and capacity are the same at first
     ///     assert_eq!(tx.max_capacity(), 5);
     ///     assert_eq!(tx.capacity(), 5);