rocketmq_common/common/message/
message_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 bytes::Bytes;
16use cheetah_string::CheetahString;
17use rocketmq_error::RocketMQError;
18use rocketmq_error::RocketMQResult;
19
20use crate::common::message::message_body::MessageBody;
21use crate::common::message::message_flag::MessageFlag;
22use crate::common::message::message_property::MessageProperties;
23use crate::common::message::message_property::MessagePropertyKey;
24use crate::common::message::MessageConst;
25
26/// Builder for constructing messages.
27///
28/// # Examples
29///
30/// ```
31/// use rocketmq_common::common::message::message_builder::MessageBuilder;
32///
33/// let msg = MessageBuilder::new()
34///     .topic("test-topic")
35///     .body_slice(b"hello world")
36///     .tags("important")
37///     .keys(vec!["key1".to_string(), "key2".to_string()])
38///     .build();
39/// ```
40#[derive(Default)]
41pub struct MessageBuilder {
42    topic: Option<CheetahString>,
43    body: MessageBody,
44    properties: MessageProperties,
45    flag: MessageFlag,
46    transaction_id: Option<CheetahString>,
47}
48
49impl MessageBuilder {
50    /// Creates a new message builder.
51    #[inline]
52    pub fn new() -> Self {
53        Self::default()
54    }
55
56    /// Sets the topic name (required).
57    #[inline]
58    pub fn topic(mut self, topic: impl Into<CheetahString>) -> Self {
59        self.topic = Some(topic.into());
60        self
61    }
62
63    /// Sets the message body from bytes.
64    #[inline]
65    pub fn body(mut self, body: impl Into<Bytes>) -> Self {
66        self.body = MessageBody::new(body);
67        self
68    }
69
70    /// Sets the message body from a byte slice.
71    #[inline]
72    pub fn body_slice(mut self, body: &[u8]) -> Self {
73        self.body = MessageBody::new(Bytes::copy_from_slice(body));
74        self
75    }
76
77    /// Sets an empty body.
78    #[inline]
79    pub fn empty_body(mut self) -> Self {
80        self.body = MessageBody::empty();
81        self
82    }
83
84    /// Sets the message tags.
85    #[inline]
86    pub fn tags(mut self, tags: impl Into<CheetahString>) -> Self {
87        self.properties.insert(MessagePropertyKey::Tags, tags.into());
88        self
89    }
90
91    /// Sets the message keys from a vector.
92    #[inline]
93    pub fn keys(mut self, keys: Vec<String>) -> Self {
94        if !keys.is_empty() {
95            let keys_str = keys.join(MessageConst::KEY_SEPARATOR);
96            self.properties.insert(MessagePropertyKey::Keys, keys_str);
97        }
98        self
99    }
100
101    /// Sets a single message key.
102    #[inline]
103    pub fn key(mut self, key: impl Into<CheetahString>) -> Self {
104        self.properties.insert(MessagePropertyKey::Keys, key.into());
105        self
106    }
107
108    /// Sets the message flag.
109    #[inline]
110    pub fn flag(mut self, flag: MessageFlag) -> Self {
111        self.flag = flag;
112        self
113    }
114
115    /// Sets the message flag from raw bits.
116    #[inline]
117    pub fn flag_bits(mut self, bits: i32) -> Self {
118        self.flag = MessageFlag::from_bits(bits);
119        self
120    }
121
122    /// Sets the delay time level (1-18 for predefined levels).
123    #[inline]
124    pub fn delay_level(mut self, level: i32) -> Self {
125        self.properties
126            .insert(MessagePropertyKey::DelayTimeLevel, level.to_string());
127        self
128    }
129
130    /// Sets the delay time in seconds.
131    #[inline]
132    pub fn delay_secs(mut self, secs: u64) -> Self {
133        self.properties
134            .insert(MessagePropertyKey::DelayTimeSec, secs.to_string());
135        self
136    }
137
138    /// Sets the delay time in milliseconds.
139    #[inline]
140    pub fn delay_millis(mut self, millis: u64) -> Self {
141        self.properties
142            .insert(MessagePropertyKey::DelayTimeMs, millis.to_string());
143        self
144    }
145
146    /// Sets the delivery time in milliseconds.
147    #[inline]
148    pub fn deliver_time_ms(mut self, time_ms: u64) -> Self {
149        self.properties
150            .insert(MessagePropertyKey::DeliverTimeMs, time_ms.to_string());
151        self
152    }
153
154    /// Sets the message priority.
155    #[inline]
156    pub fn priority(mut self, priority: u32) -> Self {
157        self.properties
158            .insert(MessagePropertyKey::Priority, priority.to_string());
159        self
160    }
161
162    /// Tries to set the message priority with Java-compatible non-negative validation.
163    pub fn try_priority(mut self, priority: i32) -> RocketMQResult<Self> {
164        if priority < 0 {
165            return Err(RocketMQError::illegal_argument(
166                "The priority must be greater than or equal to 0",
167            ));
168        }
169        self.properties
170            .insert(MessagePropertyKey::Priority, priority.to_string());
171        Ok(self)
172    }
173
174    /// Sets whether to wait for store confirmation.
175    #[inline]
176    pub fn wait_store_msg_ok(mut self, wait: bool) -> Self {
177        self.properties
178            .insert(MessagePropertyKey::WaitStoreMsgOk, if wait { "true" } else { "false" });
179        self
180    }
181
182    /// Sets the transaction ID.
183    #[inline]
184    pub fn transaction_id(mut self, id: impl Into<CheetahString>) -> Self {
185        self.transaction_id = Some(id.into());
186        self
187    }
188
189    /// Sets the buyer ID.
190    #[inline]
191    pub fn buyer_id(mut self, buyer_id: impl Into<CheetahString>) -> Self {
192        self.properties.insert(MessagePropertyKey::BuyerId, buyer_id.into());
193        self
194    }
195
196    /// Sets the instance ID.
197    #[inline]
198    pub fn instance_id(mut self, instance_id: impl Into<CheetahString>) -> Self {
199        self.properties
200            .insert(MessagePropertyKey::InstanceId, instance_id.into());
201        self
202    }
203
204    /// Sets the correlation ID.
205    #[inline]
206    pub fn correlation_id(mut self, correlation_id: impl Into<CheetahString>) -> Self {
207        self.properties
208            .insert(MessagePropertyKey::CorrelationId, correlation_id.into());
209        self
210    }
211
212    /// Sets the sharding key.
213    #[inline]
214    pub fn sharding_key(mut self, sharding_key: impl Into<CheetahString>) -> Self {
215        self.properties
216            .insert(MessagePropertyKey::ShardingKey, sharding_key.into());
217        self
218    }
219
220    /// Sets the trace switch.
221    #[inline]
222    pub fn trace_switch(mut self, enabled: bool) -> Self {
223        self.properties
224            .insert(MessagePropertyKey::TraceSwitch, if enabled { "true" } else { "false" });
225        self
226    }
227
228    /// Sets a custom property directly using MessagePropertyKey.
229    #[inline]
230    pub fn property(mut self, key: MessagePropertyKey, value: impl Into<CheetahString>) -> Self {
231        self.properties.insert(key, value.into());
232        self
233    }
234
235    /// Sets a raw property with string key (for advanced use).
236    ///
237    /// # Errors
238    ///
239    /// Returns an error if the property name is reserved by the system.
240    pub fn raw_property(
241        mut self,
242        key: impl Into<CheetahString>,
243        value: impl Into<CheetahString>,
244    ) -> RocketMQResult<Self> {
245        let key = key.into();
246        let value = value.into();
247
248        // Validate not a reserved key
249        if crate::common::message::STRING_HASH_SET.contains(key.as_str()) {
250            return Err(RocketMQError::InvalidProperty(format!(
251                "The Property<{key}> is used by system, input another please"
252            )));
253        }
254
255        self.properties.as_map_mut().insert(key, value);
256        Ok(self)
257    }
258
259    /// Builds the message.
260    ///
261    /// # Errors
262    ///
263    /// Returns an error if required fields (topic) are not set.
264    pub fn build(mut self) -> RocketMQResult<super::message_single::Message> {
265        let topic = self
266            .topic
267            .ok_or_else(|| RocketMQError::InvalidProperty("Topic is required for message".to_string()))?;
268        if self.properties.get(MessagePropertyKey::WaitStoreMsgOk).is_none() {
269            self.properties.insert(MessagePropertyKey::WaitStoreMsgOk, "true");
270        }
271
272        Ok(super::message_single::Message::from_builder(
273            topic,
274            self.body,
275            self.properties,
276            self.flag,
277            self.transaction_id,
278        ))
279    }
280
281    /// Builds the message, panicking if validation fails.
282    ///
283    /// This is convenient for test code where you know the message is valid.
284    #[track_caller]
285    pub fn build_unchecked(self) -> super::message_single::Message {
286        self.build().expect("message validation failed")
287    }
288}
289
290#[cfg(test)]
291mod tests {
292    use super::*;
293
294    #[test]
295    fn test_builder_basic() {
296        let msg = MessageBuilder::new()
297            .topic("test-topic")
298            .body_slice(b"hello world")
299            .build_unchecked();
300
301        assert_eq!(msg.topic().as_str(), "test-topic");
302        assert_eq!(msg.body_slice(), b"hello world");
303        assert_eq!(msg.properties().get(MessagePropertyKey::WaitStoreMsgOk), Some("true"));
304    }
305
306    #[test]
307    fn test_builder_with_tags_and_keys() {
308        let msg = MessageBuilder::new()
309            .topic("test-topic")
310            .body_slice(b"test")
311            .tags("tag1")
312            .keys(vec!["key1".to_string(), "key2".to_string()])
313            .build_unchecked();
314
315        assert_eq!(msg.tags(), Some("tag1"));
316        assert_eq!(msg.keys(), Some(vec!["key1".to_string(), "key2".to_string()]));
317    }
318
319    #[test]
320    fn test_builder_with_delay() {
321        let msg = MessageBuilder::new()
322            .topic("test-topic")
323            .body_slice(b"test")
324            .delay_level(3)
325            .build_unchecked();
326
327        assert_eq!(msg.delay_time_level(), 3);
328    }
329
330    #[test]
331    fn test_builder_with_priority_matches_java_property() {
332        let msg = MessageBuilder::new()
333            .topic("test-topic")
334            .body_slice(b"test")
335            .priority(7)
336            .build_unchecked();
337
338        assert_eq!(msg.priority(), 7);
339        assert_eq!(msg.properties().get(MessagePropertyKey::Priority), Some("7"));
340    }
341
342    #[test]
343    fn test_builder_try_priority_rejects_negative_like_java() {
344        match MessageBuilder::new().topic("test-topic").try_priority(-1) {
345            Ok(_) => panic!("negative priority should be rejected"),
346            Err(error) => assert!(error.to_string().contains("greater than or equal to 0")),
347        }
348    }
349
350    #[test]
351    fn test_builder_missing_topic() {
352        let result = MessageBuilder::new().body_slice(b"test").build();
353
354        assert!(result.is_err());
355    }
356
357    #[test]
358    fn test_builder_with_transaction_id() {
359        let msg = MessageBuilder::new()
360            .topic("test-topic")
361            .body_slice(b"test")
362            .transaction_id("tx-123")
363            .build_unchecked();
364
365        assert_eq!(msg.transaction_id(), Some("tx-123"));
366    }
367
368    #[test]
369    fn test_builder_wait_store_msg_ok_writes_java_boolean_property() {
370        let wait_msg = MessageBuilder::new()
371            .topic("test-topic")
372            .body_slice(b"test")
373            .wait_store_msg_ok(true)
374            .build_unchecked();
375        assert!(wait_msg.wait_store_msg_ok());
376        assert_eq!(
377            wait_msg.properties().get(MessagePropertyKey::WaitStoreMsgOk),
378            Some("true")
379        );
380
381        let no_wait_msg = MessageBuilder::new()
382            .topic("test-topic")
383            .body_slice(b"test")
384            .wait_store_msg_ok(false)
385            .build_unchecked();
386        assert!(!no_wait_msg.wait_store_msg_ok());
387        assert_eq!(
388            no_wait_msg.properties().get(MessagePropertyKey::WaitStoreMsgOk),
389            Some("false")
390        );
391    }
392}
rocketmq_common/common/message/message_builder.rs

rocketmq_common/common/message/
message_builder.rs
