1use 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#[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#[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}