use crate::buffer::{Header, Slot};
use std::sync::atomic::Ordering;
use std::thread;
use std::time::{Duration, Instant};
pub struct CpuProducer {
header: *mut Header,
slots: *mut Slot,
#[allow(dead_code)]
#[allow(dead_code)]
n_slots: usize,
delay_ms: Option<u64>,
}
unsafe impl Send for CpuProducer {}
impl CpuProducer {
pub fn new(
header: *mut Header,
slots: *mut Slot,
n_slots: usize,
delay_ms: Option<u64>,
) -> Self {
Self {
header,
slots,
n_slots,
delay_ms,
}
}
pub fn produce_messages(&mut self, n_msgs: u64) {
let start = Instant::now();
for i in 0..n_msgs {
let seq = unsafe {
(*self.header)
.producer
.write_idx
.fetch_add(1, Ordering::AcqRel)
};
let slot_idx = (seq & unsafe { (*self.header).config.wrap_mask }) as usize;
unsafe {
let slot = &mut *self.slots.add(slot_idx);
let payload_str = format!("CPU Message #{:04} [seq={}]", i, seq);
let payload_bytes = payload_str.as_bytes();
let len = payload_bytes.len().min(192);
slot.payload[..len].copy_from_slice(&payload_bytes[..len]);
std::ptr::write_volatile(&mut slot.len, len as u32);
std::ptr::write_volatile(&mut slot.flags, 0);
std::sync::atomic::fence(Ordering::Release);
std::ptr::write_volatile(&mut slot.seq, seq);
}
if let Some(delay) = self.delay_ms {
thread::sleep(Duration::from_millis(delay));
}
if i > 0 && i % 1000 == 0 {
let elapsed = start.elapsed();
let rate = i as f64 / elapsed.as_secs_f64();
println!("CPU Producer: {} messages, {:.0} msg/s", i, rate);
}
}
let elapsed = start.elapsed();
println!(
"CPU Producer: Completed {} messages in {:.2}s ({:.0} msg/s)",
n_msgs,
elapsed.as_secs_f64(),
n_msgs as f64 / elapsed.as_secs_f64()
);
}
}
pub fn start_cpu_producer(
header: *mut Header,
slots: *mut Slot,
n_slots: usize,
n_msgs: u64,
delay_ms: Option<u64>,
) -> thread::JoinHandle<()> {
let header_addr = header as usize;
let slots_addr = slots as usize;
thread::spawn(move || {
let header = header_addr as *mut Header;
let slots = slots_addr as *mut Slot;
let mut producer = CpuProducer::new(header, slots, n_slots, delay_ms);
producer.produce_messages(n_msgs);
})
}
pub struct BatchCpuProducer {
header: *mut Header,
slots: *mut Slot,
_n_slots: usize,
batch_size: usize,
}
unsafe impl Send for BatchCpuProducer {}
impl BatchCpuProducer {
pub fn new(header: *mut Header, slots: *mut Slot, n_slots: usize, batch_size: usize) -> Self {
Self {
header,
slots,
_n_slots: n_slots,
batch_size,
}
}
pub fn produce_batch(&mut self, n_msgs: u64) {
let start = Instant::now();
let mut messages_produced = 0u64;
while messages_produced < n_msgs {
let batch_count = (n_msgs - messages_produced).min(self.batch_size as u64);
let base_seq = unsafe {
(*self.header)
.producer
.write_idx
.fetch_add(batch_count, Ordering::AcqRel)
};
for i in 0..batch_count {
let seq = base_seq + i;
let slot_idx = (seq & unsafe { (*self.header).config.wrap_mask }) as usize;
unsafe {
let slot = &mut *self.slots.add(slot_idx);
let payload_str =
format!("Batch CPU #{:04} [seq={}]", messages_produced + i, seq);
let payload_bytes = payload_str.as_bytes();
let len = payload_bytes.len().min(192);
slot.payload[..len].copy_from_slice(&payload_bytes[..len]);
std::ptr::write_volatile(&mut slot.len, len as u32);
std::ptr::write_volatile(&mut slot.flags, 0x100);
std::sync::atomic::fence(Ordering::Release);
std::ptr::write_volatile(&mut slot.seq, seq);
}
}
messages_produced += batch_count;
if messages_produced % 10000 == 0 {
let elapsed = start.elapsed();
let rate = messages_produced as f64 / elapsed.as_secs_f64();
println!(
"Batch CPU Producer: {} messages, {:.0} msg/s",
messages_produced, rate
);
}
}
let elapsed = start.elapsed();
println!(
"Batch CPU Producer: Completed {} messages in {:.2}s ({:.0} msg/s)",
n_msgs,
elapsed.as_secs_f64(),
n_msgs as f64 / elapsed.as_secs_f64()
);
}
}
pub fn start_parallel_cpu_producer(
header: *mut Header,
slots: *mut Slot,
n_slots: usize,
n_msgs: u64,
n_threads: usize,
) -> Vec<thread::JoinHandle<()>> {
let msgs_per_thread = n_msgs / n_threads as u64;
let remainder = n_msgs % n_threads as u64;
let header_addr = header as usize;
let slots_addr = slots as usize;
(0..n_threads)
.map(|thread_id| {
let header_addr = header_addr;
let slots_addr = slots_addr;
let n_slots = n_slots;
let thread_msgs = if thread_id == 0 {
msgs_per_thread + remainder
} else {
msgs_per_thread
};
thread::spawn(move || {
let header = header_addr as *mut Header;
let slots = slots_addr as *mut Slot;
let mut producer = CpuProducer::new(header, slots, n_slots, None);
println!(
"CPU Producer thread {} starting, producing {} messages",
thread_id, thread_msgs
);
producer.produce_messages(thread_msgs);
println!("CPU Producer thread {} completed", thread_id);
})
})
.collect()
}
pub fn benchmark_cpu_producer(
header: *mut Header,
slots: *mut Slot,
n_slots: usize,
n_msgs: u64,
) -> Duration {
let start = Instant::now();
let mut producer = CpuProducer::new(header, slots, n_slots, None);
producer.produce_messages(n_msgs);
start.elapsed()
}
pub fn benchmark_batch_producer(
header: *mut Header,
slots: *mut Slot,
n_slots: usize,
n_msgs: u64,
batch_size: usize,
) -> Duration {
let start = Instant::now();
let mut producer = BatchCpuProducer::new(header, slots, n_slots, batch_size);
producer.produce_batch(n_msgs);
start.elapsed()
}