ktrace_rs/
lib.rs

1// Copyright 2026 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5#![no_std]
6
7use arch_rs::{InterruptDisableGuard, curr_cpu_num, ints_disabled};
8use core::cell::UnsafeCell;
9use core::mem::{MaybeUninit, size_of};
10use core::sync::atomic::{AtomicBool, AtomicPtr, AtomicU32, Ordering};
11use core::{ffi, ptr, slice};
12use kalloc::Box;
13use kernel::thread::{FxtRef, ThreadPtr};
14pub use kstring::declare_interned_category;
15use kstring::declare_interned_string;
16use kstring::interned_category::InternedCategory;
17pub use kstring::interned_string::InternedString;
18pub use platform_rs::{InstantBootTicks, timer_current_boot_ticks};
19use spsc_buffer::{Buffer, NoOpAllocator, Reservation};
20use zx_status::Status;
21
22// Re-export the ktrace macros from the sub-crate.
23pub use ktrace_macro::*;
24#[allow(unused_extern_crates)]
25extern crate self as ktrace_rs;
26
27// LINT.IfChange(KTraceState)
28#[repr(C)]
29pub struct KTraceState {
30    pub categories_bitmask: AtomicU32,
31    pub writes_enabled: AtomicBool,
32}
33const _: () = assert!(size_of::<KTraceState>() == 8);
34// LINT.ThenChange(//zircon/kernel/include/lib/ktrace.h:KTraceState)
35
36declare_interned_category!(META_CAT, "kernel:meta", extern);
37declare_interned_string!(DROP_STATS_REF, "drop_stats", extern);
38declare_interned_string!(NUM_RECORDS_REF, "num_records", extern);
39declare_interned_string!(NUM_BYTES_REF, "num_bytes", extern);
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum Context {
43    Thread = 0,
44    Cpu = 1,
45}
46
47#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub enum EventType {
49    Instant = 0,
50    Counter = 1,
51    DurationBegin = 2,
52    DurationEnd = 3,
53    DurationComplete = 4,
54    FlowBegin = 8,
55    FlowStep = 9,
56    FlowEnd = 10,
57}
58
59/// The value of a trace argument.
60pub enum ArgValue<'a> {
61    Null,
62    Bool(bool),
63    Int32(i32),
64    Uint32(u32),
65    Int64(i64),
66    Uint64(u64),
67    Double(f64),
68    String(&'a str),
69    Pointer(usize),
70    Koid(u64),
71}
72
73impl<'a> From<bool> for ArgValue<'a> {
74    fn from(v: bool) -> Self {
75        ArgValue::Bool(v)
76    }
77}
78impl<'a> From<i32> for ArgValue<'a> {
79    fn from(v: i32) -> Self {
80        ArgValue::Int32(v)
81    }
82}
83impl<'a> From<u32> for ArgValue<'a> {
84    fn from(v: u32) -> Self {
85        ArgValue::Uint32(v)
86    }
87}
88impl<'a> From<i64> for ArgValue<'a> {
89    fn from(v: i64) -> Self {
90        ArgValue::Int64(v)
91    }
92}
93impl<'a> From<u64> for ArgValue<'a> {
94    fn from(v: u64) -> Self {
95        ArgValue::Uint64(v)
96    }
97}
98impl<'a> From<f64> for ArgValue<'a> {
99    fn from(v: f64) -> Self {
100        ArgValue::Double(v)
101    }
102}
103impl<'a> From<&'a str> for ArgValue<'a> {
104    fn from(v: &'a str) -> Self {
105        ArgValue::String(v)
106    }
107}
108
109pub struct Argument<'a> {
110    name: &'static InternedString,
111    value: ArgValue<'a>,
112}
113
114impl<'a> Argument<'a> {
115    pub fn new(name: &'static InternedString, value: impl Into<ArgValue<'a>>) -> Self {
116        Self { name, value: value.into() }
117    }
118
119    /// Returns the size of this argument in 64-bit words.
120    fn size_words(&self) -> usize {
121        match &self.value {
122            ArgValue::Null | ArgValue::Bool(_) | ArgValue::Int32(_) | ArgValue::Uint32(_) => 1,
123            ArgValue::Int64(_)
124            | ArgValue::Uint64(_)
125            | ArgValue::Double(_)
126            | ArgValue::Pointer(_)
127            | ArgValue::Koid(_) => 2,
128            ArgValue::String(s) => 1 + (s.len() + 7) / 8,
129        }
130    }
131
132    fn write(&self, res: &mut KTraceReservation<'_>) -> Result<(), Status> {
133        let name_id = self.name.id() as u64;
134        let mut header = 0u64;
135        header |= (name_id & 0xffff) << 16; // NameRef
136
137        match &self.value {
138            ArgValue::Null => {
139                header |= 0u64; // ArgumentType::kNull (0)
140                res.write_word(header)?;
141            }
142            ArgValue::Int32(v) => {
143                header |= 1u64; // ArgumentType::kInt32 (1)
144                header |= ((*v as u32) as u64) << 32;
145                res.write_word(header)?;
146            }
147            ArgValue::Uint32(v) => {
148                header |= 2u64; // ArgumentType::kUint32 (2)
149                header |= (*v as u64) << 32;
150                res.write_word(header)?;
151            }
152            ArgValue::Int64(v) => {
153                header |= 3u64; // ArgumentType::kInt64 (3)
154                res.write_word(header)?;
155                res.write_word(*v as u64)?;
156            }
157            ArgValue::Uint64(v) => {
158                header |= 4u64; // ArgumentType::kUint64 (4)
159                res.write_word(header)?;
160                res.write_word(*v)?;
161            }
162            ArgValue::Double(v) => {
163                header |= 5u64; // ArgumentType::kDouble (5)
164                res.write_word(header)?;
165                res.write_word(v.to_bits())?;
166            }
167            ArgValue::String(s) => {
168                header |= 6u64; // ArgumentType::kString (6)
169                let string_len = s.len();
170                header |= (((string_len + 7) / 8) as u64) << 4; // ArgumentSize in words
171                header |= (string_len as u64) << 32; // String length in bytes
172                res.write_word(header)?;
173                res.write_bytes(s.as_bytes())?;
174            }
175            ArgValue::Pointer(v) => {
176                header |= 7u64; // ArgumentType::kPointer (7)
177                res.write_word(header)?;
178                res.write_word(*v as u64)?;
179            }
180            ArgValue::Koid(v) => {
181                header |= 8u64; // ArgumentType::kKoid (8)
182                res.write_word(header)?;
183                res.write_word(*v)?;
184            }
185            ArgValue::Bool(v) => {
186                header |= 9u64; // ArgumentType::kBool (9)
187                header |= ((*v as u64) & 1) << 32;
188                res.write_word(header)?;
189            }
190        }
191        Ok(())
192    }
193}
194
195pub struct KTraceScope<'a> {
196    category: &'static InternedCategory,
197    name: &'static InternedString,
198    timestamp: InstantBootTicks,
199    context: Context,
200    args: &'a [Argument<'a>],
201}
202
203impl<'a> KTraceScope<'a> {
204    #[inline(never)]
205    #[cold]
206    pub fn begin(
207        category: &'static InternedCategory,
208        name: &'static InternedString,
209        context: Context,
210        args: &'a [Argument<'a>],
211    ) -> Self {
212        let timestamp = timer_current_boot_ticks();
213        Self { category, name, timestamp, context, args }
214    }
215}
216
217impl<'a> Drop for KTraceScope<'a> {
218    #[inline(never)]
219    #[cold]
220    fn drop(&mut self) {
221        let end_time = timer_current_boot_ticks();
222        let ktrace = KTrace::get_instance();
223        ktrace.emit_event(
224            EventType::DurationComplete,
225            self.category,
226            self.name,
227            self.timestamp,
228            self.context,
229            Some(end_time.0 as u64),
230            self.args,
231        );
232    }
233}
234
235// LINT.IfChange(DroppedRecordDurationEvent)
236#[repr(C)]
237struct DroppedRecordDurationEvent {
238    header: u64,
239    start: InstantBootTicks,
240    process_id: u64,
241    thread_id: u64,
242    num_dropped_arg: u64,
243    bytes_dropped_arg: u64,
244    end: InstantBootTicks,
245}
246const _: () = assert!(size_of::<DroppedRecordDurationEvent>() == 56);
247// LINT.ThenChange(//zircon/kernel/lib/percpu_writer/include/lib/percpu_writer/buffer.h:DroppedRecordDurationEvent)
248
249// LINT.IfChange(DroppedRecordStats)
250/// This struct keeps track of the duration, number, and size of trace records dropped when the
251/// buffer is full. These statistics are emitted to the trace buffer as a duration as soon as
252/// space is available to do so, at which point the values are reset to 0, or false in the case
253/// of has_dropped.
254#[repr(C)]
255#[derive(Default, Debug, Clone)]
256pub struct DroppedRecordStats {
257    pub first_dropped: InstantBootTicks,
258    pub last_dropped: InstantBootTicks,
259
260    /// By storing num_dropped and bytes_dropped in 32-bit values, we ensure that they can each
261    /// be stored in a single 64-bit word in the FXT record we emit when space is available.
262    pub num_dropped: u32,
263    pub bytes_dropped: u32,
264    pub has_dropped: bool,
265}
266const _: () = assert!(size_of::<DroppedRecordStats>() == 32);
267// LINT.ThenChange(//zircon/kernel/lib/percpu_writer/include/lib/percpu_writer/buffer.h:DroppedRecordStats)
268
269impl DroppedRecordStats {
270    pub fn reset(&mut self) {
271        self.first_dropped = InstantBootTicks(0);
272        self.last_dropped = InstantBootTicks(0);
273        self.num_dropped = 0;
274        self.bytes_dropped = 0;
275        self.has_dropped = false;
276    }
277
278    pub fn track(&mut self, now: InstantBootTicks, size: u32) {
279        if !self.has_dropped {
280            self.first_dropped = now;
281            self.has_dropped = true;
282        }
283        self.last_dropped = now;
284        self.num_dropped = self.num_dropped.wrapping_add(1);
285        self.bytes_dropped = self.bytes_dropped.wrapping_add(size);
286    }
287
288    pub fn has_dropped(&self) -> bool {
289        self.has_dropped
290    }
291}
292
293/// A Rust implementation of `percpu_writer::Buffer` that wraps a static reference to an existing
294/// `spsc_buffer::Buffer` and tracks dropped trace records.
295pub struct KTraceBuffer {
296    buffer: &'static mut Buffer<NoOpAllocator>,
297    pub drop_stats: &'static mut DroppedRecordStats,
298    size: u32,
299    cpu_ref_header_entry: u16,
300    pub process_koid: u64,
301    pub thread_koid: u64,
302}
303
304// SAFETY: Access to the per-CPU buffers is synchronized by the caller (typically via disabling
305// interrupts).
306unsafe impl Send for KTraceBuffer {}
307unsafe impl Sync for KTraceBuffer {}
308
309impl KTraceBuffer {
310    /// Constructs a `KTraceBuffer` from a static reference to an existing `spsc_buffer::Buffer`,
311    /// a static reference to `DroppedRecordStats`, and its metadata.
312    pub fn new(
313        buffer: &'static mut Buffer<NoOpAllocator>,
314        drop_stats: &'static mut DroppedRecordStats,
315        cpu_ref_header_entry: u16,
316        process_koid: u64,
317        thread_koid: u64,
318    ) -> Self {
319        let size = buffer.size();
320        Self { buffer, drop_stats, size, cpu_ref_header_entry, process_koid, thread_koid }
321    }
322
323    /// Returns the size of the backing storage.
324    pub fn size(&self) -> u32 {
325        self.size
326    }
327
328    /// Drains the underlying Buffer.
329    pub fn drain(&self) -> Result<(), Status> {
330        self.buffer.drain()
331    }
332
333    /// Copies `len` bytes out of the buffer using the provided `copy_fn`.
334    pub fn read<F>(&self, copy_fn: F, len: u32) -> Result<u32, Status>
335    where
336        F: FnMut(u32, &[u8]) -> Result<(), Status>,
337    {
338        self.buffer.read(copy_fn, len)
339    }
340
341    /// Reserves a block of the given size in the buffer, interposing to write dropped record
342    /// statistics if any were tracked.
343    pub fn reserve(&mut self, header: u64) -> Result<KTraceReservation<'_>, Status> {
344        debug_assert!(ints_disabled());
345        // Compute the number of bytes we need to reserve from the provided fxt header.
346        let record_type = (header & 0xf) as u32;
347        let num_words = if record_type == 15 {
348            // Large record
349            ((header >> 4) & 0xffffffff) as u32
350        } else {
351            // Normal record
352            ((header >> 4) & 0xfff) as u32
353        };
354        let size = num_words * 8;
355
356        let mut total_size = size;
357        let event = if self.drop_stats.has_dropped() {
358            total_size += size_of::<DroppedRecordDurationEvent>() as u32;
359            Some(self.serialize_drop_stats())
360        } else {
361            None
362        };
363
364        match self.buffer.reserve(total_size) {
365            Err(status) => {
366                let now = timer_current_boot_ticks();
367                self.drop_stats.track(now, size);
368                Err(status)
369            }
370            Ok(mut res) => {
371                if let Some(event) = event {
372                    // SAFETY: DroppedRecordDurationEvent is repr(C) and has a defined binary
373                    // layout.
374                    let event_bytes = unsafe {
375                        slice::from_raw_parts(
376                            ptr::from_ref(&event).cast::<u8>(),
377                            size_of::<DroppedRecordDurationEvent>(),
378                        )
379                    };
380                    res.write(event_bytes)?;
381                    self.drop_stats.reset();
382                }
383                Ok(KTraceReservation::new(res, header))
384            }
385        }
386    }
387
388    /// Emit the dropped record stats to the trace buffer if we're currently tracking them.
389    pub fn emit_drop_stats(&mut self) -> Result<(), Status> {
390        debug_assert!(ints_disabled());
391        if !self.drop_stats.has_dropped() {
392            return Ok(());
393        }
394
395        // Serialize the event first so we release the borrow on self before calling reserve.
396        let event = self.serialize_drop_stats();
397
398        let mut res = self.buffer.reserve(size_of::<DroppedRecordDurationEvent>() as u32)?;
399        // SAFETY: DroppedRecordDurationEvent is repr(C) and has a defined binary layout.
400        let event_bytes = unsafe {
401            slice::from_raw_parts(
402                ptr::from_ref(&event).cast::<u8>(),
403                size_of::<DroppedRecordDurationEvent>(),
404            )
405        };
406        res.write(event_bytes)?;
407        res.commit()?;
408
409        // Reset the fields directly rather than calling reset_drop_stats() to avoid
410        // borrowing the entire self while res is in scope.
411        self.drop_stats.reset();
412        Ok(())
413    }
414
415    /// Resets the dropped records statistics to their initial values.
416    pub fn reset_drop_stats(&mut self) {
417        self.drop_stats.reset();
418    }
419
420    fn serialize_drop_stats(&self) -> DroppedRecordDurationEvent {
421        let mut header = 4u64; // RecordType::kEvent (4)
422        let record_size_words = (size_of::<DroppedRecordDurationEvent>() / 8) as u64;
423        header |= record_size_words << 4; // RecordSize
424        header |= 4u64 << 16; // EventType::kDurationComplete (4)
425        header |= 2u64 << 20; // ArgumentCount = 2
426        header |= (self.cpu_ref_header_entry as u64) << 24;
427        header |= (META_CAT.label().id() as u64) << 32;
428        header |= (DROP_STATS_REF.id() as u64) << 48;
429
430        // Pack the arguments.
431        // In FXT:
432        // ArgumentType::kUint32 is 2
433        // ArgumentSize is 1 word (8 bytes)
434        // NameRef is packed into bits 16..31
435        // Value is packed into bits 32..63
436        let mut num_dropped_arg = 2u64;
437        num_dropped_arg |= 1u64 << 4;
438        num_dropped_arg |= (NUM_RECORDS_REF.id() as u64) << 16;
439        num_dropped_arg |= (self.drop_stats.num_dropped as u64) << 32;
440
441        let mut bytes_dropped_arg = 2u64;
442        bytes_dropped_arg |= 1u64 << 4;
443        bytes_dropped_arg |= (NUM_BYTES_REF.id() as u64) << 16;
444        bytes_dropped_arg |= (self.drop_stats.bytes_dropped as u64) << 32;
445
446        DroppedRecordDurationEvent {
447            header,
448            start: self.drop_stats.first_dropped,
449            process_id: self.process_koid,
450            thread_id: self.thread_koid,
451            num_dropped_arg,
452            bytes_dropped_arg,
453            end: self.drop_stats.last_dropped,
454        }
455    }
456}
457
458/// KTraceReservation encapsulates a pending write to the buffer.
459#[derive(Debug)]
460pub struct KTraceReservation<'a> {
461    reservation: Reservation<'a>,
462}
463
464impl<'a> KTraceReservation<'a> {
465    fn new(reservation: Reservation<'a>, header: u64) -> Self {
466        let mut this = Self { reservation };
467        let _ = this.write_word(header);
468        this
469    }
470
471    /// Writes a single 64-bit word into the reservation.
472    pub fn write_word(&mut self, word: u64) -> Result<(), Status> {
473        debug_assert!(ints_disabled());
474        self.reservation.write(&word.to_ne_bytes())
475    }
476
477    /// Writes a byte slice into the reservation, padding to an 8-byte boundary.
478    pub fn write_bytes(&mut self, bytes: &[u8]) -> Result<(), Status> {
479        debug_assert!(ints_disabled());
480        self.reservation.write(bytes)?;
481        let num_bytes = bytes.len();
482        let aligned_bytes = (num_bytes + 7) & !7;
483        let num_zeros_to_write = aligned_bytes - num_bytes;
484        if num_zeros_to_write > 0 {
485            let zero = [0u8; 8];
486            self.reservation.write(&zero[..num_zeros_to_write])?;
487        }
488        Ok(())
489    }
490
491    /// Commits the reservation, making it visible to the reader.
492    pub fn commit(self) -> Result<(), Status> {
493        debug_assert!(ints_disabled());
494        self.reservation.commit()
495    }
496}
497
498/// A pure Rust implementation of KTrace.
499pub struct KTrace {
500    /// Reference to the shared C++ KTraceState.
501    // TODO(https://fxbug.dev/517305548): This should be made a direct allocation, and not a
502    // reference, once the C++ implementation is removed.
503    state: &'static KTraceState,
504
505    /// A heap-allocated slice of atomic pointers to per-CPU buffers.
506    /// Allocated once at boot time, and accessed completely lock-free on the hot path by CPU ID.
507    buffers: Box<[AtomicPtr<KTraceBuffer>]>,
508}
509
510// SAFETY: KTrace is a global singleton. Access to the per-CPU buffers is synchronized.
511unsafe impl Sync for KTrace {}
512unsafe impl Send for KTrace {}
513
514#[repr(transparent)]
515struct KTraceSingleton(UnsafeCell<MaybeUninit<KTrace>>);
516
517unsafe impl Sync for KTraceSingleton {}
518unsafe impl Send for KTraceSingleton {}
519
520static INSTANCE: KTraceSingleton = KTraceSingleton(UnsafeCell::new(MaybeUninit::uninit()));
521
522impl KTrace {
523    /// Retrieve the global instance of KTrace.
524    pub fn get_instance() -> &'static Self {
525        // SAFETY: KTrace must be initialized during kernel boot.
526        unsafe { &*INSTANCE.0.get().cast::<KTrace>() }
527    }
528
529    /// Returns the raw pointer to the KTraceBuffer for the given CPU.
530    pub fn get_buffer(&self, cpu: usize) -> *mut KTraceBuffer {
531        self.buffers[cpu].load(Ordering::Acquire)
532    }
533
534    /// Reserves a slot of memory to write a record into.
535    ///
536    /// # Safety
537    ///
538    /// This method MUST be invoked with interrupts disabled to enforce the single-writer invariant.
539    pub unsafe fn reserve(&self, header: u64) -> Result<KTraceReservation<'_>, Status> {
540        debug_assert!(ints_disabled());
541        if !self.writes_enabled() {
542            return Err(Status::BAD_STATE);
543        }
544
545        // Direct, lock-free indexing into the slice, followed by loading the atomic pointer!
546        let ptr = self.buffers[curr_cpu_num() as usize].load(Ordering::Acquire);
547        if ptr.is_null() {
548            return Err(Status::BAD_STATE);
549        }
550
551        let buf = unsafe { &mut *ptr };
552        buf.reserve(header)
553    }
554
555    /// Returns true if writes are currently enabled.
556    pub fn writes_enabled(&self) -> bool {
557        self.state.writes_enabled.load(Ordering::Acquire)
558    }
559
560    /// Returns the categories bitmask.
561    pub fn categories_bitmask(&self) -> u32 {
562        self.state.categories_bitmask.load(Ordering::Acquire)
563    }
564
565    /// Returns true if the given category is enabled.
566    pub fn is_category_enabled(&self, category: &InternedCategory) -> bool {
567        let category_index = category.index();
568        if category_index == InternedCategory::INVALID_INDEX {
569            return false;
570        }
571        let bitmask = self.categories_bitmask();
572        (bitmask & (1 << category_index)) != 0
573    }
574
575    /// Low-level helper to write an FXT kernel object record.
576    ///
577    /// This is not inlined to reduce code size at the instrumentation sites.
578    #[inline(never)]
579    #[cold]
580    pub fn emit_kernel_object_outlined(
581        &self,
582        koid: u64,
583        obj_type: u32,
584        name: &InternedString,
585        args: &[Argument<'_>],
586    ) {
587        let _guard = InterruptDisableGuard::new();
588        if !self.writes_enabled() {
589            return;
590        }
591
592        let base_size = 2; // Header, KOID
593        let args_size: usize = args.iter().map(|a| a.size_words()).sum();
594        let total_size_words = base_size + args_size;
595
596        if total_size_words > 0xfff {
597            return;
598        }
599
600        let mut header = 7u64; // RecordType::kKernelObject (7)
601        header |= (total_size_words as u64) << 4; // RecordSize
602        header |= (obj_type as u64) << 16; // ObjectType
603        header |= (name.id() as u64) << 24; // NameStringRef
604        header |= (args.len() as u64) << 40; // ArgumentCount
605
606        if let Ok(mut res) = unsafe { self.reserve(header) } {
607            let _ = res.write_word(koid);
608            for arg in args {
609                let _ = arg.write(&mut res);
610            }
611            let _ = res.commit();
612        }
613    }
614
615    /// Low-level helper to write a generic FXT event record.
616    ///
617    /// This is not inlined to reduce code size at the instrumentation sites.
618    #[inline(never)]
619    #[cold]
620    pub fn emit_event(
621        &self,
622        event_type: EventType,
623        category: &InternedCategory,
624        name: &InternedString,
625        timestamp: InstantBootTicks,
626        context: Context,
627        content: Option<u64>,
628        args: &[Argument<'_>],
629    ) {
630        let _guard = InterruptDisableGuard::new();
631        if !self.writes_enabled() {
632            return;
633        }
634
635        // 1. Get the process/thread KOIDs for the context.
636        let (process_koid, thread_koid) = match context {
637            Context::Thread => {
638                // SAFETY: ktrace is initialized and running after threading has been initialized.
639                let FxtRef { pid, tid } = unsafe { ThreadPtr::current().fxt_ref() };
640                (pid, tid)
641            }
642            Context::Cpu => {
643                let cpu = curr_cpu_num() as usize;
644                let ptr = self.buffers[cpu].load(Ordering::Acquire);
645                if ptr.is_null() {
646                    return;
647                }
648                let buf = unsafe { &*ptr };
649                (buf.process_koid, buf.thread_koid)
650            }
651        };
652
653        // 2. Calculate the record size.
654        let base_size = 4; // Header, Timestamp, Process KOID, Thread KOID
655        let content_size = if content.is_some() { 1 } else { 0 };
656        let args_size: usize = args.iter().map(|a| a.size_words()).sum();
657        let total_size_words = base_size + content_size + args_size;
658
659        if total_size_words > 0xfff {
660            return;
661        }
662
663        // 3. Construct the header.
664        let mut header = 4u64; // RecordType::kEvent (4)
665        header |= (total_size_words as u64) << 4; // RecordSize
666        header |= (event_type as u32 as u64) << 16; // EventType
667        header |= (args.len() as u64) << 20; // ArgumentCount
668        header |= (category.label().id() as u64) << 32; // CategoryStringRef
669        header |= (name.id() as u64) << 48; // NameStringRef
670
671        // 4. Reserve space and write the record.
672        if let Ok(mut res) = unsafe { self.reserve(header) } {
673            let _ = res.write_word(timestamp.0 as u64);
674            let _ = res.write_word(process_koid);
675            let _ = res.write_word(thread_koid);
676
677            for arg in args {
678                let _ = arg.write(&mut res);
679            }
680
681            if let Some(c) = content {
682                let _ = res.write_word(c);
683            }
684
685            let _ = res.commit();
686        }
687    }
688}
689
690/// Initializes the global KTrace instance with the given number of CPU buffers.
691///
692/// # Safety
693///
694/// This must be called at most once during kernel boot.
695#[unsafe(no_mangle)]
696pub unsafe extern "C" fn rust_ktrace_init(num_buffers: u32, state_ptr: *mut ffi::c_void) -> i32 {
697    if num_buffers == 0 || state_ptr.is_null() {
698        return Status::INVALID_ARGS.into_raw();
699    }
700
701    // SAFETY: The caller guarantees that state_ptr points to a valid KTraceState instance
702    // which has static storage duration (lives forever) and is safe to access concurrently
703    // (since its fields are atomic).
704    let state = unsafe { &*state_ptr.cast::<KTraceState>() };
705
706    let buffers = match Box::<[AtomicPtr<KTraceBuffer>]>::try_new_zeroed_slice(num_buffers as usize)
707    {
708        Ok(b) => b,
709        Err(_) => return Status::NO_MEMORY.into_raw(),
710    };
711
712    let ktrace = KTrace { state, buffers };
713
714    unsafe {
715        let slot = INSTANCE.0.get();
716        ptr::write(slot.cast::<KTrace>(), ktrace);
717    }
718
719    Status::OK.into_raw()
720}
721
722/// Initializes the KTraceBuffer for a specific CPU using a pointer to the C++ SpscBuffer.
723///
724/// # Safety
725///
726/// - `spsc_buffer_ptr` must point to a valid C++ `SpscBuffer` instance
727///   which is binary-compatible with `spsc_buffer::Buffer<NoOpAllocator>`.
728#[unsafe(no_mangle)]
729pub unsafe extern "C" fn rust_ktrace_init_cpu_buffer(
730    cpu_num: u32,
731    spsc_buffer_ptr: *mut ffi::c_void,
732    drop_stats_ptr: *mut ffi::c_void,
733    process_koid: u64,
734    thread_koid: u64,
735    cpu_ref_header_entry: u16,
736) -> i32 {
737    let ktrace = KTrace::get_instance();
738    if cpu_num >= ktrace.buffers.len() as u32 {
739        return Status::INVALID_ARGS.into_raw();
740    }
741
742    // SAFETY: The caller guarantees the pointers are valid, 'static, and binary-compatible with
743    // their respective Rust types.
744    let (buffer, drop_stats) = unsafe {
745        (
746            &mut *spsc_buffer_ptr.cast::<Buffer<NoOpAllocator>>(),
747            &mut *drop_stats_ptr.cast::<DroppedRecordStats>(),
748        )
749    };
750
751    // Allocate the KTraceBuffer on the heap.
752    let buf =
753        KTraceBuffer::new(buffer, drop_stats, cpu_ref_header_entry, process_koid, thread_koid);
754
755    let boxed_buf = match Box::try_new(buf) {
756        Ok(b) => b,
757        Err(_) => return Status::NO_MEMORY.into_raw(),
758    };
759
760    let raw_ptr = Box::into_raw(boxed_buf);
761
762    // Store the pointer atomically.
763    let old_ptr = ktrace.buffers[cpu_num as usize].swap(raw_ptr, Ordering::AcqRel);
764    if !old_ptr.is_null() {
765        // If there was a previous buffer, reclaim and drop it.
766        unsafe {
767            let _ = Box::from_raw(old_ptr);
768        }
769    }
770
771    Status::OK.into_raw()
772}
ktrace_rs/lib.rs

ktrace_rs/
lib.rs
