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

meta_client/client/
store.rs
