ktrace_rs_tests/
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::{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 kstring::interned_category::InternedCategory;
14use kstring::{declare_interned_category, declare_interned_string};
15use platform_rs::timer_current_boot_ticks;
16use spsc_buffer::{Buffer, NoOpAllocator, Reservation};
17use zx_status::Status;
18
19// LINT.IfChange(KTraceState)
20#[repr(C)]
21struct KTraceState {
22    categories_bitmask: AtomicU32,
23    writes_enabled: AtomicBool,
24}
25const _: () = assert!(size_of::<KTraceState>() == 8);
26// LINT.ThenChange(//zircon/kernel/include/lib/ktrace.h:KTraceState)
27
28declare_interned_category!(META_CAT, "kernel:meta", extern);
29declare_interned_string!(DROP_STATS_REF, "drop_stats", extern);
30declare_interned_string!(NUM_RECORDS_REF, "num_records", extern);
31declare_interned_string!(NUM_BYTES_REF, "num_bytes", extern);
32
33// LINT.IfChange(DroppedRecordDurationEvent)
34#[repr(C)]
35struct DroppedRecordDurationEvent {
36    header: u64,
37    start: i64,
38    process_id: u64,
39    thread_id: u64,
40    num_dropped_arg: u64,
41    bytes_dropped_arg: u64,
42    end: i64,
43}
44const _: () = assert!(size_of::<DroppedRecordDurationEvent>() == 56);
45// LINT.ThenChange(//zircon/kernel/lib/percpu_writer/include/lib/percpu_writer/buffer.h:DroppedRecordDurationEvent)
46
47// LINT.IfChange(DroppedRecordStats)
48/// This struct keeps track of the duration, number, and size of trace records dropped when the
49/// buffer is full. These statistics are emitted to the trace buffer as a duration as soon as
50/// space is available to do so, at which point the values are reset to 0, or false in the case
51/// of has_dropped.
52#[repr(C)]
53#[derive(Default, Debug, Clone)]
54struct DroppedRecordStats {
55    first_dropped: i64,
56    last_dropped: i64,
57
58    /// By storing num_dropped and bytes_dropped in 32-bit values, we ensure that they can each
59    /// be stored in a single 64-bit word in the FXT record we emit when space is available.
60    num_dropped: u32,
61    bytes_dropped: u32,
62    has_dropped: bool,
63}
64const _: () = assert!(size_of::<DroppedRecordStats>() == 32);
65// LINT.ThenChange(//zircon/kernel/lib/percpu_writer/include/lib/percpu_writer/buffer.h:DroppedRecordStats)
66
67impl DroppedRecordStats {
68    fn reset(&mut self) {
69        self.first_dropped = 0;
70        self.last_dropped = 0;
71        self.num_dropped = 0;
72        self.bytes_dropped = 0;
73        self.has_dropped = false;
74    }
75
76    fn track(&mut self, now: i64, size: u32) {
77        if !self.has_dropped {
78            self.first_dropped = now;
79            self.has_dropped = true;
80        }
81        self.last_dropped = now;
82        self.num_dropped = self.num_dropped.wrapping_add(1);
83        self.bytes_dropped = self.bytes_dropped.wrapping_add(size);
84    }
85
86    fn has_dropped(&self) -> bool {
87        self.has_dropped
88    }
89}
90
91/// A Rust implementation of `percpu_writer::Buffer` that wraps a static reference to an existing
92/// `spsc_buffer::Buffer` and tracks dropped trace records.
93pub struct KTraceBuffer {
94    buffer: &'static mut Buffer<NoOpAllocator>,
95    drop_stats: &'static mut DroppedRecordStats,
96    size: u32,
97    cpu_ref_header_entry: u16,
98    process_koid: u64,
99    thread_koid: u64,
100}
101
102// SAFETY: Access to the per-CPU buffers is synchronized by the caller (typically via disabling
103// interrupts).
104unsafe impl Send for KTraceBuffer {}
105unsafe impl Sync for KTraceBuffer {}
106
107impl KTraceBuffer {
108    /// Constructs a `KTraceBuffer` from a static reference to an existing `spsc_buffer::Buffer`,
109    /// a static reference to `DroppedRecordStats`, and its metadata.
110    fn new(
111        buffer: &'static mut Buffer<NoOpAllocator>,
112        drop_stats: &'static mut DroppedRecordStats,
113        cpu_ref_header_entry: u16,
114        process_koid: u64,
115        thread_koid: u64,
116    ) -> Self {
117        let size = buffer.size();
118        Self { buffer, drop_stats, size, cpu_ref_header_entry, process_koid, thread_koid }
119    }
120
121    /// Returns the size of the backing storage.
122    pub fn size(&self) -> u32 {
123        self.size
124    }
125
126    /// Drains the underlying Buffer.
127    pub fn drain(&self) -> Result<(), Status> {
128        self.buffer.drain()
129    }
130
131    /// Copies `len` bytes out of the buffer using the provided `copy_fn`.
132    pub fn read<F>(&self, copy_fn: F, len: u32) -> Result<u32, Status>
133    where
134        F: FnMut(u32, &[u8]) -> Result<(), Status>,
135    {
136        self.buffer.read(copy_fn, len)
137    }
138
139    /// Reserves a block of the given size in the buffer, interposing to write dropped record
140    /// statistics if any were tracked.
141    pub fn reserve(&mut self, header: u64) -> Result<KTraceReservation<'_>, Status> {
142        debug_assert!(ints_disabled());
143        // Compute the number of bytes we need to reserve from the provided fxt header.
144        let record_type = (header & 0xf) as u32;
145        let num_words = if record_type == 15 {
146            // Large record
147            ((header >> 4) & 0xffffffff) as u32
148        } else {
149            // Normal record
150            ((header >> 4) & 0xfff) as u32
151        };
152        let size = num_words * 8;
153
154        let mut total_size = size;
155        let event = if self.drop_stats.has_dropped() {
156            total_size += size_of::<DroppedRecordDurationEvent>() as u32;
157            Some(self.serialize_drop_stats())
158        } else {
159            None
160        };
161
162        match self.buffer.reserve(total_size) {
163            Err(status) => {
164                let now = timer_current_boot_ticks();
165                self.drop_stats.track(now, size);
166                Err(status)
167            }
168            Ok(mut res) => {
169                if let Some(event) = event {
170                    // SAFETY: DroppedRecordDurationEvent is repr(C) and has a defined binary
171                    // layout.
172                    let event_bytes = unsafe {
173                        slice::from_raw_parts(
174                            ptr::from_ref(&event).cast::<u8>(),
175                            size_of::<DroppedRecordDurationEvent>(),
176                        )
177                    };
178                    res.write(event_bytes)?;
179                    self.drop_stats.reset();
180                }
181                Ok(KTraceReservation::new(res, header))
182            }
183        }
184    }
185
186    /// Emit the dropped record stats to the trace buffer if we're currently tracking them.
187    pub fn emit_drop_stats(&mut self) -> Result<(), Status> {
188        debug_assert!(ints_disabled());
189        if !self.drop_stats.has_dropped() {
190            return Ok(());
191        }
192
193        // Serialize the event first so we release the borrow on self before calling reserve.
194        let event = self.serialize_drop_stats();
195
196        let mut res = self.buffer.reserve(size_of::<DroppedRecordDurationEvent>() as u32)?;
197        // SAFETY: DroppedRecordDurationEvent is repr(C) and has a defined binary layout.
198        let event_bytes = unsafe {
199            slice::from_raw_parts(
200                ptr::from_ref(&event).cast::<u8>(),
201                size_of::<DroppedRecordDurationEvent>(),
202            )
203        };
204        res.write(event_bytes)?;
205        res.commit()?;
206
207        // Reset the fields directly rather than calling reset_drop_stats() to avoid
208        // borrowing the entire self while res is in scope.
209        self.drop_stats.reset();
210        Ok(())
211    }
212
213    /// Resets the dropped records statistics to their initial values.
214    pub fn reset_drop_stats(&mut self) {
215        self.drop_stats.reset();
216    }
217
218    fn serialize_drop_stats(&self) -> DroppedRecordDurationEvent {
219        let mut header = 4u64; // RecordType::kEvent (4)
220        let record_size_words = (size_of::<DroppedRecordDurationEvent>() / 8) as u64;
221        header |= record_size_words << 4; // RecordSize
222        header |= 4u64 << 16; // EventType::kDurationComplete (4)
223        header |= 2u64 << 20; // ArgumentCount = 2
224        header |= (self.cpu_ref_header_entry as u64) << 24;
225        header |= (META_CAT.label().id() as u64) << 32;
226        header |= (DROP_STATS_REF.id() as u64) << 48;
227
228        // Pack the arguments.
229        // In FXT:
230        // ArgumentType::kUint32 is 2
231        // ArgumentSize is 1 word (8 bytes)
232        // NameRef is packed into bits 16..31
233        // Value is packed into bits 32..63
234        let mut num_dropped_arg = 2u64;
235        num_dropped_arg |= 1u64 << 4;
236        num_dropped_arg |= (NUM_RECORDS_REF.id() as u64) << 16;
237        num_dropped_arg |= (self.drop_stats.num_dropped as u64) << 32;
238
239        let mut bytes_dropped_arg = 2u64;
240        bytes_dropped_arg |= 1u64 << 4;
241        bytes_dropped_arg |= (NUM_BYTES_REF.id() as u64) << 16;
242        bytes_dropped_arg |= (self.drop_stats.bytes_dropped as u64) << 32;
243
244        DroppedRecordDurationEvent {
245            header,
246            start: self.drop_stats.first_dropped,
247            process_id: self.process_koid,
248            thread_id: self.thread_koid,
249            num_dropped_arg,
250            bytes_dropped_arg,
251            end: self.drop_stats.last_dropped,
252        }
253    }
254}
255
256/// KTraceReservation encapsulates a pending write to the buffer.
257#[derive(Debug)]
258pub struct KTraceReservation<'a> {
259    reservation: Reservation<'a>,
260}
261
262impl<'a> KTraceReservation<'a> {
263    fn new(reservation: Reservation<'a>, header: u64) -> Self {
264        let mut this = Self { reservation };
265        let _ = this.write_word(header);
266        this
267    }
268
269    /// Writes a single 64-bit word into the reservation.
270    pub fn write_word(&mut self, word: u64) -> Result<(), Status> {
271        debug_assert!(ints_disabled());
272        self.reservation.write(&word.to_ne_bytes())
273    }
274
275    /// Writes a byte slice into the reservation, padding to an 8-byte boundary.
276    pub fn write_bytes(&mut self, bytes: &[u8]) -> Result<(), Status> {
277        debug_assert!(ints_disabled());
278        self.reservation.write(bytes)?;
279        let num_bytes = bytes.len();
280        let aligned_bytes = (num_bytes + 7) & !7;
281        let num_zeros_to_write = aligned_bytes - num_bytes;
282        if num_zeros_to_write > 0 {
283            let zero = [0u8; 8];
284            self.reservation.write(&zero[..num_zeros_to_write])?;
285        }
286        Ok(())
287    }
288
289    /// Commits the reservation, making it visible to the reader.
290    pub fn commit(self) -> Result<(), Status> {
291        debug_assert!(ints_disabled());
292        self.reservation.commit()
293    }
294}
295
296/// A pure Rust implementation of KTrace.
297pub struct KTrace {
298    /// Reference to the shared C++ KTraceState.
299    // TODO(https://fxbug.dev/517305548): This should be made a direct allocation, and not a
300    // reference, once the C++ implementation is removed.
301    state: &'static KTraceState,
302
303    /// A heap-allocated slice of atomic pointers to per-CPU buffers.
304    /// Allocated once at boot time, and accessed completely lock-free on the hot path by CPU ID.
305    buffers: Box<[AtomicPtr<KTraceBuffer>]>,
306}
307
308// SAFETY: KTrace is a global singleton. Access to the per-CPU buffers is synchronized.
309unsafe impl Sync for KTrace {}
310unsafe impl Send for KTrace {}
311
312#[repr(transparent)]
313struct KTraceSingleton(UnsafeCell<MaybeUninit<KTrace>>);
314
315unsafe impl Sync for KTraceSingleton {}
316unsafe impl Send for KTraceSingleton {}
317
318static INSTANCE: KTraceSingleton = KTraceSingleton(UnsafeCell::new(MaybeUninit::uninit()));
319
320impl KTrace {
321    /// Retrieve the global instance of KTrace.
322    pub fn get_instance() -> &'static Self {
323        // SAFETY: KTrace must be initialized during kernel boot.
324        unsafe { &*INSTANCE.0.get().cast::<KTrace>() }
325    }
326
327    /// Reserves a slot of memory to write a record into.
328    ///
329    /// # Safety
330    ///
331    /// This method MUST be invoked with interrupts disabled to enforce the single-writer invariant.
332    pub unsafe fn reserve(&self, header: u64) -> Result<KTraceReservation<'_>, Status> {
333        debug_assert!(ints_disabled());
334        if !self.writes_enabled() {
335            return Err(Status::BAD_STATE);
336        }
337
338        // Direct, lock-free indexing into the slice, followed by loading the atomic pointer!
339        let ptr = self.buffers[curr_cpu_num() as usize].load(Ordering::Acquire);
340        if ptr.is_null() {
341            return Err(Status::BAD_STATE);
342        }
343
344        let buf = unsafe { &mut *ptr };
345        buf.reserve(header)
346    }
347
348    /// Returns true if writes are currently enabled.
349    pub fn writes_enabled(&self) -> bool {
350        self.state.writes_enabled.load(Ordering::Acquire)
351    }
352
353    /// Returns the categories bitmask.
354    pub fn categories_bitmask(&self) -> u32 {
355        self.state.categories_bitmask.load(Ordering::Acquire)
356    }
357
358    /// Returns true if the given category is enabled.
359    pub fn is_category_enabled(&self, category: &InternedCategory) -> bool {
360        let category_index = category.index();
361        if category_index == InternedCategory::INVALID_INDEX {
362            return false;
363        }
364        let bitmask = self.categories_bitmask();
365        (bitmask & (1 << category_index)) != 0
366    }
367}
368
369/// Initializes the global KTrace instance with the given number of CPU buffers.
370///
371/// # Safety
372///
373/// This must be called at most once during kernel boot.
374#[unsafe(no_mangle)]
375pub unsafe extern "C" fn rust_ktrace_init(num_buffers: u32, state_ptr: *mut ffi::c_void) -> i32 {
376    if num_buffers == 0 || state_ptr.is_null() {
377        return Status::INVALID_ARGS.into_raw();
378    }
379
380    // SAFETY: The caller guarantees that state_ptr points to a valid KTraceState instance
381    // which has static storage duration (lives forever) and is safe to access concurrently
382    // (since its fields are atomic).
383    let state = unsafe { &*state_ptr.cast::<KTraceState>() };
384
385    let buffers = match Box::<[AtomicPtr<KTraceBuffer>]>::try_new_zeroed_slice(num_buffers as usize)
386    {
387        Ok(b) => b,
388        Err(_) => return Status::NO_MEMORY.into_raw(),
389    };
390
391    let ktrace = KTrace { state, buffers };
392
393    unsafe {
394        let slot = INSTANCE.0.get();
395        ptr::write(slot.cast::<KTrace>(), ktrace);
396    }
397
398    Status::OK.into_raw()
399}
400
401/// Initializes the KTraceBuffer for a specific CPU using a pointer to the C++ SpscBuffer.
402///
403/// # Safety
404///
405/// - `spsc_buffer_ptr` must point to a valid C++ `SpscBuffer` instance
406///   which is binary-compatible with `spsc_buffer::Buffer<NoOpAllocator>`.
407#[unsafe(no_mangle)]
408pub unsafe extern "C" fn rust_ktrace_init_cpu_buffer(
409    cpu_num: u32,
410    spsc_buffer_ptr: *mut ffi::c_void,
411    drop_stats_ptr: *mut ffi::c_void,
412    process_koid: u64,
413    thread_koid: u64,
414    cpu_ref_header_entry: u16,
415) -> i32 {
416    let ktrace = KTrace::get_instance();
417    if cpu_num >= ktrace.buffers.len() as u32 {
418        return Status::INVALID_ARGS.into_raw();
419    }
420
421    // SAFETY: The caller guarantees the pointers are valid, 'static, and binary-compatible with
422    // their respective Rust types.
423    let (buffer, drop_stats) = unsafe {
424        (
425            &mut *spsc_buffer_ptr.cast::<Buffer<NoOpAllocator>>(),
426            &mut *drop_stats_ptr.cast::<DroppedRecordStats>(),
427        )
428    };
429
430    // Allocate the KTraceBuffer on the heap.
431    let buf =
432        KTraceBuffer::new(buffer, drop_stats, cpu_ref_header_entry, process_koid, thread_koid);
433
434    let boxed_buf = match Box::try_new(buf) {
435        Ok(b) => b,
436        Err(_) => return Status::NO_MEMORY.into_raw(),
437    };
438
439    let raw_ptr = Box::into_raw(boxed_buf);
440
441    // Store the pointer atomically.
442    let old_ptr = ktrace.buffers[cpu_num as usize].swap(raw_ptr, Ordering::AcqRel);
443    if !old_ptr.is_null() {
444        // If there was a previous buffer, reclaim and drop it.
445        unsafe {
446            let _ = Box::from_raw(old_ptr);
447        }
448    }
449
450    Status::OK.into_raw()
451}
452
453#[cfg(test)]
454mod tests {
455    use super::*;
456    use arch_rs::{InterruptDisableGuard, max_num_cpus};
457
458    declare_interned_category!(META_CAT, "kernel:meta", extern);
459    declare_interned_category!(MEMORY_CAT, "kernel:memory", extern);
460    declare_interned_category!(SCHED_CAT, "kernel:sched", extern);
461    declare_interned_category!(CONTENTION_CAT, "kernel:contention", extern);
462    declare_interned_category!(IPC_CAT, "kernel:ipc", extern);
463    declare_interned_category!(IRQ_CAT, "kernel:irq", extern);
464    declare_interned_string!(DROP_STATS_REF, "drop_stats", extern);
465    declare_interned_string!(NUM_RECORDS_REF, "num_records", extern);
466    declare_interned_string!(NUM_BYTES_REF, "num_bytes", extern);
467
468    //
469    // Zircon kernel-tests FFI entry points.
470    //
471
472    /// Test-only FFI helper to write a single word record from Rust.
473    ///
474    /// # Safety
475    ///
476    /// This must be called with interrupts disabled.
477    #[unsafe(no_mangle)]
478    pub unsafe extern "C" fn rust_ktrace_test_interop(header: u64, val: u64) -> i32 {
479        let ktrace = KTrace::get_instance();
480        // SAFETY: The caller guarantees interrupts are disabled.
481        if let Ok(mut res) = unsafe { ktrace.reserve(header) } {
482            let _ = res.write_word(val);
483            let _ = res.commit();
484            0
485        } else {
486            -1
487        }
488    }
489
490    /// Verifies KTraceBuffer initialization and size/metadata attributes.
491    #[unsafe(no_mangle)]
492    pub extern "C" fn rust_ktrace_test_init_and_size() -> bool {
493        let mut storage = [0u8; 256];
494        let mut inner_buf = unsafe { Buffer::from_raw_parts(storage.as_mut_ptr(), storage.len()) };
495        let leaked_ref = unsafe { &mut *ptr::from_mut(&mut inner_buf) };
496        let mut stats = DroppedRecordStats::default();
497        let leaked_stats = unsafe { &mut *ptr::from_mut(&mut stats) };
498        let kbuf = KTraceBuffer::new(leaked_ref, leaked_stats, 1, 100, 200);
499
500        if kbuf.size() != 256 {
501            return false;
502        }
503        if kbuf.process_koid != 100 {
504            return false;
505        }
506        if kbuf.thread_koid != 200 {
507            return false;
508        }
509        true
510    }
511
512    /// Verifies KTraceBuffer reservation, writing words, and committing.
513    #[unsafe(no_mangle)]
514    pub extern "C" fn rust_ktrace_test_write() -> bool {
515        let _guard = InterruptDisableGuard::new();
516        let mut storage = [0u8; 256];
517        let mut inner_buf = unsafe { Buffer::from_raw_parts(storage.as_mut_ptr(), storage.len()) };
518        let leaked_ref = unsafe { &mut *ptr::from_mut(&mut inner_buf) };
519        let mut stats = DroppedRecordStats::default();
520        let leaked_stats = unsafe { &mut *ptr::from_mut(&mut stats) };
521        let mut kbuf = KTraceBuffer::new(leaked_ref, leaked_stats, 1, 100, 200);
522
523        // Reserve 16 bytes (2 words).
524        let header = 4u64 | (2u64 << 4);
525        let mut res = match kbuf.reserve(header) {
526            Ok(r) => r,
527            Err(_) => return false,
528        };
529
530        // Write a word (8 bytes).
531        if res.write_word(0xabcdef0123456789).is_err() {
532            return false;
533        }
534        if res.commit().is_err() {
535            return false;
536        }
537
538        // Read it back.
539        let mut read_bytes = [0u8; 16];
540        let read_len = match kbuf.read(
541            |offset, src| {
542                read_bytes[offset as usize..offset as usize + src.len()].copy_from_slice(src);
543                Ok(())
544            },
545            16,
546        ) {
547            Ok(l) => l,
548            Err(_) => return false,
549        };
550
551        if read_len != 16 {
552            return false;
553        }
554        if u64::from_ne_bytes(read_bytes[0..8].try_into().unwrap()) != header {
555            return false;
556        }
557        if u64::from_ne_bytes(read_bytes[8..16].try_into().unwrap()) != 0xabcdef0123456789 {
558            return false;
559        }
560        true
561    }
562
563    /// Verifies tracking of dropped records and their subsequent serialization when space becomes
564    /// available.
565    #[unsafe(no_mangle)]
566    pub extern "C" fn rust_ktrace_test_dropped_record_tracking() -> bool {
567        let _guard = InterruptDisableGuard::new();
568        let mut storage = [0u8; 128]; // small buffer
569        let mut inner_buf = unsafe { Buffer::from_raw_parts(storage.as_mut_ptr(), storage.len()) };
570        let leaked_ref = unsafe { &mut *ptr::from_mut(&mut inner_buf) };
571        let mut stats = DroppedRecordStats::default();
572        let leaked_stats = unsafe { &mut *ptr::from_mut(&mut stats) };
573        let mut kbuf = KTraceBuffer::new(leaked_ref, leaked_stats, 1, 100, 200);
574
575        // Reserve almost all space.
576        // 128 bytes total. Let's reserve 96 bytes (12 words).
577        let header = 4u64 | (12u64 << 4);
578        let mut res = match kbuf.reserve(header) {
579            Ok(r) => r,
580            Err(_) => return false,
581        };
582        if res.write_bytes(&[0u8; 88]).is_err() {
583            return false;
584        }
585        if res.commit().is_err() {
586            return false;
587        }
588
589        // Now, try to reserve 64 bytes (8 words). This should fail because there are only 32 bytes
590        // left.
591        let header2 = 4u64 | (8u64 << 4);
592        if kbuf.reserve(header2).err() != Some(Status::NO_SPACE) {
593            return false;
594        }
595
596        // This failed reservation should have been tracked!
597        if !kbuf.drop_stats.has_dropped() {
598            return false;
599        }
600        if kbuf.drop_stats.num_dropped != 1 {
601            return false;
602        }
603        if kbuf.drop_stats.bytes_dropped != 64 {
604            return false;
605        }
606
607        // Now, drain the buffer to free all space.
608        if kbuf.drain().is_err() {
609            return false;
610        }
611
612        // Now, reserve 16 bytes (2 words).
613        // Since first_dropped was set, this reservation should successfully write the 56-byte
614        // dropped records duration event first!
615        let header3 = 4u64 | (2u64 << 4);
616        let mut res3 = match kbuf.reserve(header3) {
617            Ok(r) => r,
618            Err(_) => return false,
619        };
620        if res3.write_bytes(&[0u8; 8]).is_err() {
621            return false;
622        }
623        if res3.commit().is_err() {
624            return false;
625        }
626
627        // The dropped stats should have been reset!
628        if kbuf.drop_stats.has_dropped() {
629            return false;
630        }
631        if kbuf.drop_stats.num_dropped != 0 {
632            return false;
633        }
634        if kbuf.drop_stats.bytes_dropped != 0 {
635            return false;
636        }
637
638        // Let's read the buffer content.
639        let mut read_bytes = [0u8; 72];
640        let read_len = match kbuf.read(
641            |offset, src| {
642                read_bytes[offset as usize..offset as usize + src.len()].copy_from_slice(src);
643                Ok(())
644            },
645            72,
646        ) {
647            Ok(l) => l,
648            Err(_) => return false,
649        };
650
651        if read_len != 72 {
652            return false;
653        }
654
655        // Verify the DroppedRecordDurationEvent header:
656        let event_header = u64::from_ne_bytes(read_bytes[0..8].try_into().unwrap());
657        if (event_header & 0xf) != 4 {
658            return false;
659        }
660        if ((event_header >> 4) & 0xfff) != 7 {
661            return false;
662        }
663        if ((event_header >> 16) & 0xf) != 4 {
664            return false;
665        }
666        if ((event_header >> 20) & 0xf) != 2 {
667            return false;
668        }
669        if ((event_header >> 24) & 0xff) != 1 {
670            return false;
671        }
672        if ((event_header >> 32) & 0xffff) != u64::from(META_CAT.label().id()) {
673            return false;
674        }
675        if ((event_header >> 48) & 0xffff) != u64::from(DROP_STATS_REF.id()) {
676            return false;
677        }
678
679        // Verify process_koid (100) and thread_koid (200)
680        if u64::from_ne_bytes(read_bytes[16..24].try_into().unwrap()) != 100 {
681            return false;
682        }
683        if u64::from_ne_bytes(read_bytes[24..32].try_into().unwrap()) != 200 {
684            return false;
685        }
686
687        // Verify the two arguments:
688        let num_dropped_arg = u64::from_ne_bytes(read_bytes[32..40].try_into().unwrap());
689        if (num_dropped_arg & 0xf) != 2 {
690            return false;
691        }
692        if ((num_dropped_arg >> 4) & 0xfff) != 1 {
693            return false;
694        }
695        if ((num_dropped_arg >> 16) & 0xffff) != u64::from(NUM_RECORDS_REF.id()) {
696            return false;
697        }
698        if ((num_dropped_arg >> 32) & 0xffffffff) != 1 {
699            return false;
700        }
701
702        let bytes_dropped_arg = u64::from_ne_bytes(read_bytes[40..48].try_into().unwrap());
703        if (bytes_dropped_arg & 0xf) != 2 {
704            return false;
705        }
706        if ((bytes_dropped_arg >> 16) & 0xffff) != u64::from(NUM_BYTES_REF.id()) {
707            return false;
708        }
709        if ((bytes_dropped_arg >> 32) & 0xffffffff) != 64 {
710            return false;
711        }
712
713        // Verify the reservation header3:
714        let res_header = u64::from_ne_bytes(read_bytes[56..64].try_into().unwrap());
715        if res_header != header3 {
716            return false;
717        }
718
719        true
720    }
721
722    /// Verifies direct invocation of KTraceBuffer::emit_drop_stats.
723    #[unsafe(no_mangle)]
724    pub extern "C" fn rust_ktrace_test_emit_drop_stats() -> bool {
725        let _guard = InterruptDisableGuard::new();
726        let mut storage = [0u8; 128];
727        let mut inner_buf = unsafe { Buffer::from_raw_parts(storage.as_mut_ptr(), storage.len()) };
728        let leaked_ref = unsafe { &mut *ptr::from_mut(&mut inner_buf) };
729        let mut stats = DroppedRecordStats::default();
730        let leaked_stats = unsafe { &mut *ptr::from_mut(&mut stats) };
731        let mut kbuf = KTraceBuffer::new(leaked_ref, leaked_stats, 1, 100, 200);
732
733        // 1. Force a failed reservation to track a dropped record.
734        let header = 4u64 | (32u64 << 4);
735        if kbuf.reserve(header).err() != Some(Status::NO_SPACE) {
736            return false;
737        }
738
739        if !kbuf.drop_stats.has_dropped() {
740            return false;
741        }
742        if kbuf.drop_stats.num_dropped != 1 {
743            return false;
744        }
745        if kbuf.drop_stats.bytes_dropped != 256 {
746            return false;
747        }
748
749        // 2. Call emit_drop_stats directly.
750        if kbuf.emit_drop_stats().is_err() {
751            return false;
752        }
753
754        if kbuf.drop_stats.has_dropped() {
755            return false;
756        }
757        if kbuf.drop_stats.num_dropped != 0 {
758            return false;
759        }
760        if kbuf.drop_stats.bytes_dropped != 0 {
761            return false;
762        }
763
764        // 3. Read and verify the event.
765        let mut read_bytes = [0u8; 56];
766        let read_len = match kbuf.read(
767            |offset, src| {
768                read_bytes[offset as usize..offset as usize + src.len()].copy_from_slice(src);
769                Ok(())
770            },
771            56,
772        ) {
773            Ok(l) => l,
774            Err(_) => return false,
775        };
776
777        if read_len != 56 {
778            return false;
779        }
780
781        let event_header = u64::from_ne_bytes(read_bytes[0..8].try_into().unwrap());
782        if (event_header & 0xf) != 4 {
783            return false;
784        }
785        if ((event_header >> 4) & 0xfff) != 7 {
786            return false;
787        }
788        if ((event_header >> 16) & 0xf) != 4 {
789            return false;
790        }
791        if ((event_header >> 20) & 0xf) != 2 {
792            return false;
793        }
794        if ((event_header >> 24) & 0xff) != 1 {
795            return false;
796        }
797        if ((event_header >> 32) & 0xffff) != u64::from(META_CAT.label().id()) {
798            return false;
799        }
800        if ((event_header >> 48) & 0xffff) != u64::from(DROP_STATS_REF.id()) {
801            return false;
802        }
803
804        if u64::from_ne_bytes(read_bytes[16..24].try_into().unwrap()) != 100 {
805            return false;
806        }
807        if u64::from_ne_bytes(read_bytes[24..32].try_into().unwrap()) != 200 {
808            return false;
809        }
810
811        let num_dropped_arg = u64::from_ne_bytes(read_bytes[32..40].try_into().unwrap());
812        if (num_dropped_arg & 0xf) != 2 {
813            return false;
814        }
815        if ((num_dropped_arg >> 16) & 0xffff) != u64::from(NUM_RECORDS_REF.id()) {
816            return false;
817        }
818        if ((num_dropped_arg >> 32) & 0xffffffff) != 1 {
819            return false;
820        }
821
822        let bytes_dropped_arg = u64::from_ne_bytes(read_bytes[40..48].try_into().unwrap());
823        if (bytes_dropped_arg & 0xf) != 2 {
824            return false;
825        }
826        if ((bytes_dropped_arg >> 16) & 0xffff) != u64::from(NUM_BYTES_REF.id()) {
827            return false;
828        }
829        if ((bytes_dropped_arg >> 32) & 0xffffffff) != 256 {
830            return false;
831        }
832
833        true
834    }
835
836    /// Verifies the full global lifecycle of KTrace: initialization, category bitmasks,
837    /// CPU buffer allocation, and high-level tracing macro execution.
838    #[unsafe(no_mangle)]
839    pub extern "C" fn rust_ktrace_test_global_lifecycle() -> bool {
840        let _guard = InterruptDisableGuard::new();
841        // Initialize indices of the categories we're testing.
842        META_CAT.set_index(0, InternedCategory::INVALID_INDEX);
843        MEMORY_CAT.set_index(1, InternedCategory::INVALID_INDEX);
844        SCHED_CAT.set_index(2, InternedCategory::INVALID_INDEX);
845        CONTENTION_CAT.set_index(3, InternedCategory::INVALID_INDEX);
846        IPC_CAT.set_index(4, InternedCategory::INVALID_INDEX);
847        IRQ_CAT.set_index(5, InternedCategory::INVALID_INDEX);
848
849        // 1. Initialize the global KTrace instance with the system CPU count buffers and a local
850        // mock state.
851        let num_cpus = max_num_cpus();
852        let mut local_state = KTraceState {
853            categories_bitmask: AtomicU32::new(0),
854            writes_enabled: AtomicBool::new(false),
855        };
856        let local_state_ptr = ptr::from_mut(&mut local_state).cast::<ffi::c_void>();
857
858        let status = unsafe { rust_ktrace_init(num_cpus, local_state_ptr) };
859        if status != 0 {
860            return false;
861        }
862
863        let ktrace = KTrace::get_instance();
864
865        // 2. Verify initial states.
866        if ktrace.writes_enabled() {
867            return false;
868        }
869        if ktrace.categories_bitmask() != 0 {
870            return false;
871        }
872        if ktrace.is_category_enabled(&META_CAT) {
873            return false;
874        }
875        if ktrace.is_category_enabled(&IRQ_CAT) {
876            return false;
877        }
878
879        // 3. Test writes_enabled.
880        local_state.writes_enabled.store(true, Ordering::Release);
881        if !ktrace.writes_enabled() {
882            return false;
883        }
884        local_state.writes_enabled.store(false, Ordering::Release);
885        if ktrace.writes_enabled() {
886            return false;
887        }
888
889        // 4. Test categories_bitmask.
890        let mask = (1 << MEMORY_CAT.index()) | (1 << CONTENTION_CAT.index());
891        local_state.categories_bitmask.store(mask, Ordering::Release);
892        if ktrace.categories_bitmask() != mask {
893            return false;
894        }
895        if ktrace.is_category_enabled(&META_CAT) {
896            return false;
897        }
898        if !ktrace.is_category_enabled(&MEMORY_CAT) {
899            return false;
900        }
901        if ktrace.is_category_enabled(&SCHED_CAT) {
902            return false;
903        }
904        if !ktrace.is_category_enabled(&CONTENTION_CAT) {
905            return false;
906        }
907        if ktrace.is_category_enabled(&IPC_CAT) {
908            return false;
909        }
910
911        // 5. Test CPU buffer initialization and reserve.
912        let mut storage = [0u8; 256];
913        let mut inner_buf = unsafe { Buffer::from_raw_parts(storage.as_mut_ptr(), storage.len()) };
914        let inner_buf_ptr = ptr::from_mut(&mut inner_buf).cast::<ffi::c_void>();
915        let mut stats = DroppedRecordStats::default();
916        let stats_ptr = ptr::from_mut(&mut stats).cast::<ffi::c_void>();
917
918        // Initialize current CPU buffer.
919        let cpu = curr_cpu_num();
920        let init_status = unsafe {
921            rust_ktrace_init_cpu_buffer(
922                cpu, // cpu_num
923                inner_buf_ptr,
924                stats_ptr,
925                100, // process_koid
926                200, // thread_koid
927                1,   // cpu_ref_header_entry
928            )
929        };
930        if init_status != 0 {
931            return false;
932        }
933
934        // If writes are disabled, reserving should fail.
935        let header = 4u64 | (2u64 << 4); // Event with 2 words
936        if unsafe { ktrace.reserve(header) }.err() != Some(Status::BAD_STATE) {
937            return false;
938        }
939
940        // Enable writes.
941        local_state.writes_enabled.store(true, Ordering::Release);
942
943        // Now reserve should succeed!
944        let mut res = match unsafe { ktrace.reserve(header) } {
945            Ok(r) => r,
946            Err(_) => return false,
947        };
948        if res.write_word(0x1234567890abcdef).is_err() {
949            return false;
950        }
951        if res.commit().is_err() {
952            return false;
953        }
954
955        // Read and verify the written record from the current CPU buffer.
956        let ptr = ktrace.buffers[cpu as usize].load(Ordering::Acquire);
957        if ptr.is_null() {
958            return false;
959        }
960        let buf = unsafe { &*ptr };
961
962        let mut read_bytes = [0u8; 16];
963        let read_len = match buf.read(
964            |offset, src| {
965                read_bytes[offset as usize..offset as usize + src.len()].copy_from_slice(src);
966                Ok(())
967            },
968            16,
969        ) {
970            Ok(l) => l,
971            Err(_) => return false,
972        };
973
974        if read_len != 16 {
975            return false;
976        }
977        if u64::from_ne_bytes(read_bytes[0..8].try_into().unwrap()) != header {
978            return false;
979        }
980        if u64::from_ne_bytes(read_bytes[8..16].try_into().unwrap()) != 0x1234567890abcdef {
981            return false;
982        }
983
984        true
985    }
986}
ktrace_rs_tests/lib.rs

ktrace_rs_tests/
lib.rs
