From a79398c1c6d098a251ddbb4b56ff8a6577dd2222 Mon Sep 17 00:00:00 2001 From: Styly Date: Sun, 25 Oct 2026 21:19:20 +0000 Subject: [PATCH 1/2] perf: zero-copy server wire path, allocation-free hot paths Server allocations drop from ~17 to ~0.1 per task in bench; CPU per task roughly halves. Wire format unchanged. - Error messages are Cow<'static, str> - TaskId hex encode/parse on the stack (TaskId::to_hex) - complete_refs merges borrowed results into the finished record - submit_*_ids return TaskIds; Insert replies encode them directly - server moves zmq::Messages end to end, decodes bodies in place (RequestIn) and encodes replies from leases (Reply/LeasedView) - owner check once per lease; bytes serde for hot Vec fields Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 20 +- docs/protocol.md | 9 +- enginelib/src/api.rs | 4 +- enginelib/src/chunk.rs | 177 +++++-- enginelib/src/error.rs | 18 +- enginelib/src/nucleus/complete.rs | 186 ++++--- enginelib/src/nucleus/submit.rs | 44 +- enginelib/src/transport/server.rs | 435 ++++++++++------- enginelib/src/wire.rs | 782 +++++++++++++++++++++++++++++- 9 files changed, 1365 insertions(+), 310 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 90ea1da..06263f6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -52,7 +52,16 @@ Workspace members: - Server I/O: one thread owns the ROUTER socket and only moves bytes; handlers run on tokio (DB work on the blocking pool) and send replies back over an inproc PUSH/PULL pair. Never touch the ROUTER from a handler (zmq sockets are - not thread-safe). + not thread-safe). Frames travel as `zmq::Message`s end to end (no copies). +- Zero-copy wire path (`wire.rs` "Server-side views"): the server reads only the + header + variant first (`RequestKind`), then decodes the body as `RequestIn` + (Complete results / Insert payloads borrow from the received message; bulk + requests are decoded on the blocking pool), and encodes replies as `Reply` + (`LeasedView` writes leased tasks straight from the chunk leases, `Ids` from + `TaskId`s). These mirror `Request`/`Query`/`Response` positionally (postcard): + **any change to those types must be mirrored there**; the `wire::tests` + byte-equality tests catch drift. Hot `Vec` fields use `#[serde(with = + "bytes")]` (same bytes on the wire, one memcpy instead of per-byte). - Transport-level guards (`transport/server.rs`): at most `max_inflight` handlers run at once (the I/O thread stops reading the ROUTER when full); `Renew`/`Complete`/`CancelLease` reject a lease held by another session uid @@ -134,6 +143,8 @@ trait method) > engine default (`chunk_size` 1024, no size limits). Resolved in ### Task lifecycle (`enginelib/src/nucleus/*`) - `submit(bytes, type) -> task_id` / `submit_batch(payloads, type) -> [Result]` + (`submit_batch_ids` / `submit_derived_ids` return `chunk::TaskId`s; the + transport uses those so no id `String` is built per task) — verifies each payload (+ `max_payload`), packs the valid ones of *this call* into fresh chunks of `chunk_size`, reserves seqs, writes all chunks in one `WriteBatch`, and only then returns ids. Never merges across calls. Bumps the @@ -147,7 +158,7 @@ trait method) > engine default (`chunk_size` 1024, no size limits). Resolved in Requests a refill below the 25% watermark, both before waiting and again after taking a chunk (so workers still parked on a small queue aren't left waiting). - `renew` / `renew_batch` — extends the chunk lease named by `(task_id, token)`. -- `complete` / `complete_batch(items)` — verifies each result (+ `max_result`), +- `complete` / `complete_batch(items)` / `complete_refs(borrowed items)` — verifies each result (+ `max_result`), groups items per (type, chunk, token) and merges each group into the `f:` record in one `WriteBatch` (partial completes allowed). When the chunk's last open task is done the same batch deletes the `t:` record and the lease + dedup @@ -290,6 +301,11 @@ Report any warnings/errors that indicate risk (panic paths, serialization, auth ## Notes for future agents +- Hot paths are kept allocation-free per task (server: ~0.1 allocations/task + in `bench`). `Error` messages are `Cow<'static, str>`: use string literals for + fixed messages; task ids format via `TaskId::to_hex()` (stack buffer). + Measure with `bench` (see the lifecycle playbook) before and after changes. + - Existing code has several warnings and `unwrap()` calls in runtime paths. - Prefer minimal, targeted edits. - Preserve current external behavior unless task explicitly requests behavior change. diff --git a/docs/protocol.md b/docs/protocol.md index 31168bb..d8479eb 100644 --- a/docs/protocol.md +++ b/docs/protocol.md @@ -32,9 +32,12 @@ Envelope { req_id: u64, session: u128, body: Request | Response | ProxyRequest | IPv6 is enabled on every socket. - Messages larger than `max_message_bytes` drop the connection. -Server internals: one I/O thread owns the ROUTER and only moves bytes. Each -message is decoded and handled on the tokio runtime (DB work on the blocking -pool); replies return to the I/O thread over an inproc PUSH/PULL pair. At most +Server internals: one I/O thread owns the ROUTER and only moves bytes (frames +are handed over as ZeroMQ messages, never copied). Each message is decoded and +handled on the tokio runtime (DB work, including decoding `Complete` and +`Query` bodies, on the blocking pool); payloads and results are read in place +from the received message and leased tasks are encoded straight from the chunk +leases. Replies return to the I/O thread over an inproc PUSH/PULL pair. At most `max_inflight` requests (default 65 536, parked `Lease` waits included) are handled at once; at the cap the I/O thread stops reading the ROUTER until a handler finishes, so ZeroMQ's high-water marks and TCP push back on clients. diff --git a/enginelib/src/api.rs b/enginelib/src/api.rs index a0c1ca5..ac0ecea 100644 --- a/enginelib/src/api.rs +++ b/enginelib/src/api.rs @@ -871,7 +871,9 @@ impl Registry for EngineTaskRegistry { ); self.tasks.insert(identifier, task); } - #[instrument(skip_all, fields(namespace = %identifier.0, name = %identifier.1))] + // Runs on every insert / complete / typed query: TRACE, so release builds + // (INFO) don't build and format a span per request. + #[instrument(level = "trace", skip_all, fields(namespace = %identifier.0, name = %identifier.1))] fn get(&self, identifier: &Identifier) -> Option> { self.tasks.get(identifier).map(|obj| obj.clone_box()) } diff --git a/enginelib/src/chunk.rs b/enginelib/src/chunk.rs index fe7f4e9..a3eeced 100644 --- a/enginelib/src/chunk.rs +++ b/enginelib/src/chunk.rs @@ -45,24 +45,75 @@ impl TaskId { /// Parses a 24-char hex task id. Anything else is InvalidArgument /// `"InvalidTaskId"`. pub fn parse(s: &str) -> Result { - let invalid = || Error::invalid_argument(format!("InvalidTaskId: {s}")); - if s.len() != TASK_ID_HEX_LEN || !s.bytes().all(|b| b.is_ascii_hexdigit()) { - return Err(invalid()); + let bytes: &[u8; TASK_ID_HEX_LEN] = + s.as_bytes().try_into().map_err(|_| Self::invalid(s))?; + let mut raw = [0u8; 12]; + for (out, pair) in raw.iter_mut().zip(bytes.as_chunks::<2>().0) { + let (Some(hi), Some(lo)) = (hex_value(pair[0]), hex_value(pair[1])) else { + return Err(Self::invalid(s)); + }; + *out = hi << 4 | lo; + } + Ok(Self::new( + u16::from_be_bytes([raw[0], raw[1]]), + u64::from_be_bytes(raw[2..10].try_into().expect("8 bytes")), + u16::from_be_bytes([raw[10], raw[11]]), + )) + } + + #[cold] + fn invalid(s: &str) -> Error { + Error::invalid_argument(format!("InvalidTaskId: {s}")) + } + + /// The wire form, hex-encoded into a stack buffer (no allocation). + pub fn to_hex(&self) -> TaskIdHex { + const DIGITS: &[u8; 16] = b"0123456789abcdef"; + let mut raw = [0u8; 12]; + raw[0..2].copy_from_slice(&self.server_id.to_be_bytes()); + raw[2..10].copy_from_slice(&self.chunk_seq.to_be_bytes()); + raw[10..12].copy_from_slice(&self.idx.to_be_bytes()); + let mut hex = [0u8; TASK_ID_HEX_LEN]; + for (pair, byte) in hex.as_chunks_mut::<2>().0.iter_mut().zip(raw) { + pair[0] = DIGITS[(byte >> 4) as usize]; + pair[1] = DIGITS[(byte & 0xf) as usize]; } - let server_id = u16::from_str_radix(&s[0..4], 16).map_err(|_| invalid())?; - let chunk_seq = u64::from_str_radix(&s[4..20], 16).map_err(|_| invalid())?; - let idx = u16::from_str_radix(&s[20..24], 16).map_err(|_| invalid())?; - Ok(Self::new(server_id, chunk_seq, idx)) + TaskIdHex(hex) + } +} + +fn hex_value(digit: u8) -> Option { + match digit { + b'0'..=b'9' => Some(digit - b'0'), + b'a'..=b'f' => Some(digit - b'a' + 10), + b'A'..=b'F' => Some(digit - b'A' + 10), + _ => None, } } impl fmt::Display for TaskId { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!( - f, - "{:04x}{:016x}{:04x}", - self.server_id, self.chunk_seq, self.idx - ) + f.write_str(&self.to_hex()) + } +} + +/// A task id's 24 lowercase hex chars, on the stack. Derefs to `str`. +#[derive(Clone, Copy)] +pub struct TaskIdHex([u8; TASK_ID_HEX_LEN]); + +impl std::ops::Deref for TaskIdHex { + type Target = str; + + fn deref(&self) -> &str { + // Only ever filled with ASCII hex digits. + std::str::from_utf8(&self.0).expect("hex digits are ASCII") + } +} + +/// Serializes as the wire `task_id` string, without building a `String`. +impl Serialize for TaskId { + fn serialize(&self, serializer: S) -> Result { + serializer.serialize_str(&self.to_hex()) } } @@ -340,50 +391,17 @@ impl FinishedChunk { } pub fn encode(&self) -> Vec { - let n = self.results.len(); - let size: usize = self.results.iter().flatten().map(|r| r.len() + 4).sum(); - let mut out = Vec::with_capacity(5 + n.div_ceil(8) + size); - out.push(FORMAT_V1); - put_u32(&mut out, n); - let start = out.len(); - out.resize(start + n.div_ceil(8), 0); - for (idx, result) in self.results.iter().enumerate() { - if result.is_some() { - out[start + idx / 8] |= 1 << (idx % 8); - } - } - for result in self.results.iter().flatten() { - put_u32(&mut out, result.len()); - out.extend_from_slice(result); - } - out + encode_finished(self.results.iter().map(Option::as_deref)) } pub fn decode(data: &[u8]) -> Result { - let mut at = 0; - if data.first() != Some(&FORMAT_V1) { - return Err(corrupt("finished")); - } - at += 1; - let n = read_u32(data, &mut at).ok_or_else(|| corrupt("finished"))?; - if n > MAX_CHUNK_SIZE as usize { - return Err(corrupt("finished")); - } - let bitmap = data - .get(at..at + n.div_ceil(8)) - .ok_or_else(|| corrupt("finished"))?; - at += bitmap.len(); - let mut results = vec![None; n]; - for (idx, result) in results.iter_mut().enumerate() { - if bitmap[idx / 8] & (1 << (idx % 8)) == 0 { - continue; - } - let len = read_u32(data, &mut at).ok_or_else(|| corrupt("finished"))?; - let bytes = data.get(at..at + len).ok_or_else(|| corrupt("finished"))?; - *result = Some(bytes.to_vec()); - at += len; - } - Ok(Self { results }) + let slots = decode_finished(data)?; + Ok(Self { + results: slots + .into_iter() + .map(|slot| slot.map(<[u8]>::to_vec)) + .collect(), + }) } /// Walks the results of an encoded record in index order without copying @@ -441,6 +459,59 @@ impl FinishedChunk { } } +/// Encodes a finished chunk record from its result slots (`Some` = done), in +/// index order. Takes borrowed results so callers can merge stored and new +/// results without copying each one into its own buffer first. +pub fn encode_finished<'a>( + slots: impl ExactSizeIterator> + Clone, +) -> Vec { + let n = slots.len(); + let size: usize = slots.clone().flatten().map(|r| r.len() + 4).sum(); + let mut out = Vec::with_capacity(5 + n.div_ceil(8) + size); + out.push(FORMAT_V1); + put_u32(&mut out, n); + let start = out.len(); + out.resize(start + n.div_ceil(8), 0); + for (idx, slot) in slots.clone().enumerate() { + if slot.is_some() { + out[start + idx / 8] |= 1 << (idx % 8); + } + } + for result in slots.flatten() { + put_u32(&mut out, result.len()); + out.extend_from_slice(result); + } + out +} + +/// Decodes a finished chunk record into result slots borrowing from `data`. +pub fn decode_finished(data: &[u8]) -> Result>, Error> { + let mut at = 0; + if data.first() != Some(&FORMAT_V1) { + return Err(corrupt("finished")); + } + at += 1; + let n = read_u32(data, &mut at).ok_or_else(|| corrupt("finished"))?; + if n > MAX_CHUNK_SIZE as usize { + return Err(corrupt("finished")); + } + let bitmap = data + .get(at..at + n.div_ceil(8)) + .ok_or_else(|| corrupt("finished"))?; + at += bitmap.len(); + let mut slots = vec![None; n]; + for (idx, slot) in slots.iter_mut().enumerate() { + if bitmap[idx / 8] & (1 << (idx % 8)) == 0 { + continue; + } + let len = read_u32(data, &mut at).ok_or_else(|| corrupt("finished"))?; + let end = at.checked_add(len).ok_or_else(|| corrupt("finished"))?; + *slot = Some(data.get(at..end).ok_or_else(|| corrupt("finished"))?); + at = end; + } + Ok(slots) +} + // --------------------------------------------------------------------------- // Chunk sequence allocator // --------------------------------------------------------------------------- diff --git a/enginelib/src/error.rs b/enginelib/src/error.rs index 9f54cfc..c831075 100644 --- a/enginelib/src/error.rs +++ b/enginelib/src/error.rs @@ -1,3 +1,4 @@ +use std::borrow::Cow; use std::fmt; use serde::{Deserialize, Serialize}; @@ -11,36 +12,39 @@ pub enum ErrorKind { Unknown, } +/// An engine error. The message is a `Cow` so the common static messages +/// (`"TaskNotFound"`, `"LeaseTokenMismatch"`, ...) never allocate, and cloning +/// one error into every item of a batch is free. #[derive(Debug, Clone, PartialEq, Eq)] pub struct Error { kind: ErrorKind, - message: String, + message: Cow<'static, str>, } impl Error { pub fn new(message: String) -> Error { Error { kind: ErrorKind::Unknown, - message, + message: Cow::Owned(message), } } - pub fn with_kind(kind: ErrorKind, message: impl Into) -> Error { + pub fn with_kind(kind: ErrorKind, message: impl Into>) -> Error { Error { kind, message: message.into(), } } - pub fn not_found(message: impl Into) -> Error { + pub fn not_found(message: impl Into>) -> Error { Self::with_kind(ErrorKind::NotFound, message) } - pub fn invalid_argument(message: impl Into) -> Error { + pub fn invalid_argument(message: impl Into>) -> Error { Self::with_kind(ErrorKind::InvalidArgument, message) } - pub fn io_error(message: impl Into) -> Error { + pub fn io_error(message: impl Into>) -> Error { Self::with_kind(ErrorKind::IOError, message) } @@ -57,7 +61,7 @@ impl AsRef for Error { impl From for String { fn from(e: Error) -> String { - e.message + e.message.into_owned() } } diff --git a/enginelib/src/nucleus/complete.rs b/enginelib/src/nucleus/complete.rs index 7d9fe87..f4e6903 100644 --- a/enginelib/src/nucleus/complete.rs +++ b/enginelib/src/nucleus/complete.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use rust_rocksdb::WriteBatch; use crate::api::{TaskSettings, finished_task_key, lock_lease, task_key}; -use crate::chunk::{Bitmap, FinishedChunk}; +use crate::chunk::{Bitmap, decode_finished, encode_finished}; use crate::task::Task; use crate::{Identifier, Registry, api::ServerAPI, error::Error}; @@ -17,6 +17,28 @@ pub struct CompleteItem { pub bytes: Vec, } +impl CompleteItem { + pub fn as_ref(&self) -> CompleteRef<'_> { + CompleteRef { + task_type: (&self.task_type.0, &self.task_type.1), + task_id: &self.task_id, + lease_token: self.lease_token, + bytes: &self.bytes, + } + } +} + +/// Borrowed form of [`CompleteItem`], e.g. decoded straight out of a request +/// buffer: completing never copies a result before it is encoded into the +/// chunk's finished record. +#[derive(Debug, Clone, Copy)] +pub struct CompleteRef<'a> { + pub task_type: (&'a str, &'a str), + pub task_id: &'a str, + pub lease_token: u64, + pub bytes: &'a [u8], +} + /// Finalizes one leased task with its result bytes. See [`complete_batch`]. pub fn complete( api: Arc, @@ -25,13 +47,13 @@ pub fn complete( lease_token: u64, task_bytes: &[u8], ) -> Result<(), Error> { - complete_batch( + complete_refs( &api, - vec![CompleteItem { - task_type, - task_id, + [CompleteRef { + task_type: (&task_type.0, &task_type.1), + task_id: &task_id, lease_token, - bytes: task_bytes.to_vec(), + bytes: task_bytes, }], ) .pop() @@ -54,42 +76,66 @@ pub fn complete( /// NotFound `"TaskNotLeased"` (or `"LeaseTokenMismatch"` if someone else holds /// the chunk's lease). pub fn complete_batch(api: &Arc, items: Vec) -> Vec> { - let mut results: Vec> = Vec::with_capacity(items.len()); - // Groups keyed by (type, chunk, token), in first-seen order. + complete_refs(api, items.iter().map(CompleteItem::as_ref)) +} + +/// [`complete_batch`] over borrowed items. +pub fn complete_refs<'a>( + api: &Arc, + items: impl IntoIterator>, +) -> Vec> { + let items = items.into_iter(); + let mut results: Vec> = Vec::with_capacity(items.size_hint().0); + // Task types seen in this call; a batch rarely has more than one, so a + // linear scan beats hashing two strings per item. + let mut types: Vec<(Identifier, Option)> = Vec::new(); + // Groups keyed by (type, chunk, token), in first-seen order. Items of one + // chunk usually arrive together, so the last group is checked first. let mut index: HashMap = HashMap::new(); - let mut groups: Vec<(GroupKey, Vec)> = Vec::new(); - let mut types: HashMap> = HashMap::new(); + let mut groups: Vec<(GroupKey, Vec>)> = Vec::new(); - for (pos, item) in items.into_iter().enumerate() { + for (pos, item) in items.enumerate() { let checked = (|| { - let typed = types - .entry(item.task_type.clone()) - .or_insert_with_key(|task_type| { - api.task_registry - .get(task_type) - .map(|task| (task, api.settings(task_type))) - }); - let (task, settings) = typed.as_ref().ok_or(Error::not_found("TaskTypeNotFound"))?; - let id = api.parse_task_id(&item.task_id)?; + let (a, b) = item.task_type; + let ty = match types.iter().position(|(t, _)| t.0 == a && t.1 == b) { + Some(ty) => ty, + None => { + let task_type = (a.to_owned(), b.to_owned()); + let info = api + .task_registry + .get(&task_type) + .map(|task| (task, api.settings(&task_type))); + types.push((task_type, info)); + types.len() - 1 + } + }; + let (task, settings) = types[ty] + .1 + .as_ref() + .ok_or(Error::not_found("TaskTypeNotFound"))?; + let id = api.parse_task_id(item.task_id)?; if settings .max_result .is_some_and(|max| item.bytes.len() > max as usize) { return Err(Error::invalid_argument("ResultTooLarge")); } - if !task.verify(&item.bytes) { + if !task.verify(item.bytes) { return Err(Error::invalid_argument("Failed to verify task bytes")); } - Ok(id) + Ok((ty, id)) })(); match checked { - Ok(id) => { + Ok((ty, id)) => { results.push(Ok(())); - let key = (item.task_type, id.chunk_seq, item.lease_token); - let group = *index.entry(key.clone()).or_insert_with(|| { - groups.push((key, Vec::new())); - groups.len() - 1 - }); + let key = (ty, id.chunk_seq, item.lease_token); + let group = match groups.last() { + Some((last, _)) if *last == key => groups.len() - 1, + _ => *index.entry(key).or_insert_with(|| { + groups.push((key, Vec::new())); + groups.len() - 1 + }), + }; groups[group].1.push(Entry { pos, idx: id.idx as usize, @@ -100,20 +146,31 @@ pub fn complete_batch(api: &Arc, items: Vec) -> Vec, TaskSettings); -struct Entry { +struct Entry<'a> { pos: usize, idx: usize, - bytes: Vec, + bytes: &'a [u8], +} + +/// Reads the raw finished record of a chunk (pinned, not copied). +fn read_finished_raw<'db>( + api: &'db ServerAPI, + task_type: &Identifier, + chunk_seq: u64, +) -> Result>, Error> { + api.db + .get_pinned_cf(api.cf_finished(), finished_task_key(task_type, chunk_seq)) + .map_err(|err| Error::io_error(format!("Failed to read finished chunk: {err}"))) } fn complete_chunk( @@ -121,7 +178,7 @@ fn complete_chunk( task_type: &Identifier, chunk_seq: u64, token: u64, - entries: Vec, + entries: Vec>, results: &mut [Result<(), Error>], ) { let Some(cell) = api.leased_tasks.cell(task_type, chunk_seq) else { @@ -164,10 +221,11 @@ fn complete_chunk( } // The stored results are only needed to merge into, or to check retries. + // They stay in the pinned record; `slots` borrows them. let needs_stored = lease.chunk.has_finished || retries.iter().any(|e| !newly.get(e.idx)); - let mut finished = if needs_stored { - match api.read_finished(task_type, chunk_seq) { - Ok(finished) => finished.unwrap_or_else(|| FinishedChunk::new(len)), + let stored = if needs_stored { + match read_finished_raw(api, task_type, chunk_seq) { + Ok(stored) => stored, Err(err) => { for entry in fresh.iter().chain(&retries) { results[entry.pos] = Err(err.clone()); @@ -176,11 +234,18 @@ fn complete_chunk( } } } else { - FinishedChunk::new(len) + None }; - if finished.results.len() != len { - finished.results.resize(len, None); - } + let mut slots = match stored.as_deref().map(decode_finished).transpose() { + Ok(slots) => slots.unwrap_or_default(), + Err(err) => { + for entry in fresh.iter().chain(&retries) { + results[entry.pos] = Err(err.clone()); + } + return; + } + }; + slots.resize(len, None); // Retries of tasks done before this call: compare with what is stored. let mut in_call: Vec = Vec::new(); @@ -188,8 +253,8 @@ fn complete_chunk( if newly.get(entry.idx) { in_call.push(entry); } else { - results[entry.pos] = match &finished.results[entry.idx] { - Some(stored) if *stored == entry.bytes => Ok(()), + results[entry.pos] = match slots[entry.idx] { + Some(stored) if stored == entry.bytes => Ok(()), _ => Err(Error::not_found("TaskNotLeased")), }; } @@ -200,11 +265,9 @@ fn complete_chunk( } let mut open_after = lease.chunk.open.clone(); - let mut positions = Vec::with_capacity(fresh.len()); - for entry in fresh { + for entry in &fresh { open_after.clear(entry.idx); - positions.push(entry.pos); - finished.results[entry.idx] = Some(entry.bytes); + slots[entry.idx] = Some(entry.bytes); } let all_done = open_after.none(); @@ -214,7 +277,7 @@ fn complete_chunk( batch.put_cf( api.cf_finished(), finished_task_key(task_type, chunk_seq), - finished.encode(), + encode_finished(slots.iter().copied()), ); if all_done { batch.delete_cf(api.cf_pending(), task_key(task_type, chunk_seq)); @@ -222,13 +285,13 @@ fn complete_chunk( if let Err(err) = api.db.write(&batch) { // Nothing changed: the lease is intact and the worker may retry. let err = Error::io_error(format!("Failed to persist results: {err}")); - for pos in positions.into_iter().chain(in_call.iter().map(|e| e.pos)) { - results[pos] = Err(err.clone()); + for entry in fresh.iter().chain(&in_call) { + results[entry.pos] = Err(err.clone()); } return; } for entry in in_call { - if finished.results[entry.idx].as_deref() != Some(entry.bytes.as_slice()) { + if slots[entry.idx] != Some(entry.bytes) { results[entry.pos] = Err(Error::not_found("TaskNotLeased")); } } @@ -252,7 +315,7 @@ fn already_completed( api: &ServerAPI, task_type: &Identifier, chunk_seq: u64, - entries: Vec, + entries: Vec>, results: &mut [Result<(), Error>], otherwise: Option, ) { @@ -260,8 +323,13 @@ fn already_completed( return; } let otherwise = otherwise.unwrap_or_else(|| Error::not_found("TaskNotLeased")); - let finished = match api.read_finished(task_type, chunk_seq) { - Ok(finished) => finished, + let stored = read_finished_raw(api, task_type, chunk_seq); + let slots = match &stored { + Ok(stored) => stored.as_deref().map(decode_finished).transpose(), + Err(err) => Err(err.clone()), + }; + let slots = match &slots { + Ok(slots) => slots.as_deref().unwrap_or_default(), Err(err) => { for entry in entries { results[entry.pos] = Err(err.clone()); @@ -270,12 +338,8 @@ fn already_completed( } }; for entry in entries { - let stored = finished - .as_ref() - .and_then(|finished| finished.results.get(entry.idx)) - .and_then(Option::as_deref); - results[entry.pos] = match stored { - Some(stored) if stored == entry.bytes.as_slice() => Ok(()), + results[entry.pos] = match slots.get(entry.idx) { + Some(Some(stored)) if *stored == entry.bytes => Ok(()), _ => Err(otherwise.clone()), }; } diff --git a/enginelib/src/nucleus/submit.rs b/enginelib/src/nucleus/submit.rs index 19edb84..d1e5854 100644 --- a/enginelib/src/nucleus/submit.rs +++ b/enginelib/src/nucleus/submit.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use rust_rocksdb::WriteBatch; use crate::api::task_key; -use crate::chunk::{DerivedRange, encode_derived, encode_pending}; +use crate::chunk::{DerivedRange, TaskId, encode_derived, encode_pending}; use crate::{Identifier, Registry, api::ServerAPI, error::Error}; // t:namespace:task_name: -> pending chunk record (payloads) @@ -35,6 +35,18 @@ pub fn submit_batch>( payloads: &[P], task_type: Identifier, ) -> Result>, Error> { + Ok(submit_batch_ids(api, payloads, task_type)? + .into_iter() + .map(|result| result.map(|id| id.to_string())) + .collect()) +} + +/// [`submit_batch`], returning parsed [`TaskId`]s instead of their strings. +pub fn submit_batch_ids>( + api: Arc, + payloads: &[P], + task_type: Identifier, +) -> Result>, Error> { let task = api .task_registry .get(&task_type) @@ -44,7 +56,9 @@ pub fn submit_batch>( return Err(Error::invalid_argument("DerivedTaskType")); } - let mut results: Vec> = Vec::with_capacity(payloads.len()); + // Placeholder for valid payloads until their chunk is written. + let unwritten = TaskId::new(0, 0, 0); + let mut results: Vec> = Vec::with_capacity(payloads.len()); let mut valid: Vec = Vec::with_capacity(payloads.len()); for (i, payload) in payloads.iter().enumerate() { let payload = payload.as_ref(); @@ -57,7 +71,7 @@ pub fn submit_batch>( results.push(Err(Error::invalid_argument("Failed to verify task bytes"))); } else { valid.push(i); - results.push(Ok(String::new())); + results.push(Ok(unwritten)); } } if valid.is_empty() { @@ -72,7 +86,11 @@ pub fn submit_batch>( Ok(start) => { for (offset, members) in chunks.iter().enumerate() { for (idx, &i) in members.iter().enumerate() { - results[i] = Ok(api.task_id(start + offset as u64, idx as u16)); + results[i] = Ok(TaskId::new( + api.server_id(), + start + offset as u64, + idx as u16, + )); } } } @@ -106,6 +124,20 @@ pub fn submit_derived( count: u64, task_type: Identifier, ) -> Result, Error> { + Ok(submit_derived_ids(api, base, start, count, task_type)? + .into_iter() + .map(|id| id.to_string()) + .collect()) +} + +/// [`submit_derived`], returning parsed [`TaskId`]s instead of their strings. +pub fn submit_derived_ids( + api: Arc, + base: &[u8], + start: u64, + count: u64, + task_type: Identifier, +) -> Result, Error> { let task = api .task_registry .get(&task_type) @@ -145,7 +177,9 @@ pub fn submit_derived( encode_derived(base, range) }); let first = write_chunks(&api, &task_type, n as usize, records)?; - Ok((0..n).map(|k| api.task_id(first + k, 0)).collect()) + Ok((0..n) + .map(|k| TaskId::new(api.server_id(), first + k, 0)) + .collect()) } /// Reserves `n` fresh seqs, writes `records` (exactly `n` of them) as the diff --git a/enginelib/src/transport/server.rs b/enginelib/src/transport/server.rs index 42b8348..0fdf539 100644 --- a/enginelib/src/transport/server.rs +++ b/enginelib/src/transport/server.rs @@ -16,6 +16,7 @@ use tokio::runtime::Handle; use tracing::{debug, error, info, warn}; use super::{Background, POLL_MS, constant_time_eq, context, sleep_unless_stopped, zmq_err}; +use crate::Identifier; use crate::api::{ServerAPI, lock_lease}; use crate::config::ConfigTomlServer; use crate::error::{Error, ErrorKind}; @@ -26,17 +27,16 @@ use crate::nucleus::query::{ }; use crate::nucleus::{ cancel::cancel_batch, - complete::{CompleteItem, complete_batch}, + complete::{CompleteRef, complete_refs}, lease::lease_many, renew::renew_batch, - submit::{submit_batch, submit_derived}, + submit::{submit_batch_ids, submit_derived_ids}, }; use crate::wire::{ - self, Auth, LeaseRef, Leased, LeasedItem, MAX_LEASE_BATCH, MAX_LEASE_WAIT_MS, MAX_SCAN_LIMIT, - PROTO_VERSION, ProxyRequest, ProxyResponse, Query, QueryResponse, Request, Response, - ServerInfo, Stats, TaskResult, TaskTypeInfo, WireError, + self, Auth, LeaseRef, LeaseRefIn, LeasedView, MAX_LEASE_BATCH, MAX_LEASE_WAIT_MS, + MAX_SCAN_LIMIT, PROTO_VERSION, ProxyRequest, ProxyResponse, QueryIn, QueryResponse, Reply, + RequestIn, RequestKind, Response, ServerInfo, Stats, TaskResultIn, TaskTypeInfo, WireError, }; -use crate::{Identifier, Registry}; /// Socket high-water marks, in messages. Generous so bursts of small messages /// queue instead of being dropped. @@ -245,13 +245,7 @@ fn io_loop( let (outbound, inbound) = (items[0].is_readable(), items[1].is_readable()); if outbound { - // Forward every pending reply. A peer that has gone away (or whose - // queue is full) just loses the reply; the ROUTER drops it. - while let Ok(frames) = replies_in.recv_multipart(zmq::DONTWAIT) { - if let Err(err) = router.send_multipart(frames, zmq::DONTWAIT) { - debug!("Dropped reply: {err}"); - } - } + forward_replies(replies_in, router); } if inbound { @@ -259,19 +253,15 @@ fn io_loop( if shared.saturated() { break; } - let frames = match router.recv_multipart(zmq::DONTWAIT) { - Ok(frames) => frames, + let (identity, payload) = match recv_routed(router) { + Ok(Some(frames)) => frames, + Ok(None) => continue, Err(zmq::Error::EAGAIN) => break, Err(err) => { warn!("ROUTER recv failed: {err}"); break; } }; - // DEALER clients send a single frame: [identity, payload]. - let mut frames = frames.into_iter(); - let (Some(identity), Some(payload)) = (frames.next(), frames.next()) else { - continue; - }; dispatch(shared, runtime, identity, payload); } } @@ -308,13 +298,7 @@ fn drain(shared: &Shared, router: &zmq::Socket, replies_in: &zmq::Socket) { if zmq::poll(&mut items, 10).is_err() { return; } - let mut forwarded = false; - while let Ok(frames) = replies_in.recv_multipart(zmq::DONTWAIT) { - forwarded = true; - if let Err(err) = router.send_multipart(frames, zmq::DONTWAIT) { - debug!("Dropped reply: {err}"); - } - } + let forwarded = forward_replies(replies_in, router); if shared.inflight.load(Ordering::SeqCst) == 0 { let since = *quiet_since.get_or_insert_with(Instant::now); if forwarded { @@ -333,29 +317,77 @@ fn drain(shared: &Shared, router: &zmq::Socket, replies_in: &zmq::Socket) { } } -fn dispatch(shared: &Arc, runtime: &Handle, identity: Vec, payload: Vec) { +/// Receives one ROUTER message as `(identity, payload)`, moving both frames +/// out of ZeroMQ without copying them. DEALER clients send a single frame, so +/// anything after the payload is dropped; a message without one is `None`. +fn recv_routed(router: &zmq::Socket) -> Result, zmq::Error> { + let identity = router.recv_msg(zmq::DONTWAIT)?; + if !identity.get_more() { + return Ok(None); + } + // The rest of a multipart message is delivered with its first frame. + let payload = router.recv_msg(0)?; + let mut more = payload.get_more(); + while more { + more = router.recv_msg(0)?.get_more(); + } + Ok(Some((identity, payload))) +} + +/// Forwards every pending `[identity, reply]` from the handlers to the ROUTER, +/// frames moved as-is. A peer that has gone away (or whose queue is full) just +/// loses the reply; the ROUTER drops it. Returns whether anything was pending. +fn forward_replies(replies_in: &zmq::Socket, router: &zmq::Socket) -> bool { + let mut forwarded = false; + while let Ok(Some((identity, reply))) = recv_routed(replies_in) { + forwarded = true; + let sent = router + .send(identity, zmq::SNDMORE | zmq::DONTWAIT) + .and_then(|()| router.send(reply, zmq::DONTWAIT)); + if let Err(err) = sent { + debug!("Dropped reply: {err}"); + } + } + forwarded +} + +fn dispatch(shared: &Arc, runtime: &Handle, identity: zmq::Message, payload: zmq::Message) { let slot = InflightSlot::acquire(shared); let shared = shared.clone(); runtime.spawn(async move { - let (req_id, response) = match wire::decode::(&payload) { - Ok(envelope) => ( - envelope.req_id, + // Only the header here: where the body is decoded depends on the kind. + let (req_id, reply) = match wire::decode::(&payload) { + Ok(head) => ( + head.req_id, shared - .handle(&identity, envelope.session, envelope.body) + .handle(&identity, head.session, head.body, payload) .await, ), - Err(err) => (0, Response::Error(err.into())), + Err(err) => (0, Response::Error(err.into()).into()), }; // Free the slot before handing over the reply, so the I/O thread sees // the capacity as soon as it wakes for the reply. drop(slot); - match wire::encode(req_id, &response) { + let capacity = match &reply { + Reply::Leased(leased) => leased.size_hint(), + Reply::Completed { results, next } => { + 64 + results.len() * 2 + next.as_ref().map_or(0, LeasedView::size_hint) + } + Reply::Ids(ids) => 64 + ids.len() * (2 + crate::chunk::TASK_ID_HEX_LEN), + Reply::Response(_) => 256, + }; + match wire::encode_reply(req_id, &reply, capacity) { Ok(bytes) => shared.replies.send(identity, bytes), Err(err) => error!("Failed to encode response: {err}"), } }); } +/// Decodes the body of a request whose kind was already read. +fn decode_body(payload: &[u8]) -> Result, Error> { + wire::decode::(payload).map(|envelope| envelope.body) +} + /// One handler's share of [`ServerOptions::max_inflight`]; released on drop, /// including when the handler panics. struct InflightSlot(Arc); @@ -444,7 +476,7 @@ struct ReplyPool { } impl ReplyPool { - fn send(&self, identity: Vec, bytes: Vec) { + fn send(&self, identity: zmq::Message, bytes: Vec) { let socket = self .sockets .lock() @@ -460,7 +492,11 @@ impl ReplyPool { } }, }; - if let Err(err) = socket.send_multipart([identity, bytes], 0) { + // `Vec` -> `Message` hands the buffer to ZeroMQ without a copy. + let sent = socket + .send(identity, zmq::SNDMORE) + .and_then(|()| socket.send(bytes, 0)); + if let Err(err) = sent { // Only times out if the I/O thread is gone (shutting down). debug!("Dropped reply: {err}"); } @@ -552,53 +588,77 @@ impl Shared { /// changes, so this check can't be raced by reaping or re-leasing. Leases /// that are gone or held under another token pass through: the nucleus op /// rejects them (or treats a `complete` as an idempotent retry). - fn check_owner(&self, lease: &LeaseRef, uid: &str) -> Result<(), Error> { + fn check_owner(&self, task_type: (&str, &str), chunk_seq: u64, token: u64, uid: &str) -> bool { // Leases are per chunk: every task of a chunk lease shares its holder. - // Malformed / foreign ids pass through for the nucleus op to reject. - let Ok(id) = self.api.parse_task_id(&lease.task_id) else { - return Ok(()); - }; - let not_owned = self + // Find the type's lease map without building an owned `Identifier`. + let leases = self .api .leased_tasks - .cell(&lease.task_type, id.chunk_seq) - .is_some_and(|cell| { - let held = lock_lease(&cell); - !held.gone && held.token == lease.token && held.user_id != uid - }); - if not_owned { - return Err(Error::invalid_argument("LeaseNotOwned")); - } - Ok(()) + .tasks + .iter() + .find(|entry| entry.key().0 == task_type.0 && entry.key().1 == task_type.1) + .map(|entry| entry.value().clone()); + let cell = leases.and_then(|leases| leases.get(&chunk_seq).map(|cell| cell.clone())); + !cell.is_some_and(|cell| { + let held = lock_lease(&cell); + !held.gone && held.token == token && held.user_id != uid + }) } /// Runs the batch op `op` on the items whose lease `uid` may act on and /// answers the rest with `LeaseNotOwned`; one result per item, in order. + /// The check is per chunk lease, so consecutive items of the same lease + /// share one lookup. fn owned_batch( &self, uid: &str, items: Vec, - lease_of: impl Fn(&I) -> &LeaseRef, + lease_of: impl Fn(&I) -> LeaseRefIn<'_>, op: impl FnOnce(Vec) -> Vec>, ) -> Vec> { - let checks: Vec> = items + let mut last: Option<((&str, &str), u64, u64, bool)> = None; + let checks: Vec = items .iter() - .map(|item| self.check_owner(lease_of(item), uid)) + .map(|item| { + let lease = lease_of(item); + // Malformed / foreign ids pass through for the nucleus op to reject. + let Ok(id) = self.api.parse_task_id(lease.task_id) else { + return true; + }; + match last { + Some((task_type, chunk_seq, token, owned)) + if task_type == lease.task_type + && chunk_seq == id.chunk_seq + && token == lease.token => + { + owned + } + _ => { + let owned = + self.check_owner(lease.task_type, id.chunk_seq, lease.token, uid); + last = Some((lease.task_type, id.chunk_seq, lease.token, owned)); + owned + } + } + }) .collect(); let allowed: Vec = items .into_iter() .zip(&checks) - .filter(|(_, check)| check.is_ok()) + .filter(|(_, owned)| **owned) .map(|(item, _)| item) .collect(); let mut results = op(allowed).into_iter(); checks .into_iter() - .map(|check| match check { - Ok(()) => results - .next() - .unwrap_or_else(|| Err(Error::new("batch op returned too few results".into()))), - Err(err) => Err(err), + .map(|owned| { + if owned { + results.next().unwrap_or_else(|| { + Err(Error::new("batch op returned too few results".into())) + }) + } else { + Err(Error::invalid_argument("LeaseNotOwned")) + } }) .collect() } @@ -607,13 +667,18 @@ impl Shared { self: &Arc, identity: &[u8], session_key: u128, - request: Request, - ) -> Response { - if let Request::Hello { - proto_ver, auth, .. - } = request - { - return self.hello(identity, proto_ver, auth); + kind: RequestKind, + payload: zmq::Message, + ) -> Reply { + if kind == RequestKind::Hello { + return match decode_body(&payload) { + Ok(RequestIn::Hello { + proto_ver, auth, .. + }) => self.hello(identity, proto_ver, auth), + Ok(_) => Response::Error(Error::new("Unexpected request".into()).into()), + Err(err) => Response::Error(err.into()), + } + .into(); } // Every other request needs a session; refresh its idle timer. The @@ -626,18 +691,31 @@ impl Shared { session.last_seen = Instant::now(); (session.uid.clone(), session.admin, session.inflight.clone()) } - _ => return not_authenticated(), + _ => return not_authenticated().into(), }; let Some(_slot) = SessionSlot::acquire(&inflight, self.max_inflight_per_session) else { return Response::Error(WireError { kind: ErrorKind::InvalidArgument, message: "TooManyRequests".into(), - }); + }) + .into(); }; + // Bulk, DB-bound requests are decoded on the blocking pool, straight + // from the received message: their payloads and results are never + // copied before they reach RocksDB. + match kind { + RequestKind::Complete => return self.complete_request(payload, uid).await, + RequestKind::Query => return self.query_request(payload, admin).await, + _ => {} + } + + let request = match decode_body(&payload) { + Ok(request) => request, + Err(err) => return Response::Error(err.into()).into(), + }; match request { - Request::Hello { .. } => unreachable!("handled above"), - Request::Lease { + RequestIn::Lease { task_type, max_tasks, wait_ms, @@ -650,51 +728,24 @@ impl Shared { ) .await { - Ok(leased) => Response::Leased(leased), - Err(err) => Response::Error(err.into()), + Ok(leased) => Reply::Leased(leased), + Err(err) => Response::Error(err.into()).into(), }, - Request::Renew { leases } => Response::Renewed { + RequestIn::Renew { leases } => Response::Renewed { results: self - .owned_batch( - &uid, - leases, - |lease| lease, - |leases| renew_batch(&self.api, leases.into_iter().map(lease_ref_parts)), - ) + .owned_batch(&uid, leases, LeaseRef::as_in, |leases| { + renew_batch(&self.api, leases.into_iter().map(lease_ref_parts)) + }) .into_iter() .map(|result| result.map(wire::unix_ms).map_err(Into::into)) .collect(), - }, - Request::Complete { - results, - then_lease, - } => { - let this = self.clone(); - let owner = uid.clone(); - let results = - match tokio::task::spawn_blocking(move || this.complete(results, &owner)).await - { - Ok(results) => results, - Err(err) => return Response::Error(Error::new(err.to_string()).into()), - }; - // The next lease never waits: an empty queue just means no `next`. - let next = match then_lease { - Some(req) => self - .lease(req.task_type, uid, req.max_tasks, Duration::ZERO) - .await - .ok(), - None => None, - }; - Response::Completed { results, next } } - Request::CancelLease { leases } => { + .into(), + RequestIn::CancelLease { leases } => { let results: Vec> = self - .owned_batch( - &uid, - leases, - |lease| lease, - |leases| cancel_batch(&self.api, leases.into_iter().map(lease_ref_parts)), - ) + .owned_batch(&uid, leases, LeaseRef::as_in, |leases| { + cancel_batch(&self.api, leases.into_iter().map(lease_ref_parts)) + }) .into_iter() .map(|result| result.map_err(Into::into)) .collect(); @@ -702,19 +753,55 @@ impl Shared { &self.counters.cancelled, results.iter().filter(|result| result.is_ok()).count(), ); - Response::Cancelled { results } + Response::Cancelled { results }.into() } - Request::Query(query) => { - let this = self.clone(); - match tokio::task::spawn_blocking(move || this.query(query, admin)).await { - Ok(Ok(response)) => Response::Query(response), - Ok(Err(err)) => Response::Error(err.into()), - Err(err) => Response::Error(Error::new(err.to_string()).into()), - } + RequestIn::Hello { .. } | RequestIn::Complete { .. } | RequestIn::Query(_) => { + unreachable!("handled above") } } } + async fn complete_request(self: &Arc, payload: zmq::Message, uid: String) -> Reply { + let this = self.clone(); + let owner = uid.clone(); + let completed = tokio::task::spawn_blocking(move || match decode_body(&payload)? { + RequestIn::Complete { + results, + then_lease, + } => Ok((this.complete(results, &owner), then_lease)), + _ => Err(Error::new("Unexpected request".into())), + }) + .await; + let (results, then_lease) = match completed { + Ok(Ok(completed)) => completed, + Ok(Err(err)) => return Response::Error(err.into()).into(), + Err(err) => return Response::Error(Error::new(err.to_string()).into()).into(), + }; + // The next lease never waits: an empty queue just means no `next`. + let next = match then_lease { + Some(req) => self + .lease(req.task_type, uid, req.max_tasks, Duration::ZERO) + .await + .ok(), + None => None, + }; + Reply::Completed { results, next } + } + + async fn query_request(self: &Arc, payload: zmq::Message, admin: bool) -> Reply { + let this = self.clone(); + let reply = tokio::task::spawn_blocking(move || match decode_body(&payload)? { + RequestIn::Query(query) => this.query(query, admin), + _ => Err(Error::new("Unexpected request".into())), + }) + .await; + match reply { + Ok(Ok(reply)) => reply, + Ok(Err(err)) => Response::Error(err.into()).into(), + Err(err) => Response::Error(Error::new(err.to_string()).into()).into(), + } + } + fn hello(&self, identity: &[u8], proto_ver: u16, auth: Auth) -> Response { if proto_ver != PROTO_VERSION { return Response::Error(WireError { @@ -791,7 +878,7 @@ impl Shared { uid: String, max_tasks: u32, wait: Duration, - ) -> Result { + ) -> Result { let max_tasks = max_tasks.clamp(1, MAX_LEASE_BATCH) as usize; // Shutdown can cut the wait short: `lease_many` only awaits the first // chunk's cancel-safe `recv`; the rest is synchronous `try_lease`. @@ -805,42 +892,37 @@ impl Shared { } leases = lease_many(self.api.clone(), task_type.clone(), uid, max_tasks, wait) => leases?, }; - let mut expires_at = chrono::DateTime::::MAX_UTC; - let mut tasks = Vec::with_capacity(leases.iter().map(|lease| lease.len()).sum()); - for lease in &leases { - expires_at = expires_at.min(lease.expires_at); - tasks.extend(lease.tasks().map(|(task_id, bytes)| LeasedItem { - task_id, - token: lease.token, - bytes: bytes.to_vec(), - range: lease.range(), - })); - } - add(&self.counters.leased, tasks.len()); - Ok(Leased { + let expires_at = leases + .iter() + .map(|lease| lease.expires_at) + .min() + .unwrap_or(chrono::DateTime::::MAX_UTC); + // Tasks are written into the reply straight from the chunk leases. + let leased = LeasedView { task_type, expires_at_ms: wire::unix_ms(expires_at), - tasks, - }) + leases, + }; + add(&self.counters.leased, leased.len()); + Ok(leased) } - fn complete(&self, results: Vec, uid: &str) -> Vec> { + fn complete(&self, results: Vec>, uid: &str) -> Vec> { let results: Vec> = self .owned_batch( uid, results, - |result| &result.lease, + |result| result.lease, |results| { - let items = results - .into_iter() - .map(|TaskResult { lease, bytes }| CompleteItem { - task_type: lease.task_type, - task_id: lease.task_id, - lease_token: lease.token, - bytes, - }) - .collect(); - complete_batch(&self.api, items) + complete_refs( + &self.api, + results.iter().map(|result| CompleteRef { + task_type: result.lease.task_type, + task_id: result.lease.task_id, + lease_token: result.lease.token, + bytes: result.bytes, + }), + ) }, ) .into_iter() @@ -853,39 +935,39 @@ impl Shared { results } - fn query(&self, query: Query, admin: bool) -> Result { + fn query(&self, query: QueryIn<'_>, admin: bool) -> Result { let api = &self.api; - match query { - Query::Insert { + let response = match query { + QueryIn::Insert { task_type, payloads, } => { - let results: Vec> = - submit_batch(api.clone(), &payloads, task_type)? + let ids: Vec> = + submit_batch_ids(api.clone(), &payloads, task_type)? .into_iter() .map(|result| result.map_err(Into::into)) .collect(); add( &self.counters.inserted, - results.iter().filter(|result| result.is_ok()).count(), + ids.iter().filter(|result| result.is_ok()).count(), ); - Ok(QueryResponse::Ids(results)) + return Ok(Reply::Ids(ids)); } - Query::InsertDerived { + QueryIn::InsertDerived { task_type, base, start, count, } => { - let ids = submit_derived(api.clone(), &base, start, count, task_type)?; + let ids = submit_derived_ids(api.clone(), base, start, count, task_type)?; add(&self.counters.inserted, ids.len()); - Ok(QueryResponse::Ids(ids.into_iter().map(Ok).collect())) + return Ok(Reply::Ids(ids.into_iter().map(Ok).collect())); } - Query::Get { task_type, ids } => { + QueryIn::Get { task_type, ids } => { self.require_type(&task_type)?; - Ok(QueryResponse::Records(get_tasks(api, &task_type, ids)?)) + QueryResponse::Records(get_tasks(api, &task_type, ids)?) } - Query::Scan { + QueryIn::Scan { task_type, scope, after, @@ -894,21 +976,21 @@ impl Shared { self.require_type(&task_type)?; let limit = limit.clamp(1, MAX_SCAN_LIMIT) as usize; let (records, next) = scan_tasks(api, &task_type, scope, after.as_deref(), limit)?; - Ok(QueryResponse::Page { records, next }) + QueryResponse::Page { records, next } } - Query::EstimateKeys => match db_query::query(api.clone(), DbQuery::EstimateKeys)? { - QueryResult::EstimateKeys(keys) => Ok(QueryResponse::Estimate(keys)), - _ => Err(Error::new("Unexpected estimate result".into())), + QueryIn::EstimateKeys => match db_query::query(api.clone(), DbQuery::EstimateKeys)? { + QueryResult::EstimateKeys(keys) => QueryResponse::Estimate(keys), + _ => return Err(Error::new("Unexpected estimate result".into())), }, - Query::EstimateSize { task_type, scope } => { + QueryIn::EstimateSize { task_type, scope } => { self.require_type(&task_type)?; let prefix = scope_prefix(&task_type, scope); match db_query::query(api.clone(), DbQuery::EstimateSize { prefix })? { - QueryResult::EstimateSize(size) => Ok(QueryResponse::Estimate(size)), - _ => Err(Error::new("Unexpected estimate result".into())), + QueryResult::EstimateSize(size) => QueryResponse::Estimate(size), + _ => return Err(Error::new("Unexpected estimate result".into())), } } - Query::Delete { + QueryIn::Delete { task_type, scope, ids, @@ -916,26 +998,27 @@ impl Shared { require_admin(admin)?; self.require_type(&task_type)?; delete_tasks(api, &task_type, scope, &ids)?; - Ok(QueryResponse::Done) + QueryResponse::Done } - Query::Purge { task_type, scope } => { + QueryIn::Purge { task_type, scope } => { require_admin(admin)?; self.require_type(&task_type)?; purge_tasks(api, &task_type, scope)?; - Ok(QueryResponse::Done) + QueryResponse::Done } - Query::Stats => Ok(QueryResponse::Stats(self.stats())), - } + QueryIn::Stats => QueryResponse::Stats(self.stats()), + }; + Ok(Response::Query(response).into()) } /// Typed queries only accept registered task types, so a crafted namespace /// or name can't build a key prefix that reaches another type's records. fn require_type(&self, task_type: &Identifier) -> Result<(), Error> { - self.api - .task_registry - .get(task_type) - .map(|_| ()) - .ok_or(Error::not_found("TaskTypeNotFound")) + if self.api.task_registry.tasks.contains_key(task_type) { + Ok(()) + } else { + Err(Error::not_found("TaskTypeNotFound")) + } } fn task_types(&self) -> Vec { diff --git a/enginelib/src/wire.rs b/enginelib/src/wire.rs index 7366bd9..59d8ba0 100644 --- a/enginelib/src/wire.rs +++ b/enginelib/src/wire.rs @@ -6,10 +6,13 @@ //! order. The request set mirrors `nucleus/*` one-to-one. See //! `docs/protocol.md`. -use serde::{Deserialize, Serialize, de::DeserializeOwned}; +use serde::ser::{SerializeSeq, SerializeStruct, SerializeStructVariant}; +use serde::{Deserialize, Serialize, Serializer}; use crate::Identifier; +use crate::api::Lease; pub use crate::chunk::DerivedRange; +use crate::chunk::TaskId; use crate::error::{Error, ErrorKind}; /// v2: authenticated proxy `Register { info, token }`. @@ -53,7 +56,8 @@ pub fn encode_in_session( .map_err(|err| Error::invalid_argument(format!("Failed to encode message: {err}"))) } -pub fn decode(bytes: &[u8]) -> Result, Error> { +/// Decodes an envelope. `T` may borrow from `bytes` (see [`RequestIn`]). +pub fn decode<'a, T: Deserialize<'a>>(bytes: &'a [u8]) -> Result, Error> { postcard::from_bytes(bytes) .map_err(|err| Error::invalid_argument(format!("Failed to decode message: {err}"))) } @@ -134,6 +138,7 @@ pub struct LeaseReq { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct TaskResult { pub lease: LeaseRef, + #[serde(with = "bytes")] pub bytes: Vec, } @@ -142,6 +147,7 @@ pub struct LeasedItem { pub task_id: String, pub token: u64, /// The task payload; for a derived chunk, the base. + #[serde(with = "bytes")] pub bytes: Vec, /// Derived chunk: the task indexes to derive from `bytes`. Complete the /// whole range with one result under `task_id`. @@ -217,6 +223,7 @@ pub enum Query { /// Submit tasks. Every payload is verified; ids come back in payload order. Insert { task_type: Identifier, + #[serde(with = "bytes::seq")] payloads: Vec>, }, /// Submit the indexes `start..start + count` of a derived type. `base` is @@ -224,6 +231,7 @@ pub enum Query { /// Returns one id per chunk, in range order. InsertDerived { task_type: Identifier, + #[serde(with = "bytes")] base: Vec, start: u64, count: u64, @@ -276,6 +284,7 @@ pub struct TaskRecord { pub task_id: String, pub status: TaskStatus, /// The pending payload, or the result once finished. + #[serde(with = "bytes::option")] pub bytes: Option>, /// Derived chunk not finished yet: its index range (`bytes` is the base). pub range: Option, @@ -284,6 +293,7 @@ pub struct TaskRecord { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PageRecord { pub task_id: String, + #[serde(with = "bytes")] pub bytes: Vec, /// Pending derived chunk: its index range (`bytes` is the base). pub range: Option, @@ -397,3 +407,771 @@ pub fn endpoint(addr: &str) -> String { format!("tcp://{addr}") } } + +// --------------------------------------------------------------------------- +// Byte fields +// --------------------------------------------------------------------------- + +/// `serde(with)` helpers for `Vec` fields. Postcard encodes a byte string +/// exactly like a sequence of `u8` (varint length, then the bytes), so these +/// change nothing on the wire; they just copy the bytes in one go instead of +/// one element at a time. +pub mod bytes { + use std::fmt; + + use serde::de::{Deserializer, SeqAccess, Visitor}; + use serde::ser::Serializer; + + pub fn serialize(bytes: &[u8], serializer: S) -> Result { + serializer.serialize_bytes(bytes) + } + + pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + deserializer.deserialize_byte_buf(ByteBuf) + } + + struct ByteBuf; + + impl<'de> Visitor<'de> for ByteBuf { + type Value = Vec; + + fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.write_str("bytes") + } + + fn visit_bytes(self, bytes: &[u8]) -> Result, E> { + Ok(bytes.to_vec()) + } + + fn visit_byte_buf(self, bytes: Vec) -> Result, E> { + Ok(bytes) + } + + fn visit_seq>(self, mut seq: A) -> Result, A::Error> { + let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(0).min(4096)); + while let Some(byte) = seq.next_element()? { + out.push(byte); + } + Ok(out) + } + } + + /// For `Option>`. + pub mod option { + use serde::{Deserialize, Deserializer, Serializer}; + + #[derive(Deserialize)] + #[serde(transparent)] + struct Bytes(#[serde(with = "super")] Vec); + + pub fn serialize( + bytes: &Option>, + serializer: S, + ) -> Result { + match bytes { + Some(bytes) => serializer.serialize_some(&super::Slice(bytes)), + None => serializer.serialize_none(), + } + } + + pub fn deserialize<'de, D: Deserializer<'de>>( + deserializer: D, + ) -> Result>, D::Error> { + Ok(Option::::deserialize(deserializer)?.map(|bytes| bytes.0)) + } + } + + /// For `Vec>`. + pub mod seq { + use serde::{Deserialize, Deserializer, Serializer}; + + #[derive(Deserialize)] + #[serde(transparent)] + struct Bytes(#[serde(with = "super")] Vec); + + pub fn serialize( + items: &[Vec], + serializer: S, + ) -> Result { + serializer.collect_seq(items.iter().map(|bytes| super::Slice(bytes))) + } + + pub fn deserialize<'de, D: Deserializer<'de>>( + deserializer: D, + ) -> Result>, D::Error> { + let items = Vec::::deserialize(deserializer)?; + Ok(items.into_iter().map(|bytes| bytes.0).collect()) + } + } + + /// A borrowed byte string that serializes as bytes. + pub(crate) struct Slice<'a>(pub &'a [u8]); + + impl serde::Serialize for Slice<'_> { + fn serialize(&self, serializer: S) -> Result { + serializer.serialize_bytes(self.0) + } + } +} + +// --------------------------------------------------------------------------- +// Server-side views +// --------------------------------------------------------------------------- +// +// Postcard is positional, so a type with the same variants and fields in the +// same order (and compatible field types) shares the encoding. The server uses +// that to decode requests and encode replies without a heap copy per task. +// Each view must mirror its public type exactly; the tests below check it. + +/// Just the envelope header and the request variant, without decoding the +/// body (postcard encodes a variant as its index, then its fields, so a unit +/// variant reads only the index). Lets the server pick where to decode. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +pub(crate) enum RequestKind { + Hello, + Lease, + Renew, + Complete, + CancelLease, + Query, +} + +/// [`Request`] whose bulk fields (`Complete` results, `Insert` payloads) +/// borrow from the received message. +#[derive(Debug, Deserialize)] +pub(crate) enum RequestIn<'a> { + Hello { + proto_ver: u16, + auth: Auth, + #[allow(dead_code)] + caps: Caps, + }, + Lease { + task_type: Identifier, + max_tasks: u32, + wait_ms: u32, + }, + Renew { + leases: Vec, + }, + Complete { + #[serde(borrow)] + results: Vec>, + then_lease: Option, + }, + CancelLease { + leases: Vec, + }, + Query(#[serde(borrow)] QueryIn<'a>), +} + +/// Borrowed [`LeaseRef`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +pub(crate) struct LeaseRefIn<'a> { + #[serde(borrow)] + pub task_type: (&'a str, &'a str), + pub task_id: &'a str, + pub token: u64, +} + +impl LeaseRef { + pub(crate) fn as_in(&self) -> LeaseRefIn<'_> { + LeaseRefIn { + task_type: (&self.task_type.0, &self.task_type.1), + task_id: &self.task_id, + token: self.token, + } + } +} + +/// Borrowed [`TaskResult`]. +#[derive(Debug, Clone, Copy, Deserialize)] +pub(crate) struct TaskResultIn<'a> { + #[serde(borrow)] + pub lease: LeaseRefIn<'a>, + pub bytes: &'a [u8], +} + +/// [`Query`] whose insert payloads borrow from the received message. +#[derive(Debug, Deserialize)] +pub(crate) enum QueryIn<'a> { + Insert { + task_type: Identifier, + #[serde(borrow)] + payloads: Vec<&'a [u8]>, + }, + InsertDerived { + task_type: Identifier, + base: &'a [u8], + start: u64, + count: u64, + }, + Get { + task_type: Identifier, + ids: Vec, + }, + Scan { + task_type: Identifier, + scope: Scope, + after: Option, + limit: u32, + }, + EstimateKeys, + EstimateSize { + task_type: Identifier, + scope: Scope, + }, + Delete { + task_type: Identifier, + scope: Scope, + ids: Vec, + }, + Purge { + task_type: Identifier, + scope: Scope, + }, + Stats, +} + +/// A server reply: any [`Response`], or one of the bulk responses written +/// straight from engine state. Encodes exactly like the matching `Response`. +pub(crate) enum Reply { + Response(Response), + /// `Response::Leased`. + Leased(LeasedView), + /// `Response::Completed`. + Completed { + results: Vec>, + next: Option, + }, + /// `Response::Query(QueryResponse::Ids(..))`. + Ids(Vec>), +} + +impl From for Reply { + fn from(response: Response) -> Self { + Reply::Response(response) + } +} + +/// [`Leased`] built from the chunk leases themselves: task ids are +/// hex-encoded on the stack and payloads written from the shared chunk data. +pub(crate) struct LeasedView { + pub task_type: Identifier, + pub expires_at_ms: i64, + pub leases: Vec, +} + +impl LeasedView { + /// Number of leased tasks. + pub fn len(&self) -> usize { + self.leases.iter().map(Lease::len).sum() + } + + /// Roughly the encoded size, to allocate the reply buffer once. + pub fn size_hint(&self) -> usize { + let per_task = TASK_ID_HEX_LEN_WIRE + 10 + 5 + 18; + let payloads: usize = self + .leases + .iter() + .flat_map(|lease| { + lease + .open + .iter() + .map(|idx| lease.payloads.get(idx).map_or(0, <[u8]>::len)) + }) + .sum(); + 64 + self.len() * per_task + payloads + } +} + +/// A task id on the wire: length prefix plus 24 hex chars. +const TASK_ID_HEX_LEN_WIRE: usize = 1 + crate::chunk::TASK_ID_HEX_LEN; + +impl Serialize for Reply { + fn serialize(&self, serializer: S) -> Result { + // Variant indexes of `Response` and `QueryResponse`. + const LEASED: u32 = 1; + const COMPLETED: u32 = 3; + const QUERY: u32 = 5; + const IDS: u32 = 0; + struct Ids<'a>(&'a [Result]); + impl Serialize for Ids<'_> { + fn serialize(&self, serializer: S) -> Result { + serializer.serialize_newtype_variant("QueryResponse", IDS, "Ids", self.0) + } + } + match self { + Reply::Response(response) => response.serialize(serializer), + Reply::Leased(leased) => { + serializer.serialize_newtype_variant("Response", LEASED, "Leased", leased) + } + Reply::Completed { results, next } => { + let mut variant = + serializer.serialize_struct_variant("Response", COMPLETED, "Completed", 2)?; + variant.serialize_field("results", results)?; + variant.serialize_field("next", next)?; + variant.end() + } + Reply::Ids(ids) => { + serializer.serialize_newtype_variant("Response", QUERY, "Query", &Ids(ids)) + } + } + } +} + +impl Serialize for LeasedView { + fn serialize(&self, serializer: S) -> Result { + struct Tasks<'a>(&'a LeasedView); + impl Serialize for Tasks<'_> { + fn serialize(&self, serializer: S) -> Result { + let mut seq = serializer.serialize_seq(Some(self.0.len()))?; + for lease in &self.0.leases { + for idx in lease.open.iter() { + seq.serialize_element(&Item { + task_id: TaskId::new(lease.server_id, lease.chunk_seq, idx as u16), + token: lease.token, + bytes: lease.payloads.get(idx).unwrap_or_default(), + range: lease.range(), + })?; + } + } + seq.end() + } + } + /// Mirrors [`LeasedItem`]. + struct Item<'a> { + task_id: TaskId, + token: u64, + bytes: &'a [u8], + range: Option, + } + impl Serialize for Item<'_> { + fn serialize(&self, serializer: S) -> Result { + let mut item = serializer.serialize_struct("LeasedItem", 4)?; + item.serialize_field("task_id", &self.task_id)?; + item.serialize_field("token", &self.token)?; + item.serialize_field("bytes", &bytes::Slice(self.bytes))?; + item.serialize_field("range", &self.range)?; + item.end() + } + } + let mut leased = serializer.serialize_struct("Leased", 3)?; + leased.serialize_field("task_type", &self.task_type)?; + leased.serialize_field("expires_at_ms", &self.expires_at_ms)?; + leased.serialize_field("tasks", &Tasks(self))?; + leased.end() + } +} + +/// Encodes a reply into a buffer of `capacity` bytes (grown if needed). +pub(crate) fn encode_reply(req_id: u64, reply: &Reply, capacity: usize) -> Result, Error> { + postcard::to_extend( + &Envelope { + req_id, + session: 0, + body: reply, + }, + Vec::with_capacity(capacity), + ) + .map_err(|err| Error::invalid_argument(format!("Failed to encode message: {err}"))) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use crate::chunk::{Bitmap, ChunkPayloads, encode_derived, encode_pending}; + + fn ty() -> Identifier { + ("ns".into(), "name".into()) + } + + fn lease_ref(task_id: &str, token: u64) -> LeaseRef { + LeaseRef { + task_type: ty(), + task_id: task_id.into(), + token, + } + } + + fn lease(chunk_seq: u64, record: Vec, open: &[usize], token: u64) -> Lease { + let payloads = ChunkPayloads::decode(record).unwrap(); + let mut bitmap = Bitmap::empty(payloads.len()); + for &idx in open { + bitmap.set(idx); + } + Lease { + task_type: ty(), + server_id: 7, + chunk_seq, + payloads: Arc::new(payloads), + open: bitmap, + token, + expires_at: chrono::Utc::now(), + } + } + + fn leases() -> Vec { + let regular = encode_pending([&b"aa"[..], b"", b"cccc"].into_iter()); + let derived = encode_derived(b"base", DerivedRange { start: 5, count: 9 }); + vec![ + lease(1, regular, &[0, 2], 11), + lease(1 << 40, derived, &[0], 12), + ] + } + + /// The owned `Leased` the view stands for. + fn owned(view: &LeasedView) -> Leased { + Leased { + task_type: view.task_type.clone(), + expires_at_ms: view.expires_at_ms, + tasks: view + .leases + .iter() + .flat_map(|lease| { + lease.tasks().map(|(task_id, bytes)| LeasedItem { + task_id, + token: lease.token, + bytes: bytes.to_vec(), + range: lease.range(), + }) + }) + .collect(), + } + } + + fn view() -> LeasedView { + LeasedView { + task_type: ty(), + expires_at_ms: 1_700_000_000_123, + leases: leases(), + } + } + + #[test] + fn reply_encodes_like_response() { + let view = view(); + assert_eq!(view.len(), 3); + let leased = owned(&view); + let same = |reply: Reply, response: Response| { + assert_eq!( + encode_reply(9, &reply, 0).unwrap(), + encode(9, &response).unwrap() + ); + }; + same(Reply::Leased(view), Response::Leased(leased.clone())); + + let results = vec![Ok(()), Err(Error::not_found("TaskNotLeased").into())]; + same( + Reply::Completed { + results: results.clone(), + next: Some(super::tests::view()), + }, + Response::Completed { + results: results.clone(), + next: Some(leased), + }, + ); + same( + Reply::Completed { + results: results.clone(), + next: None, + }, + Response::Completed { + results, + next: None, + }, + ); + + let id = TaskId::new(7, 42, 3); + let err: WireError = Error::invalid_argument("PayloadTooLarge").into(); + same( + Reply::Ids(vec![Ok(id), Err(err.clone())]), + Response::Query(QueryResponse::Ids(vec![Ok(id.to_string()), Err(err)])), + ); + same( + Response::Cancelled { results: vec![] }.into(), + Response::Cancelled { results: vec![] }, + ); + } + + #[test] + fn decoded_leased_view_matches() { + let view = view(); + let leased = owned(&view); + let bytes = encode_reply(1, &Reply::Leased(view), 0).unwrap(); + match decode::(&bytes).unwrap().body { + Response::Leased(decoded) => { + assert_eq!(decoded.tasks.len(), 3); + assert_eq!(decoded.tasks[1].bytes, b"cccc"); + assert_eq!( + decoded.tasks[2].range, + Some(DerivedRange { start: 5, count: 9 }) + ); + assert_eq!( + decoded.tasks.iter().map(|t| &t.task_id).collect::>(), + leased.tasks.iter().map(|t| &t.task_id).collect::>() + ); + } + other => panic!("unexpected {other:?}"), + } + } + + /// Every `Request` decodes as the matching `RequestKind` and `RequestIn`. + #[test] + fn request_in_mirrors_request() { + let id = TaskId::new(7, 42, 3).to_string(); + let requests = [ + ( + RequestKind::Hello, + Request::Hello { + proto_ver: PROTO_VERSION, + auth: Auth { + uid: "u".into(), + challenge: "c".into(), + admin_token: Some("t".into()), + }, + caps: Caps { + runners: vec![RunnerKind::Hip], + max_inflight: 3, + }, + }, + ), + ( + RequestKind::Lease, + Request::Lease { + task_type: ty(), + max_tasks: 5, + wait_ms: 6, + }, + ), + ( + RequestKind::Renew, + Request::Renew { + leases: vec![lease_ref(&id, 1)], + }, + ), + ( + RequestKind::Complete, + Request::Complete { + results: vec![ + TaskResult { + lease: lease_ref(&id, 2), + bytes: b"result".to_vec(), + }, + TaskResult { + lease: lease_ref("x", 3), + bytes: Vec::new(), + }, + ], + then_lease: Some(LeaseReq { + task_type: ty(), + max_tasks: 8, + }), + }, + ), + ( + RequestKind::CancelLease, + Request::CancelLease { + leases: vec![lease_ref(&id, 4)], + }, + ), + ( + RequestKind::Query, + Request::Query(Query::Insert { + task_type: ty(), + payloads: vec![b"p1".to_vec(), Vec::new(), vec![0xff; 300]], + }), + ), + ( + RequestKind::Query, + Request::Query(Query::InsertDerived { + task_type: ty(), + base: b"base".to_vec(), + start: 1, + count: 2, + }), + ), + ( + RequestKind::Query, + Request::Query(Query::Get { + task_type: ty(), + ids: vec![id.clone()], + }), + ), + ( + RequestKind::Query, + Request::Query(Query::Scan { + task_type: ty(), + scope: Scope::Finished, + after: Some(id.clone()), + limit: 9, + }), + ), + (RequestKind::Query, Request::Query(Query::EstimateKeys)), + ( + RequestKind::Query, + Request::Query(Query::EstimateSize { + task_type: ty(), + scope: Scope::Pending, + }), + ), + ( + RequestKind::Query, + Request::Query(Query::Delete { + task_type: ty(), + scope: Scope::Pending, + ids: vec![id.clone()], + }), + ), + ( + RequestKind::Query, + Request::Query(Query::Purge { + task_type: ty(), + scope: Scope::Finished, + }), + ), + (RequestKind::Query, Request::Query(Query::Stats)), + ]; + for (kind, request) in requests { + let bytes = encode_in_session(3, 77, &request).unwrap(); + let head = decode::(&bytes).unwrap(); + assert_eq!((head.req_id, head.session, head.body), (3, 77, kind)); + let body = decode::(&bytes).unwrap().body; + // Field-by-field: the borrowed view must see the same values. + let same = match (&request, &body) { + ( + Request::Hello { + proto_ver, auth, .. + }, + RequestIn::Hello { + proto_ver: p, + auth: a, + .. + }, + ) => proto_ver == p && auth.uid == a.uid && auth.admin_token == a.admin_token, + ( + Request::Lease { + task_type, + max_tasks, + wait_ms, + }, + RequestIn::Lease { + task_type: t, + max_tasks: m, + wait_ms: w, + }, + ) => task_type == t && max_tasks == m && wait_ms == w, + (Request::Renew { leases }, RequestIn::Renew { leases: l }) + | (Request::CancelLease { leases }, RequestIn::CancelLease { leases: l }) => { + leases == l + } + ( + Request::Complete { + results, + then_lease, + }, + RequestIn::Complete { + results: r, + then_lease: t, + }, + ) => { + results.len() == r.len() + && results + .iter() + .zip(r) + .all(|(a, b)| a.lease.as_in() == b.lease && a.bytes == b.bytes) + && then_lease.as_ref().map(|l| l.max_tasks) + == t.as_ref().map(|l| l.max_tasks) + } + (Request::Query(query), RequestIn::Query(q)) => match (query, q) { + ( + Query::Insert { + task_type, + payloads, + }, + QueryIn::Insert { + task_type: t, + payloads: p, + }, + ) => task_type == t && payloads.iter().map(Vec::as_slice).eq(p.iter().copied()), + ( + Query::InsertDerived { + base, start, count, .. + }, + QueryIn::InsertDerived { + base: b, + start: s, + count: c, + .. + }, + ) => base == b && start == s && count == c, + (Query::Get { ids, .. }, QueryIn::Get { ids: i, .. }) + | (Query::Delete { ids, .. }, QueryIn::Delete { ids: i, .. }) => ids == i, + ( + Query::Scan { + scope, + after, + limit, + .. + }, + QueryIn::Scan { + scope: s, + after: a, + limit: l, + .. + }, + ) => scope == s && after == a && limit == l, + (Query::EstimateSize { scope, .. }, QueryIn::EstimateSize { scope: s, .. }) + | (Query::Purge { scope, .. }, QueryIn::Purge { scope: s, .. }) => scope == s, + (Query::EstimateKeys, QueryIn::EstimateKeys) + | (Query::Stats, QueryIn::Stats) => true, + _ => false, + }, + _ => false, + }; + assert!(same, "{request:?} decoded as {body:?}"); + } + } + + /// The `bytes` helpers leave the encoding of a plain `Vec` unchanged. + #[test] + fn byte_fields_encode_like_vec() { + #[derive(Serialize, Deserialize, PartialEq, Debug)] + struct Plain { + one: Vec, + opt: Option>, + none: Option>, + many: Vec>, + } + #[derive(Serialize, Deserialize, PartialEq, Debug)] + struct Fast { + #[serde(with = "bytes")] + one: Vec, + #[serde(with = "bytes::option")] + opt: Option>, + #[serde(with = "bytes::option")] + none: Option>, + #[serde(with = "bytes::seq")] + many: Vec>, + } + let big: Vec = (0..=255).cycle().take(1000).collect(); + let plain = Plain { + one: big.clone(), + opt: Some(b"x".to_vec()), + none: None, + many: vec![Vec::new(), big.clone()], + }; + let fast = Fast { + one: big.clone(), + opt: Some(b"x".to_vec()), + none: None, + many: vec![Vec::new(), big], + }; + let encoded = postcard::to_allocvec(&plain).unwrap(); + assert_eq!(postcard::to_allocvec(&fast).unwrap(), encoded); + assert_eq!(postcard::from_bytes::(&encoded).unwrap(), fast); + } +} From 3c66a37891ee7faaf7d71ae25ac85663267cb048 Mon Sep 17 00:00:00 2001 From: Styly Date: Mon, 26 Oct 2026 00:01:24 +0000 Subject: [PATCH 2/2] perf: keyed task-type lookups in complete_refs and owner checks Addresses review on the linear type scans: - complete_refs: previous item's type first, then a map keyed by the borrowed names (no owned key per item); mixed-type batches stay O(1) per item. Adds a mixed-type regression test. - owned_batch: resolve each named type's lease map once per request with a keyed lookup instead of iterating every DashMap shard per chunk. Co-Authored-By: Claude Opus 5.5 --- enginelib/macros/src/lib.rs | 8 ++-- enginelib/src/nucleus/complete.rs | 33 ++++++++------ enginelib/src/transport/server.rs | 35 ++++++++++----- enginelib/tests/nucleus.rs | 71 +++++++++++++++++++++++++++++++ 4 files changed, 119 insertions(+), 28 deletions(-) diff --git a/enginelib/macros/src/lib.rs b/enginelib/macros/src/lib.rs index 7890edb..78a5f7a 100644 --- a/enginelib/macros/src/lib.rs +++ b/enginelib/macros/src/lib.rs @@ -47,10 +47,10 @@ pub fn module(_attr: TokenStream, item: TokenStream) -> TokenStream { env!("CARGO_PKG_NAME"), );), ); - item_fn.block.stmts.insert( - 0, - parse_quote!(::enginelib::api::init_tracing();), - ); + item_fn + .block + .stmts + .insert(0, parse_quote!(::enginelib::api::init_tracing();)); item_fn .attrs .push(parse_quote!(#[unsafe(export_name="run")])); diff --git a/enginelib/src/nucleus/complete.rs b/enginelib/src/nucleus/complete.rs index f4e6903..330f90c 100644 --- a/enginelib/src/nucleus/complete.rs +++ b/enginelib/src/nucleus/complete.rs @@ -86,9 +86,12 @@ pub fn complete_refs<'a>( ) -> Vec> { let items = items.into_iter(); let mut results: Vec> = Vec::with_capacity(items.size_hint().0); - // Task types seen in this call; a batch rarely has more than one, so a - // linear scan beats hashing two strings per item. + // Task types seen in this call, looked up by the borrowed names (no owned + // key per item). A batch rarely has more than one, so the previous item's + // type is checked first and the map is only hashed when the type changes. let mut types: Vec<(Identifier, Option)> = Vec::new(); + let mut type_index: HashMap<(&'a str, &'a str), usize> = HashMap::new(); + let mut last_type: Option<((&'a str, &'a str), usize)> = None; // Groups keyed by (type, chunk, token), in first-seen order. Items of one // chunk usually arrive together, so the last group is checked first. let mut index: HashMap = HashMap::new(); @@ -96,17 +99,21 @@ pub fn complete_refs<'a>( for (pos, item) in items.enumerate() { let checked = (|| { - let (a, b) = item.task_type; - let ty = match types.iter().position(|(t, _)| t.0 == a && t.1 == b) { - Some(ty) => ty, - None => { - let task_type = (a.to_owned(), b.to_owned()); - let info = api - .task_registry - .get(&task_type) - .map(|task| (task, api.settings(&task_type))); - types.push((task_type, info)); - types.len() - 1 + let ty = match last_type { + Some((name, ty)) if name == item.task_type => ty, + _ => { + let ty = *type_index.entry(item.task_type).or_insert_with(|| { + let (a, b) = item.task_type; + let task_type = (a.to_owned(), b.to_owned()); + let info = api + .task_registry + .get(&task_type) + .map(|task| (task, api.settings(&task_type))); + types.push((task_type, info)); + types.len() - 1 + }); + last_type = Some((item.task_type, ty)); + ty } }; let (task, settings) = types[ty] diff --git a/enginelib/src/transport/server.rs b/enginelib/src/transport/server.rs index 0fdf539..fc4dfed 100644 --- a/enginelib/src/transport/server.rs +++ b/enginelib/src/transport/server.rs @@ -7,6 +7,7 @@ //! peer. ZeroMQ sockets are not thread-safe, so handlers never touch the ROUTER //! directly. +use std::collections::HashMap; use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -17,7 +18,7 @@ use tracing::{debug, error, info, warn}; use super::{Background, POLL_MS, constant_time_eq, context, sleep_unless_stopped, zmq_err}; use crate::Identifier; -use crate::api::{ServerAPI, lock_lease}; +use crate::api::{LeaseCell, ServerAPI, lock_lease}; use crate::config::ConfigTomlServer; use crate::error::{Error, ErrorKind}; use crate::events::Events; @@ -565,6 +566,9 @@ struct Shared { max_inflight_per_session: usize, } +/// One task type's chunk leases (`LeasedTaskQueue::tasks` value). +type TypeLeases = Arc>; + fn not_authenticated() -> Response { Response::Error(WireError { kind: ErrorKind::InvalidArgument, @@ -588,16 +592,10 @@ impl Shared { /// changes, so this check can't be raced by reaping or re-leasing. Leases /// that are gone or held under another token pass through: the nucleus op /// rejects them (or treats a `complete` as an idempotent retry). - fn check_owner(&self, task_type: (&str, &str), chunk_seq: u64, token: u64, uid: &str) -> bool { + /// + /// `leases` is the type's lease map (see [`Shared::type_leases`]). + fn check_owner(leases: Option<&TypeLeases>, chunk_seq: u64, token: u64, uid: &str) -> bool { // Leases are per chunk: every task of a chunk lease shares its holder. - // Find the type's lease map without building an owned `Identifier`. - let leases = self - .api - .leased_tasks - .tasks - .iter() - .find(|entry| entry.key().0 == task_type.0 && entry.key().1 == task_type.1) - .map(|entry| entry.value().clone()); let cell = leases.and_then(|leases| leases.get(&chunk_seq).map(|cell| cell.clone())); !cell.is_some_and(|cell| { let held = lock_lease(&cell); @@ -605,6 +603,16 @@ impl Shared { }) } + /// The lease map of one task type, if it has ever been leased. + fn type_leases(&self, (a, b): (&str, &str)) -> Option { + let task_type: Identifier = (a.to_owned(), b.to_owned()); + self.api + .leased_tasks + .tasks + .get(&task_type) + .map(|entry| entry.value().clone()) + } + /// Runs the batch op `op` on the items whose lease `uid` may act on and /// answers the rest with `LeaseNotOwned`; one result per item, in order. /// The check is per chunk lease, so consecutive items of the same lease @@ -616,6 +624,8 @@ impl Shared { lease_of: impl Fn(&I) -> LeaseRefIn<'_>, op: impl FnOnce(Vec) -> Vec>, ) -> Vec> { + // Each named type's lease map is resolved once per request. + let mut types: HashMap<(&str, &str), Option> = HashMap::new(); let mut last: Option<((&str, &str), u64, u64, bool)> = None; let checks: Vec = items .iter() @@ -634,8 +644,11 @@ impl Shared { owned } _ => { + let leases = types + .entry(lease.task_type) + .or_insert_with(|| self.type_leases(lease.task_type)); let owned = - self.check_owner(lease.task_type, id.chunk_seq, lease.token, uid); + Self::check_owner(leases.as_ref(), id.chunk_seq, lease.token, uid); last = Some((lease.task_type, id.chunk_seq, lease.token, owned)); owned } diff --git a/enginelib/tests/nucleus.rs b/enginelib/tests/nucleus.rs index b53de8e..6549c9b 100644 --- a/enginelib/tests/nucleus.rs +++ b/enginelib/tests/nucleus.rs @@ -1275,3 +1275,74 @@ async fn derived_chunks_survive_restart() { ); close(api).await; } + +/// One batch interleaving several task types (and an unknown one): each item +/// is routed to its own type's chunk, whatever the order. +#[tokio::test] +async fn complete_batch_mixes_task_types() { + let dir = TempDir::new("mixed-types"); + let mut api = ServerAPI::open(&dir.0, Config::default()); + let types: Vec = (0..3).map(|i| ("m".into(), format!("t{i}"))).collect(); + for id in &types { + api.task_registry.register( + Arc::new(TestTask { + id: id.clone(), + limits: TaskLimits::default(), + }), + id.clone(), + ); + } + let api = Arc::new(api); + ServerAPI::populate(&api); + + let mut leases = Vec::new(); + for id in &types { + submit_batch(api.clone(), &[&b"a"[..], b"b"], id.clone()).unwrap(); + leases.push( + lease_timeout(api.clone(), id.clone(), "w".into(), WAIT) + .await + .unwrap(), + ); + } + + // Round-robin over the types, task by task, with an unknown type mixed in. + let mut items = Vec::new(); + for idx in 0..2 { + for lease in &leases { + items.push(CompleteItem { + task_type: lease.task_type.clone(), + task_id: lease.task_ids()[idx].clone(), + lease_token: lease.token, + bytes: vec![idx as u8 + 1], + }); + } + items.push(CompleteItem { + task_type: ("m".into(), "missing".into()), + task_id: leases[0].task_ids()[idx].clone(), + lease_token: leases[0].token, + bytes: vec![9], + }); + } + let results = complete_batch(&api, items); + assert_eq!(results.len(), 8); + for (pos, result) in results.iter().enumerate() { + if pos % 4 == 3 { + assert_eq!(result.as_ref().unwrap_err().as_ref(), "TaskTypeNotFound"); + } else { + assert!(result.is_ok(), "item {pos}: {result:?}"); + } + } + for lease in &leases { + assert!( + api.leased_tasks + .cell(&lease.task_type, lease.chunk_seq) + .is_none() + ); + let done = api + .read_finished(&lease.task_type, lease.chunk_seq) + .unwrap() + .unwrap(); + assert_eq!(done.results, vec![Some(vec![1]), Some(vec![2])]); + } + close(api).await; +}