style: prefer stack pinning (#9786)

Signed-off-by: TennyZhuang <[email protected]>

Author: TennyZhuang

grafana/generate.sh

@@ -6,4 +6,4 @@ for dashboard_name in "risingwave-dev-dashboard" "risingwave-user-dashboard"; do
     generate-dashboard -o $dashboard_name.gen.json $dashboard_name.dashboard.py
     jq -c . $dashboard_name.gen.json > $dashboard_name.json
     cp $dashboard_name.json ../docker/dashboards/
-done
+done

src/common/src/util/stream_cancel.rs

@@ -19,14 +19,16 @@
 //!
 //! Following is an example of how it works:
 //! ```rust
+//! use std::pin::pin;
+//!
 //! use futures::stream::iter;
 //! use futures::StreamExt;
 //! use risingwave_common::util::stream_cancel::{cancellable_stream, stream_tripwire};
 //!
 //! #[tokio::main]
 //! async fn main() {
 //!     let (trigger, tripwire) = stream_tripwire(|| 5);
-//!     let mut s = Box::pin(cancellable_stream(iter(vec![1i32, 2, 3, 4]), tripwire));
+//!     let mut s = pin!(cancellable_stream(iter(vec![1i32, 2, 3, 4]), tripwire));
 //!     assert_eq!(Some(1), s.next().await);
 //!     assert_eq!(Some(2), s.next().await);
 //!     trigger.abort();
@@ -113,6 +115,8 @@ where
 
 #[cfg(test)]
 mod tests {
+    use std::pin::pin;
+
     use futures::stream::iter;
     use futures::StreamExt;
 
@@ -121,7 +125,7 @@ mod tests {
     #[tokio::test]
     async fn test_cancellable_stream_aborted() {
         let (trigger, tripwire) = stream_tripwire(|| 5);
-        let mut s = Box::pin(cancellable_stream(iter(vec![1i32, 2, 3, 4]), tripwire));
+        let mut s = pin!(cancellable_stream(iter(vec![1i32, 2, 3, 4]), tripwire));
         assert_eq!(Some(1), s.next().await);
         assert_eq!(Some(2), s.next().await);
         trigger.abort();
@@ -133,7 +137,7 @@ mod tests {
     #[tokio::test]
     async fn test_cancellable_stream_not_aborted() {
         let (trigger, tripwire) = stream_tripwire(|| 5);
-        let mut s = Box::pin(cancellable_stream(iter(vec![1, 2, 3, 4]), tripwire));
+        let mut s = pin!(cancellable_stream(iter(vec![1, 2, 3, 4]), tripwire));
         assert_eq!(Some(1), s.next().await);
         assert_eq!(Some(2), s.next().await);
         assert_eq!(Some(3), s.next().await);

src/connector/src/source/filesystem/s3/source/reader.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::collections::HashMap;
+use std::pin::pin;
 
 use anyhow::{anyhow, Result};
 use async_trait::async_trait;
@@ -69,7 +70,7 @@ impl S3FileReader {
             byte_stream.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e)),
         );
 
-        let reader = Box::pin(BufReader::new(stream_reader));
+        let reader = pin!(BufReader::new(stream_reader));
 
         let stream = ReaderStream::with_capacity(reader, STREAM_READER_CAPACITY);
 
@@ -202,9 +203,9 @@ impl S3FileReader {
                 parser,
                 ByteStreamSourceParserImpl::Json(_) | ByteStreamSourceParserImpl::Csv(_)
             ) {
-                NdByteStreamWrapper::new(parser).into_stream(Box::pin(data_stream))
+                NdByteStreamWrapper::new(parser).into_stream(data_stream)
             } else {
-                parser.into_stream(Box::pin(data_stream))
+                parser.into_stream(data_stream)
             };
             #[for_await]
             for msg in msg_stream {

src/stream/src/executor/backfill.rs

@@ -14,6 +14,7 @@
 
 use std::cmp::Ordering;
 use std::ops::Bound;
+use std::pin::pin;
 use std::sync::Arc;
 
 use await_tree::InstrumentAwait;
@@ -192,10 +193,13 @@ where
 
             let left_upstream = upstream.by_ref().map(Either::Left);
 
-            let right_snapshot = Box::pin(
-                Self::snapshot_read(&self.table, snapshot_read_epoch, current_pos.clone(), true)
-                    .map(Either::Right),
-            );
+            let right_snapshot = pin!(Self::snapshot_read(
+                &self.table,
+                snapshot_read_epoch,
+                current_pos.clone(),
+                true
+            )
+            .map(Either::Right),);
 
             // Prefer to select upstream, so we can stop snapshot stream as soon as the barrier
             // comes.

src/stream/src/executor/lookup/sides.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::pin::pin;
+
 use anyhow::Context;
 use either::Either;
 use futures::stream::PollNext;
@@ -122,8 +124,8 @@ fn prefer_right(_: &mut ()) -> PollNext {
 /// available for both left and right side, instead of being combined.
 #[try_stream(ok = BarrierAlignedMessage, error = StreamExecutorError)]
 pub async fn align_barrier(left: impl MessageStream, right: impl MessageStream) {
-    let mut left = Box::pin(left);
-    let mut right = Box::pin(right);
+    let mut left = pin!(left);
+    let mut right = pin!(right);
 
     enum SideStatus {
         LeftBarrier,
@@ -201,7 +203,7 @@ pub async fn stream_lookup_arrange_prev_epoch(
     stream: Box<dyn Executor>,
     arrangement: Box<dyn Executor>,
 ) {
-    let mut input = Box::pin(align_barrier(stream.execute(), arrangement.execute()));
+    let mut input = pin!(align_barrier(stream.execute(), arrangement.execute()));
     let mut arrange_buf = vec![];
     let mut stream_side_end = false;
 
@@ -285,7 +287,7 @@ pub async fn stream_lookup_arrange_this_epoch(
     stream: Box<dyn Executor>,
     arrangement: Box<dyn Executor>,
 ) {
-    let mut input = Box::pin(align_barrier(stream.execute(), arrangement.execute()));
+    let mut input = pin!(align_barrier(stream.execute(), arrangement.execute()));
     let mut stream_buf = vec![];
     let mut arrange_buf = vec![];
 

src/stream/src/executor/rearranged_chain.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::pin::pin;
+
 use futures::channel::{mpsc, oneshot};
 use futures::stream::select_with_strategy;
 use futures::{stream, StreamExt};
@@ -105,7 +107,7 @@ impl RearrangedChainExecutor {
 
     #[try_stream(ok = Message, error = StreamExecutorError)]
     async fn execute_inner(mut self) {
-        let mut upstream = Box::pin(self.upstream.execute());
+        let mut upstream = pin!(self.upstream.execute());
 
         // 1. Poll the upstream to get the first barrier.
         let first_barrier = expect_first_barrier(&mut upstream).await?;
@@ -133,106 +135,107 @@ impl RearrangedChainExecutor {
                 .unbounded_send(RearrangedMessage::PhantomBarrier(first_barrier))
                 .unwrap();
 
-            // 3. Rearrange stream, will yield the barriers polled from upstream to rearrange.
-            let rearranged_barrier = Box::pin(
-                Self::rearrange_barrier(&mut upstream, upstream_tx, stop_rearrange_rx)
-                    .map(|result| result.map(RearrangedMessage::RearrangedBarrier)),
-            );
-
-            // 4. Init the snapshot with reading epoch.
-            let snapshot = self.snapshot.execute_with_epoch(create_epoch.prev);
-
-            // Chain the `snapshot` and `upstream_rx` to get a unified `rearranged_chunks` stream.
-            let rearranged_chunks = snapshot
-                .map(|result| result.map(RearrangedMessage::rearranged_from))
-                .chain(upstream_rx.map(Ok));
-
-            // 5. Merge the rearranged barriers with chunks, with the priority of barrier.
-            let mut rearranged =
-                select_with_strategy(rearranged_barrier, rearranged_chunks, |_: &mut ()| {
-                    stream::PollNext::Left
-                });
-
-            // Record the epoch of the last rearranged barrier we received.
-            let mut last_rearranged_epoch = create_epoch;
-            let mut stop_rearrange_tx = Some(stop_rearrange_tx);
-
-            let mut processed_rows: u64 = 0;
-
-            // 6. Consume the merged `rearranged` stream.
-            #[for_await]
-            for rearranged_msg in &mut rearranged {
-                match rearranged_msg? {
-                    // If we received a phantom barrier, update the progress and check whether we
-                    // catches up with the progress of upstream MV.
-                    //
-                    // Note that there's no phantom barrier in the snapshot. So we must have already
-                    // consumed the whole snapshot and be on the upstream now.
-                    RearrangedMessage::PhantomBarrier(barrier) => {
-                        // Update the progress since we've consumed all chunks before this phantom.
-                        self.progress.update(
-                            last_rearranged_epoch.curr,
-                            barrier.epoch.curr,
-                            processed_rows,
-                        );
-
-                        if barrier.epoch.curr >= last_rearranged_epoch.curr {
-                            // Stop the background rearrangement task.
-                            stop_rearrange_tx.take().unwrap().send(()).map_err(|_| {
-                                StreamExecutorError::channel_closed("stop rearrange")
-                            })?;
-                            break;
+            {
+                // 3. Rearrange stream, will yield the barriers polled from upstream to rearrange.
+                let rearranged_barrier =
+                    pin!(
+                        Self::rearrange_barrier(&mut upstream, upstream_tx, stop_rearrange_rx)
+                            .map(|result| result.map(RearrangedMessage::RearrangedBarrier)),
+                    );
+
+                // 4. Init the snapshot with reading epoch.
+                let snapshot = self.snapshot.execute_with_epoch(create_epoch.prev);
+
+                // Chain the `snapshot` and `upstream_rx` to get a unified `rearranged_chunks`
+                // stream.
+                let rearranged_chunks = snapshot
+                    .map(|result| result.map(RearrangedMessage::rearranged_from))
+                    .chain(upstream_rx.map(Ok));
+
+                // 5. Merge the rearranged barriers with chunks, with the priority of barrier.
+                let mut rearranged =
+                    select_with_strategy(rearranged_barrier, rearranged_chunks, |_: &mut ()| {
+                        stream::PollNext::Left
+                    });
+
+                // Record the epoch of the last rearranged barrier we received.
+                let mut last_rearranged_epoch = create_epoch;
+                let mut stop_rearrange_tx = Some(stop_rearrange_tx);
+
+                let mut processed_rows: u64 = 0;
+
+                #[for_await]
+                for rearranged_msg in &mut rearranged {
+                    match rearranged_msg? {
+                        // If we received a phantom barrier, update the progress and check whether
+                        // we catches up with the progress of upstream MV.
+                        //
+                        // Note that there's no phantom barrier in the snapshot. So we must have
+                        // already consumed the whole snapshot and be on the
+                        // upstream now.
+                        RearrangedMessage::PhantomBarrier(barrier) => {
+                            // Update the progress since we've consumed all chunks before this
+                            // phantom.
+                            self.progress.update(
+                                last_rearranged_epoch.curr,
+                                barrier.epoch.curr,
+                                processed_rows,
+                            );
+
+                            if barrier.epoch.curr >= last_rearranged_epoch.curr {
+                                // Stop the background rearrangement task.
+                                stop_rearrange_tx.take().unwrap().send(()).map_err(|_| {
+                                    StreamExecutorError::channel_closed("stop rearrange")
+                                })?;
+                                break;
+                            }
                         }
-                    }
 
-                    // If we received a message, yield it.
-                    RearrangedMessage::RearrangedBarrier(barrier) => {
-                        last_rearranged_epoch = barrier.epoch;
-                        yield Message::Barrier(barrier);
-                    }
-                    RearrangedMessage::Chunk(chunk) => {
-                        processed_rows += chunk.cardinality() as u64;
-                        yield Message::Chunk(chunk)
-                    }
-                    RearrangedMessage::Watermark => {
-                        // Ignore watermark during snapshot consumption.
+                        // If we received a message, yield it.
+                        RearrangedMessage::RearrangedBarrier(barrier) => {
+                            last_rearranged_epoch = barrier.epoch;
+                            yield Message::Barrier(barrier);
+                        }
+                        RearrangedMessage::Chunk(chunk) => {
+                            processed_rows += chunk.cardinality() as u64;
+                            yield Message::Chunk(chunk)
+                        }
+                        RearrangedMessage::Watermark => {
+                            // Ignore watermark during snapshot consumption.
+                        }
                     }
                 }
-            }
-
-            // 7. Rearranged task finished.
-            // The reason for finish must be that we told it to stop.
-            tracing::trace!(actor = self.actor_id, "rearranged task finished");
-            if stop_rearrange_tx.is_some() {
-                tracing::error!(actor = self.actor_id, "rearrangement finished passively");
-            }
 
-            // 8. Consume remainings.
-            let mut finish_on_barrier = |msg: &Message| {
-                if let Some(barrier) = msg.as_barrier() {
-                    self.progress.finish(barrier.epoch.curr);
+                // 7. Rearranged task finished.
+                // The reason for finish must be that we told it to stop.
+                tracing::trace!(actor = self.actor_id, "rearranged task finished");
+                if stop_rearrange_tx.is_some() {
+                    tracing::error!(actor = self.actor_id, "rearrangement finished passively");
                 }
-            };
 
-            // Note that there may still be some messages in `rearranged`. However the rearranged
-            // barriers must be ignored, we should take the phantoms.
-            #[for_await]
-            for msg in rearranged {
-                let msg: RearrangedMessage = msg?;
-                let Some(msg) = msg.phantom_into() else { continue };
-                finish_on_barrier(&msg);
-                yield msg;
+                // 8. Consume remainings.
+                // Note that there may still be some messages in `rearranged`. However the
+                // rearranged barriers must be ignored, we should take the phantoms.
+                #[for_await]
+                for msg in rearranged {
+                    let msg: RearrangedMessage = msg?;
+                    let Some(msg) = msg.phantom_into() else { continue };
+                    if let Some(barrier) = msg.as_barrier() {
+                        self.progress.finish(barrier.epoch.curr);
+                    }
+                    yield msg;
+                }
             }
 
-            let mut remaining_upstream = upstream;
-
             // Consume remaining upstream.
             tracing::trace!(actor = self.actor_id, "begin to consume remaining upstream");
 
             #[for_await]
-            for msg in &mut remaining_upstream {
+            for msg in upstream {
                 let msg: Message = msg?;
-                finish_on_barrier(&msg);
+                if let Some(barrier) = msg.as_barrier() {
+                    self.progress.finish(barrier.epoch.curr);
+                }
                 yield msg;
             }
         } else {

src/stream/src/executor/temporal_join.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::alloc::Global;
+use std::pin::pin;
 use std::sync::Arc;
 
 use either::Either;
@@ -124,8 +125,8 @@ pub async fn chunks_until_barrier(stream: impl MessageStream, expected_barrier:
 // `InternalMessage::Barrier(right_chunks, barrier)`.
 #[try_stream(ok = InternalMessage, error = StreamExecutorError)]
 async fn align_input(left: Box<dyn Executor>, right: Box<dyn Executor>) {
-    let mut left = Box::pin(left.execute());
-    let mut right = Box::pin(right.execute());
+    let mut left = pin!(left.execute());
+    let mut right = pin!(right.execute());
     // Keep producing intervals until stream exhaustion or errors.
     loop {
         let mut right_chunks = vec![];