refactor: replace more GAT-based async trait with RPITIT

src/batch/src/exchange_source.rs

@@ -16,6 +16,7 @@ use std::fmt::Debug;
 use std::future::Future;
 
 use risingwave_common::array::DataChunk;
+use risingwave_common::error::Result;
 
 use crate::execution::grpc_exchange::GrpcExchangeSource;
 use crate::execution::local_exchange::LocalExchangeSource;
@@ -24,11 +25,7 @@ use crate::task::TaskId;
 
 /// Each `ExchangeSource` maps to one task, it takes the execution result from task chunk by chunk.
 pub trait ExchangeSource: Send + Debug {
-    type TakeDataFuture<'a>: Future<Output = risingwave_common::error::Result<Option<DataChunk>>>
-        + 'a
-    where
-        Self: 'a;
-    fn take_data(&mut self) -> Self::TakeDataFuture<'_>;
+    fn take_data(&mut self) -> impl Future<Output = Result<Option<DataChunk>>> + '_;
 
     /// Get upstream task id.
     fn get_task_id(&self) -> TaskId;
@@ -42,9 +39,7 @@ pub enum ExchangeSourceImpl {
 }
 
 impl ExchangeSourceImpl {
-    pub(crate) async fn take_data(
-        &mut self,
-    ) -> risingwave_common::error::Result<Option<DataChunk>> {
+    pub(crate) async fn take_data(&mut self) -> Result<Option<DataChunk>> {
         match self {
             ExchangeSourceImpl::Grpc(grpc) => grpc.take_data().await,
             ExchangeSourceImpl::Local(local) => local.take_data().await,

src/batch/src/execution/grpc_exchange.rs

@@ -13,7 +13,6 @@
 // limitations under the License.
 
 use std::fmt::{Debug, Formatter};
-use std::future::Future;
 
 use futures::StreamExt;
 use risingwave_common::array::DataChunk;
@@ -73,26 +72,22 @@ impl Debug for GrpcExchangeSource {
 }
 
 impl ExchangeSource for GrpcExchangeSource {
-    type TakeDataFuture<'a> = impl Future<Output = Result<Option<DataChunk>>> + 'a;
-
-    fn take_data(&mut self) -> Self::TakeDataFuture<'_> {
-        async {
-            let res = match self.stream.next().await {
-                None => {
-                    return Ok(None);
-                }
-                Some(r) => r,
-            };
-            let task_data = res?;
-            let data = DataChunk::from_protobuf(task_data.get_record_batch()?)?.compact();
-            trace!(
-                "Receiver taskOutput = {:?}, data = {:?}",
-                self.task_output_id,
-                data
-            );
+    async fn take_data(&mut self) -> Result<Option<DataChunk>> {
+        let res = match self.stream.next().await {
+            None => {
+                return Ok(None);
+            }
+            Some(r) => r,
+        };
+        let task_data = res?;
+        let data = DataChunk::from_protobuf(task_data.get_record_batch()?)?.compact();
+        trace!(
+            "Receiver taskOutput = {:?}, data = {:?}",
+            self.task_output_id,
+            data
+        );
 
-            Ok(Some(data))
-        }
+        Ok(Some(data))
     }
 
     fn get_task_id(&self) -> TaskId {

src/batch/src/execution/local_exchange.rs

@@ -13,7 +13,6 @@
 // limitations under the License.
 
 use std::fmt::{Debug, Formatter};
-use std::future::Future;
 
 use risingwave_common::array::DataChunk;
 use risingwave_common::error::Result;
@@ -52,23 +51,19 @@ impl Debug for LocalExchangeSource {
 }
 
 impl ExchangeSource for LocalExchangeSource {
-    type TakeDataFuture<'a> = impl Future<Output = Result<Option<DataChunk>>> + 'a;
-
-    fn take_data(&mut self) -> Self::TakeDataFuture<'_> {
-        async {
-            let ret = self.task_output.direct_take_data().await?;
-            if let Some(data) = ret {
-                let data = data.compact();
-                trace!(
-                    "Receiver task: {:?}, source task output: {:?}, data: {:?}",
-                    self.task_id,
-                    self.task_output.id(),
-                    data
-                );
-                Ok(Some(data))
-            } else {
-                Ok(None)
-            }
+    async fn take_data(&mut self) -> Result<Option<DataChunk>> {
+        let ret = self.task_output.direct_take_data().await?;
+        if let Some(data) = ret {
+            let data = data.compact();
+            trace!(
+                "Receiver task: {:?}, source task output: {:?}, data: {:?}",
+                self.task_id,
+                self.task_output.id(),
+                data
+            );
+            Ok(Some(data))
+        } else {
+            Ok(None)
         }
     }
 

src/batch/src/executor/test_utils.rs

@@ -13,7 +13,6 @@
 // limitations under the License.
 
 use std::collections::VecDeque;
-use std::future::Future;
 
 use assert_matches::assert_matches;
 use futures::StreamExt;
@@ -246,15 +245,11 @@ impl FakeExchangeSource {
 }
 
 impl ExchangeSource for FakeExchangeSource {
-    type TakeDataFuture<'a> = impl Future<Output = Result<Option<DataChunk>>> + 'a;
-
-    fn take_data(&mut self) -> Self::TakeDataFuture<'_> {
-        async {
-            if let Some(chunk) = self.chunks.pop() {
-                Ok(chunk)
-            } else {
-                Ok(None)
-            }
+    async fn take_data(&mut self) -> Result<Option<DataChunk>> {
+        if let Some(chunk) = self.chunks.pop() {
+            Ok(chunk)
+        } else {
+            Ok(None)
         }
     }
 

src/batch/src/lib.rs

@@ -34,6 +34,7 @@
 #![feature(result_option_inspect)]
 #![feature(assert_matches)]
 #![feature(lazy_cell)]
+#![feature(return_position_impl_trait_in_trait)]
 
 mod error;
 pub mod exchange_source;

src/connector/src/source/external.rs

@@ -200,13 +200,9 @@ impl MySqlOffset {
 }
 
 pub trait ExternalTableReader {
-    type CdcOffsetFuture<'a>: Future<Output = ConnectorResult<CdcOffset>> + Send + 'a
-    where
-        Self: 'a;
-
     fn get_normalized_table_name(&self, table_name: &SchemaTableName) -> String;
 
-    fn current_cdc_offset(&self) -> Self::CdcOffsetFuture<'_>;
+    fn current_cdc_offset(&self) -> impl Future<Output = ConnectorResult<CdcOffset>> + Send + '_;
 
     fn parse_binlog_offset(&self, offset: &str) -> ConnectorResult<CdcOffset>;
 
@@ -248,32 +244,28 @@ pub struct ExternalTableConfig {
 }
 
 impl ExternalTableReader for MySqlExternalTableReader {
-    type CdcOffsetFuture<'a> = impl Future<Output = ConnectorResult<CdcOffset>> + 'a;
-
     fn get_normalized_table_name(&self, table_name: &SchemaTableName) -> String {
         format!("`{}`", table_name.table_name)
     }
 
-    fn current_cdc_offset(&self) -> Self::CdcOffsetFuture<'_> {
-        async move {
-            let mut conn = self
-                .pool
-                .get_conn()
-                .await
-                .map_err(|e| ConnectorError::Connection(anyhow!(e)))?;
-
-            let sql = "SHOW MASTER STATUS".to_string();
-            let mut rs = conn.query::<mysql_async::Row, _>(sql).await?;
-            let row = rs
-                .iter_mut()
-                .exactly_one()
-                .map_err(|e| ConnectorError::Internal(anyhow!("read binlog error: {}", e)))?;
-
-            Ok(CdcOffset::MySql(MySqlOffset {
-                filename: row.take("File").unwrap(),
-                position: row.take("Position").unwrap(),
-            }))
-        }
+    async fn current_cdc_offset(&self) -> ConnectorResult<CdcOffset> {
+        let mut conn = self
+            .pool
+            .get_conn()
+            .await
+            .map_err(|e| ConnectorError::Connection(anyhow!(e)))?;
+
+        let sql = "SHOW MASTER STATUS".to_string();
+        let mut rs = conn.query::<mysql_async::Row, _>(sql).await?;
+        let row = rs
+            .iter_mut()
+            .exactly_one()
+            .map_err(|e| ConnectorError::Internal(anyhow!("read binlog error: {}", e)))?;
+
+        Ok(CdcOffset::MySql(MySqlOffset {
+            filename: row.take("File").unwrap(),
+            position: row.take("Position").unwrap(),
+        }))
     }
 
     fn parse_binlog_offset(&self, offset: &str) -> ConnectorResult<CdcOffset> {
@@ -478,21 +470,17 @@ impl MySqlExternalTableReader {
 }
 
 impl ExternalTableReader for ExternalTableReaderImpl {
-    type CdcOffsetFuture<'a> = impl Future<Output = ConnectorResult<CdcOffset>> + 'a;
-
     fn get_normalized_table_name(&self, table_name: &SchemaTableName) -> String {
         match self {
             ExternalTableReaderImpl::MySql(mysql) => mysql.get_normalized_table_name(table_name),
             ExternalTableReaderImpl::Mock(mock) => mock.get_normalized_table_name(table_name),
         }
     }
 
-    fn current_cdc_offset(&self) -> Self::CdcOffsetFuture<'_> {
-        async move {
-            match self {
-                ExternalTableReaderImpl::MySql(mysql) => mysql.current_cdc_offset().await,
-                ExternalTableReaderImpl::Mock(mock) => mock.current_cdc_offset().await,
-            }
+    async fn current_cdc_offset(&self) -> ConnectorResult<CdcOffset> {
+        match self {
+            ExternalTableReaderImpl::MySql(mysql) => mysql.current_cdc_offset().await,
+            ExternalTableReaderImpl::Mock(mock) => mock.current_cdc_offset().await,
         }
     }
 

src/connector/src/source/mock_external_table.rs

@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::future::Future;
 use std::sync::atomic::AtomicUsize;
 
 use futures::stream::BoxStream;
@@ -91,24 +90,21 @@ impl MockExternalTableReader {
 }
 
 impl ExternalTableReader for MockExternalTableReader {
-    type CdcOffsetFuture<'a> = impl Future<Output = ConnectorResult<CdcOffset>> + 'a;
-
     fn get_normalized_table_name(&self, _table_name: &SchemaTableName) -> String {
         "`mock_table`".to_string()
     }
 
-    fn current_cdc_offset(&self) -> Self::CdcOffsetFuture<'_> {
+    async fn current_cdc_offset(&self) -> ConnectorResult<CdcOffset> {
         static IDX: AtomicUsize = AtomicUsize::new(0);
-        async move {
-            let idx = IDX.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
-            if idx < self.binlog_watermarks.len() {
-                Ok(CdcOffset::MySql(self.binlog_watermarks[idx].clone()))
-            } else {
-                Ok(CdcOffset::MySql(MySqlOffset {
-                    filename: "1.binlog".to_string(),
-                    position: u64::MAX,
-                }))
-            }
+
+        let idx = IDX.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+        if idx < self.binlog_watermarks.len() {
+            Ok(CdcOffset::MySql(self.binlog_watermarks[idx].clone()))
+        } else {
+            Ok(CdcOffset::MySql(MySqlOffset {
+                filename: "1.binlog".to_string(),
+                position: u64::MAX,
+            }))
         }
     }