meta_srv/
mocks.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::sync::Arc;
16use std::time::Duration;
17
18use api::v1::meta::cluster_server::ClusterServer;
19use api::v1::meta::config_server::ConfigServer;
20use api::v1::meta::heartbeat_server::HeartbeatServer;
21use api::v1::meta::procedure_service_server::ProcedureServiceServer;
22use api::v1::meta::store_server::StoreServer;
23use client::client_manager::NodeClients;
24use common_base::Plugins;
25use common_grpc::channel_manager::{ChannelConfig, ChannelManager};
26use common_meta::key::TableMetadataManager;
27use common_meta::kv_backend::etcd::EtcdStore;
28use common_meta::kv_backend::memory::MemoryKvBackend;
29use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef};
30use hyper_util::rt::TokioIo;
31use servers::grpc::GrpcOptions;
32use tonic::codec::CompressionEncoding;
33use tower::service_fn;
34
35use crate::add_compressed_service;
36use crate::metasrv::builder::MetasrvBuilder;
37use crate::metasrv::{Metasrv, MetasrvOptions, SelectorRef};
38
39#[derive(Clone)]
40pub struct MockInfo {
41    pub server_addr: String,
42    pub channel_manager: ChannelManager,
43    pub metasrv: Arc<Metasrv>,
44    pub kv_backend: KvBackendRef,
45    pub in_memory: Option<ResettableKvBackendRef>,
46}
47
48pub async fn mock_with_memstore() -> MockInfo {
49    let kv_backend = Arc::new(MemoryKvBackend::new());
50    let in_memory = Arc::new(MemoryKvBackend::new());
51    mock(
52        MetasrvOptions {
53            grpc: GrpcOptions {
54                server_addr: "127.0.0.1:3002".to_string(),
55                ..Default::default()
56            },
57            ..Default::default()
58        },
59        kv_backend,
60        None,
61        None,
62        Some(in_memory),
63    )
64    .await
65}
66
67pub async fn mock_with_etcdstore(addr: &str) -> MockInfo {
68    let kv_backend = EtcdStore::with_endpoints([addr], 128).await.unwrap();
69    mock(
70        MetasrvOptions {
71            grpc: GrpcOptions {
72                server_addr: "127.0.0.1:3002".to_string(),
73                ..Default::default()
74            },
75            ..Default::default()
76        },
77        kv_backend,
78        None,
79        None,
80        None,
81    )
82    .await
83}
84
85pub async fn mock(
86    opts: MetasrvOptions,
87    kv_backend: KvBackendRef,
88    selector: Option<SelectorRef>,
89    datanode_clients: Option<Arc<NodeClients>>,
90    in_memory: Option<ResettableKvBackendRef>,
91) -> MockInfo {
92    mock_inner(
93        opts,
94        kv_backend,
95        selector,
96        datanode_clients,
97        in_memory,
98        None,
99        Plugins::default(),
100    )
101    .await
102}
103
104pub async fn mock_with_client_channel_config(
105    opts: MetasrvOptions,
106    kv_backend: KvBackendRef,
107    selector: Option<SelectorRef>,
108    datanode_clients: Option<Arc<NodeClients>>,
109    in_memory: Option<ResettableKvBackendRef>,
110    client_channel_config: ChannelConfig,
111) -> MockInfo {
112    mock_inner(
113        opts,
114        kv_backend,
115        selector,
116        datanode_clients,
117        in_memory,
118        Some(client_channel_config),
119        Plugins::default(),
120    )
121    .await
122}
123
124/// Builds a mock metasrv with the given [`Plugins`].
125///
126/// Same as [`mock`], except that the plugins are passed to the metasrv builder.
127pub async fn mock_with_plugins(
128    opts: MetasrvOptions,
129    kv_backend: KvBackendRef,
130    selector: Option<SelectorRef>,
131    datanode_clients: Option<Arc<NodeClients>>,
132    in_memory: Option<ResettableKvBackendRef>,
133    plugins: Plugins,
134) -> MockInfo {
135    mock_inner(
136        opts,
137        kv_backend,
138        selector,
139        datanode_clients,
140        in_memory,
141        None,
142        plugins,
143    )
144    .await
145}
146
147async fn mock_inner(
148    opts: MetasrvOptions,
149    kv_backend: KvBackendRef,
150    selector: Option<SelectorRef>,
151    datanode_clients: Option<Arc<NodeClients>>,
152    in_memory: Option<ResettableKvBackendRef>,
153    client_channel_config: Option<ChannelConfig>,
154    plugins: Plugins,
155) -> MockInfo {
156    let server_addr = opts.grpc.server_addr.clone();
157    let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend.clone()));
158
159    table_metadata_manager.init().await.unwrap();
160
161    let grpc_config = opts.grpc.as_config();
162    let grpc_options = opts.grpc.clone();
163    let builder = MetasrvBuilder::new()
164        .options(opts)
165        .kv_backend(kv_backend.clone());
166
167    let builder = match selector {
168        Some(s) => builder.selector(s),
169        None => builder,
170    };
171
172    let builder = match datanode_clients {
173        Some(clients) => builder.node_manager(clients),
174        None => builder,
175    };
176
177    let builder = match &in_memory {
178        Some(in_memory) => builder.in_memory(in_memory.clone()),
179        None => builder,
180    };
181
182    let builder = builder.plugins(plugins);
183
184    let metasrv = builder.build().await.unwrap();
185    metasrv.try_start().await.unwrap();
186
187    let (client, server) = tokio::io::duplex(1024);
188    let metasrv = Arc::new(metasrv);
189    let service = metasrv.clone();
190
191    let _handle = tokio::spawn(async move {
192        let mut router = tonic::transport::Server::builder();
193        let router = add_compressed_service!(
194            router,
195            HeartbeatServer::from_arc(service.clone()),
196            grpc_config
197        );
198        let router =
199            add_compressed_service!(router, StoreServer::from_arc(service.clone()), grpc_config);
200        let router = add_compressed_service!(
201            router,
202            ProcedureServiceServer::from_arc(service.clone()),
203            grpc_config
204        );
205        let router = add_compressed_service!(
206            router,
207            ClusterServer::from_arc(service.clone()),
208            grpc_config
209        );
210        let router =
211            add_compressed_service!(router, ConfigServer::from_arc(service.clone()), grpc_config);
212        router
213            .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)]))
214            .await
215    });
216
217    // Keep the mock client's codec limits aligned with the server by default.
218    let config = client_channel_config.unwrap_or_else(|| {
219        let config = ChannelConfig::new()
220            // Use an long timeout to prevent test failures due to slow operations (e.g., when testing with S3).
221            .timeout(Some(Duration::from_secs(60)))
222            .connect_timeout(Duration::from_secs(10))
223            .tcp_nodelay(true);
224        ChannelConfig {
225            max_recv_message_size: grpc_options.max_recv_message_size,
226            max_send_message_size: grpc_options.max_send_message_size,
227            ..config
228        }
229    });
230    let channel_manager = ChannelManager::with_config(config, None);
231
232    // Move client to an option so we can _move_ the inner value
233    // on the first attempt to connect. All other attempts will fail.
234    let mut client = Some(client);
235    let res = channel_manager.reset_with_connector(
236        &server_addr,
237        service_fn(move |_| {
238            let client = client.take();
239
240            async move {
241                if let Some(client) = client {
242                    Ok(TokioIo::new(client))
243                } else {
244                    Err(std::io::Error::other("Client already taken"))
245                }
246            }
247        }),
248    );
249    let _ = res.unwrap();
250
251    MockInfo {
252        server_addr,
253        channel_manager,
254        metasrv,
255        kv_backend,
256        in_memory,
257    }
258}
meta_srv/mocks.rs

meta_srv/
mocks.rs
