rocketmq_client_rust/consumer/
default_mq_push_consumer_builder.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::collections::HashMap;
16use std::sync::Arc;
17
18use cheetah_string::CheetahString;
19use rocketmq_model::common::consumer::consume_from_where::ConsumeFromWhere;
20use rocketmq_protocol::protocol::heartbeat::message_model::MessageModel;
21use rocketmq_transport::api::RPCHook;
22
23use crate::base::client_config::ClientConfig;
24use crate::base::client_options::ClientOptions;
25use crate::consumer::allocate_message_queue_strategy::AllocateMessageQueueStrategy;
26use crate::consumer::default_mq_push_consumer::ConsumerConfig;
27use crate::consumer::default_mq_push_consumer::ConsumerTuningProfile;
28use crate::consumer::default_mq_push_consumer::DefaultMQPushConsumer;
29use crate::consumer::message_queue_listener::ArcMessageQueueListener;
30use crate::nameserver_discovery::NameServerDiscoveryConfig;
31use crate::runtime::ClientRuntime;
32use crate::trace::trace_dispatcher::ArcTraceDispatcher;
33
34pub struct DefaultMQPushConsumerBuilder {
35    client_runtime: Arc<ClientRuntime>,
36    client_config: ClientConfig,
37    nameserver_discovery: Option<NameServerDiscoveryConfig>,
38    consumer_group: Option<CheetahString>,
39    message_model: Option<MessageModel>,
40    consume_from_where: Option<ConsumeFromWhere>,
41    consume_timestamp: Option<CheetahString>,
42    allocate_message_queue_strategy: Option<Arc<dyn AllocateMessageQueueStrategy>>,
43    subscription: Option<Arc<HashMap<CheetahString, CheetahString>>>,
44
45    message_queue_listener: Option<ArcMessageQueueListener>,
46
47    consume_thread_min: Option<u32>,
48    consume_thread_max: Option<u32>,
49    adjust_thread_pool_nums_threshold: Option<u64>,
50    consume_concurrently_max_span: Option<u32>,
51    pull_threshold_for_queue: Option<u32>,
52    pop_threshold_for_queue: Option<u32>,
53    pull_threshold_size_for_queue: Option<u32>,
54    pull_threshold_for_topic: Option<i32>,
55    pull_threshold_size_for_topic: Option<i32>,
56    pull_interval: Option<u64>,
57    consume_message_batch_max_size: Option<u32>,
58    pull_batch_size: Option<u32>,
59    pull_batch_size_in_bytes: Option<u32>,
60    post_subscription_when_pull: Option<bool>,
61    unit_mode: Option<bool>,
62    max_reconsume_times: Option<i32>,
63    suspend_current_queue_time_millis: Option<u64>,
64    consume_timeout: Option<u64>,
65    pop_invisible_time: Option<u64>,
66    pop_batch_nums: Option<u32>,
67    await_termination_millis_when_shutdown: Option<u64>,
68    trace_dispatcher: Option<ArcTraceDispatcher>,
69    client_rebalance: Option<bool>,
70    rpc_hook: Option<Arc<dyn RPCHook>>,
71    tuning_profile: Option<ConsumerTuningProfile>,
72}
73
74impl DefaultMQPushConsumerBuilder {
75    #[must_use]
76    pub fn new(client_runtime: Arc<ClientRuntime>) -> Self {
77        Self {
78            client_runtime,
79            client_config: ClientConfig::default(),
80            nameserver_discovery: None,
81            consumer_group: None,
82            message_model: None,
83            consume_from_where: None,
84            consume_timestamp: None,
85            allocate_message_queue_strategy: None,
86            subscription: None,
87            message_queue_listener: None,
88            consume_thread_min: None,
89            consume_thread_max: None,
90            adjust_thread_pool_nums_threshold: None,
91            consume_concurrently_max_span: None,
92            pull_threshold_for_queue: None,
93            pop_threshold_for_queue: None,
94            pull_threshold_size_for_queue: None,
95            pull_threshold_for_topic: None,
96            pull_threshold_size_for_topic: None,
97            pull_interval: None,
98            consume_message_batch_max_size: None,
99            pull_batch_size: None,
100            pull_batch_size_in_bytes: None,
101            post_subscription_when_pull: None,
102            unit_mode: None,
103            max_reconsume_times: None,
104            suspend_current_queue_time_millis: None,
105            consume_timeout: None,
106            pop_invisible_time: None,
107            pop_batch_nums: None,
108            await_termination_millis_when_shutdown: None,
109            trace_dispatcher: None,
110            client_rebalance: None,
111            rpc_hook: None,
112            tuning_profile: None,
113        }
114    }
115    #[inline]
116    pub fn name_server_addr(mut self, name_server_addr: impl Into<CheetahString>) -> Self {
117        self.client_config.namesrv_addr = Some(name_server_addr.into());
118        self.client_config
119            .namespace_initialized
120            .store(false, std::sync::atomic::Ordering::Release);
121        self
122    }
123
124    #[inline]
125    pub fn client_config(mut self, client_config: ClientConfig) -> Self {
126        self.client_config = client_config;
127        self.nameserver_discovery = None;
128        self
129    }
130
131    #[inline]
132    pub fn client_options(mut self, options: ClientOptions) -> Self {
133        self.client_config = options.client_config().clone();
134        self.nameserver_discovery = options.nameserver_discovery().cloned();
135        self
136    }
137
138    #[inline]
139    pub fn nameserver_discovery(mut self, discovery: NameServerDiscoveryConfig) -> Self {
140        self.nameserver_discovery = Some(discovery);
141        self
142    }
143
144    // Methods to set each field
145    #[inline]
146    pub fn consumer_group(mut self, consumer_group: impl Into<CheetahString>) -> Self {
147        self.consumer_group = Some(consumer_group.into());
148        self
149    }
150
151    #[inline]
152    pub fn message_model(mut self, message_model: MessageModel) -> Self {
153        self.message_model = Some(message_model);
154        self
155    }
156
157    #[inline]
158    pub fn consume_from_where(mut self, consume_from_where: ConsumeFromWhere) -> Self {
159        self.consume_from_where = Some(consume_from_where);
160        self
161    }
162
163    #[inline]
164    pub fn consume_timestamp(mut self, consume_timestamp: impl Into<CheetahString>) -> Self {
165        self.consume_timestamp = Some(consume_timestamp.into());
166        self
167    }
168
169    #[inline]
170    pub fn allocate_message_queue_strategy(
171        mut self,
172        allocate_message_queue_strategy: Arc<dyn AllocateMessageQueueStrategy>,
173    ) -> Self {
174        self.allocate_message_queue_strategy = Some(allocate_message_queue_strategy);
175        self
176    }
177
178    #[inline]
179    pub fn message_queue_listener(mut self, message_queue_listener: Option<ArcMessageQueueListener>) -> Self {
180        self.message_queue_listener = message_queue_listener;
181        self
182    }
183
184    #[inline]
185    pub fn consume_thread_min(mut self, consume_thread_min: u32) -> Self {
186        self.consume_thread_min = Some(consume_thread_min);
187        self
188    }
189
190    #[inline]
191    pub fn consume_thread_max(mut self, consume_thread_max: u32) -> Self {
192        self.consume_thread_max = Some(consume_thread_max);
193        self
194    }
195
196    #[inline]
197    pub fn adjust_thread_pool_nums_threshold(mut self, adjust_thread_pool_nums_threshold: u64) -> Self {
198        self.adjust_thread_pool_nums_threshold = Some(adjust_thread_pool_nums_threshold);
199        self
200    }
201
202    #[inline]
203    pub fn consume_concurrently_max_span(mut self, consume_concurrently_max_span: u32) -> Self {
204        self.consume_concurrently_max_span = Some(consume_concurrently_max_span);
205        self
206    }
207
208    #[inline]
209    pub fn pull_threshold_for_queue(mut self, pull_threshold_for_queue: u32) -> Self {
210        self.pull_threshold_for_queue = Some(pull_threshold_for_queue);
211        self
212    }
213
214    #[inline]
215    pub fn pop_threshold_for_queue(mut self, pop_threshold_for_queue: u32) -> Self {
216        self.pop_threshold_for_queue = Some(pop_threshold_for_queue);
217        self
218    }
219
220    #[inline]
221    pub fn pull_threshold_size_for_queue(mut self, pull_threshold_size_for_queue: u32) -> Self {
222        self.pull_threshold_size_for_queue = Some(pull_threshold_size_for_queue);
223        self
224    }
225
226    #[inline]
227    pub fn pull_threshold_for_topic(mut self, pull_threshold_for_topic: i32) -> Self {
228        self.pull_threshold_for_topic = Some(pull_threshold_for_topic);
229        self
230    }
231
232    #[inline]
233    pub fn pull_threshold_size_for_topic(mut self, pull_threshold_size_for_topic: i32) -> Self {
234        self.pull_threshold_size_for_topic = Some(pull_threshold_size_for_topic);
235        self
236    }
237
238    #[inline]
239    pub fn pull_interval(mut self, pull_interval: u64) -> Self {
240        self.pull_interval = Some(pull_interval);
241        self
242    }
243
244    #[inline]
245    pub fn consume_message_batch_max_size(mut self, consume_message_batch_max_size: u32) -> Self {
246        self.consume_message_batch_max_size = Some(consume_message_batch_max_size);
247        self
248    }
249
250    #[inline]
251    pub fn pull_batch_size(mut self, pull_batch_size: u32) -> Self {
252        self.pull_batch_size = Some(pull_batch_size);
253        self
254    }
255
256    #[inline]
257    pub fn pull_batch_size_in_bytes(mut self, pull_batch_size_in_bytes: u32) -> Self {
258        self.pull_batch_size_in_bytes = Some(pull_batch_size_in_bytes);
259        self
260    }
261
262    #[inline]
263    pub fn post_subscription_when_pull(mut self, post_subscription_when_pull: bool) -> Self {
264        self.post_subscription_when_pull = Some(post_subscription_when_pull);
265        self
266    }
267
268    #[inline]
269    pub fn unit_mode(mut self, unit_mode: bool) -> Self {
270        self.unit_mode = Some(unit_mode);
271        self
272    }
273
274    #[inline]
275    pub fn max_reconsume_times(mut self, max_reconsume_times: i32) -> Self {
276        self.max_reconsume_times = Some(max_reconsume_times);
277        self
278    }
279
280    #[inline]
281    pub fn suspend_current_queue_time_millis(mut self, suspend_current_queue_time_millis: u64) -> Self {
282        self.suspend_current_queue_time_millis = Some(suspend_current_queue_time_millis);
283        self
284    }
285
286    #[inline]
287    pub fn consume_timeout(mut self, consume_timeout: u64) -> Self {
288        self.consume_timeout = Some(consume_timeout);
289        self
290    }
291
292    #[inline]
293    pub fn pop_invisible_time(mut self, pop_invisible_time: u64) -> Self {
294        self.pop_invisible_time = Some(pop_invisible_time);
295        self
296    }
297
298    #[inline]
299    pub fn pop_batch_nums(mut self, pop_batch_nums: u32) -> Self {
300        self.pop_batch_nums = Some(pop_batch_nums);
301        self
302    }
303
304    #[inline]
305    pub fn await_termination_millis_when_shutdown(mut self, await_termination_millis_when_shutdown: u64) -> Self {
306        self.await_termination_millis_when_shutdown = Some(await_termination_millis_when_shutdown);
307        self
308    }
309
310    #[inline]
311    pub fn trace_dispatcher(mut self, trace_dispatcher: Option<ArcTraceDispatcher>) -> Self {
312        self.trace_dispatcher = trace_dispatcher;
313        self
314    }
315
316    #[inline]
317    pub fn client_rebalance(mut self, client_rebalance: bool) -> Self {
318        self.client_rebalance = Some(client_rebalance);
319        self
320    }
321
322    #[inline]
323    pub fn rpc_hook(mut self, rpc_hook: Option<Arc<dyn RPCHook>>) -> Self {
324        self.rpc_hook = rpc_hook;
325        self
326    }
327
328    #[inline]
329    pub fn tuning_profile(mut self, tuning_profile: ConsumerTuningProfile) -> Self {
330        self.tuning_profile = Some(tuning_profile);
331        self
332    }
333
334    // Build method to create a DefaultMQPushConsumer instance
335    pub fn build(mut self) -> DefaultMQPushConsumer {
336        let mut consumer_config = ConsumerConfig::default();
337        if let Some(profile) = self.tuning_profile {
338            consumer_config.apply_tuning_profile(profile);
339        }
340
341        // Set optional fields with consistent pattern
342        if let Some(value) = self.consumer_group {
343            consumer_config.consumer_group = value;
344        }
345        if let Some(value) = self.message_model {
346            consumer_config.message_model = value;
347        }
348        if let Some(value) = self.consume_from_where {
349            consumer_config.consume_from_where = value;
350        }
351        if let Some(value) = self.consume_timestamp {
352            consumer_config.consume_timestamp = Some(value);
353        }
354        if let Some(value) = self.allocate_message_queue_strategy.take() {
355            consumer_config.allocate_message_queue_strategy = Some(value);
356        }
357        if let Some(value) = self.subscription {
358            consumer_config.subscription = value;
359        }
360        consumer_config.message_queue_listener = self.message_queue_listener.take();
361
362        // Thread pool configuration
363        if let Some(value) = self.consume_thread_min {
364            consumer_config.consume_thread_min = value;
365        }
366        if let Some(value) = self.consume_thread_max {
367            consumer_config.consume_thread_max = value;
368        }
369        if let Some(value) = self.adjust_thread_pool_nums_threshold {
370            consumer_config.adjust_thread_pool_nums_threshold = value;
371        }
372
373        // Pull thresholds
374        if let Some(value) = self.consume_concurrently_max_span {
375            consumer_config.consume_concurrently_max_span = value;
376        }
377        if let Some(value) = self.pull_threshold_for_queue {
378            consumer_config.pull_threshold_for_queue = value;
379        }
380        if let Some(value) = self.pop_threshold_for_queue {
381            consumer_config.pop_threshold_for_queue = value;
382        }
383        if let Some(value) = self.pull_threshold_size_for_queue {
384            consumer_config.pull_threshold_size_for_queue = value;
385        }
386        if let Some(value) = self.pull_threshold_for_topic {
387            consumer_config.pull_threshold_for_topic = value;
388        }
389        if let Some(value) = self.pull_threshold_size_for_topic {
390            consumer_config.pull_threshold_size_for_topic = value;
391        }
392        if let Some(value) = self.pull_interval {
393            consumer_config.pull_interval = value;
394        }
395
396        // Batch configuration
397        if let Some(value) = self.consume_message_batch_max_size {
398            consumer_config.consume_message_batch_max_size = value;
399        }
400        if let Some(value) = self.pull_batch_size {
401            consumer_config.pull_batch_size = value;
402        }
403        if let Some(value) = self.pull_batch_size_in_bytes {
404            consumer_config.pull_batch_size_in_bytes = value;
405        }
406
407        // Other configuration
408        if let Some(value) = self.post_subscription_when_pull {
409            consumer_config.post_subscription_when_pull = value;
410        }
411        if let Some(value) = self.unit_mode {
412            consumer_config.unit_mode = value;
413        }
414        if let Some(value) = self.max_reconsume_times {
415            consumer_config.max_reconsume_times = value;
416        }
417        if let Some(value) = self.suspend_current_queue_time_millis {
418            consumer_config.suspend_current_queue_time_millis = value;
419        }
420        if let Some(value) = self.consume_timeout {
421            consumer_config.consume_timeout = value;
422        }
423        if let Some(value) = self.pop_invisible_time {
424            consumer_config.pop_invisible_time = value;
425        }
426        if let Some(value) = self.pop_batch_nums {
427            consumer_config.pop_batch_nums = value;
428        }
429        if let Some(value) = self.await_termination_millis_when_shutdown {
430            consumer_config.await_termination_millis_when_shutdown = value;
431        }
432
433        // Trace and RPC configuration
434        consumer_config.trace_dispatcher = self.trace_dispatcher.clone();
435        if let Some(value) = self.client_rebalance {
436            consumer_config.client_rebalance = value;
437        }
438        consumer_config.rpc_hook = self.rpc_hook.clone();
439
440        DefaultMQPushConsumer::new_with_options(
441            self.client_runtime,
442            ClientOptions::from_parts(self.client_config, self.nameserver_discovery),
443            consumer_config,
444        )
445    }
446}
447
448#[cfg(test)]
449mod tests {
450    use super::*;
451    use crate::consumer::default_mq_push_consumer::ConsumerTuningProfile;
452
453    #[test]
454    fn tuning_profile_applies_before_explicit_builder_overrides() {
455        let consumer =
456            DefaultMQPushConsumerBuilder::new(crate::runtime::test_client_runtime("push-consumer-builder-test"))
457                .consumer_group("push_profile_builder_group")
458                .tuning_profile(ConsumerTuningProfile::Throughput)
459                .pull_batch_size(64)
460                .pull_batch_size_in_bytes(768 * 1024)
461                .build();
462
463        assert_eq!(consumer.consume_message_batch_max_size(), 16);
464        assert_eq!(consumer.pull_batch_size(), 64);
465        assert_eq!(consumer.pull_batch_size_in_bytes(), 768 * 1024);
466        assert_eq!(consumer.pull_threshold_for_queue(), 4096);
467        assert_eq!(consumer.pull_threshold_size_for_queue(), 256);
468    }
469}
rocketmq_client_rust/consumer/default_mq_push_consumer_builder.rs

rocketmq_client_rust/consumer/
default_mq_push_consumer_builder.rs
