mpi-rs 0.1.0

//! Shared-memory transport fast-path for same-host ranks (opt-in `shm`
//! feature). Each directed pair of same-host ranks gets a POSIX shared-memory
//! segment holding a single-producer/single-consumer ring of fixed-size slots;
//! messages larger than a slot are chunked and reassembled in order. Cross-host
//! pairs continue to use TCP.
//!
//! This module is transport-agnostic: it moves opaque framed byte blobs and
//! hands received ones to a callback, so it needs no knowledge of MPI headers.

use std::collections::HashMap;
use std::ffi::CString;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, Mutex};

const NSLOTS: usize = 32;
const SLOT_DATA: usize = 8192;
const SLOT_HDR: usize = 16; // state(4) + more(4) + len(4) + pad(4)
const SLOT_SIZE: usize = SLOT_HDR + SLOT_DATA;
const RING_SIZE: usize = NSLOTS * SLOT_SIZE;

const EMPTY: u32 = 0;
const FULL: u32 = 1;

/// A callback invoked with each fully-reassembled framed message.
pub type OnRecv = Arc<dyn Fn(Vec<u8>) + Send + Sync>;

/// One directed SPSC ring in a shared-memory segment.
struct ShmRing {
    ptr: *mut u8,
    name: CString,
    creator: bool,
    idx: usize, // next slot to use (write index for writers, read for readers)
}

// SAFETY: the mapping is shared by design; concurrent access is coordinated by
// per-slot atomic state (SPSC) and, on the writer side, a Mutex.
unsafe impl Send for ShmRing {}

impl ShmRing {
    fn slot(&self, i: usize) -> *mut u8 {
        // SAFETY: `i < NSLOTS`, region is `RING_SIZE` bytes.
        unsafe { self.ptr.add(i * SLOT_SIZE) }
    }
    fn state(&self, i: usize) -> &AtomicU32 {
        // SAFETY: slot header starts with the atomic state word.
        unsafe { &*(self.slot(i) as *const AtomicU32) }
    }
    fn set_more(&self, i: usize, more: u32) {
        unsafe { ((self.slot(i) as *mut u32).add(1)).write_volatile(more) }
    }
    fn more(&self, i: usize) -> u32 {
        unsafe { ((self.slot(i) as *const u32).add(1)).read_volatile() }
    }
    fn set_len(&self, i: usize, len: u32) {
        unsafe { ((self.slot(i) as *mut u32).add(2)).write_volatile(len) }
    }
    fn len(&self, i: usize) -> u32 {
        unsafe { ((self.slot(i) as *const u32).add(2)).read_volatile() }
    }
    fn data(&self, i: usize) -> *mut u8 {
        unsafe { self.slot(i).add(SLOT_HDR) }
    }

    /// Create (writer) or open (reader) the named segment.
    fn open(name: String, creator: bool) -> Option<ShmRing> {
        let cname = CString::new(name).ok()?;
        // SAFETY: standard POSIX shm/mmap FFI.
        unsafe {
            let fd = if creator {
                let fd = libc::shm_open(
                    cname.as_ptr(),
                    libc::O_CREAT | libc::O_RDWR,
                    0o600 as libc::c_uint,
                );
                if fd < 0 {
                    return None;
                }
                if libc::ftruncate(fd, RING_SIZE as libc::off_t) != 0 {
                    // Already sized (created earlier) — tolerate.
                }
                fd
            } else {
                // Reader: retry until the writer has created it.
                let mut tries = 0;
                loop {
                    let fd = libc::shm_open(cname.as_ptr(), libc::O_RDWR, 0o600 as libc::c_uint);
                    if fd >= 0 {
                        break fd;
                    }
                    tries += 1;
                    if tries > 5000 {
                        return None;
                    }
                    std::thread::sleep(std::time::Duration::from_millis(1));
                }
            };
            let ptr = libc::mmap(
                std::ptr::null_mut(),
                RING_SIZE,
                libc::PROT_READ | libc::PROT_WRITE,
                libc::MAP_SHARED,
                fd,
                0,
            );
            libc::close(fd);
            if ptr == libc::MAP_FAILED {
                return None;
            }
            Some(ShmRing {
                ptr: ptr as *mut u8,
                name: cname,
                creator,
                idx: 0,
            })
        }
    }

    /// Write one framed message, chunking across slots. Blocks (spinning) while
    /// the ring is full; the remote poller always drains, so this makes
    /// progress.
    fn write(&mut self, framed: &[u8]) {
        let mut off = 0;
        loop {
            let remaining = framed.len() - off;
            let n = remaining.min(SLOT_DATA);
            let more = if off + n < framed.len() { 1 } else { 0 };
            let i = self.idx;
            // Wait for the slot to be empty.
            while self.state(i).load(Ordering::Acquire) != EMPTY {
                std::thread::yield_now();
            }
            self.set_len(i, n as u32);
            self.set_more(i, more);
            // SAFETY: `n <= SLOT_DATA`.
            unsafe {
                std::ptr::copy_nonoverlapping(framed.as_ptr().add(off), self.data(i), n);
            }
            self.state(i).store(FULL, Ordering::Release);
            self.idx = (self.idx + 1) % NSLOTS;
            off += n;
            if more == 0 {
                break;
            }
        }
    }

    /// Try to drain complete messages, invoking `on_recv` for each. Returns
    /// whether any slot was consumed.
    fn drain(&mut self, acc: &mut Vec<u8>, on_recv: &OnRecv) -> bool {
        let mut progressed = false;
        loop {
            let i = self.idx;
            if self.state(i).load(Ordering::Acquire) != FULL {
                break;
            }
            let n = self.len(i) as usize;
            let more = self.more(i);
            let start = acc.len();
            acc.resize(start + n, 0);
            // SAFETY: slot holds `n <= SLOT_DATA` valid bytes.
            unsafe {
                std::ptr::copy_nonoverlapping(self.data(i), acc.as_mut_ptr().add(start), n);
            }
            self.state(i).store(EMPTY, Ordering::Release);
            self.idx = (self.idx + 1) % NSLOTS;
            progressed = true;
            if more == 0 {
                let msg = std::mem::take(acc);
                on_recv(msg);
            }
        }
        progressed
    }
}

impl Drop for ShmRing {
    fn drop(&mut self) {
        // SAFETY: unmap our mapping; the creator also unlinks the name.
        unsafe {
            libc::munmap(self.ptr as *mut libc::c_void, RING_SIZE);
            if self.creator {
                libc::shm_unlink(self.name.as_ptr());
            }
        }
    }
}

/// The shared-memory transport for one rank: writer rings to same-host peers,
/// and a poller thread draining reader rings into the delivery callback.
pub(crate) struct ShmTransport {
    writers: HashMap<i32, Mutex<ShmRing>>,
}

impl ShmTransport {
    /// Set up shm rings to/from every `same_host` peer. `on_recv` is invoked
    /// with each reassembled framed message.
    pub(crate) fn init(
        jobid: u64,
        rank: i32,
        same_host: &[i32],
        on_recv: OnRecv,
    ) -> Option<ShmTransport> {
        if same_host.is_empty() {
            return None;
        }
        // Writer rings: this rank writes to `/mpi<jobid>.<rank>.<peer>`.
        let mut writers = HashMap::new();
        for &peer in same_host {
            let name = format!("/mpi{jobid}.{rank}.{peer}");
            if let Some(ring) = ShmRing::open(name, true) {
                writers.insert(peer, Mutex::new(ring));
            }
        }

        // Reader rings: this rank reads `/mpi<jobid>.<peer>.<rank>`.
        let reader_names: Vec<String> = same_host
            .iter()
            .map(|&peer| format!("/mpi{jobid}.{peer}.{rank}"))
            .collect();
        std::thread::spawn(move || {
            let mut readers: Vec<(ShmRing, Vec<u8>)> = reader_names
                .into_iter()
                .filter_map(|n| ShmRing::open(n, false).map(|r| (r, Vec::new())))
                .collect();
            loop {
                let mut any = false;
                for (ring, acc) in readers.iter_mut() {
                    any |= ring.drain(acc, &on_recv);
                }
                if !any {
                    std::thread::yield_now();
                }
            }
        });

        Some(ShmTransport { writers })
    }

    /// Send `framed` to `dest` over shared memory if a ring exists; returns
    /// whether it was sent (false ⇒ caller should use TCP).
    pub(crate) fn try_send(&self, dest: i32, framed: &[u8]) -> bool {
        match self.writers.get(&dest) {
            Some(ring) => {
                ring.lock().unwrap().write(framed);
                true
            }
            None => false,
        }
    }
}