meta_client/client/
store.rs

1// Copyright 2023 Greptime Team
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::collections::HashSet;
16use std::sync::Arc;
17
18use api::v1::meta::store_client::StoreClient;
19use api::v1::meta::{
20    BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, BatchPutRequest,
21    BatchPutResponse, CompareAndPutRequest, CompareAndPutResponse, DeleteRangeRequest,
22    DeleteRangeResponse, PutRequest, PutResponse, RangeRequest, RangeResponse, Role,
23};
24use common_grpc::channel_manager::ChannelManager;
25use common_telemetry::tracing_context::TracingContext;
26use snafu::{OptionExt, ResultExt, ensure};
27use tokio::sync::RwLock;
28use tonic::codec::CompressionEncoding;
29use tonic::transport::Channel;
30
31use crate::client::{Id, load_balance as lb};
32use crate::error;
33use crate::error::Result;
34
35#[derive(Clone, Debug)]
36pub struct Client {
37    inner: Arc<RwLock<Inner>>,
38}
39
40impl Client {
41    pub fn new(id: Id, role: Role, channel_manager: ChannelManager) -> Self {
42        Self::new_with_read_only(id, role, channel_manager, true)
43    }
44
45    /// Builds a writable direct Store RPC client for tests.
46    #[cfg(test)]
47    pub(super) fn new_writable(id: Id, role: Role, channel_manager: ChannelManager) -> Self {
48        Self::new_with_read_only(id, role, channel_manager, false)
49    }
50
51    fn new_with_read_only(
52        id: Id,
53        role: Role,
54        channel_manager: ChannelManager,
55        read_only: bool,
56    ) -> Self {
57        let inner = Arc::new(RwLock::new(Inner {
58            id,
59            role,
60            channel_manager,
61            peers: vec![],
62            read_only,
63        }));
64
65        Self { inner }
66    }
67
68    pub async fn start<U, A>(&mut self, urls: A) -> Result<()>
69    where
70        U: AsRef<str>,
71        A: AsRef<[U]>,
72    {
73        let mut inner = self.inner.write().await;
74        inner.start(urls).await
75    }
76
77    pub async fn range(&self, req: RangeRequest) -> Result<RangeResponse> {
78        let inner = self.inner.read().await;
79        inner.range(req).await
80    }
81
82    pub async fn put(&self, req: PutRequest) -> Result<PutResponse> {
83        let inner = self.inner.read().await;
84        inner.put(req).await
85    }
86
87    pub async fn batch_get(&self, req: BatchGetRequest) -> Result<BatchGetResponse> {
88        let inner = self.inner.read().await;
89        inner.batch_get(req).await
90    }
91
92    pub async fn batch_put(&self, req: BatchPutRequest) -> Result<BatchPutResponse> {
93        let inner = self.inner.read().await;
94        inner.batch_put(req).await
95    }
96
97    pub async fn batch_delete(&self, req: BatchDeleteRequest) -> Result<BatchDeleteResponse> {
98        let inner = self.inner.read().await;
99        inner.batch_delete(req).await
100    }
101
102    pub async fn compare_and_put(
103        &self,
104        req: CompareAndPutRequest,
105    ) -> Result<CompareAndPutResponse> {
106        let inner = self.inner.read().await;
107        inner.compare_and_put(req).await
108    }
109
110    pub async fn delete_range(&self, req: DeleteRangeRequest) -> Result<DeleteRangeResponse> {
111        let inner = self.inner.read().await;
112        inner.delete_range(req).await
113    }
114}
115
116#[derive(Debug)]
117struct Inner {
118    id: Id,
119    role: Role,
120    channel_manager: ChannelManager,
121    peers: Vec<String>,
122    read_only: bool,
123}
124
125impl Inner {
126    async fn start<U, A>(&mut self, urls: A) -> Result<()>
127    where
128        U: AsRef<str>,
129        A: AsRef<[U]>,
130    {
131        ensure!(
132            !self.is_started(),
133            error::IllegalGrpcClientStateSnafu {
134                err_msg: "Store client already started",
135            }
136        );
137
138        self.peers = urls
139            .as_ref()
140            .iter()
141            .map(|url| url.as_ref().to_string())
142            .collect::<HashSet<_>>()
143            .drain()
144            .collect::<Vec<_>>();
145
146        Ok(())
147    }
148
149    async fn range(&self, mut req: RangeRequest) -> Result<RangeResponse> {
150        let mut client = self.random_client()?;
151        req.set_header(
152            self.id,
153            self.role,
154            TracingContext::from_current_span().to_w3c(),
155        );
156        let res = client.range(req).await.map_err(error::Error::from)?;
157
158        Ok(res.into_inner())
159    }
160
161    async fn put(&self, mut req: PutRequest) -> Result<PutResponse> {
162        self.ensure_writable()?;
163
164        let mut client = self.random_client()?;
165        req.set_header(
166            self.id,
167            self.role,
168            TracingContext::from_current_span().to_w3c(),
169        );
170        let res = client.put(req).await.map_err(error::Error::from)?;
171
172        Ok(res.into_inner())
173    }
174
175    async fn batch_get(&self, mut req: BatchGetRequest) -> Result<BatchGetResponse> {
176        let mut client = self.random_client()?;
177        req.set_header(
178            self.id,
179            self.role,
180            TracingContext::from_current_span().to_w3c(),
181        );
182
183        let res = client.batch_get(req).await.map_err(error::Error::from)?;
184
185        Ok(res.into_inner())
186    }
187
188    async fn batch_put(&self, mut req: BatchPutRequest) -> Result<BatchPutResponse> {
189        self.ensure_writable()?;
190
191        let mut client = self.random_client()?;
192        req.set_header(
193            self.id,
194            self.role,
195            TracingContext::from_current_span().to_w3c(),
196        );
197        let res = client.batch_put(req).await.map_err(error::Error::from)?;
198
199        Ok(res.into_inner())
200    }
201
202    async fn batch_delete(&self, mut req: BatchDeleteRequest) -> Result<BatchDeleteResponse> {
203        self.ensure_writable()?;
204
205        let mut client = self.random_client()?;
206        req.set_header(
207            self.id,
208            self.role,
209            TracingContext::from_current_span().to_w3c(),
210        );
211        let res = client.batch_delete(req).await.map_err(error::Error::from)?;
212
213        Ok(res.into_inner())
214    }
215
216    async fn compare_and_put(
217        &self,
218        mut req: CompareAndPutRequest,
219    ) -> Result<CompareAndPutResponse> {
220        self.ensure_writable()?;
221
222        let mut client = self.random_client()?;
223        req.set_header(
224            self.id,
225            self.role,
226            TracingContext::from_current_span().to_w3c(),
227        );
228        let res = client
229            .compare_and_put(req)
230            .await
231            .map_err(error::Error::from)?;
232
233        Ok(res.into_inner())
234    }
235
236    async fn delete_range(&self, mut req: DeleteRangeRequest) -> Result<DeleteRangeResponse> {
237        self.ensure_writable()?;
238
239        let mut client = self.random_client()?;
240        req.set_header(
241            self.id,
242            self.role,
243            TracingContext::from_current_span().to_w3c(),
244        );
245        let res = client.delete_range(req).await.map_err(error::Error::from)?;
246
247        Ok(res.into_inner())
248    }
249
250    fn random_client(&self) -> Result<StoreClient<Channel>> {
251        let len = self.peers.len();
252        let peer = lb::random_get(len, |i| Some(&self.peers[i])).context(
253            error::IllegalGrpcClientStateSnafu {
254                err_msg: "Empty peers, store client may not start yet",
255            },
256        )?;
257
258        self.make_client(peer)
259    }
260
261    fn ensure_writable(&self) -> Result<()> {
262        if self.read_only {
263            return error::ReadOnlyKvBackendSnafu {
264                name: "MetaClient Store".to_string(),
265            }
266            .fail();
267        }
268
269        Ok(())
270    }
271
272    fn make_client(&self, addr: impl AsRef<str>) -> Result<StoreClient<Channel>> {
273        let channel = self
274            .channel_manager
275            .get(addr)
276            .context(error::CreateChannelSnafu)?;
277
278        let max_decoding_message_size = self
279            .channel_manager
280            .config()
281            .max_recv_message_size
282            .as_bytes() as usize;
283
284        Ok(StoreClient::new(channel)
285            .accept_compressed(CompressionEncoding::Gzip)
286            .accept_compressed(CompressionEncoding::Zstd)
287            .send_compressed(CompressionEncoding::Zstd)
288            .max_decoding_message_size(max_decoding_message_size))
289    }
290
291    #[inline]
292    fn is_started(&self) -> bool {
293        !self.peers.is_empty()
294    }
295}
296
297#[cfg(test)]
298mod test {
299    use super::*;
300
301    #[tokio::test]
302    async fn test_already_start() {
303        let mut client = Client::new(0, Role::Frontend, ChannelManager::default());
304        client
305            .start(&["127.0.0.1:1000", "127.0.0.1:1001"])
306            .await
307            .unwrap();
308        let res = client.start(&["127.0.0.1:1002"]).await;
309        assert!(res.is_err());
310        assert!(matches!(
311            res.err(),
312            Some(error::Error::IllegalGrpcClientState { .. })
313        ));
314    }
315
316    #[tokio::test]
317    async fn test_start_with_duplicate_peers() {
318        let mut client = Client::new(0, Role::Frontend, ChannelManager::default());
319        client
320            .start(&["127.0.0.1:1000", "127.0.0.1:1000", "127.0.0.1:1000"])
321            .await
322            .unwrap();
323        assert_eq!(1, client.inner.write().await.peers.len());
324    }
325
326    #[tokio::test]
327    async fn test_read_only_store_rejects_writes_before_rpc() {
328        let client = Client::new(0, Role::Frontend, ChannelManager::default());
329
330        fn assert_read_only<T>(result: Result<T>) {
331            assert!(matches!(
332                result,
333                Err(error::Error::ReadOnlyKvBackend { .. })
334            ));
335        }
336
337        assert_read_only(client.put(PutRequest::default()).await);
338        assert_read_only(client.batch_put(BatchPutRequest::default()).await);
339        assert_read_only(client.batch_delete(BatchDeleteRequest::default()).await);
340        assert_read_only(
341            client
342                .compare_and_put(CompareAndPutRequest::default())
343                .await,
344        );
345        assert_read_only(client.delete_range(DeleteRangeRequest::default()).await);
346    }
347}
meta_client/client/store.rs

meta_client/client/
store.rs
