Skip to main content

vyre_runtime/
tenant.rs

1//! Multi-tenant megakernel multiplexing.
2//!
3//! A single persistent megakernel per GPU can service many producer
4//! tools without each one paying the dispatch-setup cost. The
5//! `tenant_id` field already lives in the ring-slot protocol
6//! (`protocol::TENANT_WORD`); this module owns the host-side
7//! bookkeeping that hands each producer a stable id, reserves an
8//! opcode-range per producer, and gates publish operations against a
9//! per-tenant mask so one producer cannot accidentally drive another
10//! producer's opcodes.
11//!
12//! ## Tenants and opcodes
13//!
14//! Every tenant owns an opcode range `[base, base + cap)` where the
15//! whole range sits inside the user-extension space reserved by
16//! `vyre_runtime::megakernel::protocol::opcode` (≥ `0x4000_0000`).
17//! When [`TenantRegistry::register`] returns a [`TenantHandle`],
18//! callers publish into slot args `[rule_local_opcode, ...]` and
19//! the registry maps that to `(tenant_base + rule_local_opcode)`
20//! before writing into the ring. A tenant that tries to publish an
21//! opcode outside its own range fails with a structured error.
22//!
23//! ## Draining
24//!
25//! Unregistering a tenant revokes future publishes but does NOT
26//! revoke in-flight slots  -  the GPU is still going to execute any
27//! slot it already CAS-claimed. Callers that need hard draining
28//! drive [`TenantHandle::quiesce`] which spins on the megakernel
29//! DONE_COUNT until every slot the tenant published has been
30//! acknowledged.
31//!
32//! ## Daemon surface
33//!
34//! The registry is the reusable piece. A full `MegakernelDaemon`
35//! (listening on a Unix socket, vending handles over RPC) is a thin
36//! wrapper that we can ship alongside the runtime  -  the registry
37//! here already handles the interesting concurrency.
38
39use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
40use std::sync::Arc;
41use std::time::{Duration, Instant};
42
43use dashmap::DashMap;
44
45use crate::resident_work_queue::protocol::opcode::SHUTDOWN;
46use crate::resident_work_queue::ResidentWorkQueue;
47use crate::PipelineError;
48
49/// First opcode the tenant registry hands out. Sits inside the
50/// user-extension range reserved by the megakernel protocol so fused
51/// rule documents compose with tenant allocation without colliding
52/// with built-in opcodes.
53pub const TENANT_OPCODE_BASE: u32 = 0x4000_0000;
54
55/// Upper bound on the tenant-id space. `tenant_id == TENANT_ID_MAX`
56/// is reserved as an invalid / revoked sentinel.
57pub const TENANT_ID_MAX: u32 = u32::MAX - 1;
58
59/// Size of the opcode window reserved per tenant. 1 << 20 = 1 MiB
60/// of opcodes  -  well over any realistic rule count per producer
61/// while still allowing ~4094 simultaneous tenants inside the u32
62/// opcode range.
63pub const OPCODE_RANGE_PER_TENANT: u32 = 1 << 20;
64
65const QUIESCE_SPIN_POLLS: u64 = 64;
66const QUIESCE_MIN_PARK: Duration = Duration::from_micros(2);
67const QUIESCE_MAX_PARK: Duration = Duration::from_micros(50);
68const QUIESCE_BACKOFF_SHIFT_CAP: u64 = 5;
69
70#[allow(clippy::unnecessary_min_or_max)]
71fn quiesce_backoff_duration(poll: u64) -> Duration {
72    let parked_poll = poll.saturating_sub(QUIESCE_SPIN_POLLS);
73    let shift = parked_poll.min(QUIESCE_BACKOFF_SHIFT_CAP) as u32;
74    let multiplier = 1_u32 << shift;
75    QUIESCE_MIN_PARK
76        .checked_mul(multiplier)
77        .unwrap_or(QUIESCE_MAX_PARK)
78        .min(QUIESCE_MAX_PARK)
79}
80
81fn quiesce_idle(poll: u64) {
82    if poll < QUIESCE_SPIN_POLLS {
83        std::hint::spin_loop();
84    } else {
85        std::thread::park_timeout(quiesce_backoff_duration(poll));
86    }
87}
88
89fn tenant_registry_retry_idle(retry: u64) {
90    if retry < QUIESCE_SPIN_POLLS {
91        std::hint::spin_loop();
92    } else {
93        std::thread::park_timeout(quiesce_backoff_duration(retry));
94    }
95}
96
97/// Errors surfaced by the tenant registry.
98#[derive(Debug, thiserror::Error)]
99#[non_exhaustive]
100pub enum TenantError {
101    /// The registry ran out of tenant ids. Unregister unused tenants
102    /// or raise the range per tenant.
103    #[error("tenant registry exhausted after {issued} registrations. Fix: shrink OPCODE_RANGE_PER_TENANT or recycle tenants.")]
104    RegistryFull {
105        /// Number of tenants already issued when exhaustion hit.
106        issued: u32,
107    },
108    /// Tried to publish an opcode outside the tenant's reserved
109    /// range. Almost always a caller bug.
110    #[error(
111        "tenant {tenant_id} published local opcode {local_opcode}; out of range [0, {cap}). \
112         Fix: caller must stay inside the opcode window returned by `register()`."
113    )]
114    OpcodeOutOfRange {
115        /// Tenant id that tripped.
116        tenant_id: u32,
117        /// Local opcode the caller supplied.
118        local_opcode: u32,
119        /// Cap on the tenant's local opcode range.
120        cap: u32,
121    },
122    /// Tenant was unregistered concurrently; its handle is stale.
123    #[error("tenant {tenant_id} was revoked; handle is stale. Fix: acquire a fresh handle from the registry.")]
124    Revoked {
125        /// Tenant id that was revoked.
126        tenant_id: u32,
127    },
128    /// Quiesce timed out with inflight slots still outstanding.
129    #[error(
130        "tenant {tenant_id} quiesce timed out with {outstanding} inflight slots. \
131         Fix: ensure the megakernel is making progress (check DONE_COUNT) or raise the timeout."
132    )]
133    QuiesceTimeout {
134        /// Tenant id whose quiesce tripped.
135        tenant_id: u32,
136        /// Number of slots still inflight at timeout.
137        outstanding: u64,
138    },
139    /// Tenant has reached its configured outstanding-slot cap.
140    #[error(
141        "tenant {tenant_id} has {outstanding} outstanding slots, cap {cap}. \
142         Fix: wait for drain progress or register the tenant with a larger bounded backlog."
143    )]
144    Backpressure {
145        /// Tenant id whose backlog is full.
146        tenant_id: u32,
147        /// Current host-visible outstanding slots.
148        outstanding: u64,
149        /// Configured outstanding-slot cap.
150        cap: u64,
151    },
152    /// Tenant has reached its configured staging-byte cap.
153    #[error(
154        "tenant {tenant_id} requested {requested} staging bytes with {used} already reserved, cap {cap}. \
155         Fix: release staging reservations after publish/readback progress or register the tenant with a larger bounded staging budget."
156    )]
157    StagingBackpressure {
158        /// Tenant id whose staging byte budget is full.
159        tenant_id: u32,
160        /// New bytes requested.
161        requested: u64,
162        /// Current reserved staging bytes.
163        used: u64,
164        /// Configured staging byte cap.
165        cap: u64,
166    },
167    /// Tenant has reached its configured resident-handle cap.
168    #[error(
169        "tenant {tenant_id} requested {requested} resident handles with {used} already reserved, cap {cap}. \
170         Fix: release resident handles when backend ownership ends or register the tenant with a larger bounded resident-handle budget."
171    )]
172    ResidentHandleBackpressure {
173        /// Tenant id whose resident handle budget is full.
174        tenant_id: u32,
175        /// New handles requested.
176        requested: u64,
177        /// Current reserved resident handles.
178        used: u64,
179        /// Configured resident handle cap.
180        cap: u64,
181    },
182    /// Tenant resource accounting would underflow.
183    #[error(
184        "tenant {tenant_id} released {requested} {resource} with only {used} reserved. \
185         Fix: pair every tenant resource release with a successful reservation."
186    )]
187    ResourceUnderflow {
188        /// Tenant id whose counter would underflow.
189        tenant_id: u32,
190        /// Resource counter being released.
191        resource: &'static str,
192        /// Release count requested.
193        requested: u64,
194        /// Current reserved count.
195        used: u64,
196    },
197    /// Protocol error bubbled up from [`ResidentWorkQueue::publish_slot`].
198    #[error("{0}")]
199    Pipeline(#[from] PipelineError),
200}
201
202/// Per-tenant resource quota.
203#[derive(Debug, Clone, Copy, PartialEq, Eq)]
204pub struct TenantQuota {
205    /// Maximum host-visible ring slots the tenant may keep outstanding.
206    pub max_outstanding_slots: u64,
207    /// Maximum staging bytes the tenant may reserve for pending work.
208    pub max_staging_bytes: u64,
209    /// Maximum resident handles the tenant may hold at once.
210    pub max_resident_handles: u64,
211}
212
213impl TenantQuota {
214    /// Unbounded tenant quota for compatibility with the legacy registration
215    /// API. Individual fields are still normalized to at least one resource
216    /// slot during registration.
217    #[must_use]
218    pub const fn unbounded() -> Self {
219        Self {
220            max_outstanding_slots: u64::MAX,
221            max_staging_bytes: u64::MAX,
222            max_resident_handles: u64::MAX,
223        }
224    }
225
226    /// Build a bounded tenant quota.
227    #[must_use]
228    pub const fn bounded(
229        max_outstanding_slots: u64,
230        max_staging_bytes: u64,
231        max_resident_handles: u64,
232    ) -> Self {
233        Self {
234            max_outstanding_slots,
235            max_staging_bytes,
236            max_resident_handles,
237        }
238    }
239}
240
241/// One tenant's accounting state. Lives inside an `Arc` so handles
242/// stay valid after the registry borrow drops.
243struct TenantState {
244    id: u32,
245    base_opcode: u32,
246    opcode_cap: u32,
247    /// Number of slots this tenant has ever published.
248    published_count: AtomicU64,
249    /// Maximum host-visible slots this tenant may keep outstanding.
250    max_outstanding_slots: u64,
251    /// Number of staging bytes currently reserved by this tenant.
252    staging_bytes: AtomicU64,
253    /// Maximum staging bytes this tenant may reserve.
254    max_staging_bytes: u64,
255    /// Number of resident handles currently reserved by this tenant.
256    resident_handles: AtomicU64,
257    /// Maximum resident handles this tenant may reserve.
258    max_resident_handles: u64,
259    /// Number of slots the GPU has reported DONE for this tenant.
260    /// Advanced by [`TenantHandle::note_drained`].
261    drained_count: AtomicU64,
262    /// Number of quiesce calls completed or timed out for this tenant.
263    quiesce_calls: AtomicU64,
264    /// Number of quiesce calls that timed out before the tenant drained.
265    quiesce_timeouts: AtomicU64,
266    /// Cumulative host-observed drain wait across quiesce calls.
267    quiesce_wait_ns: AtomicU64,
268    /// Set to 1 on `unregister`; publishes reject afterwards.
269    revoked: AtomicU32,
270    /// Stable label for diagnostics (for example, `"scanner-a"`, `"scanner-b"`).
271    label: String,
272}
273
274/// Stable handle returned by [`TenantRegistry::register`]. Clones
275/// share the same underlying state, so multiple producer threads
276/// inside one tenant can publish through their own handles.
277#[derive(Clone)]
278pub struct TenantHandle {
279    state: Arc<TenantState>,
280}
281
282/// Host-visible tenant runtime counters.
283#[derive(Debug, Clone, Copy, PartialEq, Eq)]
284pub struct TenantRuntimeCounters {
285    /// Tenant id.
286    pub tenant_id: u32,
287    /// Number of slots ever published by this tenant.
288    pub published_count: u64,
289    /// Number of slots observed drained for this tenant.
290    pub drained_count: u64,
291    /// Current host-visible backlog (`published_count - drained_count`).
292    pub outstanding_slots: u64,
293    /// Configured outstanding-slot cap for this tenant.
294    pub max_outstanding_slots: u64,
295    /// Number of quiesce calls recorded for this tenant.
296    pub quiesce_calls: u64,
297    /// Number of quiesce calls that timed out.
298    pub quiesce_timeouts: u64,
299    /// Cumulative nanoseconds spent waiting for this tenant to drain.
300    pub quiesce_wait_ns: u64,
301}
302
303/// Host-visible tenant quota counters.
304#[derive(Debug, Clone, Copy, PartialEq, Eq)]
305pub struct TenantQuotaCounters {
306    /// Tenant id.
307    pub tenant_id: u32,
308    /// Current reserved staging bytes.
309    pub staging_bytes: u64,
310    /// Configured staging byte cap.
311    pub max_staging_bytes: u64,
312    /// Current reserved resident handle count.
313    pub resident_handles: u64,
314    /// Configured resident handle cap.
315    pub max_resident_handles: u64,
316}
317
318impl std::fmt::Debug for TenantHandle {
319    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
320        f.debug_struct("TenantHandle")
321            .field("id", &self.state.id)
322            .field("label", &self.state.label)
323            .field("base_opcode", &self.state.base_opcode)
324            .field(
325                "published_count",
326                &self.state.published_count.load(Ordering::Relaxed),
327            )
328            .field("max_outstanding_slots", &self.state.max_outstanding_slots)
329            .field(
330                "staging_bytes",
331                &self.state.staging_bytes.load(Ordering::Relaxed),
332            )
333            .field("max_staging_bytes", &self.state.max_staging_bytes)
334            .field(
335                "resident_handles",
336                &self.state.resident_handles.load(Ordering::Relaxed),
337            )
338            .field("max_resident_handles", &self.state.max_resident_handles)
339            .field(
340                "drained_count",
341                &self.state.drained_count.load(Ordering::Relaxed),
342            )
343            .field(
344                "revoked",
345                &(self.state.revoked.load(Ordering::Acquire) != 0),
346            )
347            .finish()
348    }
349}
350
351impl TenantHandle {
352    /// Stable tenant id; maps onto the ring-slot `TENANT_WORD`.
353    #[must_use]
354    pub fn id(&self) -> u32 {
355        self.state.id
356    }
357
358    /// Human-readable label supplied at registration time.
359    #[must_use]
360    pub fn label(&self) -> &str {
361        &self.state.label
362    }
363
364    /// First opcode this tenant owns.
365    #[must_use]
366    pub fn base_opcode(&self) -> u32 {
367        self.state.base_opcode
368    }
369
370    /// Convert a tenant-local opcode to the global opcode used in
371    /// the ring slot. Caller enforces `local < opcode_cap()`.
372    ///
373    /// # Errors
374    ///
375    /// Returns [`TenantError::OpcodeOutOfRange`] when the local
376    /// value is outside the reserved window.
377    pub fn global_opcode(&self, local: u32) -> Result<u32, TenantError> {
378        self.ensure_not_revoked()?;
379        if local >= self.state.opcode_cap {
380            return Err(TenantError::OpcodeOutOfRange {
381                tenant_id: self.id(),
382                local_opcode: local,
383                cap: self.state.opcode_cap,
384            });
385        }
386        let global = self.state.base_opcode + local;
387        if let Err(e) = crate::resident_work_queue::protocol::opcode::validate_user_opcode(global) {
388            return Err(TenantError::Pipeline(PipelineError::Backend(format!(
389                "tenant registry produced invalid global opcode {global}: {e}. Fix: repair tenant opcode window allocation before publishing."
390            ))));
391        }
392        Ok(global)
393    }
394
395    /// Publish a tenant-local opcode through [`ResidentWorkQueue::publish_slot`].
396    ///
397    /// # Errors
398    ///
399    /// - [`TenantError::Revoked`] if the tenant was unregistered.
400    /// - [`TenantError::OpcodeOutOfRange`] if `local_opcode` is
401    ///   outside the tenant's window.
402    /// - [`TenantError::Pipeline`] when the underlying
403    ///   `publish_slot` rejects (e.g., slot still in-flight).
404    pub fn publish_slot(
405        &self,
406        ring_bytes: &mut [u8],
407        slot_idx: u32,
408        local_opcode: u32,
409        args: &[u32],
410    ) -> Result<(), TenantError> {
411        self.ensure_not_revoked()?;
412        let global = self.global_opcode(local_opcode)?;
413        self.reserve_publish_slot()?;
414        if let Err(error) =
415            ResidentWorkQueue::publish_slot(ring_bytes, slot_idx, self.state.id, global, args)
416        {
417            saturating_atomic_sub_u64(&self.state.published_count, 1, "tenant published rollback");
418            return Err(error.into());
419        }
420        Ok(())
421    }
422
423    fn ensure_not_revoked(&self) -> Result<(), TenantError> {
424        if self.state.revoked.load(Ordering::Acquire) != 0 {
425            return Err(TenantError::Revoked {
426                tenant_id: self.state.id,
427            });
428        }
429        Ok(())
430    }
431
432    fn reserve_publish_slot(&self) -> Result<(), TenantError> {
433        let cap = self.state.max_outstanding_slots;
434        vyre_driver::accounting::checked_atomic_update_u64_with_order(
435            &self.state.published_count,
436            Ordering::Acquire,
437            Ordering::AcqRel,
438            Ordering::Acquire,
439            |published| {
440                let drained = self.state.drained_count.load(Ordering::Acquire);
441                let outstanding = vyre_driver::accounting::checked_sub_u64_lazy(
442                    published,
443                    drained,
444                    || {
445                        TenantError::Pipeline(PipelineError::QueueFull {
446                            queue: "tenant",
447                            fix: "tenant drained_count exceeded published_count; rebuild tenant accounting state",
448                        })
449                    },
450                )?;
451                if outstanding >= cap {
452                    return Err(TenantError::Backpressure {
453                        tenant_id: self.state.id,
454                        outstanding,
455                        cap,
456                    });
457                }
458                vyre_driver::accounting::checked_add_u64_lazy(published, 1, || {
459                    TenantError::Pipeline(PipelineError::QueueFull {
460                        queue: "tenant",
461                        fix: "tenant published_count overflowed u64; quiesce or recreate the tenant before publishing more slots",
462                    })
463                })
464            },
465            |_, _| Ok(()),
466        )?;
467        Ok(())
468    }
469
470    /// Number of slots this tenant has ever published.
471    #[must_use]
472    pub fn published_count(&self) -> u64 {
473        self.state.published_count.load(Ordering::Relaxed)
474    }
475
476    /// Number of slots this tenant has observed drained (via
477    /// [`note_drained`](Self::note_drained)).
478    #[must_use]
479    pub fn drained_count(&self) -> u64 {
480        self.state.drained_count.load(Ordering::Relaxed)
481    }
482
483    /// Maximum host-visible slots this tenant may keep outstanding.
484    #[must_use]
485    pub fn max_outstanding_slots(&self) -> u64 {
486        self.state.max_outstanding_slots
487    }
488
489    /// Reserve staging bytes against this tenant's quota.
490    pub fn reserve_staging_bytes(&self, byte_count: u64) -> Result<(), TenantError> {
491        self.ensure_not_revoked()?;
492        reserve_resource_quota(
493            &self.state.staging_bytes,
494            byte_count,
495            self.state.max_staging_bytes,
496            || {
497                TenantError::StagingBackpressure {
498                    tenant_id: self.state.id,
499                    requested: byte_count,
500                    used: self.state.staging_bytes.load(Ordering::Acquire),
501                    cap: self.state.max_staging_bytes,
502                }
503            },
504            "tenant staging byte reservation overflowed u64; release staging reservations or recreate the tenant before reserving more bytes",
505        )
506    }
507
508    /// Release staging bytes previously reserved by this tenant.
509    pub fn release_staging_bytes(&self, byte_count: u64) -> Result<(), TenantError> {
510        release_resource_quota(
511            &self.state.staging_bytes,
512            byte_count,
513            self.state.id,
514            "staging bytes",
515        )
516    }
517
518    /// Reserve resident handles against this tenant's quota.
519    pub fn reserve_resident_handles(&self, handle_count: u64) -> Result<(), TenantError> {
520        self.ensure_not_revoked()?;
521        reserve_resource_quota(
522            &self.state.resident_handles,
523            handle_count,
524            self.state.max_resident_handles,
525            || {
526                TenantError::ResidentHandleBackpressure {
527                    tenant_id: self.state.id,
528                    requested: handle_count,
529                    used: self.state.resident_handles.load(Ordering::Acquire),
530                    cap: self.state.max_resident_handles,
531                }
532            },
533            "tenant resident handle reservation overflowed u64; release resident handles or recreate the tenant before reserving more handles",
534        )
535    }
536
537    /// Release resident handles previously reserved by this tenant.
538    pub fn release_resident_handles(&self, handle_count: u64) -> Result<(), TenantError> {
539        release_resource_quota(
540            &self.state.resident_handles,
541            handle_count,
542            self.state.id,
543            "resident handles",
544        )
545    }
546
547    /// Snapshot quota counters for this tenant.
548    #[must_use]
549    pub fn quota_counters(&self) -> TenantQuotaCounters {
550        TenantQuotaCounters {
551            tenant_id: self.state.id,
552            staging_bytes: self.state.staging_bytes.load(Ordering::Acquire),
553            max_staging_bytes: self.state.max_staging_bytes,
554            resident_handles: self.state.resident_handles.load(Ordering::Acquire),
555            max_resident_handles: self.state.max_resident_handles,
556        }
557    }
558
559    fn release_all_resource_reservations(&self) {
560        self.state.staging_bytes.store(0, Ordering::Release);
561        self.state.resident_handles.store(0, Ordering::Release);
562    }
563
564    /// Snapshot host-visible runtime counters for this tenant.
565    #[must_use]
566    pub fn runtime_counters(&self) -> TenantRuntimeCounters {
567        let published_count = self.state.published_count.load(Ordering::Acquire);
568        let drained_count = self.state.drained_count.load(Ordering::Acquire);
569        TenantRuntimeCounters {
570            tenant_id: self.state.id,
571            published_count,
572            drained_count,
573            outstanding_slots: published_count.saturating_sub(drained_count),
574            max_outstanding_slots: self.state.max_outstanding_slots,
575            quiesce_calls: self.state.quiesce_calls.load(Ordering::Acquire),
576            quiesce_timeouts: self.state.quiesce_timeouts.load(Ordering::Acquire),
577            quiesce_wait_ns: self.state.quiesce_wait_ns.load(Ordering::Acquire),
578        }
579    }
580
581    /// Mark `count` slots as drained. The host pump that observes
582    /// DONE_COUNT calls this when it sees the global counter
583    /// advance past the tenant's last-published cursor.
584    pub fn note_drained(&self, count: u64) {
585        saturating_atomic_add_u64(&self.state.drained_count, count, "tenant drained_count");
586    }
587
588    /// Block-style quiesce: bounded backoff until every published
589    /// slot has been drained or `max_spins` polls elapse.
590    ///
591    /// # Errors
592    ///
593    /// Returns [`TenantError::QuiesceTimeout`] when `max_spins`
594    /// iterations pass without full drain. The outstanding count
595    /// at timeout is included for diagnostics.
596    pub fn quiesce(&self, max_spins: u64) -> Result<(), TenantError> {
597        let started = Instant::now();
598        for poll in 0..max_spins {
599            let pub_count = self.state.published_count.load(Ordering::Acquire);
600            let drained = self.state.drained_count.load(Ordering::Acquire);
601            if drained >= pub_count {
602                self.record_quiesce(started, false);
603                return Ok(());
604            }
605            quiesce_idle(poll);
606        }
607        let pub_count = self.state.published_count.load(Ordering::Acquire);
608        let drained = self.state.drained_count.load(Ordering::Acquire);
609        self.record_quiesce(started, true);
610        Err(TenantError::QuiesceTimeout {
611            tenant_id: self.state.id,
612            outstanding: vyre_driver::accounting::checked_sub_u64_lazy(pub_count, drained, || {
613                TenantError::Pipeline(PipelineError::QueueFull {
614                    queue: "tenant",
615                    fix: "tenant drained_count exceeded published_count during quiesce; rebuild tenant accounting state",
616                })
617            })?,
618        })
619    }
620
621    fn record_quiesce(&self, started: Instant, timed_out: bool) {
622        saturating_atomic_add_u64(&self.state.quiesce_calls, 1, "tenant quiesce_calls");
623        if timed_out {
624            saturating_atomic_add_u64(&self.state.quiesce_timeouts, 1, "tenant quiesce_timeouts");
625        }
626        let elapsed_ns = match u64::try_from(started.elapsed().as_nanos()) {
627            Ok(elapsed_ns) => elapsed_ns,
628            Err(_) => u64::MAX,
629        };
630        saturating_atomic_add_u64(
631            &self.state.quiesce_wait_ns,
632            elapsed_ns,
633            "tenant quiesce_wait_ns",
634        );
635    }
636}
637
638/// Thread-safe tenant registry. One per megakernel instance.
639pub struct TenantRegistry {
640    tenants: DashMap<u32, TenantHandle>,
641    next_id: AtomicU32,
642}
643
644impl Default for TenantRegistry {
645    fn default() -> Self {
646        Self {
647            tenants: DashMap::new(),
648            next_id: AtomicU32::new(0),
649        }
650    }
651}
652
653/// Caller-owned scratch for repeated concurrent-tenant selection.
654#[derive(Debug, Default)]
655pub struct TenantSelectionScratch {
656    active_ids: Vec<u32>,
657    selected_indices: Vec<usize>,
658}
659
660impl TenantSelectionScratch {
661    /// Construct empty tenant-selection scratch.
662    #[must_use]
663    pub const fn new() -> Self {
664        Self {
665            active_ids: Vec::new(),
666            selected_indices: Vec::new(),
667        }
668    }
669}
670
671fn saturating_atomic_add_u64(counter: &AtomicU64, value: u64, _label: &'static str) {
672    let mut current = counter.load(Ordering::Acquire);
673    loop {
674        let next = current.saturating_add(value);
675        match counter.compare_exchange_weak(current, next, Ordering::AcqRel, Ordering::Acquire) {
676            Ok(_) => return,
677            Err(observed) => current = observed,
678        }
679    }
680}
681
682fn saturating_atomic_sub_u64(counter: &AtomicU64, value: u64, _label: &'static str) {
683    let mut current = counter.load(Ordering::Acquire);
684    loop {
685        let next = current.saturating_sub(value);
686        match counter.compare_exchange_weak(current, next, Ordering::AcqRel, Ordering::Acquire) {
687            Ok(_) => return,
688            Err(observed) => current = observed,
689        }
690    }
691}
692
693fn reserve_resource_quota(
694    counter: &AtomicU64,
695    value: u64,
696    cap: u64,
697    backpressure: impl Fn() -> TenantError,
698    overflow_fix: &'static str,
699) -> Result<(), TenantError> {
700    vyre_driver::accounting::checked_atomic_update_u64_with_order(
701        counter,
702        Ordering::Acquire,
703        Ordering::AcqRel,
704        Ordering::Acquire,
705        |used| {
706            let next = vyre_driver::accounting::checked_add_u64_lazy(used, value, || {
707                TenantError::Pipeline(PipelineError::QueueFull {
708                    queue: "tenant resource quota",
709                    fix: overflow_fix,
710                })
711            })?;
712            if next > cap {
713                return Err(backpressure());
714            }
715            Ok(next)
716        },
717        |_, _| Ok(()),
718    )?;
719    Ok(())
720}
721
722fn release_resource_quota(
723    counter: &AtomicU64,
724    value: u64,
725    tenant_id: u32,
726    resource: &'static str,
727) -> Result<(), TenantError> {
728    vyre_driver::accounting::checked_atomic_update_u64_with_order(
729        counter,
730        Ordering::Acquire,
731        Ordering::AcqRel,
732        Ordering::Acquire,
733        |used| {
734            used.checked_sub(value)
735                .ok_or(TenantError::ResourceUnderflow {
736                    tenant_id,
737                    resource,
738                    requested: value,
739                    used,
740                })
741        },
742        |_, _| Ok(()),
743    )?;
744    Ok(())
745}
746
747impl TenantRegistry {
748    /// Fresh registry with no tenants.
749    #[must_use]
750    pub fn new() -> Self {
751        Self::default()
752    }
753
754    /// Register a new tenant with the given diagnostic label.
755    /// Returns a handle whose opcode range is reserved until
756    /// [`unregister`](Self::unregister) is called.
757    ///
758    /// # Errors
759    ///
760    /// Returns [`TenantError::RegistryFull`] when the tenant id or
761    /// opcode space is exhausted.
762    pub fn register(&self, label: impl Into<String>) -> Result<TenantHandle, TenantError> {
763        self.register_with_backpressure(label, u64::MAX)
764    }
765
766    /// Register a new tenant with a bounded outstanding-slot budget.
767    ///
768    /// # Errors
769    ///
770    /// Returns [`TenantError::RegistryFull`] when the tenant id or opcode space
771    /// is exhausted.
772    pub fn register_with_backpressure(
773        &self,
774        label: impl Into<String>,
775        max_outstanding_slots: u64,
776    ) -> Result<TenantHandle, TenantError> {
777        self.register_with_quotas(
778            label,
779            TenantQuota {
780                max_outstanding_slots,
781                ..TenantQuota::unbounded()
782            },
783        )
784    }
785
786    /// Register a tenant with explicit ring-slot, staging-byte, and
787    /// resident-handle quotas.
788    ///
789    /// # Errors
790    ///
791    /// Returns [`TenantError::RegistryFull`] when the tenant id or opcode space
792    /// is exhausted.
793    pub fn register_with_quotas(
794        &self,
795        label: impl Into<String>,
796        quota: TenantQuota,
797    ) -> Result<TenantHandle, TenantError> {
798        let mut registration_retries = 0u64;
799        let issued = vyre_driver::accounting::checked_atomic_update_u32_with_order(
800            &self.next_id,
801            Ordering::Relaxed,
802            Ordering::SeqCst,
803            Ordering::Relaxed,
804            |current| {
805                if current >= TENANT_ID_MAX {
806                    return Err(TenantError::RegistryFull { issued: current });
807                }
808                let id = current.max(1);
809                id.checked_add(1)
810                    .ok_or(TenantError::RegistryFull { issued: current })
811            },
812            |_, _| {
813                tenant_registry_retry_idle(registration_retries);
814                registration_retries = vyre_driver::accounting::checked_add_u64_lazy(
815                    registration_retries,
816                    1,
817                    || {
818                        TenantError::Pipeline(PipelineError::QueueFull {
819                            queue: "tenant",
820                            fix: "tenant registration retry counter overflowed u64; retry registration later",
821                        })
822                    },
823                )?;
824                Ok(())
825            },
826        )?;
827        let id = issued.max(1);
828
829        let tenant_offset = vyre_driver::accounting::checked_mul_u32_value(
830            id,
831            OPCODE_RANGE_PER_TENANT,
832            TenantError::RegistryFull { issued },
833        )?;
834        let base_opcode = vyre_driver::accounting::checked_add_u32_value(
835            TENANT_OPCODE_BASE,
836            tenant_offset,
837            TenantError::RegistryFull { issued },
838        )?;
839        let top_opcode = vyre_driver::accounting::checked_add_u32_value(
840            base_opcode,
841            OPCODE_RANGE_PER_TENANT,
842            TenantError::RegistryFull { issued },
843        )?;
844        if top_opcode == SHUTDOWN {
845            return Err(TenantError::RegistryFull { issued });
846        }
847        let handle = TenantHandle {
848            state: Arc::new(TenantState {
849                id,
850                base_opcode,
851                opcode_cap: OPCODE_RANGE_PER_TENANT,
852                published_count: AtomicU64::new(0),
853                max_outstanding_slots: quota.max_outstanding_slots.max(1),
854                staging_bytes: AtomicU64::new(0),
855                max_staging_bytes: quota.max_staging_bytes.max(1),
856                resident_handles: AtomicU64::new(0),
857                max_resident_handles: quota.max_resident_handles.max(1),
858                drained_count: AtomicU64::new(0),
859                quiesce_calls: AtomicU64::new(0),
860                quiesce_timeouts: AtomicU64::new(0),
861                quiesce_wait_ns: AtomicU64::new(0),
862                revoked: AtomicU32::new(0),
863                label: label.into(),
864            }),
865        };
866        self.tenants.insert(id, handle.clone());
867        Ok(handle)
868    }
869
870    /// Unregister a tenant. Future publishes on the handle fail
871    /// with [`TenantError::Revoked`]. In-flight slots already on
872    /// the GPU still execute  -  the host is responsible for
873    /// quiescing before unregister if it needs that guarantee.
874    pub fn unregister(&self, tenant_id: u32) -> Option<TenantHandle> {
875        let (_, handle) = self.tenants.remove(&tenant_id)?;
876        handle.state.revoked.store(1, Ordering::Release);
877        handle.release_all_resource_reservations();
878        Some(handle)
879    }
880
881    /// Snapshot of active tenants for observability / diagnostics.
882    #[must_use]
883    pub fn active_tenants(&self) -> Vec<TenantHandle> {
884        let mut out = Vec::with_capacity(self.tenants.len());
885        out.extend(self.tenants.iter().map(|entry| entry.value().clone()));
886        out.sort_by_key(TenantHandle::id);
887        out
888    }
889
890    /// Snapshot active tenants into caller-owned storage.
891    pub fn active_tenants_into(&self, out: &mut Vec<TenantHandle>) {
892        out.clear();
893        out.reserve(self.tenants.len());
894        self.tenants
895            .iter()
896            .for_each(|entry| out.push(entry.value().clone()));
897        out.sort_by_key(TenantHandle::id);
898    }
899
900    /// Look up a tenant by id. Returns `None` if the id was
901    /// unregistered.
902    #[must_use]
903    pub fn lookup(&self, tenant_id: u32) -> Option<TenantHandle> {
904        self.tenants
905            .get(&tenant_id)
906            .map(|entry| entry.value().clone())
907    }
908
909    /// Snapshot runtime counters for every active tenant.
910    #[must_use]
911    pub fn runtime_counters(&self) -> Vec<TenantRuntimeCounters> {
912        let mut out = Vec::with_capacity(self.tenants.len());
913        self.tenants
914            .iter()
915            .map(|entry| entry.value().runtime_counters())
916            .for_each(|counters| out.push(counters));
917        out.sort_by_key(|counters| counters.tenant_id);
918        out
919    }
920
921    /// Snapshot runtime counters into caller-owned storage.
922    pub fn runtime_counters_into(&self, out: &mut Vec<TenantRuntimeCounters>) {
923        out.clear();
924        out.reserve(self.tenants.len());
925        self.tenants
926            .iter()
927            .map(|entry| entry.value().runtime_counters())
928            .for_each(|counters| out.push(counters));
929        out.sort_by_key(|counters| counters.tenant_id);
930    }
931
932    /// Select a maximal independent subset of tenants for a fair
933    /// schedule slot.
934    ///
935    /// `conflict_adj[i*n+j] != 0` means tenants `i` and `j` cannot
936    /// share the same dispatch slot (e.g., both pinned to the same
937    /// queue, or both holding mutually-exclusive opcode locks). The
938    /// Returns a Vec of tenant ids in selection order. Empty if no
939    /// tenants are active.
940    #[must_use]
941    pub fn select_concurrent_tenants(&self, conflict_adj: &[u32]) -> Vec<u32> {
942        let mut out = Vec::new();
943        let mut scratch = TenantSelectionScratch::new();
944        self.select_concurrent_tenants_into(conflict_adj, &mut out, &mut scratch);
945        out
946    }
947
948    /// Select a maximal independent tenant subset into caller-owned storage.
949    pub fn select_concurrent_tenants_into(
950        &self,
951        conflict_adj: &[u32],
952        out: &mut Vec<u32>,
953        scratch: &mut TenantSelectionScratch,
954    ) {
955        out.clear();
956        scratch.active_ids.clear();
957        scratch.active_ids.reserve(self.tenants.len());
958        self.tenants
959            .iter()
960            .map(|entry| entry.value().id())
961            .for_each(|id| scratch.active_ids.push(id));
962        scratch.active_ids.sort_unstable();
963        let n = scratch.active_ids.len();
964        if n == 0 {
965            return;
966        }
967        if vyre_driver::accounting::checked_mul_usize_lazy(n, n, || ()).ok()
968            != Some(conflict_adj.len())
969        {
970            // Degenerate: caller didn't supply a matching adjacency.
971            // Default to all-tenants-can-run (no conflicts).
972            out.reserve(n);
973            out.extend(scratch.active_ids.iter().copied());
974            return;
975        }
976        if conflict_adj.iter().all(|conflict| *conflict == 0) {
977            out.reserve(n);
978            out.extend(scratch.active_ids.iter().copied());
979            return;
980        }
981        scratch.selected_indices.clear();
982        scratch.selected_indices.reserve(n);
983        'candidate: for candidate_idx in 0..n {
984            for &selected_idx in &scratch.selected_indices {
985                if conflict_adj[candidate_idx * n + selected_idx] != 0
986                    || conflict_adj[selected_idx * n + candidate_idx] != 0
987                {
988                    continue 'candidate;
989                }
990            }
991            scratch.selected_indices.push(candidate_idx);
992        }
993        out.reserve(scratch.selected_indices.len());
994        for &index in &scratch.selected_indices {
995            if let Some(&id) = scratch.active_ids.get(index) {
996                out.push(id);
997            }
998        }
999    }
1000}
1001
1002#[cfg(test)]
1003mod tests {
1004    use super::*;
1005
1006    #[test]
1007    fn two_tenants_get_distinct_id_and_opcode_ranges() {
1008        let reg = TenantRegistry::new();
1009        let a = reg
1010            .register("scanner-a")
1011            .expect("Fix: register a; restore this invariant before continuing.");
1012        let b = reg
1013            .register("scanner-b")
1014            .expect("Fix: register b; restore this invariant before continuing.");
1015        assert_ne!(a.id(), b.id());
1016        assert!(a.base_opcode() + OPCODE_RANGE_PER_TENANT <= b.base_opcode());
1017        assert_eq!(a.label(), "scanner-a");
1018        assert_eq!(b.label(), "scanner-b");
1019    }
1020
1021    #[test]
1022    fn global_opcode_rejects_out_of_range_local() {
1023        let reg = TenantRegistry::new();
1024        let t = reg.register("soleno").unwrap();
1025        let err = t
1026            .global_opcode(OPCODE_RANGE_PER_TENANT)
1027            .expect_err("oversized local opcode must reject");
1028        assert!(matches!(err, TenantError::OpcodeOutOfRange { .. }));
1029
1030        let ok = t
1031            .global_opcode(42)
1032            .expect("Fix: 42 < cap; restore this invariant before continuing.");
1033        assert_eq!(ok, t.base_opcode() + 42);
1034    }
1035
1036    #[test]
1037    fn publish_slot_writes_with_tenant_id_and_bumps_counter() {
1038        let reg = TenantRegistry::new();
1039        let t = reg.register("warpscan").unwrap();
1040        let mut ring = ResidentWorkQueue::try_encode_empty_ring(4).unwrap();
1041
1042        t.publish_slot(
1043            &mut ring,
1044            /* slot = */ 0,
1045            /* local = */ 7,
1046            &[1, 2, 3],
1047        )
1048        .expect("Fix: publish; restore this invariant before continuing.");
1049        assert_eq!(t.published_count(), 1);
1050
1051        // Slot 0 should carry tenant=t.id(), opcode=t.base_opcode()+7.
1052        let tenant_off = super::super::resident_work_queue::protocol::TENANT_WORD as usize * 4;
1053        let opcode_off = super::super::resident_work_queue::protocol::OPCODE_WORD as usize * 4;
1054        let stored_tenant =
1055            u32::from_le_bytes(ring[tenant_off..tenant_off + 4].try_into().unwrap());
1056        let stored_opcode =
1057            u32::from_le_bytes(ring[opcode_off..opcode_off + 4].try_into().unwrap());
1058        assert_eq!(stored_tenant, t.id());
1059        assert_eq!(stored_opcode, t.base_opcode() + 7);
1060    }
1061
1062    #[test]
1063    fn unregister_blocks_future_publishes() {
1064        let reg = TenantRegistry::new();
1065        let t = reg.register("vein").unwrap();
1066        let tenant_id = t.id();
1067        let mut ring = ResidentWorkQueue::try_encode_empty_ring(2).unwrap();
1068        t.publish_slot(&mut ring, 0, 0, &[0, 0, 0])
1069            .expect("Fix: first publish ok; restore this invariant before continuing.");
1070        reg.unregister(tenant_id)
1071            .expect("Fix: unregister; restore this invariant before continuing.");
1072        let err = t
1073            .publish_slot(&mut ring, 1, 0, &[0, 0, 0])
1074            .expect_err("publish after unregister must reject");
1075        assert!(matches!(err, TenantError::Revoked { .. }));
1076        assert!(reg.lookup(tenant_id).is_none());
1077    }
1078
1079    #[test]
1080    fn quiesce_returns_when_drained_catches_up() {
1081        let reg = TenantRegistry::new();
1082        let t = reg.register("t1").unwrap();
1083        let mut ring = ResidentWorkQueue::try_encode_empty_ring(2).unwrap();
1084        t.publish_slot(&mut ring, 0, 0, &[1, 2, 3]).unwrap();
1085        t.publish_slot(&mut ring, 1, 0, &[4, 5, 6]).unwrap();
1086        assert_eq!(t.published_count(), 2);
1087        t.note_drained(2);
1088        t.quiesce(1)
1089            .expect("Fix: drained == published after note_drained; restore this invariant before continuing.");
1090        let counters = t.runtime_counters();
1091        assert_eq!(counters.published_count, 2);
1092        assert_eq!(counters.drained_count, 2);
1093        assert_eq!(counters.outstanding_slots, 0);
1094        assert_eq!(counters.quiesce_calls, 1);
1095        assert_eq!(counters.quiesce_timeouts, 0);
1096    }
1097
1098    #[test]
1099    fn quiesce_times_out_when_drain_stalled() {
1100        let reg = TenantRegistry::new();
1101        let t = reg.register("t2").unwrap();
1102        let mut ring = ResidentWorkQueue::try_encode_empty_ring(1).unwrap();
1103        t.publish_slot(&mut ring, 0, 0, &[0, 0, 0]).unwrap();
1104        // Never note_drained → quiesce must time out.
1105        let err = t.quiesce(4).expect_err("stalled quiesce must time out");
1106        assert!(matches!(
1107            err,
1108            TenantError::QuiesceTimeout { outstanding: 1, .. }
1109        ));
1110        let counters = t.runtime_counters();
1111        assert_eq!(counters.outstanding_slots, 1);
1112        assert_eq!(counters.quiesce_calls, 1);
1113        assert_eq!(counters.quiesce_timeouts, 1);
1114    }
1115
1116    #[test]
1117    fn bounded_tenant_backpressure_rejects_unbounded_publish_backlog() {
1118        let reg = TenantRegistry::new();
1119        let t = reg.register_with_backpressure("bounded", 2).unwrap();
1120        let mut ring = ResidentWorkQueue::try_encode_empty_ring(4).unwrap();
1121
1122        t.publish_slot(&mut ring, 0, 0, &[1]).unwrap();
1123        t.publish_slot(&mut ring, 1, 0, &[2]).unwrap();
1124        let err = t
1125            .publish_slot(&mut ring, 2, 0, &[3])
1126            .expect_err("third outstanding publish must hit tenant backpressure");
1127        assert!(matches!(
1128            err,
1129            TenantError::Backpressure {
1130                outstanding: 2,
1131                cap: 2,
1132                ..
1133            }
1134        ));
1135        assert_eq!(t.published_count(), 2);
1136        let counters = t.runtime_counters();
1137        assert_eq!(counters.max_outstanding_slots, 2);
1138        assert_eq!(counters.outstanding_slots, 2);
1139    }
1140
1141    #[test]
1142    fn tenant_backpressure_reopens_after_drain_progress() {
1143        let reg = TenantRegistry::new();
1144        let t = reg.register_with_backpressure("bounded", 1).unwrap();
1145        let mut ring = ResidentWorkQueue::try_encode_empty_ring(2).unwrap();
1146
1147        t.publish_slot(&mut ring, 0, 0, &[1]).unwrap();
1148        assert!(matches!(
1149            t.publish_slot(&mut ring, 1, 0, &[2]).unwrap_err(),
1150            TenantError::Backpressure { .. }
1151        ));
1152        t.note_drained(1);
1153        t.publish_slot(&mut ring, 1, 0, &[2])
1154            .expect("Fix: drain progress must reopen the bounded tenant queue; restore this invariant before continuing.");
1155        assert_eq!(t.published_count(), 2);
1156        assert_eq!(t.runtime_counters().outstanding_slots, 1);
1157    }
1158
1159    #[test]
1160    fn tenant_resource_quotas_reject_overcommit_and_cleanup_on_unregister() {
1161        let reg = TenantRegistry::new();
1162        let t = reg
1163            .register_with_quotas("quota", TenantQuota::bounded(2, 16, 1))
1164            .unwrap();
1165
1166        t.reserve_staging_bytes(8).unwrap();
1167        let staging_error = t
1168            .reserve_staging_bytes(9)
1169            .expect_err("staging byte quota must reject overcommit");
1170        assert!(matches!(
1171            staging_error,
1172            TenantError::StagingBackpressure {
1173                requested: 9,
1174                cap: 16,
1175                ..
1176            }
1177        ));
1178        assert_eq!(t.quota_counters().staging_bytes, 8);
1179
1180        t.release_staging_bytes(4).unwrap();
1181        t.reserve_staging_bytes(12).unwrap();
1182        assert_eq!(t.quota_counters().staging_bytes, 16);
1183        let underflow = t
1184            .release_staging_bytes(17)
1185            .expect_err("staging release must reject underflow");
1186        assert!(matches!(
1187            underflow,
1188            TenantError::ResourceUnderflow {
1189                resource: "staging bytes",
1190                requested: 17,
1191                used: 16,
1192                ..
1193            }
1194        ));
1195
1196        t.reserve_resident_handles(1).unwrap();
1197        let handle_error = t
1198            .reserve_resident_handles(1)
1199            .expect_err("resident handle quota must reject overcommit");
1200        assert!(matches!(
1201            handle_error,
1202            TenantError::ResidentHandleBackpressure {
1203                requested: 1,
1204                cap: 1,
1205                ..
1206            }
1207        ));
1208        assert_eq!(t.quota_counters().resident_handles, 1);
1209
1210        let removed = reg.unregister(t.id()).unwrap();
1211        assert_eq!(removed.quota_counters().staging_bytes, 0);
1212        assert_eq!(removed.quota_counters().resident_handles, 0);
1213        assert!(matches!(
1214            t.reserve_staging_bytes(1).unwrap_err(),
1215            TenantError::Revoked { .. }
1216        ));
1217        assert!(matches!(
1218            t.reserve_resident_handles(1).unwrap_err(),
1219            TenantError::Revoked { .. }
1220        ));
1221    }
1222
1223    #[test]
1224    fn tenant_registry_registration_retry_uses_adaptive_idle_not_unbounded_spin() {
1225        for retry in [0, 1, 2, QUIESCE_SPIN_POLLS - 1, QUIESCE_SPIN_POLLS] {
1226            tenant_registry_retry_idle(retry);
1227        }
1228        assert_eq!(
1229            quiesce_backoff_duration(QUIESCE_SPIN_POLLS),
1230            QUIESCE_MIN_PARK
1231        );
1232        assert_eq!(quiesce_backoff_duration(u64::MAX), QUIESCE_MAX_PARK);
1233    }
1234
1235    #[test]
1236    fn quiesce_backoff_is_bounded_and_monotonic() {
1237        let samples = [
1238            quiesce_backoff_duration(0),
1239            quiesce_backoff_duration(1),
1240            quiesce_backoff_duration(2),
1241            quiesce_backoff_duration(8),
1242            quiesce_backoff_duration(64),
1243        ];
1244        assert_eq!(samples[0], QUIESCE_MIN_PARK);
1245        for pair in samples.windows(2) {
1246            assert!(pair[0] <= pair[1], "quiesce backoff must not shrink");
1247            assert!(pair[1] <= QUIESCE_MAX_PARK, "quiesce backoff must cap");
1248        }
1249        assert_eq!(quiesce_backoff_duration(u64::MAX), QUIESCE_MAX_PARK);
1250    }
1251
1252    #[test]
1253    fn active_tenants_tracks_registrations() {
1254        let reg = TenantRegistry::new();
1255        let a = reg.register("a").unwrap();
1256        let b = reg.register("b").unwrap();
1257        let active: Vec<u32> = reg.active_tenants().iter().map(|t| t.id()).collect();
1258        assert!(active.contains(&a.id()));
1259        assert!(active.contains(&b.id()));
1260        reg.unregister(a.id());
1261        let after: Vec<u32> = reg.active_tenants().iter().map(|t| t.id()).collect();
1262        assert!(!after.contains(&a.id()));
1263        assert!(after.contains(&b.id()));
1264        let counters: Vec<u32> = reg
1265            .runtime_counters()
1266            .iter()
1267            .map(|tenant| tenant.tenant_id)
1268            .collect();
1269        assert_eq!(counters, vec![b.id()]);
1270    }
1271
1272    #[test]
1273    fn tenant_snapshots_reuse_caller_storage() {
1274        let reg = TenantRegistry::new();
1275        let a = reg.register("a").unwrap();
1276        let b = reg.register("b").unwrap();
1277        let mut active = Vec::with_capacity(2);
1278        let mut counters = Vec::with_capacity(2);
1279
1280        reg.active_tenants_into(&mut active);
1281        reg.runtime_counters_into(&mut counters);
1282        let active_ptr = active.as_ptr();
1283        let counters_ptr = counters.as_ptr();
1284        reg.active_tenants_into(&mut active);
1285        reg.runtime_counters_into(&mut counters);
1286
1287        assert_eq!(active.as_ptr(), active_ptr);
1288        assert_eq!(counters.as_ptr(), counters_ptr);
1289        assert!(active.iter().any(|tenant| tenant.id() == a.id()));
1290        assert!(active.iter().any(|tenant| tenant.id() == b.id()));
1291        assert!(counters.iter().any(|tenant| tenant.tenant_id == a.id()));
1292        assert!(counters.iter().any(|tenant| tenant.tenant_id == b.id()));
1293    }
1294
1295    #[test]
1296    fn concurrent_tenant_selection_reuses_scratch_and_output() {
1297        let reg = TenantRegistry::new();
1298        let a = reg.register("a").unwrap();
1299        let b = reg.register("b").unwrap();
1300        let c = reg.register("c").unwrap();
1301        let n = 3;
1302        let mut conflicts = vec![0_u32; n * n];
1303        conflicts[0 * n + 1] = 1;
1304        conflicts[1 * n + 0] = 1;
1305        let mut out = Vec::with_capacity(3);
1306        let mut scratch = TenantSelectionScratch::new();
1307
1308        reg.select_concurrent_tenants_into(&conflicts, &mut out, &mut scratch);
1309        let out_ptr = out.as_ptr();
1310        let active_ids_ptr = scratch.active_ids.as_ptr();
1311        let selected_ptr = scratch.selected_indices.as_ptr();
1312        reg.select_concurrent_tenants_into(&conflicts, &mut out, &mut scratch);
1313
1314        assert_eq!(out.as_ptr(), out_ptr);
1315        assert_eq!(scratch.active_ids.as_ptr(), active_ids_ptr);
1316        assert_eq!(scratch.selected_indices.as_ptr(), selected_ptr);
1317        assert!(out.contains(&a.id()) || out.contains(&b.id()));
1318        assert!(!(out.contains(&a.id()) && out.contains(&b.id())));
1319        assert!(out.contains(&c.id()));
1320    }
1321
1322    #[test]
1323    fn concurrent_tenant_selection_fast_paths_all_zero_conflicts() {
1324        let reg = TenantRegistry::new();
1325        let a = reg.register("a").unwrap();
1326        let b = reg.register("b").unwrap();
1327        let c = reg.register("c").unwrap();
1328        let mut out = Vec::with_capacity(8);
1329        let mut scratch = TenantSelectionScratch::new();
1330        let conflicts = vec![0_u32; 9];
1331        let out_ptr = out.as_ptr();
1332
1333        reg.select_concurrent_tenants_into(&conflicts, &mut out, &mut scratch);
1334
1335        assert_eq!(out, vec![a.id(), b.id(), c.id()]);
1336        assert_eq!(
1337            out.as_ptr(),
1338            out_ptr,
1339            "all-zero conflict fast path must reuse caller-owned output storage"
1340        );
1341        assert!(
1342            scratch.selected_indices.is_empty(),
1343            "all-zero conflict fast path must not populate pairwise selection scratch"
1344        );
1345    }
1346
1347    #[test]
1348    fn concurrent_tenant_selection_respects_conflicts() {
1349        let reg = TenantRegistry::new();
1350        let a = reg.register("a").unwrap();
1351        let b = reg.register("b").unwrap();
1352        let c = reg.register("c").unwrap();
1353        let n = 3;
1354        let mut conflicts = vec![0_u32; n * n];
1355        conflicts[0 * n + 1] = 1;
1356        conflicts[1 * n + 0] = 1;
1357
1358        let selected = reg.select_concurrent_tenants(&conflicts);
1359
1360        assert!(selected.contains(&a.id()) || selected.contains(&b.id()));
1361        assert!(!(selected.contains(&a.id()) && selected.contains(&b.id())));
1362        assert!(selected.contains(&c.id()));
1363    }
1364
1365    #[test]
1366    fn concurrent_registration_assigns_unique_ids() {
1367        use std::thread;
1368        let reg = Arc::new(TenantRegistry::new());
1369        let mut handles = Vec::new();
1370        for i in 0..32 {
1371            let reg = Arc::clone(&reg);
1372            handles.push(thread::spawn(move || {
1373                reg.register(format!("t{i}")).unwrap().id()
1374            }));
1375        }
1376        let ids: Vec<u32> = handles.into_iter().map(|h| h.join().unwrap()).collect();
1377        let mut sorted = ids.clone();
1378        sorted.sort();
1379        sorted.dedup();
1380        assert_eq!(sorted.len(), ids.len(), "concurrent ids must be unique");
1381    }
1382}