rocketmq_protocol/trace/
trace_transfer_bean.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;
16
17use cheetah_string::CheetahString;
18
19/// Container for batched trace data ready for transmission.
20///
21/// Holds encoded trace data and associated message keys for dispatch to the trace backend.
22#[derive(Debug, Clone, Default)]
23pub struct TraceTransferBean {
24    /// Encoded trace data
25    pub trans_data: CheetahString,
26    /// Message keys associated with this batch
27    pub trans_key: HashSet<CheetahString>,
28}
29
30impl TraceTransferBean {
31    /// Creates an empty trace transfer bean.
32    #[inline]
33    pub fn new() -> Self {
34        Self::default()
35    }
36
37    /// Creates a trace transfer bean with encoded data and keys.
38    #[inline]
39    pub fn with_data(trans_data: CheetahString, trans_key: HashSet<CheetahString>) -> Self {
40        Self { trans_data, trans_key }
41    }
42
43    /// Sets the encoded trace data.
44    #[inline]
45    pub fn set_trans_data(&mut self, trans_data: CheetahString) {
46        self.trans_data = trans_data;
47    }
48
49    /// Returns a reference to the encoded trace data.
50    #[inline]
51    pub fn trans_data(&self) -> &CheetahString {
52        &self.trans_data
53    }
54
55    /// Sets the message keys.
56    #[inline]
57    pub fn set_trans_key(&mut self, trans_key: HashSet<CheetahString>) {
58        self.trans_key = trans_key;
59    }
60
61    /// Returns a reference to the message keys.
62    #[inline]
63    pub fn trans_key(&self) -> &HashSet<CheetahString> {
64        &self.trans_key
65    }
66
67    /// Returns a mutable reference to the message keys.
68    #[inline]
69    pub fn trans_key_mut(&mut self) -> &mut HashSet<CheetahString> {
70        &mut self.trans_key
71    }
72
73    /// Inserts a key into the set.
74    ///
75    /// Returns `true` if the key was newly inserted, `false` if it already existed.
76    #[inline]
77    pub fn add_key(&mut self, key: CheetahString) -> bool {
78        self.trans_key.insert(key)
79    }
80
81    /// Returns `true` if both trace data and keys are empty.
82    #[inline]
83    pub fn is_empty(&self) -> bool {
84        self.trans_data.is_empty() && self.trans_key.is_empty()
85    }
86
87    /// Returns the number of keys.
88    #[inline]
89    pub fn key_count(&self) -> usize {
90        self.trans_key.len()
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97    use cheetah_string::CheetahString;
98
99    #[test]
100    fn test_new_creates_empty_bean() {
101        let bean = TraceTransferBean::new();
102        assert!(bean.trans_data.is_empty());
103        assert!(bean.trans_key.is_empty());
104        assert!(bean.is_empty());
105    }
106
107    #[test]
108    fn test_add_key_logic() {
109        let mut bean = TraceTransferBean::new();
110        let key = CheetahString::from("unique_key");
111
112        assert!(bean.add_key(key.clone()));
113        assert_eq!(bean.key_count(), 1);
114
115        assert!(!bean.add_key(key));
116        assert_eq!(bean.key_count(), 1);
117    }
118
119    #[test]
120    fn test_trans_key_mut() {
121        let mut bean = TraceTransferBean::new();
122        bean.trans_key_mut().insert(CheetahString::from("manual_key"));
123
124        assert!(bean.trans_key().contains(&CheetahString::from("manual_key")));
125        assert_eq!(bean.key_count(), 1);
126    }
127
128    #[test]
129    fn test_is_empty_conditions() {
130        let mut bean = TraceTransferBean::new();
131        assert!(bean.is_empty(), "Should be empty initially");
132
133        bean.set_trans_data(CheetahString::from("data"));
134        assert!(!bean.is_empty());
135
136        bean.set_trans_data(CheetahString::from(""));
137        bean.add_key(CheetahString::from("key"));
138        assert!(!bean.is_empty());
139    }
140}
rocketmq_protocol/trace/trace_transfer_bean.rs

rocketmq_protocol/trace/
trace_transfer_bean.rs
