1use 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
49pub const TENANT_OPCODE_BASE: u32 = 0x4000_0000;
54
55pub const TENANT_ID_MAX: u32 = u32::MAX - 1;
58
59pub 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#[derive(Debug, thiserror::Error)]
99#[non_exhaustive]
100pub enum TenantError {
101 #[error("tenant registry exhausted after {issued} registrations. Fix: shrink OPCODE_RANGE_PER_TENANT or recycle tenants.")]
104 RegistryFull {
105 issued: u32,
107 },
108 #[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: u32,
117 local_opcode: u32,
119 cap: u32,
121 },
122 #[error("tenant {tenant_id} was revoked; handle is stale. Fix: acquire a fresh handle from the registry.")]
124 Revoked {
125 tenant_id: u32,
127 },
128 #[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: u32,
136 outstanding: u64,
138 },
139 #[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: u32,
147 outstanding: u64,
149 cap: u64,
151 },
152 #[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: u32,
160 requested: u64,
162 used: u64,
164 cap: u64,
166 },
167 #[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: u32,
175 requested: u64,
177 used: u64,
179 cap: u64,
181 },
182 #[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: u32,
190 resource: &'static str,
192 requested: u64,
194 used: u64,
196 },
197 #[error("{0}")]
199 Pipeline(#[from] PipelineError),
200}
201
202#[derive(Debug, Clone, Copy, PartialEq, Eq)]
204pub struct TenantQuota {
205 pub max_outstanding_slots: u64,
207 pub max_staging_bytes: u64,
209 pub max_resident_handles: u64,
211}
212
213impl TenantQuota {
214 #[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 #[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
241struct TenantState {
244 id: u32,
245 base_opcode: u32,
246 opcode_cap: u32,
247 published_count: AtomicU64,
249 max_outstanding_slots: u64,
251 staging_bytes: AtomicU64,
253 max_staging_bytes: u64,
255 resident_handles: AtomicU64,
257 max_resident_handles: u64,
259 drained_count: AtomicU64,
262 quiesce_calls: AtomicU64,
264 quiesce_timeouts: AtomicU64,
266 quiesce_wait_ns: AtomicU64,
268 revoked: AtomicU32,
270 label: String,
272}
273
274#[derive(Clone)]
278pub struct TenantHandle {
279 state: Arc<TenantState>,
280}
281
282#[derive(Debug, Clone, Copy, PartialEq, Eq)]
284pub struct TenantRuntimeCounters {
285 pub tenant_id: u32,
287 pub published_count: u64,
289 pub drained_count: u64,
291 pub outstanding_slots: u64,
293 pub max_outstanding_slots: u64,
295 pub quiesce_calls: u64,
297 pub quiesce_timeouts: u64,
299 pub quiesce_wait_ns: u64,
301}
302
303#[derive(Debug, Clone, Copy, PartialEq, Eq)]
305pub struct TenantQuotaCounters {
306 pub tenant_id: u32,
308 pub staging_bytes: u64,
310 pub max_staging_bytes: u64,
312 pub resident_handles: u64,
314 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 #[must_use]
354 pub fn id(&self) -> u32 {
355 self.state.id
356 }
357
358 #[must_use]
360 pub fn label(&self) -> &str {
361 &self.state.label
362 }
363
364 #[must_use]
366 pub fn base_opcode(&self) -> u32 {
367 self.state.base_opcode
368 }
369
370 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 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 #[must_use]
472 pub fn published_count(&self) -> u64 {
473 self.state.published_count.load(Ordering::Relaxed)
474 }
475
476 #[must_use]
479 pub fn drained_count(&self) -> u64 {
480 self.state.drained_count.load(Ordering::Relaxed)
481 }
482
483 #[must_use]
485 pub fn max_outstanding_slots(&self) -> u64 {
486 self.state.max_outstanding_slots
487 }
488
489 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 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 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 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 #[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 #[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 pub fn note_drained(&self, count: u64) {
585 saturating_atomic_add_u64(&self.state.drained_count, count, "tenant drained_count");
586 }
587
588 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
638pub 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#[derive(Debug, Default)]
655pub struct TenantSelectionScratch {
656 active_ids: Vec<u32>,
657 selected_indices: Vec<usize>,
658}
659
660impl TenantSelectionScratch {
661 #[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 #[must_use]
750 pub fn new() -> Self {
751 Self::default()
752 }
753
754 pub fn register(&self, label: impl Into<String>) -> Result<TenantHandle, TenantError> {
763 self.register_with_backpressure(label, u64::MAX)
764 }
765
766 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 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 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 #[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 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 #[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 #[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 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 #[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 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 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 0,
1045 7,
1046 &[1, 2, 3],
1047 )
1048 .expect("Fix: publish; restore this invariant before continuing.");
1049 assert_eq!(t.published_count(), 1);
1050
1051 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 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(®);
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}