tcontract_exec/workspace.rs
1//! Owning workspace: page-aligned, grow-only, reused per worker.
2//!
3//! A workspace belongs to exactly one owner — a `tcontract_exec::Pool` or a caller
4//! that wants reuse — and is *lent* to the driver for one call at a time. There
5//! is deliberately no process-global arena: storage handed to a team is
6//! returned to the provider that issued it, and a second owner can never see
7//! it.
8//!
9//! Two payload shapes:
10//!
11//! * **worker buffers** — the packed A block, the accumulator tile and the
12//! scratch vector of one thread. They are allocated by the thread that will
13//! write them, which is what makes their first touch local to the node that
14//! keeps using them. The owner maps each thread to its slot, so a lookup is
15//! O(1) and the mapping dies with the owner: the storage is the owner's, the
16//! entry is only an index, and a slot that has been trimmed is simply
17//! re-created.
18//! * **team sets** — the shared packed B panel, the team's scatter vectors and
19//! its barriers. They are taken exclusively for the duration of a call and
20//! returned on drop, so two concurrent executes can never share one.
21//!
22//! # Exclusivity
23//!
24//! A slot that is *in the owner's map* is touched only while the map lock is
25//! held. A slot that is *checked out* is not in the map at all: the thread that
26//! checked it out is its only owner until it is returned, so a re-entrant call
27//! on that thread finds nothing and allocates fresh call-local buffers instead
28//! of aliasing the ones its outer call is using (waiting for itself would
29//! deadlock). `trim` and the accounting snapshot therefore never observe a slot
30//! another thread is writing.
31//!
32//! `ArenaProvider::traced()` additionally records every growth's thread and
33//! size, which is how the tests prove that a worker's buffers were allocated by
34//! the worker. Tracing is off by default and only covers buffers the owner
35//! keeps; the fresh re-entrant path is not traced.
36
37use std::collections::HashMap;
38use std::sync::atomic::{AtomicU64, Ordering};
39use std::sync::{Barrier, Mutex};
40
41/// The 4096-byte alignment every buffer in the workspace has.
42const PAGE: usize = 4096;
43
44/// Bytes and counts one execute needs. Buffer sizes are bytes because the
45/// provider is element-type erased: the driver knows the element type and
46/// converts its real counts once.
47///
48/// A zero entry never allocates.
49#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
50pub struct WorkspaceReq {
51 /// This thread's packed A block.
52 pub a_bytes: usize,
53 /// This thread's accumulator tile, including any induced-method scratch.
54 pub tile_bytes: usize,
55 /// Elements reserved for this thread's scatter scratch.
56 pub worker_scatter: usize,
57 /// The team's shared packed B panel.
58 pub b_bytes: usize,
59 /// Elements reserved for the team's block-scatter vectors.
60 pub team_scatter: usize,
61 /// Barriers the team needs, or zero when nothing is shared.
62 pub barriers: usize,
63}
64
65impl WorkspaceReq {
66 /// The payload this request needs, in bytes: the buffers, the `i64` scatter
67 /// elements and the barriers. It is the value a [`Workspace`] is sized to and
68 /// the amount [`Workspace::have`] reports.
69 ///
70 /// # Examples
71 ///
72 /// ```
73 /// use tcontract_exec::WorkspaceReq;
74 /// let req = WorkspaceReq { a_bytes: 16, team_scatter: 2, ..Default::default() };
75 /// assert_eq!(req.bytes(), 16 + 2 * core::mem::size_of::<i64>());
76 /// ```
77 pub fn bytes(&self) -> usize {
78 self.a_bytes
79 + self.tile_bytes
80 + self.b_bytes
81 + (self.worker_scatter + self.team_scatter) * core::mem::size_of::<i64>()
82 + self.barriers * core::mem::size_of::<Barrier>()
83 }
84
85 /// Whether `have` covers every component of this request.
86 ///
87 /// Capacity is checked component-wise, never by the byte total: a large A
88 /// block does not stand in for a missing B panel, and the driver would grow
89 /// the arena for the shortage rather than refuse it. [`Workspace::reserve`]
90 /// keeps the same component-wise maxima.
91 pub fn covered_by(&self, have: &WorkspaceReq) -> bool {
92 self.a_bytes <= have.a_bytes
93 && self.tile_bytes <= have.tile_bytes
94 && self.worker_scatter <= have.worker_scatter
95 && self.b_bytes <= have.b_bytes
96 && self.team_scatter <= have.team_scatter
97 && self.barriers <= have.barriers
98 }
99}
100
101/// What an owner currently holds and currently lends out.
102///
103/// Exactness: both values are settled at every checkout, return, team take,
104/// lease drop and `trim`, so `ArenaProvider::stats` is exact whenever no
105/// execution is in flight. While one is, a checked-out slot is counted at the
106/// capacity it had when it was checked out, so `retained_bytes` is an unsettled
107/// estimate that can lag growth or shrinkage until the call returns.
108#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
109pub struct WorkspaceStats {
110 /// Capacity the owner holds: worker A/tile/scatter and team
111 /// panel/scatter/barriers, idle or in use.
112 pub retained_bytes: usize,
113 /// The part of `retained_bytes` lent to a call right now.
114 pub leased_bytes: usize,
115}
116
117/// Page-aligned, grow-only, never zeroed storage.
118///
119/// The caller writes before reading, which is the contract the packing and
120/// kernel code already has; zeroing would only double the traffic.
121#[derive(Default, Debug)]
122struct PageBuf {
123 ptr: *mut u8,
124 cap: usize,
125 /// Owner key this buffer reports growths under, or zero when untraced.
126 traced: u64,
127}
128
129// SAFETY: the buffer owns its allocation and hands out pointers only through
130// `&mut self` or an explicit, contracted call.
131unsafe impl Send for PageBuf {}
132
133impl PageBuf {
134 /// A buffer that reports its growths under `owner`, for tests that prove
135 /// where storage was first touched. Zero means untraced.
136 fn traced(owner: u64) -> Self {
137 Self {
138 traced: owner,
139 ..Self::default()
140 }
141 }
142
143 /// A pointer to at least `bytes` writable, page-aligned bytes, growing the
144 /// allocation if needed and reusing it otherwise. `ensure(0)` returns a
145 /// dangling but well-aligned pointer and allocates nothing.
146 fn ensure(&mut self, bytes: usize) -> *mut u8 {
147 if bytes == 0 {
148 return std::ptr::NonNull::<u8>::dangling().as_ptr();
149 }
150 if bytes > self.cap {
151 // Grow by doubling, so a sequence of slightly larger calls does not
152 // reallocate every time.
153 let want = bytes.next_power_of_two();
154 let layout = std::alloc::Layout::from_size_align(want, PAGE).expect("workspace layout");
155 // SAFETY: `want` is nonzero, and any old block was allocated with
156 // this same alignment and a power-of-two size.
157 let new = unsafe {
158 match self.ptr.is_null() {
159 true => std::alloc::alloc(layout),
160 false => {
161 let old = std::alloc::Layout::from_size_align_unchecked(self.cap, PAGE);
162 std::alloc::realloc(self.ptr, old, want)
163 }
164 }
165 };
166 assert!(
167 !new.is_null(),
168 "workspace allocation of {want} bytes failed"
169 );
170 self.ptr = new;
171 self.cap = want;
172 if self.traced != 0 {
173 record(self.traced, want);
174 }
175 }
176 self.ptr
177 }
178
179 /// Pointer to the current allocation, or a dangling pointer when empty.
180 fn as_ptr(&self) -> *mut u8 {
181 if self.ptr.is_null() {
182 std::ptr::NonNull::<u8>::dangling().as_ptr()
183 } else {
184 self.ptr
185 }
186 }
187
188 /// Bytes currently allocated.
189 fn cap_bytes(&self) -> usize {
190 self.cap
191 }
192
193 /// Release the allocation. Callers must hold the buffer exclusively.
194 fn trim(&mut self) {
195 if !self.ptr.is_null() {
196 // SAFETY: allocated with this alignment and a power-of-two size.
197 unsafe {
198 std::alloc::dealloc(
199 self.ptr,
200 std::alloc::Layout::from_size_align_unchecked(self.cap, PAGE),
201 )
202 };
203 self.ptr = std::ptr::null_mut();
204 self.cap = 0;
205 }
206 }
207}
208
209impl Drop for PageBuf {
210 fn drop(&mut self) {
211 self.trim();
212 }
213}
214
215/// One thread's buffers, as the owner keeps them between calls.
216#[derive(Debug)]
217struct WorkerSlot {
218 a: PageBuf,
219 tile: PageBuf,
220 scratch: Vec<i64>,
221}
222
223impl WorkerSlot {
224 /// `owner` is the tracing key, or zero for an untraced owner.
225 fn new(owner: u64) -> Self {
226 Self {
227 a: PageBuf::traced(owner),
228 tile: PageBuf::traced(owner),
229 scratch: Vec::new(),
230 }
231 }
232
233 fn bytes(&self) -> usize {
234 self.a.cap_bytes()
235 + self.tile.cap_bytes()
236 + self.scratch.capacity() * core::mem::size_of::<i64>()
237 }
238
239 fn trim(&mut self) {
240 self.a.trim();
241 self.tile.trim();
242 self.scratch = Vec::new();
243 }
244}
245
246/// A team's shared storage: one packed B panel, the scatter vectors and the
247/// barriers that publish the panel between its threads.
248#[derive(Debug, Default)]
249struct TeamSet {
250 /// The packed B panel.
251 b: PageBuf,
252 /// Block-scatter vectors for this call, capacity-reused.
253 scatter: Vec<i64>,
254 /// One barrier per column group when `pm > 1`, else empty.
255 barriers: Vec<Barrier>,
256 pm: usize,
257 pn: usize,
258}
259
260impl TeamSet {
261 fn prepare(&mut self, req: &WorkspaceReq, pm: usize, pn: usize) {
262 self.scatter.clear();
263 self.scatter.reserve(req.team_scatter);
264 if (self.pm, self.pn) != (pm, pn) {
265 // A barrier cannot be reset, so a differently shaped team needs new
266 // ones; that is why the free list keeps sets and rebuilds here.
267 self.barriers = (0..req.barriers).map(|_| Barrier::new(pm)).collect();
268 self.pm = pm;
269 self.pn = pn;
270 }
271 }
272
273 fn bytes(&self) -> usize {
274 self.b.cap_bytes()
275 + self.scatter.capacity() * core::mem::size_of::<i64>()
276 + self.barriers.capacity() * core::mem::size_of::<Barrier>()
277 }
278}
279
280/// An exclusive loan of a team set, returned to its owner on drop.
281///
282/// The set is deliberately opaque: the panel grows only through
283/// [`TeamLease::panel`], so every byte the owner hands out is accounted.
284#[derive(Debug)]
285pub struct TeamLease<'a> {
286 owner: &'a ArenaProvider,
287 set: Option<TeamSet>,
288 /// Bytes this lease has added to the owner's leased total.
289 billed: usize,
290}
291
292impl<'a> TeamLease<'a> {
293 /// Size the shared panel and return its pointer, accounting the growth to
294 /// the owner. The driver sizes it because only the driver knows the element
295 /// type; the provider keeps the storage.
296 pub fn panel(&mut self, bytes: usize) -> *mut u8 {
297 let set = self.set.as_mut().expect("team lease holds its set");
298 let before = set.b.cap_bytes();
299 let ptr = set.b.ensure(bytes);
300 let after = set.b.cap_bytes();
301 if after != before {
302 let delta = after as isize - before as isize;
303 let mut st = lock(&self.owner.state);
304 st.owned = add_signed(st.owned, delta, "retained bytes");
305 st.leased = add_signed(st.leased, delta, "leased bytes");
306 st.panel = add_signed(st.panel, delta, "panel bytes");
307 self.billed = add_signed(self.billed, delta, "billed bytes");
308 }
309 ptr
310 }
311
312 /// The team's block-scatter buffer, sized by the driver each call.
313 pub fn scatter_mut(&mut self) -> &mut Vec<i64> {
314 &mut self.set.as_mut().expect("team lease holds its set").scatter
315 }
316
317 /// The team's block-scatter buffer.
318 pub fn scatter(&self) -> &[i64] {
319 &self.set.as_ref().expect("team lease holds its set").scatter
320 }
321
322 /// One barrier per column group, empty when nothing is shared.
323 pub fn barriers(&self) -> &[Barrier] {
324 &self
325 .set
326 .as_ref()
327 .expect("team lease holds its set")
328 .barriers
329 }
330}
331
332impl Drop for TeamLease<'_> {
333 fn drop(&mut self) {
334 if let Some(set) = self.set.take() {
335 let mut st = lock(&self.owner.state);
336 // The scatter vector may have been resized through `scatter_mut`
337 // without going through `panel`, so settle both directions.
338 let delta = set.bytes() as isize - self.billed as isize;
339 st.owned = add_signed(st.owned, delta, "retained bytes");
340 st.leased = sub(st.leased, self.billed, "leased bytes");
341 st.teams.push(set);
342 }
343 }
344}
345
346mod sealed {
347 /// Only the crate's own owner implements [`WorkspaceProvider`](super::WorkspaceProvider):
348 /// the driver writes through the raw pointers it hands out, so an
349 /// implementation must honour the arena's buffer and exclusivity contracts.
350 pub trait Sealed {}
351 impl Sealed for super::ArenaProvider {}
352}
353
354/// Storage an execution may borrow. Implemented by the owner; the driver only
355/// ever sees this trait, so nothing here depends on the driver.
356///
357/// Sealed: [`ArenaProvider`] is the supported way to lend reusable storage, and
358/// the trait is deliberately not implementable outside this crate.
359pub trait WorkspaceProvider: sealed::Sealed + Sync {
360 /// Run `f` with this thread's A block, tile and scatter scratch. A
361 /// re-entrant call on the same thread gets fresh buffers instead of the
362 /// ones its outer call is using.
363 fn with_worker(&self, req: &WorkspaceReq, f: &mut dyn FnMut(*mut u8, *mut u8, &mut Vec<i64>));
364
365 /// Take the team set for `pm x pn`, exclusively, until the lease drops.
366 fn take_team(&self, req: &WorkspaceReq, pm: usize, pn: usize) -> TeamLease<'_>;
367
368 /// Release idle storage. Live leases and checked-out worker slots are kept.
369 fn trim(&self);
370}
371
372/// What one owner holds and lends. One lock covers the storage maps and the
373/// counters, so a snapshot and an update share one synchronization boundary.
374#[derive(Debug, Default)]
375struct State {
376 /// The worker slot each thread uses with this owner, created on that
377 /// thread's first call and kept until the owner is dropped.
378 threads: HashMap<std::thread::ThreadId, u64>,
379 workers: HashMap<u64, WorkerSlot>,
380 teams: Vec<TeamSet>,
381 /// Capacity of everything owned, idle or in use.
382 owned: usize,
383 /// The part of `owned` lent to a call right now.
384 leased: usize,
385 /// Capacity of the team panels alone, idle or leased. Independent of which
386 /// workers happened to take part, so a test can compare two runs of one
387 /// shape.
388 panel: usize,
389}
390
391/// A workspace owner: per-thread worker slots plus a free list of team sets.
392///
393/// One per owner — never process-global — so a lease always returns to the
394/// provider that issued it and no second owner can borrow it.
395#[derive(Debug)]
396pub struct ArenaProvider {
397 /// Process-unique and never recycled, so a trace or a diagnostic can never
398 /// be confused with another provider that happens to share an address.
399 key: u64,
400 state: Mutex<State>,
401 next_worker: AtomicU64,
402 traced: u64,
403}
404
405impl Default for ArenaProvider {
406 fn default() -> Self {
407 Self::new()
408 }
409}
410
411static NEXT_PROVIDER: AtomicU64 = AtomicU64::new(1);
412static TRACE: Mutex<Vec<(u64, std::thread::ThreadId, usize)>> = Mutex::new(Vec::new());
413
414fn lock<T>(m: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
415 m.lock().unwrap_or_else(|e| e.into_inner())
416}
417
418/// `value + delta`, rejecting a decrement below zero. Every caller establishes
419/// the bound; this turns a mistake into a loud failure instead of a wrap.
420fn add_signed(value: usize, delta: isize, what: &str) -> usize {
421 value
422 .checked_add_signed(delta)
423 .unwrap_or_else(|| panic!("workspace {what} went out of range"))
424}
425
426/// `value - amount`, rejecting an underflow for the same reason.
427fn sub(value: usize, amount: usize, what: &str) -> usize {
428 value
429 .checked_sub(amount)
430 .unwrap_or_else(|| panic!("workspace {what} went out of range"))
431}
432
433fn record(owner: u64, bytes: usize) {
434 lock(&TRACE).push((owner, std::thread::current().id(), bytes));
435}
436
437impl ArenaProvider {
438 /// A fresh owner with no storage.
439 pub fn new() -> Self {
440 Self {
441 key: NEXT_PROVIDER.fetch_add(1, Ordering::Relaxed),
442 state: Mutex::new(State::default()),
443 next_worker: AtomicU64::new(0),
444 traced: 0,
445 }
446 }
447
448 /// A fresh owner that records where its storage is first touched, readable
449 /// through [`ArenaProvider::trace_take`].
450 pub fn traced() -> Self {
451 let mut arena = Self::new();
452 arena.traced = arena.key;
453 arena
454 }
455
456 /// Take this owner's recorded buffer growths, leaving its log empty.
457 ///
458 /// Empty unless the owner was built with [`ArenaProvider::traced`].
459 pub fn trace_take(&self) -> Vec<(std::thread::ThreadId, usize)> {
460 let mut mine = Vec::new();
461 lock(&TRACE).retain(|(owner, thread, bytes)| {
462 if *owner == self.key {
463 mine.push((*thread, *bytes));
464 false
465 } else {
466 true
467 }
468 });
469 mine
470 }
471
472 /// Bytes currently retained and leased, for accounting by the host.
473 pub fn stats(&self) -> WorkspaceStats {
474 let st = lock(&self.state);
475 WorkspaceStats {
476 retained_bytes: st.owned,
477 leased_bytes: st.leased,
478 }
479 }
480
481 /// Bytes currently retained, for accounting in tests and diagnostics.
482 pub fn retained_bytes(&self) -> usize {
483 self.stats().retained_bytes
484 }
485
486 /// Panel bytes retained, idle or leased. Unlike
487 /// [`ArenaProvider::retained_bytes`] this does not depend on which workers
488 /// happened to take part, so it is what a test can compare between two runs
489 /// of one shape.
490 pub fn retained_panel_bytes(&self) -> usize {
491 lock(&self.state).panel
492 }
493
494 /// Check this thread's slot out of the map, or `None` when the thread
495 /// already holds it (a re-entrant call).
496 fn checkout(&self) -> Option<(u64, WorkerSlot, usize)> {
497 let thread = std::thread::current().id();
498 let mut st = lock(&self.state);
499 let index = match st.threads.get(&thread) {
500 Some(&index) => index,
501 None => {
502 let index = self.next_worker.fetch_add(1, Ordering::Relaxed);
503 st.threads.insert(thread, index);
504 // A fresh slot owns nothing yet, so the counters do not move.
505 return Some((index, WorkerSlot::new(self.traced), 0));
506 }
507 };
508 // The handle exists; a missing slot means an outer call holds it.
509 let slot = st.workers.remove(&index)?;
510 let billed = slot.bytes();
511 st.leased += billed;
512 Some((index, slot, billed))
513 }
514
515 /// Return a checked-out slot, settling both accounting directions.
516 fn give_back(&self, index: u64, slot: WorkerSlot, billed: usize) {
517 let delta = slot.bytes() as isize - billed as isize;
518 let mut st = lock(&self.state);
519 // `owned` follows the slot's capacity; `leased` gives back exactly what
520 // checkout added, since growth was never lent to anyone else.
521 st.owned = add_signed(st.owned, delta, "retained bytes");
522 st.leased = sub(st.leased, billed, "leased bytes");
523 st.workers.insert(index, slot);
524 }
525}
526
527impl Drop for ArenaProvider {
528 fn drop(&mut self) {
529 // Trace records outlive the provider otherwise: `traced()` appends to a
530 // process-global log that only `trace_take()` prunes, so a provider
531 // that is dropped without a `trace_take()` would keep its records for
532 // the life of the process.
533 if self.traced != 0 {
534 lock(&TRACE).retain(|(owner, _, _)| *owner != self.key);
535 }
536 }
537}
538
539/// Returns a checked-out slot even when the callback unwinds.
540struct SlotGuard<'a> {
541 owner: &'a ArenaProvider,
542 index: u64,
543 slot: Option<WorkerSlot>,
544 billed: usize,
545}
546
547impl Drop for SlotGuard<'_> {
548 fn drop(&mut self) {
549 if let Some(slot) = self.slot.take() {
550 self.owner.give_back(self.index, slot, self.billed);
551 }
552 }
553}
554
555impl WorkspaceProvider for ArenaProvider {
556 fn with_worker(&self, req: &WorkspaceReq, f: &mut dyn FnMut(*mut u8, *mut u8, &mut Vec<i64>)) {
557 let Some((index, slot, billed)) = self.checkout() else {
558 return fresh_worker(req, f);
559 };
560 // Armed before any growth, so a preparation unwind returns the slot too.
561 let mut guard = SlotGuard {
562 owner: self,
563 index,
564 slot: Some(slot),
565 billed,
566 };
567 let slot = guard.slot.as_mut().expect("guard holds its slot");
568 slot.a.ensure(req.a_bytes);
569 slot.tile.ensure(req.tile_bytes);
570 slot.scratch.clear();
571 slot.scratch.reserve(req.worker_scatter);
572 let (a, tile) = (slot.a.as_ptr(), slot.tile.as_ptr());
573 f(a, tile, &mut slot.scratch);
574 }
575
576 fn take_team(&self, req: &WorkspaceReq, pm: usize, pn: usize) -> TeamLease<'_> {
577 let (mut set, before, before_panel) = {
578 let mut st = lock(&self.state);
579 let set = st.teams.pop().unwrap_or_default();
580 let (bytes, panel) = (set.bytes(), set.b.cap_bytes());
581 (set, bytes, panel)
582 };
583 // `prepare` allocates, so it can panic. A panic drops `set` and its
584 // pages, and the owner has to stop counting the capacity it had when it
585 // was popped; the guard arms before `prepare` for that reason.
586 let mut teardown = Teardown {
587 owner: self,
588 before,
589 before_panel,
590 armed: true,
591 };
592 set.prepare(req, pm, pn);
593 let after = set.bytes();
594 teardown.armed = false;
595 let mut st = lock(&self.state);
596 st.owned = add_signed(st.owned, after as isize - before as isize, "retained bytes");
597 // `prepare` touches the scatter vector and the barriers, never the panel.
598 st.leased += after;
599 drop(st);
600 TeamLease {
601 owner: self,
602 set: Some(set),
603 billed: after,
604 }
605 }
606
607 fn trim(&self) {
608 let mut st = lock(&self.state);
609 let idle_slots: usize = st.workers.values().map(|s| s.bytes()).sum();
610 let idle_teams: usize = st.teams.iter().map(|t| t.bytes()).sum();
611 let idle_panels: usize = st.teams.iter().map(|t| t.b.cap_bytes()).sum();
612 for slot in st.workers.values_mut() {
613 slot.trim();
614 }
615 st.teams.clear();
616 st.owned = sub(st.owned, idle_slots + idle_teams, "retained bytes");
617 st.panel = sub(st.panel, idle_panels, "panel bytes");
618 }
619}
620
621/// Returns a popped team set's accounting when `prepare` unwinds.
622struct Teardown<'a> {
623 owner: &'a ArenaProvider,
624 before: usize,
625 before_panel: usize,
626 armed: bool,
627}
628
629impl Drop for Teardown<'_> {
630 fn drop(&mut self) {
631 if !self.armed {
632 return;
633 }
634 let mut st = lock(&self.owner.state);
635 st.owned = sub(st.owned, self.before, "retained bytes");
636 st.panel = sub(st.panel, self.before_panel, "panel bytes");
637 }
638}
639
640/// Buffers for one call that either re-entered its owner or has none. They are
641/// freed when the call returns, which is the price of re-entrancy, and they are
642/// not traced or accounted: the owner never holds them.
643fn fresh_worker(req: &WorkspaceReq, f: &mut dyn FnMut(*mut u8, *mut u8, &mut Vec<i64>)) {
644 let mut a = PageBuf::default();
645 let mut tile = PageBuf::default();
646 let mut scratch = Vec::new();
647 a.ensure(req.a_bytes);
648 tile.ensure(req.tile_bytes);
649 scratch.reserve(req.worker_scatter);
650 f(a.as_ptr(), tile.as_ptr(), &mut scratch);
651}
652
653/// A caller-owned, bounded workspace: an [`ArenaProvider`] pre-sized to one
654/// [`WorkspaceReq`], so a steady-state execute whose requirement is within the
655/// bound allocates nothing.
656///
657/// A plan computes a conservative requirement (`Plan::workspace_req` in the
658/// owning crate) and `execute_raw_in` checks it component-wise against
659/// [`Workspace::req`] before any write. The bound is grow-only and component-wise:
660/// [`Workspace::reserve`] keeps the maximum of each field, so a large A block
661/// does not stand in for a missing B panel. A parallel executor must call
662/// [`Workspace::prepare_for`] so every worker slot is warmed before the execute.
663///
664/// # Examples
665///
666/// ```
667/// use tcontract_exec::{Workspace, WorkspaceReq};
668/// let mut ws = Workspace::with_req(&WorkspaceReq { a_bytes: 64, ..Default::default() });
669/// assert!(ws.req().a_bytes >= 64);
670/// ws.reserve(&WorkspaceReq { b_bytes: 128, ..Default::default() });
671/// // Component-wise: the A block is kept and the B panel is added.
672/// assert!(ws.req().a_bytes >= 64 && ws.req().b_bytes >= 128);
673/// let _provider = ws.as_provider();
674/// ```
675#[derive(Debug, Default)]
676pub struct Workspace {
677 arena: ArenaProvider,
678 req: WorkspaceReq,
679 /// Worker slots warmed for an executor; see [`Workspace::prepare_for`].
680 slots: usize,
681}
682
683impl Workspace {
684 /// A workspace sized and warmed for `req`: the worker slot of the calling
685 /// thread and one team set are grown to hold it, so a later execute within
686 /// the bound does not grow the arena.
687 pub fn with_req(req: &WorkspaceReq) -> Self {
688 let mut ws = Self {
689 arena: ArenaProvider::new(),
690 req: *req,
691 slots: 0,
692 };
693 ws.warm();
694 ws.slots = 1;
695 ws
696 }
697
698 /// Grow the bound to cover `req` as well, component-wise. The bound never
699 /// shrinks, and every field keeps its own maximum.
700 pub fn reserve(&mut self, req: &WorkspaceReq) {
701 let before = self.req;
702 self.req.a_bytes = self.req.a_bytes.max(req.a_bytes);
703 self.req.tile_bytes = self.req.tile_bytes.max(req.tile_bytes);
704 self.req.worker_scatter = self.req.worker_scatter.max(req.worker_scatter);
705 self.req.b_bytes = self.req.b_bytes.max(req.b_bytes);
706 self.req.team_scatter = self.req.team_scatter.max(req.team_scatter);
707 self.req.barriers = self.req.barriers.max(req.barriers);
708 if self.req != before {
709 self.warm();
710 }
711 }
712
713 /// The bound this workspace was sized to, in bytes.
714 pub fn have(&self) -> usize {
715 self.req.bytes()
716 }
717
718 /// The per-component bound this workspace was sized to.
719 pub fn req(&self) -> &WorkspaceReq {
720 &self.req
721 }
722
723 /// How many worker slots have been warmed. [`Workspace::prepare_for`] raises
724 /// it to the executor's budget.
725 pub fn slots(&self) -> usize {
726 self.slots
727 }
728
729 /// Warm one worker slot per worker the executor may run on, so a parallel
730 /// execute within the bound does not allocate a slot per lane. `broadcast`
731 /// runs the closure on `budget` distinct workers, one call to `with_worker`
732 /// each; a context that cannot co-schedule (a serial one or a same-pool
733 /// worker) keeps the calling thread's slot only.
734 pub fn prepare_for(&mut self, exec: &crate::Exec<'_>) {
735 let slots = exec.budget().max(1);
736 if slots <= self.slots {
737 return;
738 }
739 let (arena, req) = (&self.arena, &self.req);
740 let warmed = exec.broadcast(slots, &|_| arena.with_worker(req, &mut |_, _, _| {}));
741 if warmed.is_ok() {
742 self.slots = slots;
743 }
744 }
745
746 /// The provider a plan lends to the driver for one call.
747 pub fn as_provider(&self) -> &dyn WorkspaceProvider {
748 &self.arena
749 }
750
751 /// Bytes the owner currently retains, idle or in use. A test observes it to
752 /// prove an execute did not grow the arena.
753 pub fn retained_bytes(&self) -> usize {
754 self.arena.retained_bytes()
755 }
756
757 /// First-touch the worker buffers (this thread's slot) and one team set so
758 /// the first execute within the bound allocates only the arena's bookkeeping.
759 fn warm(&self) {
760 self.arena.with_worker(&self.req, &mut |_, _, _| {});
761 let mut lease = self.arena.take_team(&self.req, 1, 1);
762 if self.req.b_bytes > 0 {
763 lease.panel(self.req.b_bytes);
764 }
765 drop(lease);
766 }
767}
768
769#[cfg(test)]
770mod thread_entries {
771 use super::*;
772
773 fn thread_slots(provider: &ArenaProvider) -> usize {
774 lock(&provider.state).threads.len()
775 }
776
777 fn trace_records(owner: u64) -> usize {
778 lock(&TRACE)
779 .iter()
780 .filter(|(key, _, _)| *key == owner)
781 .count()
782 }
783
784 #[test]
785 fn one_entry_per_thread_and_none_carried_over() {
786 let req = WorkspaceReq::default();
787 let provider = ArenaProvider::default();
788 assert_eq!(thread_slots(&provider), 0);
789 provider.with_worker(&req, &mut |_, _, _| {});
790 assert_eq!(thread_slots(&provider), 1);
791 provider.with_worker(&req, &mut |_, _, _| {});
792 assert_eq!(thread_slots(&provider), 1);
793
794 // The map belongs to the provider, so a new one never inherits an
795 // entry: nothing outlives the owner that created it.
796 let fresh = ArenaProvider::default();
797 assert_eq!(thread_slots(&fresh), 0);
798 fresh.with_worker(&req, &mut |_, _, _| {});
799 assert_eq!(thread_slots(&fresh), 1);
800 assert_eq!(thread_slots(&provider), 1);
801 }
802
803 #[test]
804 fn two_threads_get_two_entries() {
805 let provider = ArenaProvider::default();
806 let barrier = std::sync::Barrier::new(2);
807 std::thread::scope(|scope| {
808 for _ in 0..2 {
809 scope.spawn(|| {
810 provider.with_worker(&WorkspaceReq::default(), &mut |_, _, _| {});
811 barrier.wait();
812 });
813 }
814 });
815 assert_eq!(thread_slots(&provider), 2);
816 }
817
818 #[test]
819 fn dropping_a_traced_provider_drops_its_trace_records() {
820 let provider = ArenaProvider::traced();
821 let key = provider.key;
822 provider.with_worker(
823 &WorkspaceReq {
824 a_bytes: 4096,
825 ..Default::default()
826 },
827 &mut |_, _, _| {},
828 );
829 assert!(trace_records(key) > 0);
830 drop(provider);
831 assert_eq!(
832 trace_records(key),
833 0,
834 "a dropped traced provider must not keep records in the process-global log"
835 );
836 }
837}