This commit is contained in:
0m.ax 2026-07-18 16:46:47 +02:00
parent 495e7e3b07
commit 1e93198c31

View file

@ -6,30 +6,35 @@ use std::time::Instant;
use crate::raw_socket::RawSender;
const QUEUE_LEN: usize = 1000;
/// Per-thread packet queue capacity.
const QUEUE_LEN: usize = 200;
/// Max pixels per packet. Constrained by MTU 1500:
/// 1500 (IP payload max) - 20 (IP hdr) - 8 (UDP hdr) - 2 (pixel hdr) = 1470
/// 1470 / 7 bytes per pixel = 210 pixels.
const PIXELS_PER_PACKET: usize = 210;
const MSG_PAYLOAD_SIZE: usize = 7 * PIXELS_PER_PACKET;
const MSGSIZE: usize = 2 + MSG_PAYLOAD_SIZE;
const MSGSIZE: usize = 2 + 7 * PIXELS_PER_PACKET;
/// Maximum number of packets to batch into a single sendmmsg call.
const SEND_BATCH_SIZE: usize = 26;
/// Number of packer/sender thread pairs.
const NUM_THREADS: usize = 5;
// ---------------------------------------------------------------------------
// Packet + per-thread queue
// ---------------------------------------------------------------------------
/// A filled buffer ready to be sent, carrying its valid data length.
struct Packet {
buf: Box<[u8; MSGSIZE]>,
len: usize,
}
/// Shared state between all pixel writers and the sender thread.
struct SharedState {
/// Filled packets waiting to be sent.
/// Per-thread queue shared between a packer and its sender.
struct SenderQueue {
pending: Mutex<VecDeque<Packet>>,
/// Wakes the sender when packets are available.
condvar: Condvar,
/// Pool of empty, reusable buffers.
pool: Mutex<Vec<Box<[u8; MSGSIZE]>>>,
}
@ -41,147 +46,221 @@ fn alloc_buf() -> Box<[u8; MSGSIZE]> {
buf
}
/// A pixel writer that fills buffers and submits them for sending.
/// Grabs an empty buffer: try pool first, steal oldest pending, or allocate.
fn grab_buf(queue: &SenderQueue) -> Box<[u8; MSGSIZE]> {
if let Some(buf) = queue.pool.lock().unwrap().pop() {
return buf;
}
if let Some(oldest) = queue.pending.lock().unwrap().pop_front() {
return oldest.buf;
}
alloc_buf()
}
/// Submits a filled packet to the sender queue and wakes the sender.
fn submit_packet(queue: &SenderQueue, buf: Box<[u8; MSGSIZE]>, pixel_count: usize) {
let len = 2 + pixel_count * 7;
queue.pending.lock().unwrap().push_back(Packet { buf, len });
queue.condvar.notify_one();
}
// ---------------------------------------------------------------------------
// Simple PRNG for deterministic permutation generation
// ---------------------------------------------------------------------------
struct Xorshift64(u64);
impl Xorshift64 {
fn next(&mut self) -> u64 {
let mut x = self.0;
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
self.0 = x;
x
}
}
fn fisher_yates_shuffle(arr: &mut [usize], rng: &mut Xorshift64) {
for i in (1..arr.len()).rev() {
let j = (rng.next() as usize) % (i + 1);
arr.swap(i, j);
}
}
// ---------------------------------------------------------------------------
// Display
// ---------------------------------------------------------------------------
/// Collects pixels into a frame buffer and publishes completed frames for
/// packer threads to consume.
///
/// Multiple writers can exist concurrently, each with its own buffer,
/// all sharing the same pending queue and buffer pool. This enables
/// multi-threaded drawing.
///
/// Writers never block. If the pool is empty, they steal the oldest
/// unsent packet or allocate a fresh buffer.
/// Architecture:
/// main thread -> write_raw_pixels() -> flush_frame()
/// |
/// publishes Arc<Vec<[u8;7]>> frame snapshot
/// |
/// 5 packer threads (each with unique noise permutation)
/// continuously iterate the frame, packing pixels into packets
/// |
/// 5 sender threads (one per packer, own packet queue)
/// send via RawSender (AF_PACKET) or fallback UDP
pub struct Display {
shared: Arc<SharedState>,
/// The buffer currently being filled with pixel data.
current_buf: Box<[u8; MSGSIZE]>,
/// How many pixels have been written into `current_buf`.
pos_in_buf: usize,
/// Frame being assembled by the main thread.
building_frame: Vec<[u8; 7]>,
/// Shared pointer to the latest completed frame.
/// Packer threads clone the Arc to read it without holding the lock.
current_frame: Arc<Mutex<Arc<Vec<[u8; 7]>>>>,
}
impl Display {
/// Creates a new Display and spawns the sender thread.
///
/// The returned Display can create additional writers via
/// `create_writer()` for multi-threaded drawing.
/// Creates a new Display and spawns packer + sender thread pairs.
pub fn new(host: &str, port: u16, interface: Option<&str>) -> Self {
let shared = Arc::new(SharedState {
pending: Mutex::new(VecDeque::with_capacity(QUEUE_LEN)),
condvar: Condvar::new(),
pool: Mutex::new(Vec::with_capacity(QUEUE_LEN)),
});
let current_frame = Arc::new(Mutex::new(Arc::new(Vec::new())));
// Pre-allocate the buffer pool. Keep one out as the initial current_buf.
{
let mut pool = shared.pool.lock().unwrap();
for _ in 1..QUEUE_LEN {
pool.push(alloc_buf());
for thread_idx in 0..NUM_THREADS {
// Each packer-sender pair gets its own queue.
let queue = Arc::new(SenderQueue {
pending: Mutex::new(VecDeque::with_capacity(QUEUE_LEN)),
condvar: Condvar::new(),
pool: Mutex::new({
let mut v = Vec::with_capacity(QUEUE_LEN);
for _ in 0..QUEUE_LEN {
v.push(alloc_buf());
}
v
}),
});
// Spawn packer thread.
let frame_source = Arc::clone(&current_frame);
let packer_queue = Arc::clone(&queue);
// Each thread gets a unique seed for a different permutation.
let seed = (thread_idx as u64 + 1).wrapping_mul(0x517cc1b727220a95);
thread::Builder::new()
.name(format!("packer-{}", thread_idx))
.spawn(move || {
packer_loop(frame_source, packer_queue, seed);
})
.unwrap();
// Spawn sender thread.
let sender_queue = Arc::clone(&queue);
match RawSender::new(host, port, 0, interface) {
Ok(raw_sender) => {
eprintln!("[display] Thread {}: raw AF_PACKET sender", thread_idx);
thread::Builder::new()
.name(format!("sender-{}", thread_idx))
.spawn(move || {
sender_loop_raw(raw_sender, sender_queue);
})
.unwrap();
}
Err(e) => {
eprintln!(
"[display] Thread {}: raw unavailable ({}), falling back to UDP",
thread_idx, e
);
let remote_addr = (host, port)
.to_socket_addrs()
.expect("Invalid remote address")
.next()
.expect("Could not resolve host");
let socket = UdpSocket::bind("0.0.0.0:0").expect("Could not bind");
socket.connect(remote_addr).expect("Could not connect");
thread::Builder::new()
.name(format!("sender-{}", thread_idx))
.spawn(move || {
sender_loop_udp(socket, sender_queue);
})
.unwrap();
}
}
}
for _ in [1,2,3,4,5] {
// Spawn the sender thread.
let shared_sender = Arc::clone(&shared);
match RawSender::new(host, port, 0, interface) {
Ok(raw_sender) => {
eprintln!("[display] Using raw AF_PACKET sender");
thread::Builder::new()
.name("sender-raw".into())
.spawn(move || {
sender_loop_raw(raw_sender, shared_sender);
}).unwrap();
}
Err(e) => {
eprintln!("[display] Raw AF_PACKET unavailable ({}), falling back to UDP", e);
let remote_addr = (host, port)
.to_socket_addrs()
.expect("Invalid remote address")
.next()
.expect("Could not resolve host");
let socket = UdpSocket::bind("0.0.0.0:0").expect("Could not bind");
socket.connect(remote_addr).expect("Could not connect");
thread::Builder::new()
.name("sender-raw".into())
.spawn(move || {
sender_loop_udp(socket, shared_sender);
}).unwrap();
}
};
}
Display {
shared,
current_buf: alloc_buf(),
pos_in_buf: 0,
building_frame: Vec::with_capacity(60_000),
current_frame,
}
}
/// Grabs an empty buffer: try pool first, steal oldest pending, or allocate.
fn grab_buf(&self) -> Box<[u8; MSGSIZE]> {
// 1. Try the pool (fast path).
{
let mut pool = self.shared.pool.lock().unwrap();
if let Some(buf) = pool.pop() {
return buf;
}
}
// 2. Pool empty -- steal the oldest unsent packet.
{
let mut pending = self.shared.pending.lock().unwrap();
if let Some(oldest) = pending.pop_front() {
return oldest.buf;
}
}
// 3. Everything is in-flight. Allocate a fresh buffer.
alloc_buf()
}
/// Flushes the current buffer if it contains pixel data.
///
/// This method **never blocks**.
pub fn flush_frame(&mut self) {
if self.pos_in_buf == 0 {
return;
}
let len = 2 + self.pos_in_buf * 7;
let new_buf = self.grab_buf();
let filled_buf = std::mem::replace(&mut self.current_buf, new_buf);
{
let mut pending = self.shared.pending.lock().unwrap();
pending.push_back(Packet {
buf: filled_buf,
len,
});
}
self.shared.condvar.notify_one();
self.pos_in_buf = 0;
}
/// Appends raw 7-byte pixel entries to the frame being built.
#[inline]
pub fn write_raw_pixels(&mut self, pixels: &[[u8; 7]]) {
let mut remaining = pixels;
while !remaining.is_empty() {
let space = PIXELS_PER_PACKET - self.pos_in_buf;
let n = remaining.len().min(space);
let dst_offset = 2 + self.pos_in_buf * 7;
// SAFETY: [u8; 7] has align 1 and no padding, so &[[u8; 7]]
// is a contiguous block of bytes we can copy in one shot.
let src = unsafe {
std::slice::from_raw_parts(remaining.as_ptr() as *const u8, n * 7)
};
self.current_buf[dst_offset..dst_offset + n * 7].copy_from_slice(src);
self.pos_in_buf += n;
remaining = &remaining[n..];
if self.pos_in_buf == PIXELS_PER_PACKET {
self.flush_frame();
}
}
self.building_frame.extend_from_slice(pixels);
}
/// Publishes the current frame for packer threads and starts a new one.
pub fn flush_frame(&mut self) {
let frame = std::mem::take(&mut self.building_frame);
let capacity = frame.len();
let new_frame = Arc::new(frame);
*self.current_frame.lock().unwrap() = new_frame;
// Pre-allocate next frame to the same size.
self.building_frame.reserve(capacity);
}
}
// ---------------------------------------------------------------------------
// Packer thread
// ---------------------------------------------------------------------------
/// Continuously reads the current frame and packs pixels in a noise-shuffled
/// order into packets, feeding them to the paired sender thread.
///
/// Each packer has a unique fixed permutation so all 5 threads naturally
/// cover different regions of the image at any given moment.
fn packer_loop(
frame_source: Arc<Mutex<Arc<Vec<[u8; 7]>>>>,
queue: Arc<SenderQueue>,
seed: u64,
) {
let mut rng = Xorshift64(seed);
let mut permutation: Vec<usize> = Vec::new();
let mut last_frame_len: usize = 0;
loop {
// Snapshot the current frame (just an Arc clone -- very cheap).
let frame = frame_source.lock().unwrap().clone();
if frame.is_empty() {
thread::yield_now();
continue;
}
// Regenerate the permutation when the frame size changes.
if frame.len() != last_frame_len {
permutation = (0..frame.len()).collect();
fisher_yates_shuffle(&mut permutation, &mut rng);
last_frame_len = frame.len();
}
// Walk the permutation, packing pixels into packets.
let mut buf = grab_buf(&queue);
let mut pos: usize = 0;
for &idx in &permutation {
if idx >= frame.len() {
continue;
}
let dst_offset = 2 + pos * 7;
buf[dst_offset..dst_offset + 7].copy_from_slice(&frame[idx]);
pos += 1;
if pos == PIXELS_PER_PACKET {
submit_packet(&queue, buf, pos);
buf = grab_buf(&queue);
pos = 0;
}
}
// Flush any remaining partial packet.
if pos > 0 {
submit_packet(&queue, buf, pos);
}
}
}
// ---------------------------------------------------------------------------
@ -189,7 +268,7 @@ impl Display {
// ---------------------------------------------------------------------------
/// Raw AF_PACKET sender loop.
fn sender_loop_raw(raw_sender: RawSender, shared: Arc<SharedState>) {
fn sender_loop_raw(raw_sender: RawSender, queue: Arc<SenderQueue>) {
let mut batch: Vec<Packet> = Vec::with_capacity(SEND_BATCH_SIZE);
let mut frame_buf: Vec<u8> = Vec::with_capacity(SEND_BATCH_SIZE * (MSGSIZE + 42));
@ -202,9 +281,9 @@ fn sender_loop_raw(raw_sender: RawSender, shared: Arc<SharedState>) {
loop {
{
let mut pending = shared.pending.lock().unwrap();
let mut pending = queue.pending.lock().unwrap();
while pending.is_empty() {
pending = shared.condvar.wait(pending).unwrap();
pending = queue.condvar.wait(pending).unwrap();
}
let n = pending.len().min(SEND_BATCH_SIZE);
batch.extend(pending.drain(..n));
@ -232,7 +311,7 @@ fn sender_loop_raw(raw_sender: RawSender, shared: Arc<SharedState>) {
// Return buffers to the pool.
{
let mut pool = shared.pool.lock().unwrap();
let mut pool = queue.pool.lock().unwrap();
for packet in batch.drain(..) {
pool.push(packet.buf);
}
@ -260,7 +339,7 @@ fn sender_loop_raw(raw_sender: RawSender, shared: Arc<SharedState>) {
}
/// Fallback UDP sender loop.
fn sender_loop_udp(socket: UdpSocket, shared: Arc<SharedState>) {
fn sender_loop_udp(socket: UdpSocket, queue: Arc<SenderQueue>) {
let mut batch: Vec<Packet> = Vec::with_capacity(SEND_BATCH_SIZE);
let mut stats_sent: u64 = 0;
let mut stats_bytes: u64 = 0;
@ -268,9 +347,9 @@ fn sender_loop_udp(socket: UdpSocket, shared: Arc<SharedState>) {
loop {
{
let mut pending = shared.pending.lock().unwrap();
let mut pending = queue.pending.lock().unwrap();
while pending.is_empty() {
pending = shared.condvar.wait(pending).unwrap();
pending = queue.condvar.wait(pending).unwrap();
}
let n = pending.len().min(SEND_BATCH_SIZE);
batch.extend(pending.drain(..n));
@ -284,7 +363,7 @@ fn sender_loop_udp(socket: UdpSocket, shared: Arc<SharedState>) {
}
{
let mut pool = shared.pool.lock().unwrap();
let mut pool = queue.pool.lock().unwrap();
for packet in batch.drain(..) {
pool.push(packet.buf);
}