risingwave_rpc_client/
lib.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
15//! Wrapper gRPC clients, which help constructing the request and destructing the
16//! response gRPC message structs.
17
18#![feature(try_blocks)]
19#![feature(impl_trait_in_assoc_type)]
20#![feature(error_generic_member_access)]
21#![feature(panic_update_hook)]
22#![feature(negative_impls)]
23
24use std::any::type_name;
25use std::fmt::{Debug, Formatter};
26use std::future::Future;
27use std::str::FromStr;
28use std::sync::Arc;
29
30use anyhow::{Context, anyhow};
31use async_trait::async_trait;
32pub use compactor_client::{CompactorClient, GrpcCompactorProxyClient};
33pub use compute_client::{ComputeClient, ComputeClientPool, ComputeClientPoolRef};
34pub use connector_client::{SinkCoordinatorStreamHandle, SinkWriterStreamHandle};
35use error::Result;
36pub use frontend_client::{FrontendClientPool, FrontendClientPoolRef};
37use futures::future::try_join_all;
38use futures::stream::{BoxStream, Peekable, TryStreamExt};
39use futures::{Stream, StreamExt};
40pub use hummock_meta_client::{
41    CompactionEventItem, HummockMetaClient, HummockMetaClientChangeLogInfo,
42    IcebergCompactionEventItem,
43};
44pub use meta_client::{MetaClient, SinkCoordinationRpcClient};
45use moka::future::Cache;
46pub use monitor_client::{MonitorClient, MonitorClientPool, MonitorClientPoolRef};
47use rand::prelude::IndexedRandom;
48use risingwave_common::config::RpcClientConfig;
49use risingwave_common::util::addr::HostAddr;
50use risingwave_pb::common::{WorkerNode, WorkerType};
51use rw_futures_util::await_future_with_monitor_error_stream;
52pub use sink_coordinate_client::CoordinatorStreamHandle;
53pub use stream_client::{
54    StreamClient, StreamClientPool, StreamClientPoolRef, StreamingControlHandle,
55};
56use tokio::sync::mpsc::{
57    Receiver, Sender, UnboundedReceiver, UnboundedSender, channel, unbounded_channel,
58};
59
60pub mod error;
61
62mod channel;
63mod compactor_client;
64mod compute_client;
65mod connector_client;
66mod frontend_client;
67mod hummock_meta_client;
68mod meta_client;
69mod monitor_client;
70mod sink_coordinate_client;
71mod stream_client;
72
73#[async_trait]
74pub trait RpcClient: Send + Sync + 'static + Clone {
75    async fn new_client(host_addr: HostAddr, opts: &RpcClientConfig) -> Result<Self>;
76
77    async fn new_clients(
78        host_addr: HostAddr,
79        size: usize,
80        opts: &RpcClientConfig,
81    ) -> Result<Arc<Vec<Self>>> {
82        let make_clients = || {
83            std::iter::repeat_n(host_addr.clone(), size)
84                .map(|host_addr| Self::new_client(host_addr, opts))
85        };
86        let concurrency = opts.pool_setup_concurrency;
87        if concurrency == 0 || concurrency >= size {
88            try_join_all(make_clients()).await.map(Arc::new)
89        } else {
90            futures::stream::iter(make_clients())
91                .buffer_unordered(concurrency)
92                .try_collect::<Vec<_>>()
93                .await
94                .map(Arc::new)
95        }
96    }
97}
98
99#[derive(Clone)]
100pub struct RpcClientPool<S> {
101    connection_pool_size: u16,
102
103    clients: Cache<HostAddr, Arc<Vec<S>>>,
104
105    opts: RpcClientConfig,
106}
107
108impl<S> std::fmt::Debug for RpcClientPool<S> {
109    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
110        f.debug_struct("RpcClientPool")
111            .field("connection_pool_size", &self.connection_pool_size)
112            .field("type", &type_name::<S>())
113            .field("len", &self.clients.entry_count())
114            .finish()
115    }
116}
117
118/// Intentionally not implementing `Default` to let callers be explicit about the pool size.
119impl<S> !Default for RpcClientPool<S> {}
120
121impl<S> RpcClientPool<S>
122where
123    S: RpcClient,
124{
125    /// Create a new pool with the given `connection_pool_size`, which is the number of
126    /// connections to each node that will be reused.
127    pub fn new(connection_pool_size: u16, opts: RpcClientConfig) -> Self {
128        Self {
129            connection_pool_size,
130            clients: Cache::new(u64::MAX),
131            opts,
132        }
133    }
134
135    /// Create a pool for testing purposes. Same as [`Self::adhoc`].
136    pub fn for_test() -> Self {
137        Self::adhoc()
138    }
139
140    /// Create a pool for ad-hoc usage, where the number of connections to each node is 1.
141    pub fn adhoc() -> Self {
142        Self::new(1, RpcClientConfig::default())
143    }
144
145    /// Gets the RPC client for the given node. If the connection is not established, a
146    /// new client will be created and returned.
147    pub async fn get(&self, node: &WorkerNode) -> Result<S> {
148        let addr = if node.get_type().unwrap() == WorkerType::Frontend {
149            let prop = node
150                .property
151                .as_ref()
152                .expect("frontend node property is missing");
153            HostAddr::from_str(prop.internal_rpc_host_addr.as_str())?
154        } else {
155            node.get_host().unwrap().into()
156        };
157
158        self.get_by_addr(addr).await
159    }
160
161    /// Gets the RPC client for the given addr. If the connection is not established, a
162    /// new client will be created and returned.
163    pub async fn get_by_addr(&self, addr: HostAddr) -> Result<S> {
164        Ok(self
165            .clients
166            .try_get_with(
167                addr.clone(),
168                S::new_clients(addr.clone(), self.connection_pool_size as usize, &self.opts),
169            )
170            .await
171            .with_context(|| format!("failed to create RPC client to {addr}"))?
172            .choose(&mut rand::rng())
173            .unwrap()
174            .clone())
175    }
176
177    pub fn invalidate_all(&self) {
178        self.clients.invalidate_all()
179    }
180}
181
182#[macro_export]
183macro_rules! stream_rpc_client_method_impl {
184    ($( { $client:tt, $fn_name:ident, $req:ty, $resp:ty }),*) => {
185        $(
186            pub async fn $fn_name(&self, request: $req) -> $crate::Result<$resp> {
187                Ok(self
188                    .$client
189                    .to_owned()
190                    .$fn_name(request)
191                    .await
192                    .map_err($crate::error::RpcError::from_stream_status)?
193                    .into_inner())
194            }
195        )*
196    }
197}
198
199#[macro_export]
200macro_rules! meta_rpc_client_method_impl {
201    ($( { $client:tt, $fn_name:ident, $req:ty, $resp:ty }),*) => {
202        $(
203            pub async fn $fn_name(&self, request: $req) -> $crate::Result<$resp> {
204                let mut client = self.core.read().await.$client.to_owned();
205                match client.$fn_name(request).await {
206                    Ok(resp) => Ok(resp.into_inner()),
207                    Err(e) => {
208                        self.refresh_client_if_needed(e.code()).await;
209                        Err($crate::error::RpcError::from_meta_status(e))
210                    }
211                }
212            }
213        )*
214    }
215}
216
217pub const DEFAULT_BUFFER_SIZE: usize = 16;
218
219pub struct BidiStreamSender<REQ> {
220    tx: Sender<REQ>,
221}
222
223impl<REQ> BidiStreamSender<REQ> {
224    pub async fn send_request<R: Into<REQ>>(&mut self, request: R) -> Result<()> {
225        self.tx
226            .send(request.into())
227            .await
228            .map_err(|_| anyhow!("unable to send request {}", type_name::<REQ>()).into())
229    }
230}
231
232pub struct BidiStreamReceiver<RSP> {
233    pub stream: Peekable<BoxStream<'static, Result<RSP>>>,
234}
235
236impl<RSP> BidiStreamReceiver<RSP> {
237    pub async fn next_response(&mut self) -> Result<RSP> {
238        self.stream
239            .next()
240            .await
241            .ok_or_else(|| anyhow!("end of response stream"))?
242    }
243}
244
245pub struct BidiStreamHandle<REQ, RSP> {
246    pub request_sender: BidiStreamSender<REQ>,
247    pub response_stream: BidiStreamReceiver<RSP>,
248}
249
250impl<REQ, RSP> Debug for BidiStreamHandle<REQ, RSP> {
251    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
252        f.write_str(type_name::<Self>())
253    }
254}
255
256impl<REQ, RSP> BidiStreamHandle<REQ, RSP> {
257    pub fn for_test(
258        request_sender: Sender<REQ>,
259        response_stream: BoxStream<'static, Result<RSP>>,
260    ) -> Self {
261        Self {
262            request_sender: BidiStreamSender { tx: request_sender },
263            response_stream: BidiStreamReceiver {
264                stream: response_stream.peekable(),
265            },
266        }
267    }
268
269    pub async fn initialize<
270        F: FnOnce(Receiver<REQ>) -> Fut,
271        St: Stream<Item = Result<RSP>> + Send + Unpin + 'static,
272        Fut: Future<Output = Result<St>> + Send,
273        R: Into<REQ>,
274    >(
275        first_request: R,
276        init_stream_fn: F,
277    ) -> Result<(Self, RSP)> {
278        let (request_sender, request_receiver) = channel(DEFAULT_BUFFER_SIZE);
279
280        // Send initial request in case of the blocking receive call from creating streaming request
281        request_sender
282            .send(first_request.into())
283            .await
284            .map_err(|_err| anyhow!("unable to send first request of {}", type_name::<REQ>()))?;
285
286        let mut response_stream = init_stream_fn(request_receiver).await?;
287
288        let first_response = response_stream
289            .next()
290            .await
291            .ok_or_else(|| anyhow!("get empty response from first request"))??;
292
293        Ok((
294            Self {
295                request_sender: BidiStreamSender { tx: request_sender },
296                response_stream: BidiStreamReceiver {
297                    stream: response_stream.boxed().peekable(),
298                },
299            },
300            first_response,
301        ))
302    }
303
304    pub async fn next_response(&mut self) -> Result<RSP> {
305        self.response_stream.next_response().await
306    }
307
308    pub async fn send_request(&mut self, request: REQ) -> Result<()> {
309        match await_future_with_monitor_error_stream(
310            &mut self.response_stream.stream,
311            self.request_sender.send_request(request),
312        )
313        .await
314        {
315            Ok(send_result) => send_result,
316            Err(None) => Err(anyhow!("end of response stream").into()),
317            Err(Some(e)) => Err(e),
318        }
319    }
320}
321
322/// The handle of a bidi-stream started from the rpc client. It is similar to the `BidiStreamHandle`
323/// except that its sender is unbounded.
324pub struct UnboundedBidiStreamHandle<REQ, RSP> {
325    pub request_sender: UnboundedSender<REQ>,
326    pub response_stream: BoxStream<'static, Result<RSP>>,
327}
328
329impl<REQ, RSP> Debug for UnboundedBidiStreamHandle<REQ, RSP> {
330    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
331        f.write_str(type_name::<Self>())
332    }
333}
334
335impl<REQ, RSP> UnboundedBidiStreamHandle<REQ, RSP> {
336    pub async fn initialize<
337        F: FnOnce(UnboundedReceiver<REQ>) -> Fut,
338        St: Stream<Item = Result<RSP>> + Send + Unpin + 'static,
339        Fut: Future<Output = Result<St>> + Send,
340        R: Into<REQ>,
341    >(
342        first_request: R,
343        init_stream_fn: F,
344    ) -> Result<(Self, RSP)> {
345        let (request_sender, request_receiver) = unbounded_channel();
346
347        // Send initial request in case of the blocking receive call from creating streaming request
348        request_sender
349            .send(first_request.into())
350            .map_err(|_err| anyhow!("unable to send first request of {}", type_name::<REQ>()))?;
351
352        let mut response_stream = init_stream_fn(request_receiver).await?;
353
354        let first_response = response_stream
355            .next()
356            .await
357            .context("get empty response from first request")??;
358
359        Ok((
360            Self {
361                request_sender,
362                response_stream: response_stream.boxed(),
363            },
364            first_response,
365        ))
366    }
367
368    pub async fn next_response(&mut self) -> Result<RSP> {
369        self.response_stream
370            .next()
371            .await
372            .ok_or_else(|| anyhow!("end of response stream"))?
373    }
374
375    pub fn send_request(&mut self, request: REQ) -> Result<()> {
376        self.request_sender
377            .send(request)
378            .map_err(|_| anyhow!("unable to send request {}", type_name::<REQ>()).into())
379    }
380}
risingwave_rpc_client/lib.rs

risingwave_rpc_client/
lib.rs
