rocketmq_client_rust/implementation/
mq_client_manager.rs

1// Copyright 2023 The RocketMQ Rust Authors
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::atomic::AtomicBool;
16use std::sync::atomic::AtomicU64;
17use std::sync::atomic::Ordering;
18use std::sync::Arc;
19
20use cheetah_string::CheetahString;
21use dashmap::DashMap;
22use parking_lot::RwLock;
23use rocketmq_observability::metrics::client::ClientMetrics;
24use rocketmq_observability::TelemetryHandle;
25use rocketmq_protocol::protocol::remoting_command_defaults::RemotingCommandDefaults;
26use rocketmq_protocol::protocol::remoting_command_defaults::RemotingCommandFactory;
27use rocketmq_runtime::ChildServiceContext;
28use rocketmq_runtime::ResourceBudget;
29use rocketmq_runtime::ShutdownDeadline;
30use rocketmq_transport::api::RPCHook;
31use tracing::info;
32
33use crate::base::client_config::ClientConfig;
34use crate::base::client_options::ClientOptions;
35use crate::factory::mq_client_instance::MQClientInstance;
36use crate::nameserver_discovery::NameServerDiscoveryConfig;
37use crate::producer::produce_accumulator::ProduceAccumulator;
38use crate::producer::request_future_holder::RequestFutureHolder;
39
40struct ClientPoolEntry {
41    generation: u64,
42    leases: usize,
43    identity: ClientPoolIdentity,
44    instance: Arc<MQClientInstance>,
45}
46
47#[derive(Clone, Eq, PartialEq)]
48struct ClientPoolIdentity {
49    namesrv_addr: Option<CheetahString>,
50    discovery_fingerprint: Option<CheetahString>,
51    use_tls: bool,
52    vip_channel_enabled: bool,
53    api_timeout_millis: u64,
54    rpc_hook_identity: Option<usize>,
55    remoting_command_defaults: RemotingCommandDefaults,
56}
57
58impl ClientPoolIdentity {
59    fn new(
60        client_config: &ClientConfig,
61        discovery: Option<&NameServerDiscoveryConfig>,
62        rpc_hook: Option<&Arc<dyn RPCHook>>,
63        remoting_command_factory: RemotingCommandFactory,
64    ) -> Self {
65        Self {
66            namesrv_addr: client_config.namesrv_addr.clone(),
67            discovery_fingerprint: discovery.map(|config| CheetahString::from_string(config.fingerprint())),
68            use_tls: client_config.use_tls,
69            vip_channel_enabled: client_config.vip_channel_enabled,
70            api_timeout_millis: client_config.mq_client_api_timeout,
71            rpc_hook_identity: rpc_hook.map(|hook| Arc::as_ptr(hook).cast::<()>() as usize),
72            remoting_command_defaults: remoting_command_factory.defaults(),
73        }
74    }
75}
76
77type ClientInstanceHashMap = DashMap<CheetahString /* clientId */, ClientPoolEntry>;
78type AccumulatorHashMap = DashMap<CheetahString /* clientId */, Arc<ProduceAccumulator>>;
79
80#[derive(Clone)]
81pub struct ClientPool {
82    inner: Arc<ClientPoolInner>,
83}
84
85struct ClientPoolInner {
86    service_context: ChildServiceContext,
87    resource_budget: ResourceBudget,
88    telemetry_handle: TelemetryHandle,
89    client_metrics: ClientMetrics,
90    factory_table: Arc<ClientInstanceHashMap>,
91    accumulator_table: Arc<AccumulatorHashMap>,
92    request_future_holder: Arc<RequestFutureHolder>,
93    remoting_command_factory: RemotingCommandFactory,
94    next_generation: AtomicU64,
95    closed: AtomicBool,
96    admission: RwLock<()>,
97}
98
99#[derive(Clone, Debug, Eq, PartialEq)]
100pub struct ClientPoolToken {
101    client_id: CheetahString,
102    generation: u64,
103}
104
105pub struct PooledClient {
106    instance: Arc<MQClientInstance>,
107    token: ClientPoolToken,
108}
109
110impl PooledClient {
111    pub fn instance(&self) -> &Arc<MQClientInstance> {
112        &self.instance
113    }
114
115    pub fn into_parts(self) -> (Arc<MQClientInstance>, ClientPoolToken) {
116        (self.instance, self.token)
117    }
118}
119
120impl ClientPool {
121    pub(crate) fn new(
122        service_context: ChildServiceContext,
123        resource_budget: ResourceBudget,
124        telemetry_handle: TelemetryHandle,
125        client_metrics: ClientMetrics,
126        remoting_command_factory: RemotingCommandFactory,
127    ) -> Self {
128        let request_future_holder = Arc::new(RequestFutureHolder::new(service_context.component("request-futures")));
129        Self {
130            inner: Arc::new(ClientPoolInner {
131                service_context,
132                resource_budget,
133                telemetry_handle,
134                client_metrics,
135                factory_table: Arc::new(DashMap::with_capacity(128)),
136                accumulator_table: Arc::new(DashMap::with_capacity(128)),
137                request_future_holder,
138                remoting_command_factory,
139                next_generation: AtomicU64::new(0),
140                closed: AtomicBool::new(false),
141                admission: RwLock::new(()),
142            }),
143        }
144    }
145
146    pub fn instance_count(&self) -> usize {
147        self.inner.factory_table.len()
148    }
149
150    pub(crate) fn request_future_holder(&self) -> Arc<RequestFutureHolder> {
151        Arc::clone(&self.inner.request_future_holder)
152    }
153
154    pub fn get_or_create(
155        &self,
156        client_config: ClientConfig,
157        rpc_hook: Option<Arc<dyn RPCHook>>,
158    ) -> crate::ClientResult<PooledClient> {
159        self.get_or_create_with_options(ClientOptions::legacy(client_config), rpc_hook)
160    }
161
162    pub fn get_or_create_with_options(
163        &self,
164        options: ClientOptions,
165        rpc_hook: Option<Arc<dyn RPCHook>>,
166    ) -> crate::ClientResult<PooledClient> {
167        let _admission = self.inner.admission.read();
168        if self.inner.closed.load(Ordering::Acquire) {
169            return Err(mq_client_err!("ClientRuntime is shutting down"));
170        }
171        let (client_config, discovery, command_factory_override) = options.into_normalized_parts()?;
172        let remoting_command_factory = command_factory_override.unwrap_or(self.inner.remoting_command_factory);
173        let client_id = CheetahString::from_string(client_config.build_mq_client_id());
174        let identity = ClientPoolIdentity::new(
175            &client_config,
176            discovery.as_ref(),
177            rpc_hook.as_ref(),
178            remoting_command_factory,
179        );
180
181        let mut entry = self.inner.factory_table.entry(client_id.clone()).or_insert_with(|| {
182            let generation = self.inner.next_generation.fetch_add(1, Ordering::AcqRel) + 1;
183            info!(
184                client_id = %client_id,
185                generation,
186                "Created new MQClientInstance in ClientPool"
187            );
188            let instance = MQClientInstance::new_arc_with_resource_budget_and_discovery(
189                client_config,
190                discovery,
191                remoting_command_factory,
192                generation,
193                client_id.clone(),
194                rpc_hook,
195                self.inner.service_context.component(
196                    rocketmq_runtime::ScopeId::try_new(format!("instance-{generation}"))
197                        .expect("the client instance scope has a fixed nonblank prefix"),
198                ),
199                Arc::clone(&self.inner.request_future_holder),
200                self.inner.resource_budget.clone(),
201                self.inner.telemetry_handle.clone(),
202                self.inner.client_metrics.clone(),
203            );
204            ClientPoolEntry {
205                generation,
206                leases: 0,
207                identity: identity.clone(),
208                instance,
209            }
210        });
211        if entry.identity != identity {
212            return Err(mq_client_err!(
213                "ClientPool configuration conflicts with an existing client-id owner"
214            ));
215        }
216        entry.leases += 1;
217
218        Ok(PooledClient {
219            instance: Arc::clone(&entry.instance),
220            token: ClientPoolToken {
221                client_id,
222                generation: entry.generation,
223            },
224        })
225    }
226
227    pub fn get_or_create_produce_accumulator(&self, client_config: ClientConfig) -> Option<Arc<ProduceAccumulator>> {
228        let _admission = self.inner.admission.read();
229        if self.inner.closed.load(Ordering::Acquire) {
230            return None;
231        }
232        let client_id = CheetahString::from_string(client_config.build_mq_client_id());
233
234        Some(
235            self.inner
236                .accumulator_table
237                .entry(client_id.clone())
238                .or_insert_with(|| {
239                    info!("Created new ProduceAccumulator for clientId:[{}]", client_id);
240                    Arc::new(ProduceAccumulator::with_resource_budget(
241                        self.inner.service_context.component(
242                            rocketmq_runtime::ScopeId::try_new(format!("accumulator-{}", client_id))
243                                .expect("the accumulator scope has a fixed nonblank prefix"),
244                        ),
245                        client_id.as_str(),
246                        self.inner.resource_budget.clone(),
247                    ))
248                })
249                .clone(),
250        )
251    }
252
253    pub async fn release(&self, token: ClientPoolToken) -> bool {
254        let should_remove = self
255            .inner
256            .factory_table
257            .get_mut(&token.client_id)
258            .is_some_and(|mut current| {
259                if current.generation != token.generation || current.leases == 0 {
260                    return false;
261                }
262                current.leases -= 1;
263                current.leases == 0
264            });
265        if !should_remove {
266            return false;
267        }
268
269        let removed = self
270            .inner
271            .factory_table
272            .remove_if(&token.client_id, |_, current| {
273                current.generation == token.generation && current.leases == 0
274            })
275            .map(|(_, entry)| entry.instance);
276        if let Some(instance) = removed {
277            instance.shutdown().await;
278            true
279        } else {
280            false
281        }
282    }
283
284    pub(crate) async fn shutdown_until(&self, deadline: ShutdownDeadline) {
285        let instances = {
286            let _admission = self.inner.admission.write();
287            self.inner.closed.store(true, Ordering::Release);
288            let keys = self
289                .inner
290                .factory_table
291                .iter()
292                .map(|entry| entry.key().clone())
293                .collect::<Vec<_>>();
294            let mut instances = Vec::with_capacity(keys.len());
295            for key in keys {
296                if let Some((_, entry)) = self.inner.factory_table.remove(&key) {
297                    instances.push(entry.instance);
298                }
299            }
300            self.inner.accumulator_table.clear();
301            instances
302        };
303
304        for instance in instances {
305            if deadline.is_expired() {
306                break;
307            }
308            let _ = tokio::time::timeout(deadline.remaining(), instance.shutdown()).await;
309        }
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use crate::test_support::error_assertions::client_exception;
316    use cheetah_string::CheetahString;
317    use rocketmq_protocol::protocol::SerializeType;
318
319    use super::*;
320
321    #[test]
322    fn invalid_nameserver_address_is_rejected_before_pool_admission() {
323        let runtime = crate::runtime::test_client_runtime("invalid-nameserver-pool-test");
324        let config = ClientConfig {
325            namesrv_addr: Some(CheetahString::from_static_str("missing-port")),
326            ..Default::default()
327        };
328
329        let error = match runtime.pool().get_or_create(config, None) {
330            Ok(_) => panic!("invalid NameServer address must be rejected"),
331            Err(error) => error,
332        };
333
334        assert!(error.is(&rocketmq_error::CORE_CONFIGURATION_INVALID));
335        assert_eq!(runtime.pool().instance_count(), 0);
336    }
337
338    #[test]
339    fn canonical_nameserver_address_drives_pool_identity() {
340        let runtime = crate::runtime::test_client_runtime("canonical-nameserver-pool-test");
341        let first_config = ClientConfig {
342            namesrv_addr: Some(CheetahString::from_static_str(" NS-A:9876 ; ns-a:9876 ")),
343            ..Default::default()
344        };
345        let first = runtime
346            .pool()
347            .get_or_create(first_config, None)
348            .expect("first pooled client");
349
350        let equivalent_config = ClientConfig {
351            namesrv_addr: Some(CheetahString::from_static_str("ns-a:9876")),
352            ..Default::default()
353        };
354        let equivalent = runtime
355            .pool()
356            .get_or_create(equivalent_config, None)
357            .expect("equivalent canonical target must share the pool entry");
358        assert!(Arc::ptr_eq(first.instance(), equivalent.instance()));
359
360        let different_config = ClientConfig {
361            namesrv_addr: Some(CheetahString::from_static_str("ns-b:9876")),
362            ..Default::default()
363        };
364        assert!(runtime.pool().get_or_create(different_config, None).is_err());
365    }
366
367    #[test]
368    fn remoting_defaults_participate_in_pool_identity() {
369        let runtime = crate::runtime::test_client_runtime("remoting-defaults-pool-identity-test");
370        let mut config = ClientConfig::default();
371        config.set_instance_name("remoting-defaults-owner".into());
372        let json = RemotingCommandFactory::new(RemotingCommandDefaults::new(501, SerializeType::JSON));
373        let binary = RemotingCommandFactory::new(RemotingCommandDefaults::new(502, SerializeType::ROCKETMQ));
374
375        runtime
376            .pool()
377            .get_or_create_with_options(
378                ClientOptions::legacy(config.clone()).with_remoting_command_factory(json),
379                None,
380            )
381            .expect("first command owner");
382        let error = match runtime.pool().get_or_create_with_options(
383            ClientOptions::legacy(config).with_remoting_command_factory(binary),
384            None,
385        ) {
386            Ok(_) => panic!("different remoting defaults must not share one client-id owner"),
387            Err(error) => error,
388        };
389
390        assert!(client_exception(&error)
391            .to_string()
392            .contains("ClientPool configuration conflicts with an existing client-id owner"));
393    }
394}
rocketmq_client_rust/implementation/mq_client_manager.rs

rocketmq_client_rust/implementation/
mq_client_manager.rs
