risingwave_rpc_client/
stream_client.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::sync::Arc;
16use std::time::Duration;
17
18use anyhow::anyhow;
19use async_trait::async_trait;
20use futures::TryStreamExt;
21use risingwave_common::config::{MAX_CONNECTION_WINDOW_SIZE, RpcClientConfig};
22use risingwave_common::monitor::{EndpointExt, TcpConfig};
23use risingwave_common::util::addr::HostAddr;
24use risingwave_pb::stream_service::stream_service_client::StreamServiceClient;
25use risingwave_pb::stream_service::streaming_control_stream_request::PbInitRequest;
26use risingwave_pb::stream_service::streaming_control_stream_response::InitResponse;
27use risingwave_pb::stream_service::*;
28use tokio_stream::wrappers::UnboundedReceiverStream;
29use tonic::transport::Endpoint;
30
31use crate::channel::{Channel, WrappedChannelExt};
32use crate::error::{Result, RpcError};
33use crate::{RpcClient, RpcClientPool, UnboundedBidiStreamHandle, stream_rpc_client_method_impl};
34
35#[derive(Clone)]
36pub struct StreamClient(StreamServiceClient<Channel>);
37
38#[async_trait]
39impl RpcClient for StreamClient {
40    async fn new_client(host_addr: HostAddr, opts: &RpcClientConfig) -> Result<Self> {
41        Self::new(host_addr, opts).await
42    }
43}
44
45impl StreamClient {
46    async fn new(host_addr: HostAddr, opts: &RpcClientConfig) -> Result<Self> {
47        let channel = Endpoint::from_shared(format!("http://{}", host_addr))?
48            .initial_connection_window_size(MAX_CONNECTION_WINDOW_SIZE)
49            .connect_timeout(Duration::from_secs(opts.connect_timeout_secs))
50            .monitored_connect(
51                "grpc-stream-client",
52                TcpConfig {
53                    tcp_nodelay: true,
54                    ..Default::default()
55                },
56            )
57            .await?
58            .wrapped();
59
60        Ok(Self(
61            StreamServiceClient::new(channel).max_decoding_message_size(usize::MAX),
62        ))
63    }
64}
65
66pub type StreamClientPool = RpcClientPool<StreamClient>;
67pub type StreamClientPoolRef = Arc<StreamClientPool>;
68
69macro_rules! for_all_stream_rpc {
70    ($macro:ident) => {
71        $macro! {
72            { 0, get_min_uncommitted_object_id, GetMinUncommittedObjectIdRequest, GetMinUncommittedObjectIdResponse }
73            ,{ 0, warm_up_table_cache, WarmUpTableCacheRequest, WarmUpTableCacheResponse }
74        }
75    };
76}
77
78impl StreamClient {
79    for_all_stream_rpc! { stream_rpc_client_method_impl }
80}
81
82pub type StreamingControlHandle =
83    UnboundedBidiStreamHandle<StreamingControlStreamRequest, StreamingControlStreamResponse>;
84
85impl StreamClient {
86    pub async fn start_streaming_control(
87        &self,
88        init_request: PbInitRequest,
89    ) -> Result<StreamingControlHandle> {
90        let first_request = StreamingControlStreamRequest {
91            request: Some(streaming_control_stream_request::Request::Init(
92                init_request,
93            )),
94        };
95        let mut client = self.0.clone();
96        let (handle, first_rsp) =
97            UnboundedBidiStreamHandle::initialize(first_request, |rx| async move {
98                client
99                    .streaming_control_stream(UnboundedReceiverStream::new(rx))
100                    .await
101                    .map(|response| response.into_inner().map_err(RpcError::from_stream_status))
102                    .map_err(RpcError::from_stream_status)
103            })
104            .await?;
105        match first_rsp {
106            StreamingControlStreamResponse {
107                response: Some(streaming_control_stream_response::Response::Init(InitResponse {})),
108            } => {}
109            other => {
110                return Err(anyhow!("expect InitResponse but get {:?}", other).into());
111            }
112        };
113        Ok(handle)
114    }
115}
risingwave_rpc_client/stream_client.rs

risingwave_rpc_client/
stream_client.rs
