rocketmq_client_rust/stat/
consumer_stats_manager.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::future::Future;
16use std::sync::atomic::AtomicBool;
17use std::sync::atomic::Ordering;
18use std::sync::Arc;
19use std::thread;
20use std::time::Duration;
21
22use rocketmq_common::common::stats::stats_item_set::StatsItemSet;
23use rocketmq_remoting::protocol::body::consume_status::ConsumeStatus;
24use tracing::warn;
25
26const TOPIC_AND_GROUP_CONSUME_OK_TPS: &str = "CONSUME_OK_TPS";
27const TOPIC_AND_GROUP_CONSUME_FAILED_TPS: &str = "CONSUME_FAILED_TPS";
28const TOPIC_AND_GROUP_CONSUME_RT: &str = "CONSUME_RT";
29const TOPIC_AND_GROUP_PULL_TPS: &str = "PULL_TPS";
30const TOPIC_AND_GROUP_PULL_RT: &str = "PULL_RT";
31
32/// Tracks consumer-side statistics for each topic/group pair.
33///
34/// Maintains five time-windowed metric sets: consume OK TPS, consume failed
35/// TPS, consume RT, pull TPS, and pull RT. All `inc_*` methods are
36/// thread-safe and intended to be called on the hot consumption path.
37pub struct ConsumerStatsManager {
38    topic_and_group_consume_ok_tps: StatsItemSet,
39    topic_and_group_consume_rt: StatsItemSet,
40    topic_and_group_consume_failed_tps: StatsItemSet,
41    topic_and_group_pull_tps: StatsItemSet,
42    topic_and_group_pull_rt: StatsItemSet,
43    shutdown_signal: Arc<AtomicBool>,
44}
45
46impl Default for ConsumerStatsManager {
47    fn default() -> Self {
48        Self::new()
49    }
50}
51
52impl ConsumerStatsManager {
53    /// Creates a new `ConsumerStatsManager` with all metric sets initialised.
54    pub fn new() -> Self {
55        Self {
56            topic_and_group_consume_ok_tps: StatsItemSet::new(TOPIC_AND_GROUP_CONSUME_OK_TPS.to_string()),
57            topic_and_group_consume_rt: StatsItemSet::new(TOPIC_AND_GROUP_CONSUME_RT.to_string()),
58            topic_and_group_consume_failed_tps: StatsItemSet::new(TOPIC_AND_GROUP_CONSUME_FAILED_TPS.to_string()),
59            topic_and_group_pull_tps: StatsItemSet::new(TOPIC_AND_GROUP_PULL_TPS.to_string()),
60            topic_and_group_pull_rt: StatsItemSet::new(TOPIC_AND_GROUP_PULL_RT.to_string()),
61            shutdown_signal: Arc::new(AtomicBool::new(false)),
62        }
63    }
64
65    /// Starts background sampling tasks.
66    ///
67    /// Spawns three background tasks that advance the sliding-window snapshots used
68    /// by [`consume_status`](Self::consume_status):
69    ///
70    /// - every 10 s: `cs_list_minute` (drives per-minute stats)
71    /// - every 10 min: `cs_list_hour` (drives per-hour stats)
72    /// - every 1 h: `cs_list_day` (drives per-day stats)
73    pub fn start(&self) {
74        self.shutdown_signal.store(false, Ordering::Release);
75
76        // 10-second tick drives cs_list_minute on each StatsItem.
77        let sets_sec = [
78            self.topic_and_group_consume_ok_tps.clone(),
79            self.topic_and_group_consume_rt.clone(),
80            self.topic_and_group_consume_failed_tps.clone(),
81            self.topic_and_group_pull_tps.clone(),
82            self.topic_and_group_pull_rt.clone(),
83        ];
84        let shutdown_signal = self.shutdown_signal.clone();
85        spawn_stats_task("rocketmq-client-consumer-stats-seconds", async move {
86            loop {
87                if shutdown_signal.load(Ordering::Acquire) {
88                    break;
89                }
90                for set in &sets_sec {
91                    set.sampling_in_seconds();
92                }
93                if !sleep_or_shutdown(&shutdown_signal, Duration::from_secs(10)).await {
94                    break;
95                }
96            }
97        });
98
99        // 10-minute tick drives cs_list_hour on each StatsItem.
100        let sets_min = [
101            self.topic_and_group_consume_ok_tps.clone(),
102            self.topic_and_group_consume_rt.clone(),
103            self.topic_and_group_consume_failed_tps.clone(),
104            self.topic_and_group_pull_tps.clone(),
105            self.topic_and_group_pull_rt.clone(),
106        ];
107        let shutdown_signal = self.shutdown_signal.clone();
108        spawn_stats_task("rocketmq-client-consumer-stats-minutes", async move {
109            loop {
110                if shutdown_signal.load(Ordering::Acquire) {
111                    break;
112                }
113                for set in &sets_min {
114                    set.sampling_in_minutes();
115                }
116                if !sleep_or_shutdown(&shutdown_signal, Duration::from_secs(10 * 60)).await {
117                    break;
118                }
119            }
120        });
121
122        // 1-hour tick drives cs_list_day on each StatsItem.
123        let sets_hour = [
124            self.topic_and_group_consume_ok_tps.clone(),
125            self.topic_and_group_consume_rt.clone(),
126            self.topic_and_group_consume_failed_tps.clone(),
127            self.topic_and_group_pull_tps.clone(),
128            self.topic_and_group_pull_rt.clone(),
129        ];
130        let shutdown_signal = self.shutdown_signal.clone();
131        spawn_stats_task("rocketmq-client-consumer-stats-hours", async move {
132            loop {
133                if shutdown_signal.load(Ordering::Acquire) {
134                    break;
135                }
136                for set in &sets_hour {
137                    set.sampling_in_hours();
138                }
139                if !sleep_or_shutdown(&shutdown_signal, Duration::from_secs(3600)).await {
140                    break;
141                }
142            }
143        });
144    }
145
146    /// Shuts down the stats manager.
147    pub fn shutdown(&self) {
148        self.shutdown_signal.store(true, Ordering::Release);
149    }
150
151    /// Records a single pull response-time observation in milliseconds.
152    pub fn inc_pull_rt(&self, group: &str, topic: &str, rt: u64) {
153        self.topic_and_group_pull_rt
154            .add_rt_value(&stats_key(topic, group), rt as i64, 1);
155    }
156
157    /// Records `msgs` messages successfully pulled in one batch.
158    pub fn inc_pull_tps(&self, group: &str, topic: &str, msgs: u64) {
159        self.topic_and_group_pull_tps
160            .add_value(&stats_key(topic, group), msgs as i64, 1);
161    }
162
163    /// Records a single consume response-time observation in milliseconds.
164    pub fn inc_consume_rt(&self, group: &str, topic: &str, rt: u64) {
165        self.topic_and_group_consume_rt
166            .add_rt_value(&stats_key(topic, group), rt as i64, 1);
167    }
168
169    /// Records `msgs` messages consumed successfully in one batch.
170    pub fn inc_consume_ok_tps(&self, group: &str, topic: &str, msgs: u64) {
171        self.topic_and_group_consume_ok_tps
172            .add_value(&stats_key(topic, group), msgs as i64, 1);
173    }
174
175    /// Java-compatible acronym spelling for `incConsumeOKTPS`.
176    #[inline]
177    pub fn inc_consume_oktps(&self, group: &str, topic: &str, msgs: u64) {
178        self.inc_consume_ok_tps(group, topic, msgs);
179    }
180
181    /// Records `msgs` messages that failed consumption in one batch.
182    pub fn inc_consume_failed_tps(&self, group: &str, topic: &str, msgs: u64) {
183        self.topic_and_group_consume_failed_tps
184            .add_value(&stats_key(topic, group), msgs as i64, 1);
185    }
186
187    /// Returns a point-in-time [`ConsumeStatus`] snapshot for the given
188    /// consumer group and topic.
189    ///
190    /// - Pull / consume RT uses the per-minute average; consume RT falls back to the per-hour
191    ///   average when the per-minute window is empty.
192    /// - `consume_failed_msgs` accumulates the per-hour sum of failed messages.
193    pub fn consume_status(&self, group: &str, topic: &str) -> ConsumeStatus {
194        let key = stats_key(topic, group);
195
196        let pull_rt = self.topic_and_group_pull_rt.get_stats_data_in_minute(&key).get_avgpt();
197
198        let pull_tps = self.topic_and_group_pull_tps.get_stats_data_in_minute(&key).get_tps();
199
200        let consume_rt = {
201            let minute = self.topic_and_group_consume_rt.get_stats_data_in_minute(&key);
202            if minute.get_sum() == 0 {
203                self.topic_and_group_consume_rt.get_stats_data_in_hour(&key).get_avgpt()
204            } else {
205                minute.get_avgpt()
206            }
207        };
208
209        let consume_ok_tps = self
210            .topic_and_group_consume_ok_tps
211            .get_stats_data_in_minute(&key)
212            .get_tps();
213
214        let consume_failed_tps = self
215            .topic_and_group_consume_failed_tps
216            .get_stats_data_in_minute(&key)
217            .get_tps();
218
219        let consume_failed_msgs = self
220            .topic_and_group_consume_failed_tps
221            .get_stats_data_in_hour(&key)
222            .get_sum() as i64;
223
224        ConsumeStatus {
225            pull_rt,
226            pull_tps,
227            consume_rt,
228            consume_ok_tps,
229            consume_failed_tps,
230            consume_failed_msgs,
231        }
232    }
233}
234
235/// Builds the canonical stats key `"topic@group"` used by all metric sets.
236#[inline]
237fn stats_key(topic: &str, group: &str) -> String {
238    format!("{topic}@{group}")
239}
240
241fn spawn_stats_task<F>(thread_name: &'static str, task: F)
242where
243    F: Future<Output = ()> + Send + 'static,
244{
245    if let Ok(handle) = tokio::runtime::Handle::try_current() {
246        drop(handle.spawn(task));
247        return;
248    }
249
250    match tokio::runtime::Builder::new_current_thread().enable_all().build() {
251        Ok(runtime) => {
252            if let Err(error) = thread::Builder::new()
253                .name(thread_name.to_string())
254                .spawn(move || runtime.block_on(task))
255            {
256                warn!("Failed to spawn {} background thread: {}", thread_name, error);
257            }
258        }
259        Err(error) => warn!("Failed to build {} runtime: {}", thread_name, error),
260    }
261}
262
263async fn sleep_or_shutdown(shutdown_signal: &Arc<AtomicBool>, delay: Duration) -> bool {
264    let deadline = tokio::time::Instant::now() + delay;
265
266    loop {
267        if shutdown_signal.load(Ordering::Acquire) {
268            return false;
269        }
270
271        let now = tokio::time::Instant::now();
272        if now >= deadline {
273            return true;
274        }
275
276        tokio::time::sleep((deadline - now).min(Duration::from_millis(50))).await;
277    }
278}
279
280#[cfg(test)]
281mod tests {
282    use super::*;
283
284    fn make_manager() -> ConsumerStatsManager {
285        ConsumerStatsManager::new()
286    }
287
288    #[test]
289    fn stats_key_format() {
290        assert_eq!(stats_key("TopicA", "GroupA"), "TopicA@GroupA");
291    }
292
293    #[test]
294    fn smoke_inc_consume_ok_tps() {
295        let mgr = make_manager();
296        mgr.inc_consume_ok_tps("GroupA", "TopicA", 5);
297    }
298
299    #[test]
300    fn java_acronym_alias_inc_consume_oktps_records_ok_tps() {
301        let mgr = make_manager();
302        mgr.inc_consume_oktps("GroupA", "TopicA", 5);
303        let item = mgr
304            .topic_and_group_consume_ok_tps
305            .get_stats_item(&stats_key("TopicA", "GroupA"))
306            .expect("alias should record the same stats item as inc_consume_ok_tps");
307
308        assert_eq!(item.get_value(), 5);
309        assert_eq!(item.get_times(), 1);
310    }
311
312    #[test]
313    fn smoke_inc_consume_failed_tps() {
314        let mgr = make_manager();
315        mgr.inc_consume_failed_tps("GroupA", "TopicA", 3);
316    }
317
318    #[test]
319    fn smoke_inc_consume_rt() {
320        let mgr = make_manager();
321        mgr.inc_consume_rt("GroupA", "TopicA", 42);
322    }
323
324    #[test]
325    fn smoke_inc_pull_rt() {
326        let mgr = make_manager();
327        mgr.inc_pull_rt("GroupA", "TopicA", 10);
328    }
329
330    #[test]
331    fn smoke_inc_pull_tps() {
332        let mgr = make_manager();
333        mgr.inc_pull_tps("GroupA", "TopicA", 100);
334    }
335
336    #[test]
337    fn consume_status_returns_zero_for_empty_stats() {
338        let mgr = make_manager();
339        let status = mgr.consume_status("GroupA", "TopicA");
340        assert_eq!(status.pull_rt, 0.0);
341        assert_eq!(status.pull_tps, 0.0);
342        assert_eq!(status.consume_rt, 0.0);
343        assert_eq!(status.consume_ok_tps, 0.0);
344        assert_eq!(status.consume_failed_tps, 0.0);
345        assert_eq!(status.consume_failed_msgs, 0);
346    }
347
348    #[tokio::test]
349    async fn start_launches_background_tasks() {
350        let mgr = make_manager();
351        // Verifies start() does not panic in a Tokio runtime context.
352        mgr.start();
353        mgr.shutdown();
354    }
355
356    #[test]
357    fn start_without_tokio_runtime_does_not_panic() {
358        let mgr = make_manager();
359        mgr.start();
360        mgr.shutdown();
361        thread::sleep(Duration::from_millis(80));
362    }
363
364    #[test]
365    fn default_creates_valid_manager() {
366        let mgr = ConsumerStatsManager::default();
367        // Should not panic when querying uninitialised key.
368        let _ = mgr.consume_status("G", "T");
369    }
370}
rocketmq_client_rust/stat/consumer_stats_manager.rs

rocketmq_client_rust/stat/
consumer_stats_manager.rs
