Skip to main content

rocketmq_client_rust/producer/
transaction_mq_produce_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::HashSet;
16use std::sync::Arc;
17
18use cheetah_string::CheetahString;
19use rocketmq_model::common::compression::compression_type::CompressionType;
20use rocketmq_protocol::common::compression::compressor::Compressor;
21use rocketmq_transport::api::RPCHook;
22
23use crate::base::client_config::ClientConfig;
24use crate::base::client_options::ClientOptions;
25use crate::nameserver_discovery::NameServerDiscoveryConfig;
26use crate::producer::default_mq_producer::DefaultMQProducer;
27use crate::producer::produce_accumulator::ProduceAccumulator;
28use crate::producer::producer_impl::default_mq_producer_impl::DefaultMQProducerImpl;
29use crate::producer::transaction_listener::ArcTransactionListener;
30use crate::producer::transaction_listener::TransactionListener;
31use crate::producer::transaction_mq_producer::TransactionMQProducer;
32use crate::producer::transaction_mq_producer::TransactionProducerConfig;
33use crate::runtime::ClientRuntime;
34use crate::session::ClientSession;
35use crate::trace::trace_dispatcher::ArcTraceDispatcher;
36
37pub struct TransactionMQProducerBuilder {
38    client_runtime: Arc<ClientRuntime>,
39    client_config: Option<ClientConfig>,
40    nameserver_discovery: Option<NameServerDiscoveryConfig>,
41    default_mqproducer_impl: Option<DefaultMQProducerImpl>,
42    retry_response_codes: Option<HashSet<i32>>,
43    producer_group: Option<CheetahString>,
44    topics: Option<Vec<CheetahString>>,
45    create_topic_key: Option<CheetahString>,
46    default_topic_queue_nums: Option<u32>,
47    send_msg_timeout: Option<u32>,
48    compress_msg_body_over_howmuch: Option<u32>,
49    retry_times_when_send_failed: Option<u32>,
50    retry_times_when_send_async_failed: Option<u32>,
51    retry_another_broker_when_not_store_ok: Option<bool>,
52    max_message_size: Option<u32>,
53    trace_dispatcher: Option<ArcTraceDispatcher>,
54    auto_batch: Option<bool>,
55    produce_accumulator: Option<ProduceAccumulator>,
56    enable_backpressure_for_async_mode: Option<bool>,
57    back_pressure_for_async_send_num: Option<u32>,
58    back_pressure_for_async_send_size: Option<u32>,
59    rpc_hook: Option<Arc<dyn RPCHook>>,
60    compress_level: Option<i32>,
61    compress_type: Option<CompressionType>,
62    compressor: Option<&'static (dyn Compressor + Send + Sync)>,
63    transaction_listener: Option<ArcTransactionListener>,
64    check_thread_pool_min_size: Option<u32>,
65    check_thread_pool_max_size: Option<u32>,
66    check_request_hold_max: Option<u32>,
67}
68
69impl TransactionMQProducerBuilder {
70    pub fn new(client_runtime: Arc<ClientRuntime>) -> Self {
71        Self {
72            client_runtime,
73            client_config: Some(Default::default()),
74            nameserver_discovery: None,
75            default_mqproducer_impl: None,
76            retry_response_codes: None,
77            producer_group: None,
78            topics: None,
79            create_topic_key: None,
80            default_topic_queue_nums: None,
81            send_msg_timeout: None,
82            compress_msg_body_over_howmuch: None,
83            retry_times_when_send_failed: None,
84            retry_times_when_send_async_failed: None,
85            retry_another_broker_when_not_store_ok: None,
86            max_message_size: None,
87            trace_dispatcher: None,
88            auto_batch: None,
89            produce_accumulator: None,
90            enable_backpressure_for_async_mode: None,
91            back_pressure_for_async_send_num: None,
92            back_pressure_for_async_send_size: None,
93            rpc_hook: None,
94            compress_level: None,
95            compress_type: None,
96            compressor: None,
97            transaction_listener: None,
98            check_thread_pool_min_size: None,
99            check_thread_pool_max_size: None,
100            check_request_hold_max: None,
101        }
102    }
103
104    pub fn client_config(mut self, client_config: ClientConfig) -> Self {
105        self.client_config = Some(client_config);
106        self.nameserver_discovery = None;
107        self
108    }
109
110    pub fn client_options(mut self, options: ClientOptions) -> Self {
111        self.client_config = Some(options.client_config().clone());
112        self.nameserver_discovery = options.nameserver_discovery().cloned();
113        self
114    }
115
116    pub fn nameserver_discovery(mut self, discovery: NameServerDiscoveryConfig) -> Self {
117        self.nameserver_discovery = Some(discovery);
118        self
119    }
120
121    pub fn default_mqproducer_impl(mut self, default_mqproducer_impl: DefaultMQProducerImpl) -> Self {
122        self.default_mqproducer_impl = Some(default_mqproducer_impl);
123        self
124    }
125
126    pub fn retry_response_codes(mut self, retry_response_codes: HashSet<i32>) -> Self {
127        self.retry_response_codes = Some(retry_response_codes);
128        self
129    }
130
131    pub fn producer_group(mut self, producer_group: impl Into<CheetahString>) -> Self {
132        self.producer_group = Some(producer_group.into());
133        self
134    }
135
136    pub fn topics(mut self, topics: Vec<impl Into<CheetahString>>) -> Self {
137        self.topics = Some(topics.into_iter().map(|t| t.into()).collect());
138        self
139    }
140
141    pub fn name_server_addr(mut self, name_server_addr: impl Into<CheetahString>) -> Self {
142        if let Some(client_config) = self.client_config.as_mut() {
143            client_config.namesrv_addr = Some(name_server_addr.into());
144            client_config
145                .namespace_initialized
146                .store(false, std::sync::atomic::Ordering::Release);
147        }
148        self
149    }
150
151    pub fn use_tls(mut self, use_tls: bool) -> Self {
152        self.client_config
153            .get_or_insert_with(ClientConfig::default)
154            .set_use_tls(use_tls);
155        self
156    }
157
158    pub fn create_topic_key(mut self, create_topic_key: impl Into<CheetahString>) -> Self {
159        self.create_topic_key = Some(create_topic_key.into());
160        self
161    }
162
163    pub fn default_topic_queue_nums(mut self, default_topic_queue_nums: u32) -> Self {
164        self.default_topic_queue_nums = Some(default_topic_queue_nums);
165        self
166    }
167
168    pub fn send_msg_timeout(mut self, send_msg_timeout: u32) -> Self {
169        self.send_msg_timeout = Some(send_msg_timeout);
170        self
171    }
172
173    pub fn compress_msg_body_over_howmuch(mut self, compress_msg_body_over_howmuch: u32) -> Self {
174        self.compress_msg_body_over_howmuch = Some(compress_msg_body_over_howmuch);
175        self
176    }
177
178    pub fn retry_times_when_send_failed(mut self, retry_times_when_send_failed: u32) -> Self {
179        self.retry_times_when_send_failed = Some(retry_times_when_send_failed);
180        self
181    }
182
183    pub fn retry_times_when_send_async_failed(mut self, retry_times_when_send_async_failed: u32) -> Self {
184        self.retry_times_when_send_async_failed = Some(retry_times_when_send_async_failed);
185        self
186    }
187
188    pub fn retry_another_broker_when_not_store_ok(mut self, retry_another_broker_when_not_store_ok: bool) -> Self {
189        self.retry_another_broker_when_not_store_ok = Some(retry_another_broker_when_not_store_ok);
190        self
191    }
192
193    pub fn max_message_size(mut self, max_message_size: u32) -> Self {
194        self.max_message_size = Some(max_message_size);
195        self
196    }
197
198    pub fn trace_dispatcher(mut self, trace_dispatcher: ArcTraceDispatcher) -> Self {
199        self.trace_dispatcher = Some(trace_dispatcher);
200        self
201    }
202
203    pub fn auto_batch(mut self, auto_batch: bool) -> Self {
204        self.auto_batch = Some(auto_batch);
205        self
206    }
207
208    pub fn produce_accumulator(mut self, produce_accumulator: ProduceAccumulator) -> Self {
209        self.produce_accumulator = Some(produce_accumulator);
210        self
211    }
212
213    pub fn enable_backpressure_for_async_mode(mut self, enable_backpressure_for_async_mode: bool) -> Self {
214        self.enable_backpressure_for_async_mode = Some(enable_backpressure_for_async_mode);
215        self
216    }
217
218    pub fn back_pressure_for_async_send_num(mut self, back_pressure_for_async_send_num: u32) -> Self {
219        self.back_pressure_for_async_send_num = Some(back_pressure_for_async_send_num);
220        self
221    }
222
223    pub fn back_pressure_for_async_send_size(mut self, back_pressure_for_async_send_size: u32) -> Self {
224        self.back_pressure_for_async_send_size = Some(back_pressure_for_async_send_size);
225        self
226    }
227
228    pub fn rpc_hook(mut self, rpc_hook: Arc<dyn RPCHook>) -> Self {
229        self.rpc_hook = Some(rpc_hook);
230        self
231    }
232
233    pub fn compress_level(mut self, compress_level: i32) -> Self {
234        self.compress_level = Some(compress_level);
235        self
236    }
237
238    pub fn compress_type(mut self, compress_type: CompressionType) -> Self {
239        self.compress_type = Some(compress_type);
240        self
241    }
242
243    pub fn compressor(mut self, compressor: &'static (dyn Compressor + Send + Sync)) -> Self {
244        self.compressor = Some(compressor);
245        self
246    }
247
248    pub fn transaction_listener(mut self, transaction_listener: impl TransactionListener) -> Self {
249        self.transaction_listener = Some(Arc::new(transaction_listener));
250        self
251    }
252
253    /// Set maximum size of transaction check thread pool
254    ///
255    /// Note: When using default Tokio Runtime with spawn_blocking, this serves as a reference.
256    /// To control actual thread count, configure the Runtime:
257    /// ```ignore
258    /// tokio::runtime::Builder::new_multi_thread()
259    ///     .max_blocking_threads(100)
260    ///     .build()
261    /// ```
262    pub fn check_thread_pool_max_size(mut self, size: u32) -> Self {
263        self.check_thread_pool_max_size = Some(size);
264        self
265    }
266
267    /// Set minimum size of transaction check thread pool
268    pub fn check_thread_pool_min_size(mut self, size: u32) -> Self {
269        self.check_thread_pool_min_size = Some(size);
270        self
271    }
272
273    /// Set maximum capacity of transaction check request queue
274    pub fn check_request_hold_max(mut self, size: u32) -> Self {
275        self.check_request_hold_max = Some(size);
276        self
277    }
278
279    pub fn build(self) -> TransactionMQProducer {
280        let client_session = ClientSession::new(Arc::clone(&self.client_runtime));
281        let mut mq_producer = DefaultMQProducer::unbound();
282        mq_producer.set_client_session(client_session);
283        if let Some(client_config) = self.client_config {
284            mq_producer.set_client_config(client_config);
285        }
286
287        if let Some(retry_response_codes) = self.retry_response_codes {
288            mq_producer.set_retry_response_codes(retry_response_codes);
289        }
290        if let Some(producer_group) = self.producer_group {
291            mq_producer.set_producer_group(producer_group);
292        }
293        if let Some(topics) = self.topics {
294            mq_producer.set_topics(topics);
295        }
296        if let Some(create_topic_key) = self.create_topic_key {
297            mq_producer.set_create_topic_key(create_topic_key);
298        }
299        if let Some(default_topic_queue_nums) = self.default_topic_queue_nums {
300            mq_producer.set_default_topic_queue_nums(default_topic_queue_nums);
301        }
302        if let Some(send_msg_timeout) = self.send_msg_timeout {
303            mq_producer.set_send_msg_timeout(send_msg_timeout);
304        }
305        if let Some(compress_msg_body_over_howmuch) = self.compress_msg_body_over_howmuch {
306            mq_producer.set_compress_msg_body_over_howmuch(compress_msg_body_over_howmuch);
307        }
308        if let Some(retry_times_when_send_failed) = self.retry_times_when_send_failed {
309            mq_producer.set_retry_times_when_send_failed(retry_times_when_send_failed);
310        }
311        if let Some(retry_times_when_send_async_failed) = self.retry_times_when_send_async_failed {
312            mq_producer.set_retry_times_when_send_async_failed(retry_times_when_send_async_failed);
313        }
314        if let Some(retry_another_broker_when_not_store_ok) = self.retry_another_broker_when_not_store_ok {
315            mq_producer.set_retry_another_broker_when_not_store_ok(retry_another_broker_when_not_store_ok);
316        }
317        if let Some(max_message_size) = self.max_message_size {
318            mq_producer.set_max_message_size(max_message_size);
319        }
320
321        if let Some(trace_dispatcher) = self.trace_dispatcher {
322            mq_producer.set_trace_dispatcher(trace_dispatcher);
323        }
324        if let Some(auto_batch) = self.auto_batch {
325            mq_producer.set_auto_batch(auto_batch);
326        }
327        if let Some(produce_accumulator) = self.produce_accumulator {
328            mq_producer.set_produce_accumulator(produce_accumulator);
329        }
330
331        if let Some(enable_backpressure_for_async_mode) = self.enable_backpressure_for_async_mode {
332            mq_producer.set_enable_backpressure_for_async_mode(enable_backpressure_for_async_mode);
333        }
334        if let Some(back_pressure_for_async_send_num) = self.back_pressure_for_async_send_num {
335            mq_producer.set_back_pressure_for_async_send_num(back_pressure_for_async_send_num);
336        }
337        if let Some(back_pressure_for_async_send_size) = self.back_pressure_for_async_send_size {
338            mq_producer.set_back_pressure_for_async_send_size(back_pressure_for_async_send_size);
339        }
340        if let Some(rpc_hook) = self.rpc_hook {
341            mq_producer.set_rpc_hook(rpc_hook);
342        }
343        if let Some(compress_level) = self.compress_level {
344            mq_producer.set_compress_level(compress_level);
345        }
346        if let Some(compress_type) = self.compress_type {
347            mq_producer.set_compress_type(compress_type);
348        }
349        if let Some(compressor) = self.compressor {
350            mq_producer.set_compressor(Some(compressor));
351        }
352
353        if let Some(default_mqproducer_impl) = self.default_mqproducer_impl {
354            mq_producer.set_default_mqproducer_impl(default_mqproducer_impl);
355        } else {
356            let producer_impl = DefaultMQProducerImpl::new_with_options(
357                self.client_runtime,
358                ClientOptions::from_parts(
359                    mq_producer.client_config_snapshot().as_ref().clone(),
360                    self.nameserver_discovery,
361                ),
362                mq_producer.producer_config_snapshot().as_ref().clone(),
363                mq_producer.rpc_hook(),
364            );
365            mq_producer.set_default_mqproducer_impl(producer_impl);
366        }
367        let transaction_producer_config = TransactionProducerConfig {
368            transaction_listener: self.transaction_listener,
369            transaction_check_listener: None,
370            check_thread_pool_min_size: self.check_thread_pool_min_size.unwrap_or(1),
371            check_thread_pool_max_size: self.check_thread_pool_max_size.unwrap_or(1),
372            check_request_hold_max: self.check_request_hold_max.unwrap_or(2000),
373        };
374        TransactionMQProducer::new(transaction_producer_config, mq_producer)
375    }
376}