Rust allocator backed by systemd fd store. See darkcoding.net blog post for details.

systemd_alloc.rs

#![allow(dead_code)]

use anyhow::Context;
use std::alloc::AllocError;
use std::alloc::{Allocator, Layout};
use std::ffi::{c_char, c_uint, CStr, CString};
use std::io;
use std::mem::MaybeUninit;
use std::os::fd::AsRawFd;
use std::process;
use std::ptr::{self, NonNull};
use std::slice;

pub struct SystemdAlloc {
    found: bool,
    memfd: i32,
}

impl SystemdAlloc {
    pub fn new(name: String) -> anyhow::Result<Self> {
        // Later, restoring should go in the allocator factory

        let raw_names: *mut *const c_char = ptr::null_mut();
        let num_fds = unsafe { super::sd_listen_fds_with_names(0, &raw_names) };
        println!("Received {num_fds} fds from systemd");
        let names = unsafe { Vec::from_raw_parts(raw_names, num_fds as usize, num_fds as usize) };
        let mut fd_num = 3;
        let mut fd = None;
        let mut found = false;
        for n in names {
            // We now own the string's memory, libsystemd wants us to `free` it which we don't
            let candidate = unsafe { CStr::from_ptr(n).to_str()? };
            println!("systemd sent us: {name}");
            if candidate == name {
                fd = Some(fd_num);
                found = true;
                println!("Restoring");
                break;
            }
            fd_num += 1;
        }
        if fd.is_none() {
            println!("Creating new store");
            let memfd = mem_fd(&name).unwrap();
            unsafe {
                super::sd_pid_notify_with_fds(
                    process::id() as c_uint,                                     // pid
                    0, // unset_environment
                    CString::new(format!("FDSTORE=1\nFDNAME={name}"))?.as_ptr(), // state
                    [memfd.as_raw_fd()].as_ptr(), // fds
                    1, // n_fds
                );
            }

            fd = Some(memfd);
        }

        Ok(SystemdAlloc {
            memfd: fd.unwrap(),
            found,
        })
    }

    pub fn into_box<T>(self, t: T) -> anyhow::Result<Box<T, SystemdAlloc>> {
        if self.found {
            self.restore()
        } else {
            Ok(Box::new_in(t, self))
        }
    }

    fn restore<T>(self) -> anyhow::Result<Box<T, SystemdAlloc>> {
        let mut stat_buf = MaybeUninit::<libc::stat>::uninit();
        if unsafe { libc::fstat(self.memfd, stat_buf.as_mut_ptr()) } == -1 {
            return Err(io::Error::last_os_error()).context("fstat");
        }
        let stat = unsafe { stat_buf.assume_init() };

        let buf_addr = unsafe {
            libc::mmap(
                ptr::null_mut(),
                stat.st_size as usize,
                libc::PROT_READ | libc::PROT_WRITE,
                libc::MAP_SHARED,
                self.memfd, //.as_raw_fd(),
                0,
            )
        };
        if buf_addr == libc::MAP_FAILED {
            anyhow::bail!("mmap failed: {}", io::Error::last_os_error());
        }
        let p_ptr = buf_addr as *mut u8 as *mut T;
        Ok(unsafe { Box::from_raw_in(p_ptr, self) })
    }
}

unsafe impl Allocator for SystemdAlloc {
    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        if unsafe { libc::ftruncate(self.memfd, layout.size() as i64) } == -1 {
            eprintln!("ftruncate: {}", io::Error::last_os_error());
            return Err(AllocError {});
        }
        let buf_addr = unsafe {
            libc::mmap(
                ptr::null_mut(),
                layout.size(),
                libc::PROT_READ | libc::PROT_WRITE,
                libc::MAP_SHARED,
                self.memfd,
                0,
            )
        };
        if buf_addr == libc::MAP_FAILED {
            eprintln!("mmap failed: {}", io::Error::last_os_error());
            return Err(AllocError {});
        }
        let slice: &mut [u8] =
            unsafe { slice::from_raw_parts_mut(buf_addr as *mut u8, layout.size()) };
        let ptr = unsafe { NonNull::new_unchecked(slice) };
        Ok(ptr)
    }

    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
        libc::munmap(ptr.as_ptr() as *mut std::ffi::c_void, layout.size());
    }
}

fn mem_fd(name: &str) -> anyhow::Result<i32> {
    let memfd = unsafe { libc::memfd_create(CString::new(name)?.as_ptr(), 0) };
    if memfd == -1 {
        return Err(io::Error::last_os_error()).context("memfd_create");
    }
    Ok(memfd)
}