risingwave_rpc_client/
frontend_client.rs

1// Copyright 2024 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 async_trait::async_trait;
19use risingwave_common::config::{MAX_CONNECTION_WINDOW_SIZE, RpcClientConfig};
20use risingwave_common::monitor::{EndpointExt, TcpConfig};
21use risingwave_common::util::addr::HostAddr;
22use risingwave_common::util::retry::exponential_backoff;
23use risingwave_pb::frontend_service::frontend_service_client::FrontendServiceClient;
24use risingwave_pb::frontend_service::{
25    CancelRunningSqlRequest, CancelRunningSqlResponse, GetAllCursorsRequest, GetAllCursorsResponse,
26    GetAllSubCursorsRequest, GetAllSubCursorsResponse, GetRunningSqlsRequest,
27    GetRunningSqlsResponse, GetTableReplacePlanRequest, GetTableReplacePlanResponse,
28};
29use tokio_retry::strategy::jitter;
30use tonic::Response;
31use tonic::transport::Endpoint;
32
33use crate::channel::{Channel, WrappedChannelExt};
34use crate::error::Result;
35use crate::{RpcClient, RpcClientPool};
36
37const DEFAULT_RETRY_INTERVAL: u64 = 50;
38const DEFAULT_RETRY_MAX_DELAY: Duration = Duration::from_secs(5);
39const DEFAULT_RETRY_MAX_ATTEMPTS: usize = 10;
40
41#[derive(Clone)]
42struct FrontendClient(FrontendServiceClient<Channel>);
43
44impl FrontendClient {
45    async fn new(host_addr: HostAddr, opts: &RpcClientConfig) -> Result<Self> {
46        let channel = Endpoint::from_shared(format!("http://{}", host_addr))?
47            .initial_connection_window_size(MAX_CONNECTION_WINDOW_SIZE)
48            .connect_timeout(Duration::from_secs(opts.connect_timeout_secs))
49            .monitored_connect(
50                "grpc-frontend-client",
51                TcpConfig {
52                    tcp_nodelay: true,
53                    ..Default::default()
54                },
55            )
56            .await?
57            .wrapped();
58
59        Ok(Self(
60            FrontendServiceClient::new(channel).max_decoding_message_size(usize::MAX),
61        ))
62    }
63}
64
65// similar to the stream_client used in the Meta node
66pub type FrontendClientPool = RpcClientPool<FrontendRetryClient>;
67pub type FrontendClientPoolRef = Arc<FrontendClientPool>;
68
69#[async_trait]
70impl RpcClient for FrontendRetryClient {
71    async fn new_client(host_addr: HostAddr, opts: &RpcClientConfig) -> Result<Self> {
72        Self::new(host_addr, opts).await
73    }
74}
75
76#[derive(Clone)]
77pub struct FrontendRetryClient {
78    client: FrontendClient,
79}
80
81impl FrontendRetryClient {
82    async fn new(host_addr: HostAddr, opts: &RpcClientConfig) -> Result<Self> {
83        let client = FrontendClient::new(host_addr, opts).await?;
84        Ok(Self { client })
85    }
86
87    #[inline(always)]
88    fn get_retry_strategy() -> impl Iterator<Item = Duration> {
89        exponential_backoff(
90            Duration::from_millis(DEFAULT_RETRY_INTERVAL),
91            DEFAULT_RETRY_INTERVAL,
92            DEFAULT_RETRY_MAX_DELAY,
93        )
94        .take(DEFAULT_RETRY_MAX_ATTEMPTS)
95        .map(jitter)
96    }
97
98    fn should_retry(status: &tonic::Status) -> bool {
99        if status.code() == tonic::Code::Unavailable
100            || status.code() == tonic::Code::Unknown
101            || status.code() == tonic::Code::Unauthenticated
102            || status.code() == tonic::Code::Aborted
103        {
104            return true;
105        }
106        false
107    }
108
109    pub async fn get_table_replace_plan(
110        &self,
111        request: GetTableReplacePlanRequest,
112    ) -> std::result::Result<Response<GetTableReplacePlanResponse>, tonic::Status> {
113        tokio_retry::RetryIf::spawn(
114            Self::get_retry_strategy(),
115            || async {
116                self.client
117                    .clone()
118                    .0
119                    .get_table_replace_plan(request.clone())
120                    .await
121            },
122            Self::should_retry,
123        )
124        .await
125    }
126
127    pub async fn get_running_sqls(
128        &self,
129        request: GetRunningSqlsRequest,
130    ) -> std::result::Result<Response<GetRunningSqlsResponse>, tonic::Status> {
131        self.client.0.clone().get_running_sqls(request).await
132    }
133
134    pub async fn get_all_cursors(
135        &self,
136        request: GetAllCursorsRequest,
137    ) -> std::result::Result<Response<GetAllCursorsResponse>, tonic::Status> {
138        self.client.0.clone().get_all_cursors(request).await
139    }
140
141    pub async fn get_all_sub_cursors(
142        &self,
143        request: GetAllSubCursorsRequest,
144    ) -> std::result::Result<Response<GetAllSubCursorsResponse>, tonic::Status> {
145        self.client.0.clone().get_all_sub_cursors(request).await
146    }
147
148    pub async fn cancel_running_sql(
149        &self,
150        request: CancelRunningSqlRequest,
151    ) -> std::result::Result<Response<CancelRunningSqlResponse>, tonic::Status> {
152        self.client.0.clone().cancel_running_sql(request).await
153    }
154}
risingwave_rpc_client/frontend_client.rs

risingwave_rpc_client/
frontend_client.rs
