Skip to main content

rustfs_common/
trace_bus.rs

1// Copyright 2024 RustFS Team
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 smallvec::SmallVec;
16use std::{
17    sync::{
18        Arc, OnceLock,
19        atomic::{AtomicUsize, Ordering},
20    },
21    time::{Duration, SystemTime},
22};
23use tokio::sync::broadcast;
24
25const DEFAULT_TRACE_BUS_CAPACITY: usize = 1024;
26const TRACE_ATTR_INLINE_CAPACITY: usize = 8;
27
28static GLOBAL_TRACE_BUS: OnceLock<TraceBus> = OnceLock::new();
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub enum TraceKind {
32    Heal,
33    Scanner,
34}
35
36impl TraceKind {
37    pub const fn as_str(self) -> &'static str {
38        match self {
39            Self::Heal => "heal",
40            Self::Scanner => "scanner",
41        }
42    }
43}
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub enum TraceFunc {
47    HealTask,
48    HealBucket,
49    HealObject,
50    HealCheckAbandonedParts,
51    HealErasureSetPage,
52    ScannerFolder,
53    ScannerIlmAction,
54    ScannerHealCandidate,
55    Dropped,
56}
57
58impl TraceFunc {
59    pub const fn as_str(self) -> &'static str {
60        match self {
61            Self::HealTask => "heal.Task",
62            Self::HealBucket => "heal.Bucket",
63            Self::HealObject => "heal.Object",
64            Self::HealCheckAbandonedParts => "heal.CheckAbandonedParts",
65            Self::HealErasureSetPage => "heal.ErasureSetPage",
66            Self::ScannerFolder => "scanner.Folder",
67            Self::ScannerIlmAction => "scanner.IlmAction",
68            Self::ScannerHealCandidate => "scanner.HealCandidate",
69            Self::Dropped => "trace.Dropped",
70        }
71    }
72}
73
74#[derive(Debug, Clone, PartialEq, Eq)]
75pub enum TraceVal {
76    Bool(bool),
77    U64(u64),
78    I64(i64),
79    Str(Arc<str>),
80}
81
82impl From<bool> for TraceVal {
83    fn from(value: bool) -> Self {
84        Self::Bool(value)
85    }
86}
87
88impl From<u64> for TraceVal {
89    fn from(value: u64) -> Self {
90        Self::U64(value)
91    }
92}
93
94impl From<i64> for TraceVal {
95    fn from(value: i64) -> Self {
96        Self::I64(value)
97    }
98}
99
100impl From<&str> for TraceVal {
101    fn from(value: &str) -> Self {
102        Self::Str(Arc::from(value))
103    }
104}
105
106impl From<String> for TraceVal {
107    fn from(value: String) -> Self {
108        Self::Str(Arc::from(value))
109    }
110}
111
112#[derive(Debug, Clone, PartialEq, Eq)]
113pub struct TraceAttr {
114    pub key: &'static str,
115    pub value: TraceVal,
116}
117
118#[derive(Debug, Clone, PartialEq, Eq)]
119pub struct TraceEvent {
120    pub kind: TraceKind,
121    pub func: TraceFunc,
122    pub time: SystemTime,
123    pub bucket: Option<Arc<str>>,
124    pub object: Option<Arc<str>>,
125    pub duration: Duration,
126    pub bytes: u64,
127    pub attrs: SmallVec<[TraceAttr; TRACE_ATTR_INLINE_CAPACITY]>,
128}
129
130/// A telemetry operation whose published shape cannot retain request data.
131#[derive(Debug, Clone, Copy, PartialEq, Eq)]
132pub enum TelemetryTraceOperation {
133    GetObject,
134    PutObject,
135    HeadObject,
136    ListObjects,
137    InternalRpc,
138}
139
140#[derive(Debug, Clone, Copy, PartialEq, Eq)]
141pub enum TelemetryTraceStatus {
142    Ok,
143    Error,
144}
145
146/// A pre-classified S3 or internode RPC observation.
147///
148/// This event deliberately has no request path, headers, object identity,
149/// payload, or free-form attributes. Producers must classify the operation
150/// and outcome before publishing it.
151#[derive(Debug, Clone, Copy, PartialEq, Eq)]
152pub struct TelemetryTraceEvent {
153    pub operation: TelemetryTraceOperation,
154    pub duration: Duration,
155    pub status: TelemetryTraceStatus,
156}
157
158impl TelemetryTraceEvent {
159    pub const fn new(operation: TelemetryTraceOperation, duration: Duration, status: TelemetryTraceStatus) -> Self {
160        Self {
161            operation,
162            duration,
163            status,
164        }
165    }
166}
167
168impl TraceEvent {
169    pub fn new(kind: TraceKind, func: TraceFunc) -> Self {
170        Self {
171            kind,
172            func,
173            time: SystemTime::now(),
174            bucket: None,
175            object: None,
176            duration: Duration::ZERO,
177            bytes: 0,
178            attrs: SmallVec::new(),
179        }
180    }
181
182    pub fn with_bucket(mut self, bucket: impl Into<Arc<str>>) -> Self {
183        self.bucket = Some(bucket.into());
184        self
185    }
186
187    pub fn with_object(mut self, object: impl Into<Arc<str>>) -> Self {
188        self.object = Some(object.into());
189        self
190    }
191
192    pub fn with_duration(mut self, duration: Duration) -> Self {
193        self.duration = duration;
194        self
195    }
196
197    pub fn with_bytes(mut self, bytes: u64) -> Self {
198        self.bytes = bytes;
199        self
200    }
201
202    pub fn with_attr(mut self, key: &'static str, value: impl Into<TraceVal>) -> Self {
203        self.attrs.push(TraceAttr {
204            key,
205            value: value.into(),
206        });
207        self
208    }
209}
210
211#[derive(Debug)]
212pub struct TraceBus {
213    sender: broadcast::Sender<Arc<TraceEvent>>,
214    subscriber_count: Arc<AtomicUsize>,
215    telemetry_sender: broadcast::Sender<Arc<TelemetryTraceEvent>>,
216    telemetry_subscriber_count: Arc<AtomicUsize>,
217}
218
219impl TraceBus {
220    pub fn new(capacity: usize) -> Self {
221        let capacity = capacity.max(1);
222        let (sender, _receiver) = broadcast::channel(capacity);
223        let (telemetry_sender, _telemetry_receiver) = broadcast::channel(capacity);
224        Self {
225            sender,
226            subscriber_count: Arc::new(AtomicUsize::new(0)),
227            telemetry_sender,
228            telemetry_subscriber_count: Arc::new(AtomicUsize::new(0)),
229        }
230    }
231
232    pub fn subscriber_count(&self) -> usize {
233        self.subscriber_count.load(Ordering::Acquire)
234    }
235
236    pub fn subscribe(&self) -> TraceSubscription {
237        let receiver = self.sender.subscribe();
238        self.subscriber_count.fetch_add(1, Ordering::AcqRel);
239        TraceSubscription {
240            receiver,
241            subscriber_count: Arc::clone(&self.subscriber_count),
242        }
243    }
244
245    pub fn emit(&self, build: impl FnOnce() -> TraceEvent) -> bool {
246        if self.subscriber_count() == 0 {
247            return false;
248        }
249
250        self.sender.send(Arc::new(build())).is_ok()
251    }
252
253    pub fn telemetry_subscriber_count(&self) -> usize {
254        self.telemetry_subscriber_count.load(Ordering::Acquire)
255    }
256
257    pub fn subscribe_telemetry(&self) -> TelemetryTraceSubscription {
258        let receiver = self.telemetry_sender.subscribe();
259        self.telemetry_subscriber_count.fetch_add(1, Ordering::AcqRel);
260        TelemetryTraceSubscription {
261            receiver,
262            subscriber_count: Arc::clone(&self.telemetry_subscriber_count),
263        }
264    }
265
266    pub fn emit_telemetry(&self, build: impl FnOnce() -> TelemetryTraceEvent) -> bool {
267        if self.telemetry_subscriber_count() == 0 {
268            return false;
269        }
270
271        self.telemetry_sender.send(Arc::new(build())).is_ok()
272    }
273}
274
275impl Default for TraceBus {
276    fn default() -> Self {
277        Self::new(DEFAULT_TRACE_BUS_CAPACITY)
278    }
279}
280
281#[derive(Debug)]
282pub struct TraceSubscription {
283    receiver: broadcast::Receiver<Arc<TraceEvent>>,
284    subscriber_count: Arc<AtomicUsize>,
285}
286
287impl TraceSubscription {
288    pub async fn recv(&mut self) -> Result<Arc<TraceEvent>, broadcast::error::RecvError> {
289        self.receiver.recv().await
290    }
291
292    pub fn try_recv(&mut self) -> Result<Arc<TraceEvent>, broadcast::error::TryRecvError> {
293        self.receiver.try_recv()
294    }
295}
296
297impl Drop for TraceSubscription {
298    fn drop(&mut self) {
299        self.subscriber_count.fetch_sub(1, Ordering::AcqRel);
300    }
301}
302
303#[derive(Debug)]
304pub struct TelemetryTraceSubscription {
305    receiver: broadcast::Receiver<Arc<TelemetryTraceEvent>>,
306    subscriber_count: Arc<AtomicUsize>,
307}
308
309impl TelemetryTraceSubscription {
310    pub async fn recv(&mut self) -> Result<Arc<TelemetryTraceEvent>, broadcast::error::RecvError> {
311        self.receiver.recv().await
312    }
313
314    pub fn try_recv(&mut self) -> Result<Arc<TelemetryTraceEvent>, broadcast::error::TryRecvError> {
315        self.receiver.try_recv()
316    }
317}
318
319impl Drop for TelemetryTraceSubscription {
320    fn drop(&mut self) {
321        self.subscriber_count.fetch_sub(1, Ordering::AcqRel);
322    }
323}
324
325pub fn global_trace_bus() -> &'static TraceBus {
326    GLOBAL_TRACE_BUS.get_or_init(TraceBus::default)
327}
328
329pub fn subscribe_trace_events() -> TraceSubscription {
330    global_trace_bus().subscribe()
331}
332
333pub fn trace_emit(build: impl FnOnce() -> TraceEvent) -> bool {
334    global_trace_bus().emit(build)
335}
336
337pub fn trace_subscriber_count() -> usize {
338    global_trace_bus().subscriber_count()
339}
340
341pub fn subscribe_telemetry_trace_events() -> TelemetryTraceSubscription {
342    global_trace_bus().subscribe_telemetry()
343}
344
345pub fn telemetry_trace_emit(build: impl FnOnce() -> TelemetryTraceEvent) -> bool {
346    global_trace_bus().emit_telemetry(build)
347}
348
349pub fn telemetry_trace_subscriber_count() -> usize {
350    global_trace_bus().telemetry_subscriber_count()
351}
352
353#[cfg(test)]
354mod tests {
355    use super::*;
356    use std::sync::atomic::AtomicUsize;
357
358    #[test]
359    fn trace_emit_skips_builder_without_subscribers() {
360        let bus = TraceBus::new(4);
361        let built = AtomicUsize::new(0);
362
363        let sent = bus.emit(|| {
364            built.fetch_add(1, Ordering::Relaxed);
365            TraceEvent::new(TraceKind::Heal, TraceFunc::HealTask)
366        });
367
368        assert!(!sent);
369        assert_eq!(built.load(Ordering::Relaxed), 0);
370    }
371
372    #[tokio::test]
373    async fn trace_subscriber_receives_event() {
374        let bus = TraceBus::new(4);
375        let mut subscription = bus.subscribe();
376
377        assert!(bus.emit(|| {
378            TraceEvent::new(TraceKind::Heal, TraceFunc::HealObject)
379                .with_bucket("bucket")
380                .with_object("object")
381                .with_duration(Duration::from_millis(7))
382                .with_bytes(11)
383                .with_attr("dry", true)
384        }));
385
386        let event = subscription
387            .recv()
388            .await
389            .expect("subscriber should receive emitted trace event");
390
391        assert_eq!(event.kind, TraceKind::Heal);
392        assert_eq!(event.func, TraceFunc::HealObject);
393        assert_eq!(event.bucket.as_deref(), Some("bucket"));
394        assert_eq!(event.object.as_deref(), Some("object"));
395        assert_eq!(event.duration, Duration::from_millis(7));
396        assert_eq!(event.bytes, 11);
397        assert_eq!(
398            event.attrs.as_slice(),
399            &[TraceAttr {
400                key: "dry",
401                value: TraceVal::Bool(true)
402            }]
403        );
404    }
405
406    #[test]
407    fn trace_subscription_drop_decrements_count() {
408        let bus = TraceBus::new(4);
409        let subscription = bus.subscribe();
410
411        assert_eq!(bus.subscriber_count(), 1);
412        drop(subscription);
413        assert_eq!(bus.subscriber_count(), 0);
414    }
415
416    #[tokio::test]
417    async fn lagged_subscriber_drops_events_without_blocking_publishers() {
418        let bus = TraceBus::new(2);
419        let mut subscription = bus.subscribe();
420
421        for index in 0_u64..4 {
422            assert!(bus.emit(|| { TraceEvent::new(TraceKind::Scanner, TraceFunc::ScannerFolder).with_attr("index", index) }));
423        }
424
425        let err = subscription
426            .recv()
427            .await
428            .expect_err("receiver should observe lag instead of blocking publishers");
429        assert!(matches!(err, broadcast::error::RecvError::Lagged(_)));
430    }
431
432    #[tokio::test]
433    async fn telemetry_subscription_exposes_only_classified_fields() {
434        let bus = TraceBus::new(4);
435        let _unrelated_subscription = bus.subscribe();
436        let built = AtomicUsize::new(0);
437        assert!(!bus.emit_telemetry(|| {
438            built.fetch_add(1, Ordering::Relaxed);
439            TelemetryTraceEvent::new(TelemetryTraceOperation::GetObject, Duration::ZERO, TelemetryTraceStatus::Ok)
440        }));
441        assert_eq!(built.load(Ordering::Relaxed), 0);
442
443        let mut subscription = bus.subscribe_telemetry();
444        assert_eq!(bus.telemetry_subscriber_count(), 1);
445
446        assert!(bus.emit_telemetry(|| {
447            TelemetryTraceEvent::new(TelemetryTraceOperation::GetObject, Duration::from_micros(7), TelemetryTraceStatus::Ok)
448        }));
449
450        let event = subscription.recv().await.expect("classified telemetry event");
451        assert_eq!(event.operation, TelemetryTraceOperation::GetObject);
452        assert_eq!(event.duration, Duration::from_micros(7));
453        assert_eq!(event.status, TelemetryTraceStatus::Ok);
454        drop(subscription);
455        assert_eq!(bus.telemetry_subscriber_count(), 0);
456    }
457}