diff --git a/src/aio.rs b/src/aio.rs new file mode 100644 index 0000000..3dcb2be --- /dev/null +++ b/src/aio.rs @@ -0,0 +1,116 @@ +mod pool; + +use std::alloc::{Layout, alloc}; +use std::os::raw::c_void; + +use crate::error::Error; +use crate::guard; +use crate::memcached::server::{self, ResponseHandler}; +use pool::Entry; +pub use pool::{restart_after_fork, start}; + +/// Reached from both the lane queue and the callback data; `wake` frees it after the lane. +#[repr(C)] +struct Job { + /// The first field, so its address is the job's. + entry: Entry, + cookie: *const c_void, + reply: &'static [u8], +} + +impl Job { + /// Like `Box::new`, but a failed allocation is `None` instead of an abort. + #[expect( + clippy::cast_ptr_alignment, + reason = "`alloc` aligns the memory for the layout it is given" + )] + fn new(cookie: *const c_void, reply: &'static [u8]) -> Option<*mut Self> { + // SAFETY: `Job` is not zero-sized. + let raw = unsafe { alloc(Layout::new::()) }.cast::(); + if raw.is_null() { + return None; + } + let job = Self { + entry: Entry::new(run), + cookie, + reply, + }; + // SAFETY: `raw` is a fresh allocation with `Job`'s layout. + unsafe { raw.write(job) }; + Some(raw) + } +} + +/// # Safety +/// +/// `cookie` must be the connection of the running `execute`, and `handler` must accept it. +pub unsafe fn begin(cookie: *const c_void, reply: &'static [u8], handler: ResponseHandler) -> bool { + let Some(job) = Job::new(cookie, reply) else { + // SAFETY: the caller guarantees `handler` accepts `cookie`. + return unsafe { server::respond(handler, cookie, Error::OutOfMemory.reply()) }; + }; + // Before queueing, so a lane that finishes first is still counted. + // SAFETY: `cookie` is the running command's connection. + unsafe { server::waitfor_io_complete(cookie) }; + // SAFETY: `cookie` is the running command's connection, and `wake` frees `job`. + unsafe { server::on_io_complete(cookie, wake, job.cast()) }; + // SAFETY: `job` stays alive until `wake`, which runs only after the lane's notify. + unsafe { pool::submit(pool::idlest(), &raw const (*job).entry) }; + true +} + +unsafe fn run(entry: *const Entry) { + // SAFETY: `entry` is the first field of a `#[repr(C)]` `Job`, alive until `wake`. + let cookie = unsafe { (*entry.cast::()).cookie }; + // SAFETY: the connection waits for this notify, so the core keeps it alive until then. + unsafe { server::notify_io_complete(cookie) }; +} + +unsafe extern "C" fn wake( + cookie: *const c_void, + data: *mut c_void, + handler: ResponseHandler, +) -> bool { + guard( + || false, + || { + // SAFETY: `data` is the job from `begin`; the core calls this once, after the lane. + let job = unsafe { Box::from_raw(data.cast::()) }; + // SAFETY: the core passes the handler for `cookie`. + unsafe { server::respond(handler, cookie, job.reply) } + }, + ) +} + +#[cfg(test)] +mod tests { + use std::os::raw::{c_char, c_int}; + + use super::*; + + /// # Safety + /// + /// `cookie` must be an unshared `Vec`, and `dta` must hold `nbytes` bytes. + unsafe extern "C" fn record(cookie: *const c_void, nbytes: c_int, dta: *const c_char) -> bool { + // SAFETY: the caller guarantees both, per `# Safety`. + let (reply, bytes) = unsafe { + ( + &mut *cookie.cast_mut().cast::>(), + std::slice::from_raw_parts(dta.cast::(), nbytes as usize), + ) + }; + reply.extend_from_slice(bytes); + true + } + + #[test] + fn wake_writes_the_reply_and_frees_the_job() { + let mut reply: Vec = Vec::new(); + let cookie = (&raw mut reply).cast_const().cast(); + let job = Job::new(cookie, b"NOT_SUPPORTED\r\n").unwrap(); + // SAFETY: `job` is ours, and `record` gets `reply` back as its cookie. + let ok = unsafe { wake(cookie, job.cast(), Some(record)) }; + assert!(ok); + assert_eq!(reply, b"NOT_SUPPORTED\r\n"); + } +} diff --git a/src/aio/pool.rs b/src/aio/pool.rs new file mode 100644 index 0000000..5113d3c --- /dev/null +++ b/src/aio/pool.rs @@ -0,0 +1,298 @@ +use std::cell::UnsafeCell; +use std::os::raw::c_int; +use std::panic::{AssertUnwindSafe, catch_unwind}; +use std::ptr; +use std::sync::atomic::{AtomicPtr, AtomicUsize, Ordering}; +use std::sync::{Arc, Barrier, Condvar, Mutex, MutexGuard, PoisonError}; + +const WORKERS: usize = 32; + +/// A link embedded in a job, so queueing a job never allocates. +pub struct Entry { + next: UnsafeCell<*const Entry>, + run: unsafe fn(*const Entry), +} + +// SAFETY: `next` is only accessed under the lane mutex and `run` is a plain fn pointer. +unsafe impl Send for Entry {} +// SAFETY: through `&Entry` only `run` is read freely; `next` is touched only under the lane mutex. +unsafe impl Sync for Entry {} + +impl Entry { + pub const fn new(run: unsafe fn(*const Self)) -> Self { + Self { + next: UnsafeCell::new(ptr::null()), + run, + } + } +} + +struct List { + head: *const Entry, + tail: *const Entry, +} + +// SAFETY: a `List` is only reached through its lane's mutex, and its entries stay alive until run. +unsafe impl Send for List {} + +impl List { + const fn new() -> Self { + Self { + head: ptr::null(), + tail: ptr::null(), + } + } + + /// # Safety + /// + /// `entry` must stay alive until run and must not be in any list. + unsafe fn push_back(&mut self, entry: *const Entry) { + // SAFETY: `entry` is alive and unlinked, and `&mut self` holds the lane mutex. + unsafe { *(*entry).next.get() = ptr::null() }; + if self.tail.is_null() { + self.head = entry; + } else { + // SAFETY: `tail` is in this list, so it is alive, and the lane mutex is held. + unsafe { *(*self.tail).next.get() = entry }; + } + self.tail = entry; + } + + fn pop_front(&mut self) -> Option<*const Entry> { + if self.head.is_null() { + return None; + } + let entry = self.head; + // SAFETY: `entry` is this list's head, so it is alive, and the lane mutex is held. + self.head = unsafe { *(*entry).next.get() }; + if self.head.is_null() { + self.tail = ptr::null(); + } + Some(entry) + } +} + +pub struct Lane { + queue: Mutex, + wake: Condvar, + /// Queued plus running jobs, read without the lock to pick the least busy lane. + load: AtomicUsize, +} + +struct Pool { + lanes: [Lane; WORKERS], +} + +// Set at load, and again in the child of `-d`'s fork, where the lanes made at load are gone. +static POOL: AtomicPtr = AtomicPtr::new(ptr::null_mut()); + +/// Starts the lanes; `false` if a thread could not be made. +pub fn start() -> bool { + let pool: &'static Pool = Box::leak(Box::new(Pool { + lanes: std::array::from_fn(|_| Lane { + queue: Mutex::new(List::new()), + wake: Condvar::new(), + load: AtomicUsize::new(0), + }), + })); + // Wait for every lane to settle, so a fork never catches one mid-start holding a lock. + let ready = Arc::new(Barrier::new(WORKERS + 1)); + for (n, lane) in pool.lanes.iter().enumerate() { + let ready = Arc::clone(&ready); + let spawned = std::thread::Builder::new() + .name(format!("arcus-vector-{n}")) + .spawn(move || { + ready.wait(); + run(lane); + }); + if spawned.is_err() { + return false; + } + } + ready.wait(); + POOL.store(ptr::from_ref(pool).cast_mut(), Ordering::Release); + true +} + +unsafe extern "C" { + fn pthread_atfork( + prepare: Option, + parent: Option, + child: Option, + ) -> c_int; +} + +pub fn restart_after_fork() -> bool { + // SAFETY: `restart` is a plain fn that catches its own panics. + unsafe { pthread_atfork(None, None, Some(restart)) == 0 } +} + +extern "C" fn restart() { + if !crate::guard(|| false, start) { + // A job queued on a lane with no thread would leave its connection waiting forever. + std::process::exit(1); + } +} + +#[expect(clippy::expect_used, reason = "`start` runs at load, before any job")] +fn pool() -> &'static Pool { + // SAFETY: a non-null `POOL` is a leaked pool, never freed. + unsafe { POOL.load(Ordering::Acquire).as_ref() }.expect("the lanes start at load") +} + +#[cfg(test)] +pub fn start_once() { + static ONCE: std::sync::Once = std::sync::Once::new(); + ONCE.call_once(|| assert!(start())); +} + +fn queue(lane: &Lane) -> MutexGuard<'_, List> { + lane.queue.lock().unwrap_or_else(PoisonError::into_inner) +} + +pub fn idlest() -> &'static Lane { + let load = |lane: &Lane| lane.load.load(Ordering::Relaxed); + let lanes = &pool().lanes; + lanes.iter().fold(&lanes[0], |idlest, lane| { + if load(lane) < load(idlest) { + lane + } else { + idlest + } + }) +} + +/// # Safety +/// +/// `entry` must stay alive until its `run` is called and must not be in another list meanwhile. +pub unsafe fn submit(lane: &Lane, entry: *const Entry) { + let mut queue = queue(lane); + // SAFETY: the caller guarantees `entry` is alive and unlinked; the lane mutex is held. + unsafe { queue.push_back(entry) }; + lane.load.fetch_add(1, Ordering::Relaxed); + drop(queue); + lane.wake.notify_one(); +} + +fn run(lane: &Lane) { + loop { + let entry = { + let mut queue = queue(lane); + loop { + if let Some(entry) = queue.pop_front() { + break entry; + } + queue = lane + .wake + .wait(queue) + .unwrap_or_else(PoisonError::into_inner); + } + }; + // A panicking job must not kill the lane, or its later jobs would never run. + // SAFETY: `entry` was just popped, so its submitter keeps it alive until this `run`. + let _ = catch_unwind(AssertUnwindSafe(|| unsafe { ((*entry).run)(entry) })); + lane.load.fetch_sub(1, Ordering::Relaxed); + } +} + +#[cfg(test)] +mod tests { + use std::sync::mpsc::{Sender, channel}; + use std::time::Duration; + + use super::*; + + #[repr(C)] + struct Probe { + entry: Entry, + f: Box, + } + + unsafe fn run_probe(entry: *const Entry) { + // SAFETY: `entry` is the first field of a `#[repr(C)]` `Probe` leaked by `probe`, run once. + let probe = unsafe { Box::from_raw(entry.cast_mut().cast::()) }; + (probe.f)(); + } + + fn probe(lane: usize, f: impl FnOnce() + Send + 'static) { + start_once(); + let raw = Box::into_raw(Box::new(Probe { + entry: Entry::new(run_probe), + f: Box::new(f), + })); + // SAFETY: the leaked `Probe` stays alive and unlinked until `run_probe` frees it. + unsafe { submit(&pool().lanes[lane], &raw const (*raw).entry) }; + } + + fn send(tx: &Sender, n: usize) -> impl FnOnce() + Send + 'static { + let tx = tx.clone(); + move || { + let _ = tx.send(n); + } + } + + #[test] + fn jobs_on_one_lane_run_in_order() { + let (tx, rx) = channel(); + for n in 0..100 { + probe(3, send(&tx, n)); + } + let got: Vec = (0..100) + .map(|_| rx.recv_timeout(Duration::from_secs(10)).unwrap()) + .collect(); + assert_eq!(got, (0..100).collect::>()); + } + + #[test] + fn a_busy_lane_keeps_taking_jobs() { + let (hold_tx, hold_rx) = channel::<()>(); + let (started_tx, started_rx) = channel::<()>(); + probe(5, move || { + let _ = started_tx.send(()); + let _ = hold_rx.recv(); + }); + started_rx.recv_timeout(Duration::from_secs(10)).unwrap(); + + let (tx, rx) = channel(); + for n in 0..1000 { + probe(5, send(&tx, n)); + } + drop(hold_tx); + for _ in 0..1000 { + rx.recv_timeout(Duration::from_secs(10)).unwrap(); + } + } + + #[test] + fn the_idlest_lane_runs_the_job() { + let (tx, rx) = channel(); + let raw = Box::into_raw(Box::new(Probe { + entry: Entry::new(run_probe), + f: Box::new(send(&tx, 7)), + })); + start_once(); + let lane = idlest(); + // SAFETY: the leaked `Probe` stays alive and unlinked until `run_probe` frees it. + unsafe { submit(lane, &raw const (*raw).entry) }; + assert_eq!(rx.recv_timeout(Duration::from_secs(10)).unwrap(), 7); + } + + #[test] + fn the_list_pops_in_push_order() { + unsafe fn nothing(_: *const Entry) {} + let entries = [ + Entry::new(nothing), + Entry::new(nothing), + Entry::new(nothing), + ]; + let mut list = List::new(); + for entry in &entries { + // SAFETY: each entry outlives `list` and is pushed once. + unsafe { list.push_back(entry) }; + } + for entry in &entries { + assert_eq!(list.pop_front(), Some(ptr::from_ref(entry))); + } + assert_eq!(list.pop_front(), None); + } +} diff --git a/src/ascii.rs b/src/ascii.rs index b281b20..35bc013 100644 --- a/src/ascii.rs +++ b/src/ascii.rs @@ -2,10 +2,12 @@ use std::cell::Cell; use std::os::raw::{c_char, c_int, c_void}; use std::ptr; +use crate::aio; use crate::error::Error; use crate::guard; use crate::memcached::ffi::{AUTHZ_NONE, EXTENSION_ASCII_PROTOCOL_DESCRIPTOR, token_t}; -use crate::protocol::{self, Tokens}; +use crate::memcached::server::{self, ResponseHandler}; +use crate::protocol::{self, Request, Tokens}; const COMMAND: &[u8] = b"vop"; @@ -14,9 +16,6 @@ thread_local! { static ACCEPT_PANICKED: Cell = const { Cell::new(false) }; } -type ResponseHandler = - Option bool>; - pub(crate) static mut DESCRIPTOR: EXTENSION_ASCII_PROTOCOL_DESCRIPTOR = EXTENSION_ASCII_PROTOCOL_DESCRIPTOR { get_name: Some(get_name), @@ -88,17 +87,19 @@ unsafe extern "C" fn execute( if ACCEPT_PANICKED.replace(false) { return false; } - let Some(respond) = response_handler else { - return false; - }; // SAFETY: the caller guarantees `Tokens::new`'s contract. let tokens = unsafe { Tokens::new(argv, argc) }; let reply = match protocol::parse(&tokens) { + Ok(Request::Search(_)) => { + let reply = Error::NotSupported.reply(); + // SAFETY: `cookie` is the running command's connection; the handler accepts it. + return unsafe { aio::begin(cookie, reply, response_handler) }; + } Ok(_) => Error::NotSupported.reply(), Err(e) => e.reply(), }; - // SAFETY: `respond` accepts `cookie`, and `reply` is `'static`. - unsafe { respond(cookie, reply.len() as c_int, reply.as_ptr().cast()) } + // SAFETY: the caller guarantees `response_handler` accepts `cookie`. + unsafe { server::respond(response_handler, cookie, reply) } }, ) } diff --git a/src/lib.rs b/src/lib.rs index f351925..20ebe0e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,6 @@ //! arcus-vector: vector search as an extension to arcus-memcached. +mod aio; mod ascii; mod error; mod memcached; @@ -38,6 +39,10 @@ fn initialize(get_server_api: GET_SERVER_API) -> EXTENSION_ERROR_CODE { if server.is_null() { return EXTENSION_ERROR_CODE_EXTENSION_FATAL; } + memcached::server::set(server); + if !aio::start() || !aio::restart_after_fork() { + return EXTENSION_ERROR_CODE_EXTENSION_FATAL; + } // SAFETY: `server` was null-checked above and points to the core's static handle. let extension = unsafe { (*server).extension }; if extension.is_null() { diff --git a/src/memcached.rs b/src/memcached.rs index 751146f..e7c516a 100644 --- a/src/memcached.rs +++ b/src/memcached.rs @@ -1,3 +1,4 @@ //! The interface memcached gives an extension and its engine. pub mod ffi; +pub mod server; diff --git a/src/memcached/server.rs b/src/memcached/server.rs new file mode 100644 index 0000000..7ade7c0 --- /dev/null +++ b/src/memcached/server.rs @@ -0,0 +1,85 @@ +use std::os::raw::{c_char, c_int, c_void}; +use std::ptr; +use std::sync::atomic::{AtomicPtr, Ordering}; + +use super::ffi::{ENGINE_ERROR_CODE_ENGINE_SUCCESS, SERVER_CORE_API, SERVER_HANDLE_V1}; + +pub type ResponseHandler = + Option bool>; + +pub type IoCallback = unsafe extern "C" fn( + cookie: *const c_void, + data: *mut c_void, + handler: ResponseHandler, +) -> bool; + +static SERVER: AtomicPtr = AtomicPtr::new(ptr::null_mut()); + +pub fn set(server: *mut SERVER_HANDLE_V1) { + SERVER.store(server, Ordering::Release); +} + +fn core() -> Option<&'static SERVER_CORE_API> { + let server = SERVER.load(Ordering::Acquire); + if server.is_null() { + return None; + } + // SAFETY: `server` is the core's static handle, set once at initialization. + let core = unsafe { (*server).core }; + // SAFETY: a non-null `core` is the core's static API table. + unsafe { core.as_ref() } +} + +/// # Safety +/// +/// `handler`, if present, must accept `cookie`. +pub unsafe fn respond(handler: ResponseHandler, cookie: *const c_void, reply: &[u8]) -> bool { + let Some(handler) = handler else { + return false; + }; + let Ok(len) = c_int::try_from(reply.len()) else { + return false; + }; + // SAFETY: the caller guarantees `handler` accepts `cookie`; `reply` outlives the call. + unsafe { handler(cookie, len, reply.as_ptr().cast()) } +} + +/// Call before handing the work to another thread, so a notify that comes first still counts. +/// +/// # Safety +/// +/// `cookie` must be the connection of the running `execute`. +#[expect(clippy::expect_used, reason = "the core API is to be checked at load")] +pub unsafe fn waitfor_io_complete(cookie: *const c_void) { + let waitfor = core() + .and_then(|core| core.waitfor_io_complete) + .expect("the core supports waitfor_io_complete"); + // SAFETY: the caller guarantees `cookie` is a live connection. + unsafe { waitfor(cookie) }; +} + +/// Ends one wait; also withdraws a wait that turned out not to be needed. +/// +/// # Safety +/// +/// `cookie` must be a connection with a wait from [`waitfor_io_complete`]. +#[expect(clippy::expect_used, reason = "the core API is to be checked at load")] +pub unsafe fn notify_io_complete(cookie: *const c_void) { + let notify = core() + .and_then(|core| core.notify_io_complete) + .expect("the core supports notify_io_complete"); + // SAFETY: the caller guarantees `cookie` is a waiting, so live, connection. + unsafe { notify(cookie, ENGINE_ERROR_CODE_ENGINE_SUCCESS) }; +} + +/// # Safety +/// +/// `cookie` must be the running command's connection, and `data` must live until `cb` runs. +#[expect(clippy::expect_used, reason = "the core API is to be checked at load")] +pub unsafe fn on_io_complete(cookie: *const c_void, cb: IoCallback, data: *mut c_void) { + let register = core() + .and_then(|core| core.on_io_complete) + .expect("the core supports on_io_complete"); + // SAFETY: the caller guarantees `cookie` is a live connection and `data` outlives `cb`. + unsafe { register(cookie, Some(cb), data) }; +}