risingwave_rpc_client/
compactor_client.rs

1// Copyright 2023 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 risingwave_common::monitor::EndpointExt;
19use risingwave_common::util::addr::HostAddr;
20use risingwave_common::util::retry::exponential_backoff;
21use risingwave_pb::configured_monitor_service_client;
22use risingwave_pb::hummock::hummock_manager_service_client::HummockManagerServiceClient;
23use risingwave_pb::hummock::{
24    GetNewObjectIdsRequest, GetNewObjectIdsResponse, ReportCompactionTaskRequest,
25    ReportCompactionTaskResponse,
26};
27use risingwave_pb::meta::system_params_service_client::SystemParamsServiceClient;
28use risingwave_pb::meta::{GetSystemParamsRequest, GetSystemParamsResponse};
29use risingwave_pb::monitor_service::monitor_service_client::MonitorServiceClient;
30use risingwave_pb::monitor_service::{StackTraceRequest, StackTraceResponse};
31use tokio::sync::RwLock;
32use tokio_retry::strategy::jitter;
33use tonic::transport::{Channel, Endpoint};
34
35use crate::error::{Result, RpcError};
36use crate::retry_rpc;
37const ENDPOINT_KEEP_ALIVE_INTERVAL_SEC: u64 = 60;
38const ENDPOINT_KEEP_ALIVE_TIMEOUT_SEC: u64 = 60;
39
40const DEFAULT_RETRY_INTERVAL: u64 = 20;
41const DEFAULT_RETRY_MAX_DELAY: Duration = Duration::from_secs(5);
42const DEFAULT_RETRY_MAX_ATTEMPTS: usize = 3;
43#[derive(Clone)]
44pub struct CompactorClient {
45    pub monitor_client: MonitorServiceClient<Channel>,
46}
47
48impl CompactorClient {
49    pub async fn new(host_addr: HostAddr) -> Result<Self> {
50        let channel = Endpoint::from_shared(format!("http://{}", host_addr))?
51            .connect_timeout(Duration::from_secs(5))
52            .monitored_connect("grpc-compactor-client", Default::default())
53            .await?;
54        Ok(Self {
55            monitor_client: configured_monitor_service_client(MonitorServiceClient::new(channel)),
56        })
57    }
58
59    pub async fn stack_trace(&self, req: StackTraceRequest) -> Result<StackTraceResponse> {
60        Ok(self
61            .monitor_client
62            .clone()
63            .stack_trace(req)
64            .await
65            .map_err(RpcError::from_compactor_status)?
66            .into_inner())
67    }
68}
69
70#[derive(Debug, Clone)]
71pub struct GrpcCompactorProxyClientCore {
72    hummock_client: HummockManagerServiceClient<Channel>,
73    system_params_client: SystemParamsServiceClient<Channel>,
74}
75
76impl GrpcCompactorProxyClientCore {
77    pub(crate) fn new(channel: Channel) -> Self {
78        let hummock_client =
79            HummockManagerServiceClient::new(channel.clone()).max_decoding_message_size(usize::MAX);
80        let system_params_client = SystemParamsServiceClient::new(channel);
81
82        Self {
83            hummock_client,
84            system_params_client,
85        }
86    }
87}
88
89/// Client to proxy server. Cloning the instance is lightweight.
90///
91/// Todo(wcy-fdu): add refresh client interface.
92#[derive(Debug, Clone)]
93pub struct GrpcCompactorProxyClient {
94    pub core: Arc<RwLock<GrpcCompactorProxyClientCore>>,
95    endpoint: String,
96}
97
98impl GrpcCompactorProxyClient {
99    pub async fn new(endpoint: String) -> Self {
100        let channel = Self::connect_to_endpoint(endpoint.clone()).await;
101        let core = Arc::new(RwLock::new(GrpcCompactorProxyClientCore::new(channel)));
102        Self { core, endpoint }
103    }
104
105    async fn recreate_core(&self) {
106        tracing::info!("GrpcCompactorProxyClient rpc transfer failed, try to reconnect");
107        let channel = Self::connect_to_endpoint(self.endpoint.clone()).await;
108        let mut core = self.core.write().await;
109        *core = GrpcCompactorProxyClientCore::new(channel);
110    }
111
112    async fn connect_to_endpoint(endpoint: String) -> Channel {
113        let endpoint = Endpoint::from_shared(endpoint).expect("Fail to construct tonic Endpoint");
114        endpoint
115            .http2_keep_alive_interval(Duration::from_secs(ENDPOINT_KEEP_ALIVE_INTERVAL_SEC))
116            .keep_alive_timeout(Duration::from_secs(ENDPOINT_KEEP_ALIVE_TIMEOUT_SEC))
117            .connect_timeout(Duration::from_secs(5))
118            .monitored_connect("grpc-compactor-proxy-client", Default::default())
119            .await
120            .expect("Failed to create channel via proxy rpc endpoint.")
121    }
122
123    pub async fn get_new_sst_ids(
124        &self,
125        request: GetNewObjectIdsRequest,
126    ) -> std::result::Result<tonic::Response<GetNewObjectIdsResponse>, tonic::Status> {
127        retry_rpc!(self, get_new_object_ids, request, GetNewObjectIdsResponse)
128    }
129
130    pub async fn report_compaction_task(
131        &self,
132        request: ReportCompactionTaskRequest,
133    ) -> std::result::Result<tonic::Response<ReportCompactionTaskResponse>, tonic::Status> {
134        retry_rpc!(
135            self,
136            report_compaction_task,
137            request,
138            ReportCompactionTaskResponse
139        )
140    }
141
142    pub async fn get_system_params(
143        &self,
144    ) -> std::result::Result<tonic::Response<GetSystemParamsResponse>, tonic::Status> {
145        tokio_retry::RetryIf::spawn(
146            Self::get_retry_strategy(),
147            || async {
148                let mut system_params_client = self.core.read().await.system_params_client.clone();
149                let rpc_res = system_params_client
150                    .get_system_params(GetSystemParamsRequest {})
151                    .await;
152                if rpc_res.is_err() {
153                    self.recreate_core().await;
154                }
155                rpc_res
156            },
157            Self::should_retry,
158        )
159        .await
160    }
161
162    #[inline(always)]
163    fn get_retry_strategy() -> impl Iterator<Item = Duration> {
164        exponential_backoff(
165            Duration::from_millis(DEFAULT_RETRY_INTERVAL),
166            DEFAULT_RETRY_INTERVAL,
167            DEFAULT_RETRY_MAX_DELAY,
168        )
169        .take(DEFAULT_RETRY_MAX_ATTEMPTS)
170        .map(jitter)
171    }
172
173    #[inline(always)]
174    fn should_retry(status: &tonic::Status) -> bool {
175        if status.code() == tonic::Code::Unavailable
176            || status.code() == tonic::Code::Unknown
177            || (status.code() == tonic::Code::Unauthenticated
178                && status.message().contains("invalid auth token"))
179        {
180            return true;
181        }
182        false
183    }
184}
185
186#[macro_export]
187macro_rules! retry_rpc {
188    ($self:expr, $rpc_call:ident, $request:expr, $response:ty) => {
189        tokio_retry::RetryIf::spawn(
190            Self::get_retry_strategy(),
191            || async {
192                let mut hummock_client = $self.core.read().await.hummock_client.clone();
193                let rpc_res = hummock_client.$rpc_call($request.clone()).await;
194                if rpc_res.is_err() {
195                    $self.recreate_core().await;
196                }
197                rpc_res
198            },
199            Self::should_retry,
200        )
201        .await
202    };
203}
risingwave_rpc_client/compactor_client.rs

risingwave_rpc_client/
compactor_client.rs
