rocketmq_client_rust/hook/
consume_message_context.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::any::Any;
16use std::collections::HashMap;
17use std::fmt;
18use std::sync::Arc;
19
20use cheetah_string::CheetahString;
21use rocketmq_common::common::message::message_ext::MessageExt;
22use rocketmq_common::common::message::message_queue::MessageQueue;
23use rocketmq_rust::ArcMut;
24
25use crate::base::access_channel::AccessChannel;
26
27/// Context information for message consumption operations.
28///
29/// Contains metadata about a message consumption attempt, including the consumer group,
30/// message list, target queue, and consumption result. This context is passed to
31/// consumption hooks for monitoring and tracing purposes.
32///
33/// # Examples
34///
35/// ```ignore
36/// use rocketmq_client::hook::consume_message_context::ConsumeMessageContext;
37/// use cheetah_string::CheetahString;
38///
39/// let context = ConsumeMessageContext::new(
40///     CheetahString::from("test_group"),
41///     &messages
42/// ).with_success(true)
43///   .with_status("CONSUME_SUCCESS");
44/// ```
45#[derive(Default)]
46pub struct ConsumeMessageContext<'a> {
47    /// Consumer group name
48    pub consumer_group: CheetahString,
49    /// List of messages being consumed
50    pub msg_list: &'a [ArcMut<MessageExt>],
51    /// Target message queue
52    pub mq: Option<MessageQueue>,
53    /// Whether consumption succeeded
54    pub success: bool,
55    /// Consumption status description
56    pub status: CheetahString,
57    /// Trace context for distributed tracing
58    pub mq_trace_context: Option<Arc<dyn Any + Send + Sync>>,
59    /// Additional properties
60    pub props: HashMap<CheetahString, CheetahString>,
61    /// Namespace
62    pub namespace: CheetahString,
63    /// Access channel
64    pub access_channel: Option<AccessChannel>,
65}
66
67impl<'a> ConsumeMessageContext<'a> {
68    /// Creates a new consumption context.
69    #[inline]
70    pub fn new(consumer_group: CheetahString, msg_list: &'a [ArcMut<MessageExt>]) -> Self {
71        Self {
72            consumer_group,
73            msg_list,
74            ..Default::default()
75        }
76    }
77
78    /// Sets the message queue.
79    #[inline]
80    pub fn with_mq(mut self, mq: MessageQueue) -> Self {
81        self.mq = Some(mq);
82        self
83    }
84
85    /// Sets the success flag.
86    #[inline]
87    pub fn with_success(mut self, success: bool) -> Self {
88        self.success = success;
89        self
90    }
91
92    /// Sets the status description.
93    #[inline]
94    pub fn with_status(mut self, status: impl Into<CheetahString>) -> Self {
95        self.status = status.into();
96        self
97    }
98
99    /// Sets the trace context.
100    #[inline]
101    pub fn with_trace_context(mut self, context: Arc<dyn Any + Send + Sync>) -> Self {
102        self.mq_trace_context = Some(context);
103        self
104    }
105
106    /// Sets the namespace.
107    #[inline]
108    pub fn with_namespace(mut self, namespace: impl Into<CheetahString>) -> Self {
109        self.namespace = namespace.into();
110        self
111    }
112
113    /// Sets the access channel.
114    #[inline]
115    pub fn with_access_channel(mut self, channel: AccessChannel) -> Self {
116        self.access_channel = Some(channel);
117        self
118    }
119
120    /// Adds a property.
121    #[inline]
122    pub fn add_prop(&mut self, key: impl Into<CheetahString>, value: impl Into<CheetahString>) {
123        self.props.insert(key.into(), value.into());
124    }
125
126    /// Gets the number of messages in this context.
127    #[inline]
128    pub fn message_count(&self) -> usize {
129        self.msg_list.len()
130    }
131}
132
133impl<'a> fmt::Display for ConsumeMessageContext<'a> {
134    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
135        write!(
136            f,
137            "ConsumeMessageContext {{ consumer_group: {}, msg_count: {}, ",
138            self.consumer_group,
139            self.msg_list.len()
140        )?;
141
142        if let Some(ref mq) = self.mq {
143            write!(f, "mq: {}, ", mq)?;
144        }
145
146        write!(f, "success: {}, status: {}", self.success, self.status)?;
147
148        if !self.namespace.is_empty() {
149            write!(f, ", namespace: {}", self.namespace)?;
150        }
151
152        if let Some(ref channel) = self.access_channel {
153            write!(f, ", access_channel: {:?}", channel)?;
154        }
155
156        write!(f, " }}")
157    }
158}
159
160impl<'a> fmt::Debug for ConsumeMessageContext<'a> {
161    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
162        f.debug_struct("ConsumeMessageContext")
163            .field("consumer_group", &self.consumer_group)
164            .field("msg_list_len", &self.msg_list.len())
165            .field("mq", &self.mq)
166            .field("success", &self.success)
167            .field("status", &self.status)
168            .field("mq_trace_context", &self.mq_trace_context.as_ref().map(|_| "Some(_)"))
169            .field("props_count", &self.props.len())
170            .field("namespace", &self.namespace)
171            .field("access_channel", &self.access_channel)
172            .finish()
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179
180    #[test]
181    fn consume_message_context_new() {
182        let group = CheetahString::from("test_group");
183        let msg_list = vec![ArcMut::new(MessageExt::default())];
184        let context = ConsumeMessageContext::new(group.clone(), &msg_list);
185
186        assert_eq!(context.consumer_group, group);
187        assert_eq!(context.message_count(), 1);
188        assert!(context.mq.is_none());
189        assert!(!context.success);
190    }
191
192    #[test]
193    fn consume_message_context_builder_methods() {
194        let group = CheetahString::from("test_group");
195        let msg_list = vec![];
196        let mq = MessageQueue::from_parts("topic", "broker", 1);
197
198        let context = ConsumeMessageContext::new(group, &msg_list)
199            .with_mq(mq.clone())
200            .with_success(true)
201            .with_status("SUCCESS")
202            .with_namespace("test_ns")
203            .with_access_channel(AccessChannel::Cloud);
204
205        assert_eq!(context.mq.unwrap(), mq);
206        assert!(context.success);
207        assert_eq!(context.status, "SUCCESS");
208        assert_eq!(context.namespace, "test_ns");
209        assert_eq!(context.access_channel.unwrap(), AccessChannel::Cloud);
210    }
211
212    #[test]
213    fn consume_message_context_add_prop() {
214        let mut context = ConsumeMessageContext::<'static>::default();
215        context.add_prop("key1", "value1");
216        assert_eq!(context.props.get(&CheetahString::from("key1")).unwrap(), "value1");
217    }
218
219    #[test]
220    fn consume_message_context_display() {
221        let group = CheetahString::from("test_group");
222        let msg_list = vec![];
223        let context = ConsumeMessageContext::new(group, &msg_list)
224            .with_success(true)
225            .with_status("OK");
226
227        let display = format!("{}", context);
228        assert!(display.contains("consumer_group: test_group"));
229        assert!(display.contains("msg_count: 0"));
230        assert!(display.contains("success: true"));
231        assert!(display.contains("status: OK"));
232    }
233
234    #[test]
235    fn consume_message_context_debug() {
236        let context = ConsumeMessageContext::<'static>::default();
237        let debug = format!("{:?}", context);
238        assert!(debug.contains("ConsumeMessageContext"));
239        assert!(debug.contains("consumer_group"));
240    }
241
242    #[test]
243    fn consume_message_context_with_trace_context() {
244        let context =
245            ConsumeMessageContext::<'static>::default().with_trace_context(Arc::new("trace_data".to_string()));
246        assert!(context.mq_trace_context.is_some());
247    }
248}
rocketmq_client_rust/hook/consume_message_context.rs

rocketmq_client_rust/hook/
consume_message_context.rs
