risingwave_batch/execution/
local_exchange.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::fmt::{Debug, Formatter};
16
17use risingwave_common::array::DataChunk;
18
19use crate::error::Result;
20use crate::exchange_source::ExchangeSource;
21use crate::task::{BatchTaskContext, TaskId, TaskOutput, TaskOutputId};
22
23/// Exchange data from a local task execution.
24pub struct LocalExchangeSource {
25    task_output: TaskOutput,
26
27    /// Id of task which contains the `ExchangeExecutor` of this source.
28    task_id: TaskId,
29
30    take_data_span: await_tree::Span,
31}
32
33impl LocalExchangeSource {
34    pub fn create(
35        output_id: TaskOutputId,
36        context: &dyn BatchTaskContext,
37        task_id: TaskId,
38    ) -> Result<Self> {
39        let task_output = context.get_task_output(output_id)?;
40        let take_data_span = await_tree::span!(
41            "local_exchange_take_data (task_output {:?})",
42            task_output.id()
43        );
44        Ok(Self {
45            task_output,
46            task_id,
47            take_data_span,
48        })
49    }
50}
51
52impl Debug for LocalExchangeSource {
53    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
54        f.debug_struct("LocalExchangeSource")
55            .field("task_output_id", self.task_output.id())
56            .finish()
57    }
58}
59
60impl ExchangeSource for LocalExchangeSource {
61    async fn take_data(&mut self) -> Result<Option<DataChunk>> {
62        use await_tree::InstrumentAwait;
63        let ret = self
64            .task_output
65            .direct_take_data()
66            .instrument_await(self.take_data_span.clone())
67            .await?;
68        if let Some(data) = ret {
69            let data = data.compact_vis();
70            trace!(
71                "Receiver task: {:?}, source task output: {:?}, data: {:?}",
72                self.task_id,
73                self.task_output.id(),
74                data
75            );
76            Ok(Some(data))
77        } else {
78            Ok(None)
79        }
80    }
81
82    fn get_task_id(&self) -> TaskId {
83        self.task_id.clone()
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use std::net::SocketAddr;
90    use std::sync::Arc;
91    use std::sync::atomic::{AtomicBool, Ordering};
92    use std::time::Duration;
93
94    use risingwave_common::config::RpcClientConfig;
95    use risingwave_pb::batch_plan::{TaskId, TaskOutputId};
96    use risingwave_pb::data::DataChunk;
97    use risingwave_pb::task_service::batch_exchange_service_server::{
98        BatchExchangeService, BatchExchangeServiceServer,
99    };
100    use risingwave_pb::task_service::{GetDataRequest, GetDataResponse};
101    use risingwave_rpc_client::ComputeClient;
102    use tokio::time::sleep;
103    use tokio_stream::wrappers::ReceiverStream;
104    use tonic::{Request, Response, Status};
105
106    use crate::exchange_source::ExchangeSource;
107    use crate::execution::grpc_exchange::GrpcExchangeSource;
108
109    struct FakeExchangeService {
110        rpc_called: Arc<AtomicBool>,
111    }
112
113    #[async_trait::async_trait]
114    impl BatchExchangeService for FakeExchangeService {
115        type GetDataStream = ReceiverStream<Result<GetDataResponse, Status>>;
116
117        async fn get_data(
118            &self,
119            _: Request<GetDataRequest>,
120        ) -> Result<Response<Self::GetDataStream>, Status> {
121            let (tx, rx) = tokio::sync::mpsc::channel(10);
122            self.rpc_called.store(true, Ordering::SeqCst);
123            for _ in 0..3 {
124                tx.send(Ok(GetDataResponse {
125                    record_batch: Some(DataChunk::default()),
126                }))
127                .await
128                .unwrap();
129            }
130            Ok(Response::new(ReceiverStream::new(rx)))
131        }
132    }
133
134    #[tokio::test]
135    async fn test_exchange_client() {
136        let rpc_called = Arc::new(AtomicBool::new(false));
137        let server_run = Arc::new(AtomicBool::new(false));
138        let addr: SocketAddr = "127.0.0.1:12345".parse().unwrap();
139
140        // Start a server.
141        let (shutdown_send, shutdown_recv) = tokio::sync::oneshot::channel();
142        let exchange_svc = BatchExchangeServiceServer::new(FakeExchangeService {
143            rpc_called: rpc_called.clone(),
144        });
145        let cp_server_run = server_run.clone();
146        let join_handle = tokio::spawn(async move {
147            cp_server_run.store(true, Ordering::SeqCst);
148            tonic::transport::Server::builder()
149                .add_service(exchange_svc)
150                .serve_with_shutdown(addr, async move {
151                    shutdown_recv.await.unwrap();
152                })
153                .await
154                .unwrap();
155        });
156
157        sleep(Duration::from_secs(1)).await;
158        assert!(server_run.load(Ordering::SeqCst));
159
160        let client = ComputeClient::new(addr.into(), &RpcClientConfig::default())
161            .await
162            .unwrap();
163        let task_output_id = TaskOutputId {
164            task_id: Some(TaskId::default()),
165            ..Default::default()
166        };
167        let mut src = GrpcExchangeSource::create(client, task_output_id, None)
168            .await
169            .unwrap();
170        for _ in 0..3 {
171            assert!(src.take_data().await.unwrap().is_some());
172        }
173        assert!(src.take_data().await.unwrap().is_none());
174        assert!(rpc_called.load(Ordering::SeqCst));
175
176        // Gracefully terminate the server.
177        shutdown_send.send(()).unwrap();
178        join_handle.await.unwrap();
179    }
180}
risingwave_batch/execution/local_exchange.rs

risingwave_batch/execution/
local_exchange.rs
