pw_kernel/lib/pw_kernel_tracing/buffer.rs - pigweed/sandbox - Git at Google

 // Copyright 2026 The Pigweed Authors
 //
 // Licensed under the Apache License, Version 2.0 (the "License"); you may not
 // use this file except in compliance with the License. You may obtain a copy of
 // the License at
 //
 //     https://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 // License for the specific language governing permissions and limitations under
 // the License.

 use core::cell::UnsafeCell;
 use core::sync::atomic::Ordering;

 use crate::{EventType, Record};

 /// A single entry in the trace buffer.
 #[allow(dead_code)]
 #[repr(C)]
 struct RecordEntry([u32; 4]);
 impl RecordEntry {
     pub const fn new() -> Self {
         Self([0u32; 4])
     }
 }

 /// A circular buffer for storing trace records.
 ///
 /// The buffer uses atomic operations to manage a write position, allowing multiple
 /// producers to write concurrently without locking.
 ///
 /// # Type Parameters
 /// * `AtomicUsize`: The atomic type used for the position counter (e.g., `pw_atomic::AtomicUsize`).
 /// * `NUM_ENTRIES`: The number of entries in the buffer.
 pub struct Buffer<AtomicUsize: pw_atomic::AtomicUsize, const NUM_ENTRIES: usize> {
     records: UnsafeCell<[RecordEntry; NUM_ENTRIES]>,
     pos: AtomicUsize,
 }

 /// Safety: The buffer is safe to share across threads because:
 /// 1. `pos` is an `AtomicUsize` which handles concurrent updates safely.
 /// 2. `records` is accessed via `UnsafeCell`, but the protocol for writing ensures
 ///    that each slot is effectively owned by a single writer at a time (barring wraparound collisions).
 unsafe impl<AtomicUsize: pw_atomic::AtomicUsize, const NUM_ENTRIES: usize> Sync
     for Buffer<AtomicUsize, NUM_ENTRIES>
 {
 }

 impl<AtomicUsize: pw_atomic::AtomicUsize, const NUM_ENTRIES: usize>
     Buffer<AtomicUsize, NUM_ENTRIES>
 {
     /// Creates a new, empty buffer.
     #[must_use]
     pub const fn new() -> Self {
         Self {
             records: UnsafeCell::new([const { RecordEntry::new() }; NUM_ENTRIES]),
             pos: AtomicUsize::ZERO,
         }
     }

     /// Returns the total size of the buffer in bytes.
     pub const fn buffer_len(&self) -> usize {
         NUM_ENTRIES * core::mem::size_of::<RecordEntry>()
     }

     /// Returns a raw pointer to the underlying buffer storage.
     ///
     /// # Safety
     /// This provides direct access to the raw buffer storage. The caller must ensure
     /// that they do not violate any safety guarantees, especially if reading while
     /// other threads are writing.
     pub const unsafe fn buffer(&self) -> *const () {
         self.records.get() as *const ()
     }

     /// Adds a record to the buffer.
     ///
     /// This method reserves a slot in the circular buffer using atomic operations
     /// and writes the record data.
     ///
     /// # Arguments
     /// * `time_stamp`: The timestamp of the event.
     /// * `event_type`: The type of event.
     /// * `event_data`: The raw payload of the event.
     ///
     /// # Example
     /// ```
     /// use pw_kernel_tracing::{Buffer, EventType, ContextSwitchEvent};
     /// use core::sync::atomic::AtomicUsize;
     ///
     /// let buffer = Buffer::<AtomicUsize, 4>::new();
     /// let event = ContextSwitchEvent {
     ///     old_thread_id: 1,
     ///     new_thread_id: 2,
     ///     old_thread_state: 3,
     /// };
     /// buffer.add_record(0x1234, EventType::ContextSwitch, event.encode());
     /// ```
     pub fn add_record(&self, time_stamp: u32, event_type: EventType, event_data: [u32; 3]) {
         let pos = loop {
             let pos = self.pos.load(Ordering::SeqCst);
             let next_pos = (pos + 1) % NUM_ENTRIES;
             if self
                 .pos
                 .compare_exchange(pos, next_pos, Ordering::SeqCst, Ordering::SeqCst)
                 .is_ok()
             {
                 break pos;
             }
         };

         let record = Record {
             event_type,
             time_stamp,
             event_data,
         };

         // Safety:
         // * The offset into the records array is guaranteed to be in bounds
         //   because it is bounds checked above.
         // * "Reasonable" exclusive access can be assumed because entries written
         //   as soon as their storage is reserved. This makes the likelihood of
         //   the trace buffer wrapping before the entry is re-written extremely
         //   low.
         // * In the unlikely event that the buffer wraps and an entry is corrupted,
         //   the reader of the trace buffer is expected to handle the corrupted
         //   entry.

         unsafe {
             let entry = self.records.get().cast::<RecordEntry>().add(pos);
             *entry = RecordEntry(record.encode());
         }
     }
 }

 // These tests require an implementation of core::atomic.
 #[cfg(all(test, feature = "builtin_atomics"))]
 mod tests {
     use core::sync::atomic::AtomicUsize as StdAtomicUsize;

     use unittest::test;

     use super::*;
     // Helper to access private records for testing
     impl<AtomicUsize: pw_atomic::AtomicUsize, const NUM_ENTRIES: usize>
         Buffer<AtomicUsize, NUM_ENTRIES>
     {
         fn get_record_entry(&self, index: usize) -> &RecordEntry {
             unsafe { &(*self.records.get())[index] }
         }
     }

     #[test]
     fn add_record_writes_to_buffer() -> unittest::Result<()> {
         let buffer = Buffer::<StdAtomicUsize, 4>::new();

         buffer.add_record(0x123456, EventType::ContextSwitch, [1, 2, 3]);

         let entry = buffer.get_record_entry(0);
         let header = entry.0[0];
         let timestamp = header & 0x03ff_ffff;
         let event_type_val = (header >> 26) & 0x3f;

         unittest::assert_eq!(timestamp, 0x123456);
         unittest::assert_eq!(event_type_val, EventType::ContextSwitch as u32);
         unittest::assert_eq!(entry.0[1], 1);
         unittest::assert_eq!(entry.0[2], 2);
         unittest::assert_eq!(entry.0[3], 3);

         Ok(())
     }

     #[test]
     fn buffer_wraps_around() -> unittest::Result<()> {
         let buffer = Buffer::<StdAtomicUsize, 2>::new();

         buffer.add_record(0x100, EventType::ContextSwitch, [1, 1, 1]);
         buffer.add_record(0x200, EventType::ContextSwitch, [2, 2, 2]);
         buffer.add_record(0x300, EventType::ContextSwitch, [3, 3, 3]); // Should overwrite index 0

         // Check index 0 (should be [3,3,3])
         let entry0 = buffer.get_record_entry(0);
         unittest::assert_eq!(entry0.0[1], 3);

         // Check index 1 (should be [2,2,2])
         let entry1 = buffer.get_record_entry(1);
         unittest::assert_eq!(entry1.0[1], 2);

         Ok(())
     }

     #[test]
     fn buffer_allocates_slots_correctly() -> unittest::Result<()> {
         let buffer = Buffer::<StdAtomicUsize, 4>::new();

         // internal position should be 0
         unittest::assert_eq!(buffer.pos.load(Ordering::SeqCst), 0);

         buffer.add_record(0x0, EventType::ContextSwitch, [0; 3]);
         // internal position should be 1
         unittest::assert_eq!(buffer.pos.load(Ordering::SeqCst), 1);

         buffer.add_record(0x1, EventType::ContextSwitch, [0; 3]);
         buffer.add_record(0x2, EventType::ContextSwitch, [0; 3]);
         buffer.add_record(0x3, EventType::ContextSwitch, [0; 3]);
         // Should wrap to 0
         unittest::assert_eq!(buffer.pos.load(Ordering::SeqCst), 0);

         Ok(())
     }
 }
	// Copyright 2026 The Pigweed Authors
	//
	// Licensed under the Apache License, Version 2.0 (the "License"); you may not
	// use this file except in compliance with the License. You may obtain a copy of
	// the License at
	//
	// https://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
	// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
	// License for the specific language governing permissions and limitations under
	// the License.

	use core::cell::UnsafeCell;
	use core::sync::atomic::Ordering;

	use crate::{EventType, Record};

	/// A single entry in the trace buffer.
	#[allow(dead_code)]
	#[repr(C)]
	struct RecordEntry([u32; 4]);
	impl RecordEntry {
	pub const fn new() -> Self {
	Self([0u32; 4])
	}
	}

	/// A circular buffer for storing trace records.
	///
	/// The buffer uses atomic operations to manage a write position, allowing multiple
	/// producers to write concurrently without locking.
	///
	/// # Type Parameters
	/// * `AtomicUsize`: The atomic type used for the position counter (e.g., `pw_atomic::AtomicUsize`).
	/// * `NUM_ENTRIES`: The number of entries in the buffer.
	pub struct Buffer<AtomicUsize: pw_atomic::AtomicUsize, const NUM_ENTRIES: usize> {
	records: UnsafeCell<[RecordEntry; NUM_ENTRIES]>,
	pos: AtomicUsize,
	}

	/// Safety: The buffer is safe to share across threads because:
	/// 1. `pos` is an `AtomicUsize` which handles concurrent updates safely.
	/// 2. `records` is accessed via `UnsafeCell`, but the protocol for writing ensures
	/// that each slot is effectively owned by a single writer at a time (barring wraparound collisions).
	unsafe impl<AtomicUsize: pw_atomic::AtomicUsize, const NUM_ENTRIES: usize> Sync
	for Buffer<AtomicUsize, NUM_ENTRIES>
	{
	}

	impl<AtomicUsize: pw_atomic::AtomicUsize, const NUM_ENTRIES: usize>
	Buffer<AtomicUsize, NUM_ENTRIES>
	{
	/// Creates a new, empty buffer.
	#[must_use]
	pub const fn new() -> Self {
	Self {
	records: UnsafeCell::new([const { RecordEntry::new() }; NUM_ENTRIES]),
	pos: AtomicUsize::ZERO,
	}
	}

	/// Returns the total size of the buffer in bytes.
	pub const fn buffer_len(&self) -> usize {
	NUM_ENTRIES * core::mem::size_of::<RecordEntry>()
	}

	/// Returns a raw pointer to the underlying buffer storage.
	///
	/// # Safety
	/// This provides direct access to the raw buffer storage. The caller must ensure
	/// that they do not violate any safety guarantees, especially if reading while
	/// other threads are writing.
	pub const unsafe fn buffer(&self) -> *const () {
	self.records.get() as *const ()
	}

	/// Adds a record to the buffer.
	///
	/// This method reserves a slot in the circular buffer using atomic operations
	/// and writes the record data.
	///
	/// # Arguments
	/// * `time_stamp`: The timestamp of the event.
	/// * `event_type`: The type of event.
	/// * `event_data`: The raw payload of the event.
	///
	/// # Example
	/// ```
	/// use pw_kernel_tracing::{Buffer, EventType, ContextSwitchEvent};
	/// use core::sync::atomic::AtomicUsize;
	///
	/// let buffer = Buffer::<AtomicUsize, 4>::new();
	/// let event = ContextSwitchEvent {
	/// old_thread_id: 1,
	/// new_thread_id: 2,
	/// old_thread_state: 3,
	/// };
	/// buffer.add_record(0x1234, EventType::ContextSwitch, event.encode());
	/// ```
	pub fn add_record(&self, time_stamp: u32, event_type: EventType, event_data: [u32; 3]) {
	let pos = loop {
	let pos = self.pos.load(Ordering::SeqCst);
	let next_pos = (pos + 1) % NUM_ENTRIES;
	if self
	.pos
	.compare_exchange(pos, next_pos, Ordering::SeqCst, Ordering::SeqCst)
	.is_ok()
	{
	break pos;
	}
	};

	let record = Record {
	event_type,
	time_stamp,
	event_data,
	};

	// Safety:
	// * The offset into the records array is guaranteed to be in bounds
	// because it is bounds checked above.
	// * "Reasonable" exclusive access can be assumed because entries written
	// as soon as their storage is reserved. This makes the likelihood of
	// the trace buffer wrapping before the entry is re-written extremely
	// low.
	// * In the unlikely event that the buffer wraps and an entry is corrupted,
	// the reader of the trace buffer is expected to handle the corrupted
	// entry.

	unsafe {
	let entry = self.records.get().cast::<RecordEntry>().add(pos);
	*entry = RecordEntry(record.encode());
	}
	}
	}

	// These tests require an implementation of core::atomic.
	#[cfg(all(test, feature = "builtin_atomics"))]
	mod tests {
	use core::sync::atomic::AtomicUsize as StdAtomicUsize;

	use unittest::test;

	use super::*;
	// Helper to access private records for testing
	impl<AtomicUsize: pw_atomic::AtomicUsize, const NUM_ENTRIES: usize>
	Buffer<AtomicUsize, NUM_ENTRIES>
	{
	fn get_record_entry(&self, index: usize) -> &RecordEntry {
	unsafe { &(*self.records.get())[index] }
	}
	}

	#[test]
	fn add_record_writes_to_buffer() -> unittest::Result<()> {
	let buffer = Buffer::<StdAtomicUsize, 4>::new();

	buffer.add_record(0x123456, EventType::ContextSwitch, [1, 2, 3]);

	let entry = buffer.get_record_entry(0);
	let header = entry.0[0];
	let timestamp = header & 0x03ff_ffff;
	let event_type_val = (header >> 26) & 0x3f;

	unittest::assert_eq!(timestamp, 0x123456);
	unittest::assert_eq!(event_type_val, EventType::ContextSwitch as u32);
	unittest::assert_eq!(entry.0[1], 1);
	unittest::assert_eq!(entry.0[2], 2);
	unittest::assert_eq!(entry.0[3], 3);

	Ok(())
	}

	#[test]
	fn buffer_wraps_around() -> unittest::Result<()> {
	let buffer = Buffer::<StdAtomicUsize, 2>::new();

	buffer.add_record(0x100, EventType::ContextSwitch, [1, 1, 1]);
	buffer.add_record(0x200, EventType::ContextSwitch, [2, 2, 2]);
	buffer.add_record(0x300, EventType::ContextSwitch, [3, 3, 3]); // Should overwrite index 0

	// Check index 0 (should be [3,3,3])
	let entry0 = buffer.get_record_entry(0);
	unittest::assert_eq!(entry0.0[1], 3);

	// Check index 1 (should be [2,2,2])
	let entry1 = buffer.get_record_entry(1);
	unittest::assert_eq!(entry1.0[1], 2);

	Ok(())
	}

	#[test]
	fn buffer_allocates_slots_correctly() -> unittest::Result<()> {
	let buffer = Buffer::<StdAtomicUsize, 4>::new();

	// internal position should be 0
	unittest::assert_eq!(buffer.pos.load(Ordering::SeqCst), 0);

	buffer.add_record(0x0, EventType::ContextSwitch, [0; 3]);
	// internal position should be 1
	unittest::assert_eq!(buffer.pos.load(Ordering::SeqCst), 1);

	buffer.add_record(0x1, EventType::ContextSwitch, [0; 3]);
	buffer.add_record(0x2, EventType::ContextSwitch, [0; 3]);
	buffer.add_record(0x3, EventType::ContextSwitch, [0; 3]);
	// Should wrap to 0
	unittest::assert_eq!(buffer.pos.load(Ordering::SeqCst), 0);

	Ok(())
	}
	}