rocketmq_client_rust/trace/
trace_data_encoder.rs

1/*
2 * Licensed to the Apache Software Foundation (ASF) under one or more
3 * contributor license agreements.  See the NOTICE file distributed with
4 * this work for additional information regarding copyright ownership.
5 * The ASF licenses this file to You under the Apache License, Version 2.0
6 * (the "License"); you may not use this file except in compliance with
7 * the License.  You may obtain a copy of the License at
8 *
9 *     http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 */
17
18use cheetah_string::CheetahString;
19use rocketmq_model::common::message::message_enum::MessageType;
20use rocketmq_model::common::message::MessageConst;
21
22use crate::base::access_channel::AccessChannel;
23use crate::producer::local_transaction_state::LocalTransactionState;
24use crate::trace::trace_bean::TraceBean;
25use crate::trace::trace_constants::TraceConstants;
26use crate::trace::trace_context::TraceContext;
27use crate::trace::trace_transfer_bean::TraceTransferBean;
28use crate::trace::trace_type::TraceType;
29
30/// Encodes and decodes trace data for RocketMQ message tracing.
31///
32/// Provides utilities for converting structured trace contexts into wire format strings
33/// and parsing trace data strings back into trace contexts. The encoder uses pre-allocated
34/// buffers and zero-copy string parsing to minimize allocations.
35pub struct TraceDataEncoder;
36
37impl TraceDataEncoder {
38    /// Decodes trace data from a string into structured trace contexts.
39    ///
40    /// Parses a trace data string containing one or more encoded trace records separated
41    /// by field separators. The result vector is pre-allocated based on separator count
42    /// to minimize allocations. Invalid or malformed records are silently skipped.
43    ///
44    /// Returns an empty vector if the input string is empty.
45    pub fn decoder_from_trace_data_string(trace_data: &str) -> Vec<TraceContext> {
46        // Early return for empty input
47        if trace_data.is_empty() {
48            return Vec::new();
49        }
50
51        // Pre-allocate result vector based on field separator count
52        let estimated_size = trace_data.matches(TraceConstants::FIELD_SPLITOR).count();
53        let mut res_list = Vec::with_capacity(estimated_size.max(1));
54
55        // Split by field separator
56        for context_str in trace_data.split(TraceConstants::FIELD_SPLITOR) {
57            if context_str.is_empty() {
58                continue;
59            }
60
61            let line: Vec<&str> = context_str.split(TraceConstants::CONTENT_SPLITOR).collect();
62            if line.is_empty() {
63                continue;
64            }
65
66            // Match trace type and parse accordingly
67            match line[0] {
68                "Pub" => {
69                    if let Some(ctx) = Self::decode_pub_context(&line) {
70                        res_list.push(ctx);
71                    }
72                }
73                "SubBefore" => {
74                    if let Some(ctx) = Self::decode_sub_before_context(&line) {
75                        res_list.push(ctx);
76                    }
77                }
78                "SubAfter" => {
79                    if let Some(ctx) = Self::decode_sub_after_context(&line) {
80                        res_list.push(ctx);
81                    }
82                }
83                "EndTransaction" => {
84                    if let Some(ctx) = Self::decode_end_transaction_context(&line) {
85                        res_list.push(ctx);
86                    }
87                }
88                "Recall" => {
89                    if let Some(ctx) = Self::decode_recall_context(&line) {
90                        res_list.push(ctx);
91                    }
92                }
93                _ => {} // Unknown trace type, skip
94            }
95        }
96
97        res_list
98    }
99
100    /// Encodes a trace context into a transfer bean for transmission.
101    ///
102    /// Serializes the trace context into a wire format string and extracts message keys
103    /// for indexing. The string builder is pre-allocated with 256 bytes capacity to reduce
104    /// allocations during encoding.
105    ///
106    /// Returns `None` if the context has no trace beans or an invalid trace type.
107    pub fn encoder_from_context_bean(ctx: &TraceContext) -> Option<TraceTransferBean> {
108        let trace_beans = ctx.trace_beans.as_ref()?;
109        if trace_beans.is_empty() {
110            return None;
111        }
112
113        let mut transfer_bean = TraceTransferBean::new();
114        // Pre-allocate string builder with estimated capacity
115        let mut sb = String::with_capacity(256);
116
117        match ctx.trace_type? {
118            TraceType::Pub => {
119                Self::encode_pub_context(ctx, &trace_beans[0], &mut sb);
120            }
121            TraceType::SubBefore => {
122                Self::encode_sub_before_context(ctx, trace_beans, &mut sb);
123            }
124            TraceType::SubAfter => {
125                Self::encode_sub_after_context(ctx, trace_beans, &mut sb);
126            }
127            TraceType::EndTransaction => {
128                Self::encode_end_transaction_context(ctx, &trace_beans[0], &mut sb);
129            }
130            TraceType::Recall => {
131                Self::encode_recall_context(ctx, &trace_beans[0], &mut sb);
132            }
133        }
134
135        let records = rocketmq_protocol::trace::decode_records(&sb);
136        transfer_bean.set_trans_data(rocketmq_protocol::trace::encode_records(&records));
137
138        // Extract keys from trace beans
139        for bean in trace_beans {
140            transfer_bean.add_key(bean.msg_id.clone());
141
142            if !bean.keys.is_empty() {
143                let keys: Vec<&str> = bean.keys.split_str(MessageConst::KEY_SEPARATOR).collect();
144                for key in keys {
145                    if !key.is_empty() {
146                        transfer_bean.add_key(CheetahString::from_slice(key));
147                    }
148                }
149            }
150        }
151
152        Some(transfer_bean)
153    }
154
155    // ==================== Decoder Helper Methods ====================
156
157    #[inline]
158    fn decode_pub_context(line: &[&str]) -> Option<TraceContext> {
159        if line.len() < 12 {
160            return None;
161        }
162
163        let time_stamp = line[1].parse().ok()?;
164        let region_id = CheetahString::from_slice(line[2]);
165        let group_name = CheetahString::from_slice(line[3]);
166        let cost_time = line[10].parse().ok()?;
167        let body_length = line[9].parse().ok()?;
168
169        // Parse message type
170        let msg_type = if let Ok(msg_type_ordinal) = line[11].parse::<usize>() {
171            Self::message_type_from_ordinal(msg_type_ordinal)
172        } else {
173            None
174        };
175
176        // Handle different line lengths for backward compatibility
177        let (is_success, offset_msg_id, client_host) = if line.len() == 13 {
178            (
179                line[12].parse().unwrap_or(true),
180                CheetahString::default(),
181                CheetahString::default(),
182            )
183        } else if line.len() == 14 {
184            (
185                line[13].parse().unwrap_or(true),
186                CheetahString::from_slice(line[12]),
187                CheetahString::default(),
188            )
189        } else if line.len() >= 15 {
190            (
191                line[13].parse().unwrap_or(true),
192                CheetahString::from_slice(line[12]),
193                CheetahString::from_slice(line[14]),
194            )
195        } else {
196            (true, CheetahString::default(), CheetahString::default())
197        };
198
199        let bean = TraceBean {
200            topic: CheetahString::from_slice(line[4]),
201            msg_id: CheetahString::from_slice(line[5]),
202            tags: CheetahString::from_slice(line[6]),
203            keys: CheetahString::from_slice(line[7]),
204            store_host: CheetahString::from_slice(line[8]),
205            body_length,
206            msg_type,
207            offset_msg_id,
208            client_host,
209            ..Default::default()
210        };
211
212        let ctx = TraceContext {
213            trace_type: Some(TraceType::Pub),
214            time_stamp,
215            region_id,
216            group_name,
217            cost_time,
218            is_success,
219            trace_beans: Some(vec![bean]),
220            ..Default::default()
221        };
222
223        Some(ctx)
224    }
225
226    #[inline]
227    fn decode_sub_before_context(line: &[&str]) -> Option<TraceContext> {
228        if line.len() < 8 {
229            return None;
230        }
231
232        let bean = TraceBean {
233            msg_id: CheetahString::from_slice(line[5]),
234            retry_times: line[6].parse().ok()?,
235            keys: CheetahString::from_slice(line[7]),
236            ..Default::default()
237        };
238
239        let ctx = TraceContext {
240            trace_type: Some(TraceType::SubBefore),
241            time_stamp: line[1].parse().ok()?,
242            region_id: CheetahString::from_slice(line[2]),
243            group_name: CheetahString::from_slice(line[3]),
244            request_id: CheetahString::from_slice(line[4]),
245            trace_beans: Some(vec![bean]),
246            ..Default::default()
247        };
248
249        Some(ctx)
250    }
251
252    #[inline]
253    fn decode_sub_after_context(line: &[&str]) -> Option<TraceContext> {
254        if line.len() < 6 {
255            return None;
256        }
257
258        let context_code = if line.len() >= 7 {
259            line[6].parse().unwrap_or(0)
260        } else {
261            0
262        };
263
264        let (time_stamp, group_name) = if line.len() >= 9 {
265            (line[7].parse().unwrap_or(0), CheetahString::from_slice(line[8]))
266        } else {
267            (0, CheetahString::default())
268        };
269
270        let bean = TraceBean {
271            msg_id: CheetahString::from_slice(line[2]),
272            keys: CheetahString::from_slice(line[5]),
273            ..Default::default()
274        };
275
276        let ctx = TraceContext {
277            trace_type: Some(TraceType::SubAfter),
278            request_id: CheetahString::from_slice(line[1]),
279            cost_time: line[3].parse().ok()?,
280            is_success: line[4].parse().unwrap_or(false),
281            context_code,
282            time_stamp,
283            group_name,
284            trace_beans: Some(vec![bean]),
285            ..Default::default()
286        };
287
288        Some(ctx)
289    }
290
291    #[inline]
292    fn decode_end_transaction_context(line: &[&str]) -> Option<TraceContext> {
293        if line.len() < 13 {
294            return None;
295        }
296
297        let msg_type = if let Ok(msg_type_ordinal) = line[9].parse::<usize>() {
298            Self::message_type_from_ordinal(msg_type_ordinal)
299        } else {
300            None
301        };
302
303        let bean = TraceBean {
304            topic: CheetahString::from_slice(line[4]),
305            msg_id: CheetahString::from_slice(line[5]),
306            tags: CheetahString::from_slice(line[6]),
307            keys: CheetahString::from_slice(line[7]),
308            store_host: CheetahString::from_slice(line[8]),
309            msg_type,
310            transaction_id: Some(CheetahString::from_slice(line[10])),
311            transaction_state: Self::parse_transaction_state(line[11]),
312            from_transaction_check: line[12].parse().unwrap_or(false),
313            ..Default::default()
314        };
315
316        let ctx = TraceContext {
317            trace_type: Some(TraceType::EndTransaction),
318            time_stamp: line[1].parse().ok()?,
319            region_id: CheetahString::from_slice(line[2]),
320            group_name: CheetahString::from_slice(line[3]),
321            trace_beans: Some(vec![bean]),
322            ..Default::default()
323        };
324
325        Some(ctx)
326    }
327
328    #[inline]
329    fn decode_recall_context(line: &[&str]) -> Option<TraceContext> {
330        if line.len() < 7 {
331            return None;
332        }
333
334        let bean = TraceBean {
335            topic: CheetahString::from_slice(line[4]),
336            msg_id: CheetahString::from_slice(line[5]),
337            ..Default::default()
338        };
339
340        let ctx = TraceContext {
341            trace_type: Some(TraceType::Recall),
342            time_stamp: line[1].parse().ok()?,
343            region_id: CheetahString::from_slice(line[2]),
344            group_name: CheetahString::from_slice(line[3]),
345            is_success: line[6].parse().unwrap_or(false),
346            trace_beans: Some(vec![bean]),
347            ..Default::default()
348        };
349
350        Some(ctx)
351    }
352
353    // ==================== Encoder Helper Methods ====================
354
355    #[inline]
356    fn encode_pub_context(ctx: &TraceContext, bean: &TraceBean, sb: &mut String) {
357        sb.push_str("Pub");
358        sb.push(TraceConstants::CONTENT_SPLITOR);
359        sb.push_str(&ctx.time_stamp.to_string());
360        sb.push(TraceConstants::CONTENT_SPLITOR);
361        sb.push_str(&ctx.region_id);
362        sb.push(TraceConstants::CONTENT_SPLITOR);
363        sb.push_str(&ctx.group_name);
364        sb.push(TraceConstants::CONTENT_SPLITOR);
365        sb.push_str(&bean.topic);
366        sb.push(TraceConstants::CONTENT_SPLITOR);
367        sb.push_str(&bean.msg_id);
368        sb.push(TraceConstants::CONTENT_SPLITOR);
369        sb.push_str(&bean.tags);
370        sb.push(TraceConstants::CONTENT_SPLITOR);
371        sb.push_str(&bean.keys);
372        sb.push(TraceConstants::CONTENT_SPLITOR);
373        sb.push_str(&bean.store_host);
374        sb.push(TraceConstants::CONTENT_SPLITOR);
375        sb.push_str(&bean.body_length.to_string());
376        sb.push(TraceConstants::CONTENT_SPLITOR);
377        sb.push_str(&ctx.cost_time.to_string());
378        sb.push(TraceConstants::CONTENT_SPLITOR);
379
380        if let Some(msg_type) = &bean.msg_type {
381            sb.push_str(&Self::message_type_to_ordinal(*msg_type).to_string());
382        } else {
383            sb.push('0');
384        }
385        sb.push(TraceConstants::CONTENT_SPLITOR);
386        sb.push_str(&bean.offset_msg_id);
387        sb.push(TraceConstants::CONTENT_SPLITOR);
388        sb.push_str(if ctx.is_success { "true" } else { "false" });
389        sb.push(TraceConstants::CONTENT_SPLITOR);
390        sb.push_str(&bean.client_host);
391        sb.push(TraceConstants::FIELD_SPLITOR);
392    }
393
394    #[inline]
395    fn encode_sub_before_context(ctx: &TraceContext, beans: &[TraceBean], sb: &mut String) {
396        for bean in beans {
397            sb.push_str("SubBefore");
398            sb.push(TraceConstants::CONTENT_SPLITOR);
399            sb.push_str(&ctx.time_stamp.to_string());
400            sb.push(TraceConstants::CONTENT_SPLITOR);
401            sb.push_str(&ctx.region_id);
402            sb.push(TraceConstants::CONTENT_SPLITOR);
403            sb.push_str(&ctx.group_name);
404            sb.push(TraceConstants::CONTENT_SPLITOR);
405            sb.push_str(&ctx.request_id);
406            sb.push(TraceConstants::CONTENT_SPLITOR);
407            sb.push_str(&bean.msg_id);
408            sb.push(TraceConstants::CONTENT_SPLITOR);
409            sb.push_str(&bean.retry_times.to_string());
410            sb.push(TraceConstants::CONTENT_SPLITOR);
411            sb.push_str(&bean.keys);
412            sb.push(TraceConstants::FIELD_SPLITOR);
413        }
414    }
415
416    #[inline]
417    fn encode_sub_after_context(ctx: &TraceContext, beans: &[TraceBean], sb: &mut String) {
418        for bean in beans {
419            sb.push_str("SubAfter");
420            sb.push(TraceConstants::CONTENT_SPLITOR);
421            sb.push_str(&ctx.request_id);
422            sb.push(TraceConstants::CONTENT_SPLITOR);
423            sb.push_str(&bean.msg_id);
424            sb.push(TraceConstants::CONTENT_SPLITOR);
425            sb.push_str(&ctx.cost_time.to_string());
426            sb.push(TraceConstants::CONTENT_SPLITOR);
427            sb.push_str(if ctx.is_success { "true" } else { "false" });
428            sb.push(TraceConstants::CONTENT_SPLITOR);
429            sb.push_str(&bean.keys);
430            sb.push(TraceConstants::CONTENT_SPLITOR);
431            sb.push_str(&ctx.context_code.to_string());
432            sb.push(TraceConstants::CONTENT_SPLITOR);
433
434            // Only add timestamp and group name if not CLOUD access channel
435            if !matches!(ctx.access_channel, Some(AccessChannel::Cloud)) {
436                sb.push_str(&ctx.time_stamp.to_string());
437                sb.push(TraceConstants::CONTENT_SPLITOR);
438                sb.push_str(&ctx.group_name);
439            }
440            sb.push(TraceConstants::FIELD_SPLITOR);
441        }
442    }
443
444    #[inline]
445    fn encode_end_transaction_context(ctx: &TraceContext, bean: &TraceBean, sb: &mut String) {
446        sb.push_str("EndTransaction");
447        sb.push(TraceConstants::CONTENT_SPLITOR);
448        sb.push_str(&ctx.time_stamp.to_string());
449        sb.push(TraceConstants::CONTENT_SPLITOR);
450        sb.push_str(&ctx.region_id);
451        sb.push(TraceConstants::CONTENT_SPLITOR);
452        sb.push_str(&ctx.group_name);
453        sb.push(TraceConstants::CONTENT_SPLITOR);
454        sb.push_str(&bean.topic);
455        sb.push(TraceConstants::CONTENT_SPLITOR);
456        sb.push_str(&bean.msg_id);
457        sb.push(TraceConstants::CONTENT_SPLITOR);
458        sb.push_str(&bean.tags);
459        sb.push(TraceConstants::CONTENT_SPLITOR);
460        sb.push_str(&bean.keys);
461        sb.push(TraceConstants::CONTENT_SPLITOR);
462        sb.push_str(&bean.store_host);
463        sb.push(TraceConstants::CONTENT_SPLITOR);
464
465        if let Some(msg_type) = &bean.msg_type {
466            sb.push_str(&Self::message_type_to_ordinal(*msg_type).to_string());
467        } else {
468            sb.push('0');
469        }
470        sb.push(TraceConstants::CONTENT_SPLITOR);
471
472        if let Some(ref transaction_id) = bean.transaction_id {
473            sb.push_str(transaction_id);
474        }
475        sb.push(TraceConstants::CONTENT_SPLITOR);
476
477        if let Some(ref state) = bean.transaction_state {
478            sb.push_str(&state.to_string());
479        }
480        sb.push(TraceConstants::CONTENT_SPLITOR);
481        sb.push_str(if bean.from_transaction_check { "true" } else { "false" });
482        sb.push(TraceConstants::FIELD_SPLITOR);
483    }
484
485    #[inline]
486    fn encode_recall_context(ctx: &TraceContext, bean: &TraceBean, sb: &mut String) {
487        sb.push_str("Recall");
488        sb.push(TraceConstants::CONTENT_SPLITOR);
489        sb.push_str(&ctx.time_stamp.to_string());
490        sb.push(TraceConstants::CONTENT_SPLITOR);
491        sb.push_str(&ctx.region_id);
492        sb.push(TraceConstants::CONTENT_SPLITOR);
493        sb.push_str(&ctx.group_name);
494        sb.push(TraceConstants::CONTENT_SPLITOR);
495        sb.push_str(&bean.topic);
496        sb.push(TraceConstants::CONTENT_SPLITOR);
497        sb.push_str(&bean.msg_id);
498        sb.push(TraceConstants::CONTENT_SPLITOR);
499        sb.push_str(if ctx.is_success { "true" } else { "false" });
500        sb.push(TraceConstants::FIELD_SPLITOR);
501    }
502
503    // ==================== Utility Methods ====================
504
505    #[inline]
506    fn message_type_from_ordinal(ordinal: usize) -> Option<MessageType> {
507        match ordinal {
508            0 => Some(MessageType::NormalMsg),
509            1 => Some(MessageType::TransMsgHalf),
510            2 => Some(MessageType::TransMsgCommit),
511            3 => Some(MessageType::DelayMsg),
512            4 => Some(MessageType::OrderMsg),
513            _ => None,
514        }
515    }
516
517    #[inline]
518    fn message_type_to_ordinal(msg_type: MessageType) -> usize {
519        match msg_type {
520            MessageType::NormalMsg => 0,
521            MessageType::TransMsgHalf => 1,
522            MessageType::TransMsgCommit => 2,
523            MessageType::DelayMsg => 3,
524            MessageType::OrderMsg => 4,
525        }
526    }
527
528    #[inline]
529    fn parse_transaction_state(state_str: &str) -> Option<LocalTransactionState> {
530        match state_str {
531            "COMMIT_MESSAGE" => Some(LocalTransactionState::CommitMessage),
532            "ROLLBACK_MESSAGE" => Some(LocalTransactionState::RollbackMessage),
533            "UNKNOW" | "UNKNOWN" => Some(LocalTransactionState::Unknown),
534            _ => None,
535        }
536    }
537}
538
539#[cfg(test)]
540mod tests {
541    use cheetah_string::CheetahString;
542    use rocketmq_model::common::message::message_enum::MessageType;
543
544    use crate::base::access_channel::AccessChannel;
545    use crate::producer::local_transaction_state::LocalTransactionState;
546    use crate::trace::trace_bean::TraceBean;
547    use crate::trace::trace_context::TraceContext;
548    use crate::trace::trace_type::TraceType;
549
550    use super::*;
551
552    #[test]
553    fn test_decode_empty_trace_data() {
554        let result = TraceDataEncoder::decoder_from_trace_data_string("");
555        assert!(result.is_empty());
556    }
557
558    #[test]
559    fn test_encode_and_decode_pub_context() {
560        let bean = TraceBean {
561            topic: CheetahString::from("test-topic"),
562            msg_id: CheetahString::from("msg-001"),
563            tags: CheetahString::from("tagA"),
564            keys: CheetahString::from("key1 key2"),
565            store_host: CheetahString::from("127.0.0.1:10911"),
566            body_length: 1024,
567            msg_type: Some(MessageType::NormalMsg),
568            offset_msg_id: CheetahString::from("offset-001"),
569            client_host: CheetahString::from("192.168.1.1:12345"),
570            ..Default::default()
571        };
572
573        let ctx = TraceContext {
574            trace_type: Some(TraceType::Pub),
575            time_stamp: 1234567890,
576            region_id: CheetahString::from("us-east-1"),
577            group_name: CheetahString::from("test-group"),
578            cost_time: 100,
579            is_success: true,
580            trace_beans: Some(vec![bean]),
581            ..Default::default()
582        };
583
584        let transfer_bean = TraceDataEncoder::encoder_from_context_bean(&ctx).unwrap();
585        let canonical_records = rocketmq_protocol::trace::decode_records(&transfer_bean.trans_data);
586        assert_eq!(canonical_records.len(), 1);
587        assert_eq!(
588            rocketmq_protocol::trace::encode_records(&canonical_records),
589            transfer_bean.trans_data
590        );
591        let decoded_contexts = TraceDataEncoder::decoder_from_trace_data_string(&transfer_bean.trans_data);
592
593        assert_eq!(decoded_contexts.len(), 1);
594        let decoded = &decoded_contexts[0];
595        assert_eq!(decoded.trace_type, Some(TraceType::Pub));
596        assert_eq!(decoded.time_stamp, 1234567890);
597        assert_eq!(decoded.region_id, CheetahString::from("us-east-1"));
598        assert_eq!(decoded.group_name, CheetahString::from("test-group"));
599        assert_eq!(decoded.cost_time, 100);
600        assert!(decoded.is_success);
601        assert!(decoded.trace_beans.is_some());
602        let decoded_bean = &decoded.trace_beans.as_ref().unwrap()[0];
603        assert_eq!(decoded_bean.topic, CheetahString::from("test-topic"));
604        assert_eq!(decoded_bean.msg_id, CheetahString::from("msg-001"));
605        assert_eq!(decoded_bean.tags, CheetahString::from("tagA"));
606        assert_eq!(decoded_bean.keys, CheetahString::from("key1 key2"));
607        assert_eq!(decoded_bean.store_host, CheetahString::from("127.0.0.1:10911"));
608        assert_eq!(decoded_bean.body_length, 1024);
609        assert_eq!(decoded_bean.msg_type, Some(MessageType::NormalMsg));
610        assert_eq!(decoded_bean.offset_msg_id, CheetahString::from("offset-001"));
611        assert_eq!(decoded_bean.client_host, CheetahString::from("192.168.1.1:12345"));
612    }
613
614    #[test]
615    fn test_encode_and_decode_sub_before_context() {
616        let bean = TraceBean {
617            msg_id: CheetahString::from("msg-002"),
618            retry_times: 2,
619            keys: CheetahString::from("key3"),
620            ..Default::default()
621        };
622
623        let ctx = TraceContext {
624            trace_type: Some(TraceType::SubBefore),
625            time_stamp: 9876543210,
626            region_id: CheetahString::from("eu-west-1"),
627            group_name: CheetahString::from("consumer-group"),
628            request_id: CheetahString::from("req-123"),
629            trace_beans: Some(vec![bean]),
630            ..Default::default()
631        };
632
633        let transfer_bean = TraceDataEncoder::encoder_from_context_bean(&ctx).unwrap();
634        let decoded_contexts = TraceDataEncoder::decoder_from_trace_data_string(&transfer_bean.trans_data);
635
636        assert_eq!(decoded_contexts.len(), 1);
637        let decoded = &decoded_contexts[0];
638        assert_eq!(decoded.trace_type, Some(TraceType::SubBefore));
639        assert_eq!(decoded.time_stamp, 9876543210);
640        assert_eq!(decoded.region_id, CheetahString::from("eu-west-1"));
641        assert_eq!(decoded.group_name, CheetahString::from("consumer-group"));
642        assert_eq!(decoded.request_id, CheetahString::from("req-123"));
643        assert!(decoded.trace_beans.is_some());
644        let decoded_bean = &decoded.trace_beans.as_ref().unwrap()[0];
645        assert_eq!(decoded_bean.msg_id, CheetahString::from("msg-002"));
646        assert_eq!(decoded_bean.retry_times, 2);
647        assert_eq!(decoded_bean.keys, CheetahString::from("key3"));
648    }
649
650    #[test]
651    fn test_encode_and_decode_sub_after_context() {
652        let bean = TraceBean {
653            msg_id: CheetahString::from("msg-003"),
654            keys: CheetahString::from("key4 key5"),
655            ..Default::default()
656        };
657
658        let ctx = TraceContext {
659            trace_type: Some(TraceType::SubAfter),
660            request_id: CheetahString::from("req-456"),
661            cost_time: 200,
662            is_success: false,
663            context_code: 1,
664            time_stamp: 1122334455,
665            group_name: CheetahString::from("consumer-group-2"),
666            trace_beans: Some(vec![bean]),
667            ..Default::default()
668        };
669
670        let transfer_bean = TraceDataEncoder::encoder_from_context_bean(&ctx).unwrap();
671        let decoded_contexts = TraceDataEncoder::decoder_from_trace_data_string(&transfer_bean.trans_data);
672
673        assert_eq!(decoded_contexts.len(), 1);
674        let decoded = &decoded_contexts[0];
675        assert_eq!(decoded.trace_type, Some(TraceType::SubAfter));
676        assert_eq!(decoded.request_id, CheetahString::from("req-456"));
677        assert_eq!(decoded.cost_time, 200);
678        assert!(!decoded.is_success);
679        assert_eq!(decoded.context_code, 1);
680        assert_eq!(decoded.time_stamp, 1122334455);
681        assert_eq!(decoded.group_name, CheetahString::from("consumer-group-2"));
682        assert!(decoded.trace_beans.is_some());
683        let decoded_bean = &decoded.trace_beans.as_ref().unwrap()[0];
684        assert_eq!(decoded_bean.msg_id, CheetahString::from("msg-003"));
685        assert_eq!(decoded_bean.keys, CheetahString::from("key4 key5"));
686    }
687
688    #[test]
689    fn test_encode_and_decode_sub_after_cloud_channel() {
690        let bean = TraceBean {
691            msg_id: CheetahString::from("msg-004"),
692            keys: CheetahString::from("key6"),
693            ..Default::default()
694        };
695
696        let ctx = TraceContext {
697            trace_type: Some(TraceType::SubAfter),
698            request_id: CheetahString::from("req-789"),
699            cost_time: 150,
700            is_success: true,
701            context_code: 0,
702            access_channel: Some(AccessChannel::Cloud),
703            time_stamp: 1122334455,
704            group_name: CheetahString::from("consumer-group-cloud"),
705            trace_beans: Some(vec![bean]),
706            ..Default::default()
707        };
708
709        let transfer_bean = TraceDataEncoder::encoder_from_context_bean(&ctx).unwrap();
710        let decoded_contexts = TraceDataEncoder::decoder_from_trace_data_string(&transfer_bean.trans_data);
711
712        assert_eq!(decoded_contexts.len(), 1);
713        let decoded = &decoded_contexts[0];
714        assert_eq!(decoded.trace_type, Some(TraceType::SubAfter));
715        assert_eq!(decoded.request_id, CheetahString::from("req-789"));
716        assert_eq!(decoded.cost_time, 150);
717        assert!(decoded.is_success);
718        assert_eq!(decoded.context_code, 0);
719        assert_eq!(decoded.time_stamp, 0);
720        assert_eq!(decoded.group_name, CheetahString::default());
721    }
722
723    #[test]
724    fn test_encode_and_decode_end_transaction_context() {
725        let bean = TraceBean {
726            topic: CheetahString::from("tx-topic"),
727            msg_id: CheetahString::from("msg-tx-001"),
728            tags: CheetahString::from("tx-tag"),
729            keys: CheetahString::from("tx-key"),
730            store_host: CheetahString::from("127.0.0.1:10911"),
731            msg_type: Some(MessageType::TransMsgCommit),
732            transaction_id: Some(CheetahString::from("tx-id-001")),
733            transaction_state: Some(LocalTransactionState::CommitMessage),
734            from_transaction_check: true,
735            ..Default::default()
736        };
737
738        let ctx = TraceContext {
739            trace_type: Some(TraceType::EndTransaction),
740            time_stamp: 1234567890123,
741            region_id: CheetahString::from("ap-southeast-1"),
742            group_name: CheetahString::from("tx-group"),
743            trace_beans: Some(vec![bean]),
744            ..Default::default()
745        };
746
747        let transfer_bean = TraceDataEncoder::encoder_from_context_bean(&ctx).unwrap();
748        let decoded_contexts = TraceDataEncoder::decoder_from_trace_data_string(&transfer_bean.trans_data);
749
750        assert_eq!(decoded_contexts.len(), 1);
751        let decoded = &decoded_contexts[0];
752        assert_eq!(decoded.trace_type, Some(TraceType::EndTransaction));
753        assert_eq!(decoded.time_stamp, 1234567890123);
754        assert_eq!(decoded.region_id, CheetahString::from("ap-southeast-1"));
755        assert_eq!(decoded.group_name, CheetahString::from("tx-group"));
756        assert!(decoded.trace_beans.is_some());
757        let decoded_bean = &decoded.trace_beans.as_ref().unwrap()[0];
758        assert_eq!(decoded_bean.topic, CheetahString::from("tx-topic"));
759        assert_eq!(decoded_bean.msg_id, CheetahString::from("msg-tx-001"));
760        assert_eq!(decoded_bean.tags, CheetahString::from("tx-tag"));
761        assert_eq!(decoded_bean.keys, CheetahString::from("tx-key"));
762        assert_eq!(decoded_bean.store_host, CheetahString::from("127.0.0.1:10911"));
763        assert_eq!(decoded_bean.msg_type, Some(MessageType::TransMsgCommit));
764        assert_eq!(decoded_bean.transaction_id, Some(CheetahString::from("tx-id-001")));
765        assert_eq!(
766            decoded_bean.transaction_state,
767            Some(LocalTransactionState::CommitMessage)
768        );
769        assert!(decoded_bean.from_transaction_check);
770    }
771
772    #[test]
773    fn test_encode_and_decode_recall_context() {
774        let bean = TraceBean {
775            topic: CheetahString::from("recall-topic"),
776            msg_id: CheetahString::from("msg-recall-001"),
777            ..Default::default()
778        };
779
780        let ctx = TraceContext {
781            trace_type: Some(TraceType::Recall),
782            time_stamp: 9876543210987,
783            region_id: CheetahString::from("sa-east-1"),
784            group_name: CheetahString::from("recall-group"),
785            is_success: true,
786            trace_beans: Some(vec![bean]),
787            ..Default::default()
788        };
789
790        let transfer_bean = TraceDataEncoder::encoder_from_context_bean(&ctx).unwrap();
791        let decoded_contexts = TraceDataEncoder::decoder_from_trace_data_string(&transfer_bean.trans_data);
792
793        assert_eq!(decoded_contexts.len(), 1);
794        let decoded = &decoded_contexts[0];
795        assert_eq!(decoded.trace_type, Some(TraceType::Recall));
796        assert_eq!(decoded.time_stamp, 9876543210987);
797        assert_eq!(decoded.region_id, CheetahString::from("sa-east-1"));
798        assert_eq!(decoded.group_name, CheetahString::from("recall-group"));
799        assert!(decoded.is_success);
800        assert!(decoded.trace_beans.is_some());
801        let decoded_bean = &decoded.trace_beans.as_ref().unwrap()[0];
802        assert_eq!(decoded_bean.topic, CheetahString::from("recall-topic"));
803        assert_eq!(decoded_bean.msg_id, CheetahString::from("msg-recall-001"));
804    }
805
806    #[test]
807    fn test_encode_invalid_context() {
808        let ctx_no_beans = TraceContext {
809            trace_type: Some(TraceType::Pub),
810            trace_beans: None,
811            ..Default::default()
812        };
813        assert!(TraceDataEncoder::encoder_from_context_bean(&ctx_no_beans).is_none());
814
815        let ctx_empty_beans = TraceContext {
816            trace_type: Some(TraceType::Pub),
817            trace_beans: Some(vec![]),
818            ..Default::default()
819        };
820        assert!(TraceDataEncoder::encoder_from_context_bean(&ctx_empty_beans).is_none());
821        let ctx_no_trace_type = TraceContext {
822            trace_type: None,
823            trace_beans: Some(vec![TraceBean::default()]),
824            ..Default::default()
825        };
826        assert!(TraceDataEncoder::encoder_from_context_bean(&ctx_no_trace_type).is_none());
827    }
828
829    #[test]
830    fn test_transfer_bean_keys() {
831        let bean = TraceBean {
832            msg_id: CheetahString::from("msg-key-test"),
833            keys: CheetahString::from("key1 key2 key3"),
834            ..Default::default()
835        };
836
837        let ctx = TraceContext {
838            trace_type: Some(TraceType::Pub),
839            trace_beans: Some(vec![bean]),
840            ..Default::default()
841        };
842
843        let transfer_bean = TraceDataEncoder::encoder_from_context_bean(&ctx).unwrap();
844        assert!(transfer_bean.trans_key.contains(&CheetahString::from("msg-key-test")));
845        assert!(transfer_bean.trans_key.contains(&CheetahString::from("key1")));
846        assert!(transfer_bean.trans_key.contains(&CheetahString::from("key2")));
847        assert!(transfer_bean.trans_key.contains(&CheetahString::from("key3")));
848    }
849
850    #[test]
851    fn test_message_type_ordinal_conversion() {
852        assert_eq!(TraceDataEncoder::message_type_to_ordinal(MessageType::NormalMsg), 0);
853        assert_eq!(TraceDataEncoder::message_type_to_ordinal(MessageType::TransMsgHalf), 1);
854        assert_eq!(
855            TraceDataEncoder::message_type_to_ordinal(MessageType::TransMsgCommit),
856            2
857        );
858        assert_eq!(TraceDataEncoder::message_type_to_ordinal(MessageType::DelayMsg), 3);
859        assert_eq!(TraceDataEncoder::message_type_to_ordinal(MessageType::OrderMsg), 4);
860
861        assert_eq!(
862            TraceDataEncoder::message_type_from_ordinal(0),
863            Some(MessageType::NormalMsg)
864        );
865        assert_eq!(
866            TraceDataEncoder::message_type_from_ordinal(1),
867            Some(MessageType::TransMsgHalf)
868        );
869        assert_eq!(
870            TraceDataEncoder::message_type_from_ordinal(2),
871            Some(MessageType::TransMsgCommit)
872        );
873        assert_eq!(
874            TraceDataEncoder::message_type_from_ordinal(3),
875            Some(MessageType::DelayMsg)
876        );
877        assert_eq!(
878            TraceDataEncoder::message_type_from_ordinal(4),
879            Some(MessageType::OrderMsg)
880        );
881        assert_eq!(TraceDataEncoder::message_type_from_ordinal(99), None);
882    }
883
884    #[test]
885    fn test_parse_transaction_state() {
886        assert_eq!(
887            TraceDataEncoder::parse_transaction_state("COMMIT_MESSAGE"),
888            Some(LocalTransactionState::CommitMessage)
889        );
890        assert_eq!(
891            TraceDataEncoder::parse_transaction_state("ROLLBACK_MESSAGE"),
892            Some(LocalTransactionState::RollbackMessage)
893        );
894        assert_eq!(
895            TraceDataEncoder::parse_transaction_state("UNKNOW"),
896            Some(LocalTransactionState::Unknown)
897        );
898        assert_eq!(
899            TraceDataEncoder::parse_transaction_state("UNKNOWN"),
900            Some(LocalTransactionState::Unknown)
901        );
902        assert_eq!(TraceDataEncoder::parse_transaction_state("INVALID"), None);
903    }
904
905    fn raw_record(ctx: &TraceContext) -> String {
906        TraceDataEncoder::encoder_from_context_bean(ctx)
907            .unwrap()
908            .trans_data
909            .to_string()
910    }
911
912    /// Replaces one `CONTENT_SPLITOR`-delimited field of an encoded record, preserving its
913    /// trailing `FIELD_SPLITOR` terminator if present.
914    fn corrupt_field(record: &str, field_index: usize, replacement: &str) -> String {
915        let has_trailing_field_splitor = record.ends_with(TraceConstants::FIELD_SPLITOR);
916        let body = if has_trailing_field_splitor {
917            &record[..record.len() - TraceConstants::FIELD_SPLITOR.len_utf8()]
918        } else {
919            record
920        };
921        let mut fields: Vec<&str> = body.split(TraceConstants::CONTENT_SPLITOR).collect();
922        fields[field_index] = replacement;
923        let mut joined = fields.join(&TraceConstants::CONTENT_SPLITOR.to_string());
924        if has_trailing_field_splitor {
925            joined.push(TraceConstants::FIELD_SPLITOR);
926        }
927        joined
928    }
929
930    #[test]
931    fn decode_malformed_truncated_record_between_two_valid_records_keeps_both_valid_ones_in_order() {
932        let pub_ctx = TraceContext {
933            trace_type: Some(TraceType::Pub),
934            time_stamp: 111,
935            region_id: CheetahString::from("r1"),
936            group_name: CheetahString::from("g1"),
937            cost_time: 10,
938            is_success: true,
939            trace_beans: Some(vec![TraceBean {
940                topic: CheetahString::from("topic-1"),
941                msg_id: CheetahString::from("msg-1"),
942                ..Default::default()
943            }]),
944            ..Default::default()
945        };
946        let sub_before_ctx = TraceContext {
947            trace_type: Some(TraceType::SubBefore),
948            time_stamp: 222,
949            region_id: CheetahString::from("r2"),
950            group_name: CheetahString::from("g2"),
951            request_id: CheetahString::from("req-2"),
952            trace_beans: Some(vec![TraceBean {
953                msg_id: CheetahString::from("msg-2"),
954                retry_times: 1,
955                ..Default::default()
956            }]),
957            ..Default::default()
958        };
959
960        let pub_raw = raw_record(&pub_ctx);
961        let sub_before_raw = raw_record(&sub_before_ctx);
962        // Fewer than the 12 fields `decode_pub_context` requires.
963        let truncated = format!("Pub{sep}123{sep}region", sep = TraceConstants::CONTENT_SPLITOR);
964
965        let combined = format!(
966            "{pub_raw}{truncated}{sep}{sub_before_raw}",
967            sep = TraceConstants::FIELD_SPLITOR
968        );
969        let decoded = TraceDataEncoder::decoder_from_trace_data_string(&combined);
970
971        assert_eq!(decoded.len(), 2);
972        assert_eq!(decoded[0].trace_type, Some(TraceType::Pub));
973        assert_eq!(
974            decoded[0].trace_beans.as_ref().unwrap()[0].msg_id,
975            CheetahString::from("msg-1")
976        );
977        assert_eq!(decoded[1].trace_type, Some(TraceType::SubBefore));
978        assert_eq!(
979            decoded[1].trace_beans.as_ref().unwrap()[0].msg_id,
980            CheetahString::from("msg-2")
981        );
982    }
983
984    #[test]
985    fn decode_malformed_truncated_records_of_every_type_are_skipped() {
986        let sep = TraceConstants::CONTENT_SPLITOR;
987        let truncated_records = [
988            format!("Pub{sep}1{sep}region{sep}group{sep}topic"),
989            format!("SubBefore{sep}1{sep}region{sep}group"),
990            format!("SubAfter{sep}req{sep}msg"),
991            format!("EndTransaction{sep}1{sep}region{sep}group{sep}topic"),
992            format!("Recall{sep}1{sep}region"),
993        ];
994
995        for truncated in truncated_records {
996            let decoded = TraceDataEncoder::decoder_from_trace_data_string(&truncated);
997            assert!(
998                decoded.is_empty(),
999                "expected truncated record to be skipped: {truncated:?}"
1000            );
1001        }
1002    }
1003
1004    #[test]
1005    fn decode_malformed_corrupted_timestamp_is_skipped_but_following_valid_record_decodes() {
1006        let pub_ctx = TraceContext {
1007            trace_type: Some(TraceType::Pub),
1008            time_stamp: 111,
1009            region_id: CheetahString::from("r1"),
1010            group_name: CheetahString::from("g1"),
1011            cost_time: 10,
1012            is_success: true,
1013            trace_beans: Some(vec![TraceBean {
1014                topic: CheetahString::from("topic-1"),
1015                msg_id: CheetahString::from("msg-1"),
1016                ..Default::default()
1017            }]),
1018            ..Default::default()
1019        };
1020        let corrupted_pub = corrupt_field(&raw_record(&pub_ctx), 1, "not-a-number");
1021
1022        let recall_ctx = TraceContext {
1023            trace_type: Some(TraceType::Recall),
1024            time_stamp: 333,
1025            region_id: CheetahString::from("r3"),
1026            group_name: CheetahString::from("g3"),
1027            is_success: true,
1028            trace_beans: Some(vec![TraceBean {
1029                topic: CheetahString::from("topic-3"),
1030                msg_id: CheetahString::from("msg-3"),
1031                ..Default::default()
1032            }]),
1033            ..Default::default()
1034        };
1035        let recall_raw = raw_record(&recall_ctx);
1036
1037        let combined = format!("{corrupted_pub}{recall_raw}");
1038        let decoded = TraceDataEncoder::decoder_from_trace_data_string(&combined);
1039
1040        assert_eq!(decoded.len(), 1);
1041        assert_eq!(decoded[0].trace_type, Some(TraceType::Recall));
1042        assert_eq!(
1043            decoded[0].trace_beans.as_ref().unwrap()[0].msg_id,
1044            CheetahString::from("msg-3")
1045        );
1046    }
1047
1048    #[test]
1049    fn decode_malformed_corrupted_cost_time_is_skipped_but_following_valid_record_decodes() {
1050        let pub_ctx = TraceContext {
1051            trace_type: Some(TraceType::Pub),
1052            time_stamp: 111,
1053            region_id: CheetahString::from("r1"),
1054            group_name: CheetahString::from("g1"),
1055            cost_time: 10,
1056            is_success: true,
1057            trace_beans: Some(vec![TraceBean {
1058                topic: CheetahString::from("topic-1"),
1059                msg_id: CheetahString::from("msg-1"),
1060                ..Default::default()
1061            }]),
1062            ..Default::default()
1063        };
1064        // cost_time is field index 10 for a Pub record.
1065        let corrupted_pub = corrupt_field(&raw_record(&pub_ctx), 10, "not-a-number");
1066
1067        let sub_before_ctx = TraceContext {
1068            trace_type: Some(TraceType::SubBefore),
1069            time_stamp: 222,
1070            region_id: CheetahString::from("r2"),
1071            group_name: CheetahString::from("g2"),
1072            request_id: CheetahString::from("req-2"),
1073            trace_beans: Some(vec![TraceBean {
1074                msg_id: CheetahString::from("msg-2"),
1075                retry_times: 1,
1076                ..Default::default()
1077            }]),
1078            ..Default::default()
1079        };
1080        let sub_before_raw = raw_record(&sub_before_ctx);
1081
1082        let combined = format!("{corrupted_pub}{sub_before_raw}");
1083        let decoded = TraceDataEncoder::decoder_from_trace_data_string(&combined);
1084
1085        assert_eq!(decoded.len(), 1);
1086        assert_eq!(decoded[0].trace_type, Some(TraceType::SubBefore));
1087    }
1088
1089    #[test]
1090    fn decode_malformed_corrupted_retry_times_is_skipped_but_following_valid_record_decodes() {
1091        let sub_before_ctx = TraceContext {
1092            trace_type: Some(TraceType::SubBefore),
1093            time_stamp: 222,
1094            region_id: CheetahString::from("r2"),
1095            group_name: CheetahString::from("g2"),
1096            request_id: CheetahString::from("req-2"),
1097            trace_beans: Some(vec![TraceBean {
1098                msg_id: CheetahString::from("msg-2"),
1099                retry_times: 1,
1100                ..Default::default()
1101            }]),
1102            ..Default::default()
1103        };
1104        // retry_times is field index 6 for a SubBefore record.
1105        let corrupted_sub_before = corrupt_field(&raw_record(&sub_before_ctx), 6, "not-a-number");
1106
1107        let pub_ctx = TraceContext {
1108            trace_type: Some(TraceType::Pub),
1109            time_stamp: 111,
1110            region_id: CheetahString::from("r1"),
1111            group_name: CheetahString::from("g1"),
1112            cost_time: 10,
1113            is_success: true,
1114            trace_beans: Some(vec![TraceBean {
1115                topic: CheetahString::from("topic-1"),
1116                msg_id: CheetahString::from("msg-1"),
1117                ..Default::default()
1118            }]),
1119            ..Default::default()
1120        };
1121        let pub_raw = raw_record(&pub_ctx);
1122
1123        let combined = format!("{corrupted_sub_before}{pub_raw}");
1124        let decoded = TraceDataEncoder::decoder_from_trace_data_string(&combined);
1125
1126        assert_eq!(decoded.len(), 1);
1127        assert_eq!(decoded[0].trace_type, Some(TraceType::Pub));
1128    }
1129
1130    #[test]
1131    fn decode_malformed_unknown_type_and_repeated_separators_create_no_contexts() {
1132        let unknown = format!(
1133            "Bogus{content_sep}foo{content_sep}bar",
1134            content_sep = TraceConstants::CONTENT_SPLITOR
1135        );
1136        let combined = format!(
1137            "{unknown}{field_sep}{field_sep}{field_sep}",
1138            field_sep = TraceConstants::FIELD_SPLITOR
1139        );
1140
1141        let decoded = TraceDataEncoder::decoder_from_trace_data_string(&combined);
1142
1143        assert!(decoded.is_empty());
1144    }
1145}
rocketmq_client_rust/trace/trace_data_encoder.rs

rocketmq_client_rust/trace/
trace_data_encoder.rs
