1use std::fmt;
31use std::io;
32use std::str::FromStr;
33use std::sync::atomic::{AtomicU32, Ordering};
34use std::sync::{Arc, Mutex, MutexGuard};
35use std::time::{Duration, Instant};
36
37use orbit_core::{Fleet, NetId64, NodeId, OrbitEpoch};
38
39mod layout;
40#[cfg(feature = "stream")]
41mod session;
42mod policy;
43mod table;
44
45pub use layout::PENDING_RESERVATIONS;
46use layout::{
47 GENERATION_MASK, RESOURCE_DRAINING, RESOURCE_LIVE, ResourceSlot, SLOT_BITS, SLOT_MASK,
48 pack_counts, unpack_counts,
49};
50pub use policy::{Decision, Limits, LocalFirst, LocalOnly, Policy, Reason};
51use table::Table;
52#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
53pub use orbit_core::readiness::Readiness;
54pub use table::{segment_size, segment_size_for};
55
56pub const POOL_KIND: u8 = 247;
59pub const POOL_KEY_CAPACITY: usize =
65 orbit_core::compile::usize_from_env(option_env!("ORBIT_POOL_KEY_CAPACITY"), 256);
66pub const POOL_RESOURCE_LANE_CAPACITY: usize =
72 orbit_core::compile::usize_from_env(option_env!("ORBIT_POOL_RESOURCE_LANE_CAPACITY"), 256);
73
74#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
83pub struct PoolSpec {
84 pub kind: u8,
86 pub key_capacity: usize,
88 pub lane_capacity: usize,
91}
92
93impl PoolSpec {
94 pub const DEFAULT: Self =
95 Self::new(POOL_KIND, POOL_KEY_CAPACITY, POOL_RESOURCE_LANE_CAPACITY);
96
97 pub const fn new(kind: u8, key_capacity: usize, lane_capacity: usize) -> Self {
98 Self { kind, key_capacity, lane_capacity }
99 }
100
101 fn validate(self) -> Result<()> {
104 if self.key_capacity == 0
105 || self.lane_capacity == 0
106 || !self.key_capacity.is_power_of_two()
107 || !self.lane_capacity.is_power_of_two()
108 || self.lane_capacity > 1 << SLOT_BITS
109 {
110 return Err(Error::Malformed(format!(
111 "pool spec kind={} key_capacity={} lane_capacity={}: both are powers of two, and a lane holds at most {}",
112 self.kind,
113 self.key_capacity,
114 self.lane_capacity,
115 1_usize << SLOT_BITS
116 )));
117 }
118 Ok(())
119 }
120}
121
122impl Default for PoolSpec {
123 fn default() -> Self {
124 Self::DEFAULT
125 }
126}
127
128pub type Result<T, E = Error> = std::result::Result<T, E>;
129
130const PLACING: u64 = u64::MAX;
133
134#[non_exhaustive]
138#[derive(Debug)]
139pub enum Error {
140 Stale(ResourceId),
142 Busy(ResourceId),
144 Draining(ResourceId),
146 NotOwner(ResourceId),
148 NotReserved(Lease),
152 CreationBudget {
155 key: Key,
156 max_live: u32,
157 },
158 KeyFull {
160 capacity: usize,
161 },
162 Full {
164 capacity: usize,
165 },
166 Malformed(String),
167 Io(io::Error),
168}
169
170impl fmt::Display for Error {
171 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
172 match self {
173 Self::Stale(id) => write!(f, "resource {id} has ended"),
174 Self::Busy(id) => write!(f, "resource {id} has no capacity left"),
175 Self::Draining(id) => write!(f, "resource {id} is draining"),
176 Self::NotOwner(id) => write!(f, "resource {id} belongs to another node"),
177 Self::NotReserved(lease) => write!(
178 f,
179 "lease {} on {} is not an unaccepted reservation",
180 lease.fence, lease.id
181 ),
182 Self::CreationBudget { key, max_live } => {
183 write!(f, "creation budget for {key} is spent: max_live={max_live}")
184 }
185 Self::KeyFull { capacity } => write!(f, "pool key table is full: capacity={capacity}"),
186 Self::Full { capacity } => write!(f, "pool lane is full: capacity={capacity}"),
187 Self::Malformed(text) => write!(f, "not a pool resource id: {text:?}"),
188 Self::Io(error) => write!(f, "Orbit pool io error: {error}"),
189 }
190 }
191}
192
193impl std::error::Error for Error {
194 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
195 match self {
196 Self::Io(error) => Some(error),
197 _ => None,
198 }
199 }
200}
201
202impl From<io::Error> for Error {
203 fn from(value: io::Error) -> Self {
204 Self::Io(value)
205 }
206}
207
208#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
213pub struct Key(u128);
214
215impl Key {
216 pub const fn new(digest: u128) -> Self {
217 Self(digest)
218 }
219
220 pub const fn from_bytes(digest: [u8; 16]) -> Self {
221 Self(u128::from_le_bytes(digest))
222 }
223
224 pub const fn get(self) -> u128 {
225 self.0
226 }
227
228 const fn parts(self) -> (u64, u64) {
229 (self.0 as u64, (self.0 >> 64) as u64)
230 }
231}
232
233impl fmt::Display for Key {
234 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
235 write!(f, "key:{:032x}", self.0)
236 }
237}
238
239#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
243pub struct Incarnation(u64);
244
245impl Incarnation {
246 pub const fn new(value: u64) -> Self {
247 Self(value)
248 }
249
250 pub const fn get(self) -> u64 {
251 self.0
252 }
253}
254
255#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
258pub struct ResourceId(NetId64);
259
260impl ResourceId {
261 pub const fn from_net_id(id: NetId64) -> Self {
262 Self(id)
263 }
264
265 pub const fn net_id(self) -> NetId64 {
266 self.0
267 }
268
269 pub const fn kind(self) -> u8 {
270 self.0.kind()
271 }
272
273 pub const fn node(self) -> u16 {
275 self.0.node()
276 }
277
278 pub const fn slot(self) -> u32 {
279 (self.0.counter() & SLOT_MASK) as u32
280 }
281
282 pub const fn generation(self) -> u32 {
283 (self.0.counter() >> SLOT_BITS) as u32
284 }
285
286 fn make(kind: u8, node: u16, slot: u32, generation: u32) -> Self {
287 Self(NetId64::make(
288 kind,
289 node,
290 ((generation as u64) << SLOT_BITS) | (slot as u64 & SLOT_MASK),
291 ))
292 }
293}
294
295impl fmt::Display for ResourceId {
296 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
297 self.0.fmt(f)
298 }
299}
300
301impl FromStr for ResourceId {
302 type Err = Error;
303
304 fn from_str(text: &str) -> Result<Self> {
305 text.parse::<NetId64>()
306 .map(Self)
307 .map_err(|_| Error::Malformed(text.to_owned()))
308 }
309}
310
311#[derive(Clone, Copy, Debug, Eq, PartialEq)]
314pub enum State {
315 Live,
316 Draining,
317}
318
319#[non_exhaustive]
323#[derive(Clone, Copy, Debug, Eq, PartialEq)]
324pub struct Candidate {
325 pub id: ResourceId,
326 pub owner: NodeId,
327 pub owner_incarnation: Incarnation,
328 pub local: bool,
330 pub state: State,
331 pub capacity: u32,
332 pub reserved: u32,
334 pub active: u32,
336 pub last_reserve_ms: u64,
337}
338
339impl Candidate {
340 pub fn free(&self) -> u32 {
341 self.capacity.saturating_sub(self.reserved + self.active)
342 }
343}
344
345#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
348pub struct Lease {
349 pub id: ResourceId,
350 pub fence: u64,
353 pub holder: NodeId,
354 pub holder_incarnation: Incarnation,
355}
356
357#[derive(Clone)]
360pub struct Pool {
361 table: Arc<Table>,
362}
363
364impl Pool {
365 pub fn new(fleet: Arc<Fleet>, incarnation: Incarnation) -> Result<Self> {
367 Self::with_spec(fleet, incarnation, PoolSpec::DEFAULT)
368 }
369
370 pub fn with_spec(
375 fleet: Arc<Fleet>,
376 incarnation: Incarnation,
377 spec: PoolSpec,
378 ) -> Result<Self> {
379 spec.validate()?;
380 Ok(Self {
381 table: table::open(&fleet, incarnation, spec)?,
382 })
383 }
384
385 pub fn kind(&self) -> u8 {
387 self.table.kind()
388 }
389
390 pub fn readiness(&self) -> Result<Readiness> {
403 self.table.take_readiness()
404 }
405
406 pub fn watch(&self, key: Key) -> Result<()> {
413 let (lo, hi) = key.parts();
414 let key_index = self.table.key_index(lo, hi)?;
415 self.table.watch(key_index)
416 }
417
418 pub fn node(&self) -> NodeId {
419 NodeId::new(self.table.node())
420 }
421
422 pub fn incarnation(&self) -> Incarnation {
423 Incarnation::new(self.table.incarnation())
424 }
425
426 pub fn epoch(&self) -> u64 {
427 self.table.epoch()
428 }
429
430 pub fn register(&self, key: Key, capacity: u32) -> Result<ResourceId> {
433 if capacity == 0 {
434 return Err(Error::Malformed("capacity must be at least one".to_owned()));
435 }
436 let (lo, hi) = key.parts();
437 let key_index = self.table.key_index(lo, hi)?;
438 let (index, generation) = self.table.allocate((lo, hi), key_index, capacity)?;
439 self.table.members(key_index)[index / 64].fetch_or(1 << (index % 64), Ordering::SeqCst);
440 let _ = self.table.key(key_index).counts.try_update(
441 Ordering::SeqCst,
442 Ordering::SeqCst,
443 |counts| {
444 let (live, creating) = unpack_counts(counts);
445 Some(pack_counts(live + 1, creating))
446 },
447 );
448 self.table.key_changed(key_index);
449 Ok(ResourceId::make(
450 self.table.kind(),
451 self.table.node(),
452 (index % self.table.geometry().lane_capacity) as u32,
453 generation,
454 ))
455 }
456
457 pub fn unregister(&self, id: ResourceId) -> Result<()> {
459 let (index, slot) = self.owned(id)?;
460 self.table.close_resource(index, slot);
461 Ok(())
462 }
463
464 pub fn drain(&self, id: ResourceId) -> Result<()> {
466 let (_, slot) = self.owned(id)?;
467 let _ = slot.state.compare_exchange(
468 RESOURCE_LIVE,
469 RESOURCE_DRAINING,
470 Ordering::SeqCst,
471 Ordering::SeqCst,
472 );
473 Ok(())
474 }
475
476 pub fn reconcile(&self, id: ResourceId, active: u32, grace: std::time::Duration) -> Result<()> {
482 let (_, slot) = self.owned(id)?;
483 let now = OrbitEpoch::now().as_unix_ms();
484 let mut aged = 0_u32;
485 for reservation in &slot.pending {
486 let fence = reservation.fence.load(Ordering::Acquire);
487 if fence == 0 || fence == PLACING {
488 continue;
489 }
490 let since = reservation.since_ms.load(Ordering::Relaxed);
491 if now.saturating_sub(since) > grace.as_millis() as u64
492 && reservation
493 .fence
494 .compare_exchange(fence, 0, Ordering::SeqCst, Ordering::SeqCst)
495 .is_ok()
496 {
497 aged += 1;
498 }
499 }
500 let mut freed = false;
501 let _ = slot
502 .units
503 .try_update(Ordering::SeqCst, Ordering::SeqCst, |units| {
504 let (reserved, was_active) = unpack_counts(units);
505 let reserved = reserved.saturating_sub(aged);
506 freed = reserved + active < unpack_counts(units).0 + was_active;
507 Some(pack_counts(reserved, active))
508 });
509 if freed {
510 self.table
511 .key_changed(slot.key_index.load(Ordering::Acquire) as usize);
512 }
513 Ok(())
514 }
515
516 pub fn candidates(&self, key: Key) -> Vec<Candidate> {
518 let (lo, hi) = key.parts();
519 let Ok(key_index) = self.table.key_index(lo, hi) else {
520 return Vec::new();
521 };
522 let mut found = Vec::new();
523 for (word_index, word) in self.table.members(key_index).iter().enumerate() {
524 let mut bits = word.load(Ordering::SeqCst);
525 while bits != 0 {
526 let bit = bits.trailing_zeros() as usize;
527 bits &= bits - 1;
528 let index = word_index * 64 + bit;
529 if let Some(candidate) = self.candidate_at(index, (lo, hi)) {
530 found.push(candidate);
531 }
532 }
533 }
534 found
535 }
536
537 fn candidate_at(&self, index: usize, key: (u64, u64)) -> Option<Candidate> {
538 let slot = self.table.resources().get(index)?;
539 let state = match slot.state.load(Ordering::Acquire) {
540 RESOURCE_LIVE => State::Live,
541 RESOURCE_DRAINING => State::Draining,
542 _ => return None,
543 };
544 if slot.key_lo.load(Ordering::Relaxed) != key.0
545 || slot.key_hi.load(Ordering::Relaxed) != key.1
546 {
547 return None;
548 }
549 let owner = slot.owner_node.load(Ordering::Acquire);
550 Some(Candidate {
551 id: ResourceId::make(
552 self.table.kind(),
553 owner,
554 (index % self.table.geometry().lane_capacity) as u32,
555 slot.generation.load(Ordering::Relaxed),
556 ),
557 owner: NodeId::new(owner),
558 owner_incarnation: Incarnation::new(slot.owner_incarnation.load(Ordering::Relaxed)),
559 local: owner == self.table.node(),
560 state,
561 capacity: slot.capacity.load(Ordering::Relaxed),
562 reserved: unpack_counts(slot.units.load(Ordering::SeqCst)).0,
563 active: unpack_counts(slot.units.load(Ordering::SeqCst)).1,
564 last_reserve_ms: slot.last_reserve_ms.load(Ordering::Relaxed),
565 })
566 }
567
568 pub fn reserve(&self, id: ResourceId) -> Result<Lease> {
571 let (_, slot) = self.locate(id)?;
572 match slot.state.load(Ordering::Acquire) {
573 RESOURCE_LIVE => {}
574 RESOURCE_DRAINING => return Err(Error::Draining(id)),
575 _ => return Err(Error::Stale(id)),
576 }
577 let capacity = slot.capacity.load(Ordering::Relaxed);
578 slot.units
579 .try_update(Ordering::SeqCst, Ordering::SeqCst, |units| {
580 let (reserved, active) = unpack_counts(units);
581 (reserved + active < capacity).then(|| pack_counts(reserved + 1, active))
582 })
583 .map_err(|_| Error::Busy(id))?;
584 if !slot.is(id.generation()) {
587 let _ = slot
588 .units
589 .try_update(Ordering::SeqCst, Ordering::SeqCst, |units| {
590 let (reserved, active) = unpack_counts(units);
591 Some(pack_counts(reserved.saturating_sub(1), active))
592 });
593 return Err(Error::Stale(id));
594 }
595 let fence = slot.fence.fetch_add(1, Ordering::SeqCst) + 1;
596 let now = OrbitEpoch::now().as_unix_ms();
597 let mut placed = false;
598 for reservation in &slot.pending {
599 if reservation
603 .fence
604 .compare_exchange(0, PLACING, Ordering::SeqCst, Ordering::SeqCst)
605 .is_ok()
606 {
607 reservation.since_ms.store(now, Ordering::Relaxed);
608 reservation.fence.store(fence, Ordering::SeqCst);
609 placed = true;
610 break;
611 }
612 }
613 if !placed {
614 let _ = slot
616 .units
617 .try_update(Ordering::SeqCst, Ordering::SeqCst, |units| {
618 let (reserved, active) = unpack_counts(units);
619 Some(pack_counts(reserved.saturating_sub(1), active))
620 });
621 return Err(Error::Busy(id));
622 }
623 slot.last_reserve_ms.store(now, Ordering::Release);
624 Ok(Lease {
625 id,
626 fence,
627 holder: self.node(),
628 holder_incarnation: self.incarnation(),
629 })
630 }
631
632 pub fn accept(&self, lease: Lease) -> Result<Execution> {
639 let (_, slot) = self.owned(lease.id)?;
640 if slot.state.load(Ordering::Acquire) != RESOURCE_LIVE {
641 return Err(Error::Draining(lease.id));
642 }
643 let taken = slot.pending.iter().any(|reservation| {
646 reservation
647 .fence
648 .compare_exchange(lease.fence, 0, Ordering::SeqCst, Ordering::SeqCst)
649 .is_ok()
650 });
651 if !taken {
652 return Err(Error::NotReserved(lease));
653 }
654 let _ = slot
655 .units
656 .try_update(Ordering::SeqCst, Ordering::SeqCst, |units| {
657 let (reserved, active) = unpack_counts(units);
658 Some(pack_counts(reserved.saturating_sub(1), active + 1))
659 });
660 Ok(Execution {
661 table: Arc::clone(&self.table),
662 lease,
663 })
664 }
665
666 pub fn is_current(&self, lease: Lease) -> bool {
670 self.locate(lease.id)
671 .map(|(_, slot)| slot.is(lease.id.generation()))
672 .unwrap_or(false)
673 }
674
675 pub fn acquire(&self, key: Key, limits: &Limits, policy: &dyn Policy) -> Result<Plan> {
681 for _ in 0..limits.attempts.max(1) {
682 let candidates = self.candidates(key);
683 let budget = self.budget(key);
684 match policy.decide(key, &candidates, budget, limits) {
685 Decision::Reuse(id) => match self.reserve(id) {
686 Ok(lease) if id.node() == self.table.node() => {
687 return Ok(Plan::LocalReuse(lease));
688 }
689 Ok(lease) => return Ok(Plan::RemoteReuse(lease)),
690 Err(Error::Busy(_) | Error::Stale(_) | Error::Draining(_)) => continue,
691 Err(error) => return Err(error),
692 },
693 Decision::Create => match self.claim_create(key, limits.max_live) {
694 Ok(permit) => return Ok(Plan::Create(permit)),
695 Err(Error::CreationBudget { .. }) => continue,
696 Err(error) => return Err(error),
697 },
698 Decision::Wait => return Ok(Plan::Wait(self.version(key)?)),
699 Decision::Reject(reason) => return Ok(Plan::Reject(reason)),
700 }
701 }
702 Ok(Plan::Wait(self.version(key)?))
703 }
704
705 pub fn claim_create(&self, key: Key, max_live: u32) -> Result<CreationPermit> {
709 let (lo, hi) = key.parts();
710 let key_index = self.table.key_index(lo, hi)?;
711 self.table
712 .key(key_index)
713 .counts
714 .try_update(Ordering::SeqCst, Ordering::SeqCst, |counts| {
715 let (live, creating) = unpack_counts(counts);
716 (live + creating < max_live).then(|| pack_counts(live, creating + 1))
717 })
718 .map_err(|_| Error::CreationBudget { key, max_live })?;
719 self.table
720 .claims(usize::from(self.table.node()), key_index)
721 .fetch_add(1, Ordering::SeqCst);
722 Ok(CreationPermit {
723 table: Arc::clone(&self.table),
724 key_index,
725 })
726 }
727
728 pub fn budget(&self, key: Key) -> (u32, u32) {
731 let (lo, hi) = key.parts();
732 match self.table.key_index(lo, hi) {
733 Ok(key_index) => unpack_counts(self.table.key(key_index).counts.load(Ordering::SeqCst)),
734 Err(_) => (0, 0),
735 }
736 }
737
738 #[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
746 pub fn wait_capacity_timeout(
747 &self,
748 key: Key,
749 since: u32,
750 timeout: Duration,
751 ) -> Result<Option<u32>> {
752 let (lo, hi) = key.parts();
753 let key_index = self.table.key_index(lo, hi)?;
754 let slot = self.table.key(key_index);
755 let deadline = Instant::now() + timeout;
756 loop {
757 let now = slot.changes.load(Ordering::SeqCst);
758 if now != since {
759 return Ok(Some(now));
760 }
761 let left = deadline.saturating_duration_since(Instant::now());
762 if left.is_zero() {
763 return Ok(None);
764 }
765 slot.waiters.fetch_add(1, Ordering::SeqCst);
766 let outcome = if slot.changes.load(Ordering::SeqCst) == since {
767 wait_on_timeout(&slot.changes, since, left)
768 } else {
769 Ok(true)
770 };
771 slot.waiters.fetch_sub(1, Ordering::SeqCst);
772 if !outcome? {
773 let now = slot.changes.load(Ordering::SeqCst);
776 return Ok((now != since).then_some(now));
777 }
778 }
779 }
780
781 pub fn version(&self, key: Key) -> Result<u32> {
783 let (lo, hi) = key.parts();
784 let key_index = self.table.key_index(lo, hi)?;
785 Ok(self.table.key(key_index).changes.load(Ordering::SeqCst))
786 }
787
788 pub fn wait_capacity(&self, key: Key, since: u32) -> Result<u32> {
792 let (lo, hi) = key.parts();
793 let key_index = self.table.key_index(lo, hi)?;
794 let slot = self.table.key(key_index);
795 loop {
796 let now = slot.changes.load(Ordering::SeqCst);
797 if now != since {
798 return Ok(now);
799 }
800 slot.waiters.fetch_add(1, Ordering::SeqCst);
801 let outcome = if slot.changes.load(Ordering::SeqCst) == since {
802 wait_on(&slot.changes, since)
803 } else {
804 Ok(())
805 };
806 slot.waiters.fetch_sub(1, Ordering::SeqCst);
807 outcome?;
808 }
809 }
810
811 pub fn poll_capacity(
815 &self,
816 key: Key,
817 since: u32,
818 cx: &mut std::task::Context<'_>,
819 ) -> std::task::Poll<Result<u32>> {
820 let (lo, hi) = key.parts();
821 let key_index = match self.table.key_index(lo, hi) {
822 Ok(index) => index,
823 Err(error) => return std::task::Poll::Ready(Err(error)),
824 };
825 let changes = &self.table.key(key_index).changes;
826 let now = changes.load(Ordering::SeqCst);
827 if now != since {
828 return std::task::Poll::Ready(Ok(now));
829 }
830 if let Err(error) = self.table.register(key_index, cx.waker()) {
831 return std::task::Poll::Ready(Err(error));
832 }
833 let now = changes.load(Ordering::SeqCst);
834 if now != since {
835 std::task::Poll::Ready(Ok(now))
836 } else {
837 std::task::Poll::Pending
838 }
839 }
840
841 pub fn owners(&self) -> Vec<(NodeId, Incarnation)> {
856 let mut seen: Vec<(NodeId, Incarnation)> = Vec::new();
857 for slot in self.table.resources() {
858 let state = slot.state.load(Ordering::Acquire);
859 if state != crate::layout::RESOURCE_LIVE
860 && state != crate::layout::RESOURCE_DRAINING
861 {
862 continue;
863 }
864 let owner = (
865 NodeId::new(slot.owner_node.load(Ordering::Acquire)),
866 Incarnation::new(slot.owner_incarnation.load(Ordering::Acquire)),
867 );
868 if !seen.contains(&owner) {
869 seen.push(owner);
870 }
871 }
872 seen.sort_by_key(|(node, incarnation)| (node.get(), incarnation.get()));
873 seen
874 }
875
876 pub fn node_dead(&self, node: NodeId, incarnation: Incarnation) {
881 self.table.node_dead(node.get(), incarnation.get());
882 }
883
884 pub fn reset_all(&self) {
886 self.table.reset_all();
887 }
888
889 #[cfg(unix)]
890 pub fn unlink(&self) -> Result<()> {
891 self.table.unlink()
892 }
893
894 fn locate(&self, id: ResourceId) -> Result<(usize, &ResourceSlot)> {
895 let geometry = self.table.geometry();
896 if id.kind() != self.table.kind()
897 || usize::from(id.node()) >= geometry.fleet_capacity
898 || id.slot() as usize >= geometry.lane_capacity
899 || id.generation() == 0
900 || id.generation() > GENERATION_MASK
901 {
902 return Err(Error::Malformed(id.to_string()));
903 }
904 let index = usize::from(id.node()) * geometry.lane_capacity + id.slot() as usize;
905 Ok((index, &self.table.resources()[index]))
906 }
907
908 fn owned(&self, id: ResourceId) -> Result<(usize, &ResourceSlot)> {
909 let (index, slot) = self.locate(id)?;
910 if !slot.is(id.generation()) {
911 return Err(Error::Stale(id));
912 }
913 if slot.owner_node.load(Ordering::Acquire) != self.table.node()
914 || slot.owner_incarnation.load(Ordering::Acquire) != self.table.incarnation()
915 {
916 return Err(Error::NotOwner(id));
917 }
918 Ok((index, slot))
919 }
920}
921
922#[derive(Debug)]
926pub enum Plan {
927 LocalReuse(Lease),
930 RemoteReuse(Lease),
933 Create(CreationPermit),
935 Wait(u32),
937 Reject(Reason),
938}
939
940pub struct Execution {
945 table: Arc<Table>,
946 lease: Lease,
947}
948
949impl Execution {
950 pub fn lease(&self) -> Lease {
951 self.lease
952 }
953
954 pub fn complete(self) {}
955}
956
957impl Drop for Execution {
958 fn drop(&mut self) {
959 let index = usize::from(self.lease.id.node()) * self.table.geometry().lane_capacity
960 + self.lease.id.slot() as usize;
961 let slot = &self.table.resources()[index];
962 if slot.generation.load(Ordering::Acquire) != self.lease.id.generation() {
965 return;
966 }
967 let _ = slot
968 .units
969 .try_update(Ordering::SeqCst, Ordering::SeqCst, |units| {
970 let (reserved, active) = unpack_counts(units);
971 Some(pack_counts(reserved, active.saturating_sub(1)))
972 });
973 self.table
974 .key_changed(slot.key_index.load(Ordering::Acquire) as usize);
975 }
976}
977
978pub struct CreationPermit {
982 table: Arc<Table>,
983 key_index: usize,
984}
985
986impl CreationPermit {
987 pub fn finish(self) {}
989}
990
991impl fmt::Debug for CreationPermit {
992 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
993 f.debug_struct("CreationPermit")
994 .field("key_index", &self.key_index)
995 .finish_non_exhaustive()
996 }
997}
998
999impl Drop for CreationPermit {
1000 fn drop(&mut self) {
1001 let _ = self.table.key(self.key_index).counts.try_update(
1002 Ordering::SeqCst,
1003 Ordering::SeqCst,
1004 |counts| {
1005 let (live, creating) = unpack_counts(counts);
1006 Some(pack_counts(live, creating.saturating_sub(1)))
1007 },
1008 );
1009 let _ = self
1010 .table
1011 .claims(usize::from(self.table.node()), self.key_index)
1012 .try_update(Ordering::SeqCst, Ordering::SeqCst, |held| {
1013 held.checked_sub(1)
1014 });
1015 self.table.key_changed(self.key_index);
1016 }
1017}
1018
1019#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
1020#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
1024pub(crate) fn waits_supported() -> bool {
1025 orbit_core::sync::supported()
1026}
1027
1028#[cfg(not(any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
1029pub(crate) fn waits_supported() -> bool {
1030 false
1031}
1032
1033#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
1036pub(crate) fn wait_on_timeout(word: &AtomicU32, expected: u32, timeout: Duration) -> Result<bool> {
1037 orbit_core::sync::wait_word_timeout(word, expected, timeout).map_err(Error::Io)
1038}
1039
1040pub(crate) fn wait_on(word: &AtomicU32, expected: u32) -> Result<()> {
1041 orbit_core::sync::wait_word(word, expected).map_err(Error::Io)
1042}
1043
1044#[cfg(not(any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
1045pub(crate) fn wait_on(_word: &AtomicU32, _expected: u32) -> Result<()> {
1046 Err(Error::Io(std::io::Error::new(std::io::ErrorKind::Unsupported, "orbit-pool needs a platform that can wait on a shared word")))
1047}
1048
1049pub(crate) fn wake_on(word: &AtomicU32) {
1050 #[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
1051 let _ = orbit_core::sync::wake_word(word);
1052 #[cfg(not(any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
1053 let _ = word;
1054}
1055
1056pub(crate) fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
1057 mutex.lock().unwrap_or_else(|error| error.into_inner())
1058}
1059
1060#[cfg(test)]
1061mod tests {
1062 use std::sync::Arc;
1063 use std::time::Duration;
1064
1065 use orbit_core::{Fleet, NodeId};
1066
1067 use super::{
1068 Error, Incarnation, Key, Limits, LocalFirst, LocalOnly, POOL_RESOURCE_LANE_CAPACITY, Plan,
1069 Pool, State,
1070 };
1071
1072 fn pool(name: &'static str) -> Pool {
1073 Pool::new(Arc::new(Fleet::join(name, 2).unwrap()), Incarnation::new(1)).unwrap()
1074 }
1075
1076 const KEY: Key = Key::new(0xC0FFEE);
1077
1078 #[test]
1079 fn an_exclusive_resource_is_leased_once_and_freed_by_the_owner() {
1080 let pool = pool("pool-exclusive");
1081 let id = pool.register(KEY, 1).unwrap();
1082 assert_eq!(id.node(), 0);
1083 let candidates = pool.candidates(KEY);
1084 assert_eq!(candidates.len(), 1);
1085 assert!(candidates[0].local);
1086 assert_eq!(candidates[0].free(), 1);
1087
1088 let lease = pool.reserve(id).unwrap();
1089 assert_eq!(lease.fence, 1);
1090 assert!(matches!(pool.reserve(id), Err(Error::Busy(_))));
1091 assert_eq!(pool.candidates(KEY)[0].reserved, 1);
1092
1093 let execution = pool.accept(lease).unwrap();
1094 let snapshot = pool.candidates(KEY)[0];
1095 assert_eq!((snapshot.reserved, snapshot.active), (0, 1));
1096 assert!(matches!(pool.reserve(id), Err(Error::Busy(_))));
1098 execution.complete();
1099 let again = pool.reserve(id).unwrap();
1100 assert_eq!(again.fence, 2);
1101 drop(pool.accept(again).unwrap());
1102 assert_eq!(pool.candidates(KEY)[0].free(), 1);
1103 }
1104
1105 #[test]
1106 fn an_abandoned_reservation_is_aged_out_by_the_owner_not_by_time_alone() {
1107 let pool = pool("pool-abandon");
1108 let id = pool.register(KEY, 1).unwrap();
1109 let abandoned = pool.reserve(id).unwrap();
1110 assert!(matches!(pool.reserve(id), Err(Error::Busy(_))));
1111 pool.reconcile(id, 0, Duration::from_secs(60)).unwrap();
1113 assert!(matches!(pool.reserve(id), Err(Error::Busy(_))));
1114 std::thread::sleep(Duration::from_millis(5));
1117 pool.reconcile(id, 0, Duration::from_millis(1)).unwrap();
1118 assert!(matches!(pool.accept(abandoned), Err(Error::NotReserved(_))));
1119 assert!(pool.reserve(id).is_ok());
1120 }
1121
1122 #[test]
1123 fn a_lease_is_accepted_exactly_once_and_leases_are_told_apart() {
1124 let pool = pool("pool-fence");
1125 let id = pool.register(KEY, 2).unwrap();
1126 let first = pool.reserve(id).unwrap();
1127 let second = pool.reserve(id).unwrap();
1128 assert_ne!(first.fence, second.fence);
1129 let running_second = pool.accept(second).unwrap();
1131 assert!(matches!(pool.accept(second), Err(Error::NotReserved(_))));
1132 let running_first = pool.accept(first).unwrap();
1133 assert!(matches!(pool.accept(first), Err(Error::NotReserved(_))));
1134 let snapshot = pool.candidates(KEY)[0];
1135 assert_eq!((snapshot.reserved, snapshot.active), (0, 2));
1136 drop(running_first);
1137 drop(running_second);
1138 assert_eq!(pool.candidates(KEY)[0].free(), 2);
1139
1140 let old = pool.reserve(id).unwrap();
1142 std::thread::sleep(Duration::from_millis(5));
1143 let fresh = pool.reserve(id).unwrap();
1144 pool.reconcile(id, 0, Duration::from_millis(2)).unwrap();
1145 assert!(matches!(pool.accept(old), Err(Error::NotReserved(_))));
1146 assert!(pool.accept(fresh).is_ok());
1147 }
1148
1149 #[test]
1150 fn an_owner_far_behind_makes_the_resource_refuse_reservations() {
1151 let pool = pool("pool-pending");
1152 let id = pool.register(KEY, u32::MAX).unwrap();
1153 let leases = (0..super::PENDING_RESERVATIONS)
1154 .map(|_| pool.reserve(id).unwrap())
1155 .collect::<Vec<_>>();
1156 assert!(matches!(pool.reserve(id), Err(Error::Busy(_))));
1157 let running = pool.accept(leases[0]).unwrap();
1158 assert!(pool.reserve(id).is_ok());
1159 drop(running);
1160 }
1161
1162 #[test]
1163 fn capacity_counts_and_drain_refuses_new_leases() {
1164 let pool = pool("pool-capacity");
1165 let id = pool.register(KEY, 3).unwrap();
1166 let leases = (0..3)
1167 .map(|_| pool.reserve(id).unwrap())
1168 .collect::<Vec<_>>();
1169 assert!(matches!(pool.reserve(id), Err(Error::Busy(_))));
1170 let running = pool.accept(leases[0]).unwrap();
1171 pool.drain(id).unwrap();
1172 assert!(matches!(pool.reserve(id), Err(Error::Draining(_))));
1173 assert!(matches!(pool.accept(leases[1]), Err(Error::Draining(_))));
1174 assert_eq!(pool.candidates(KEY)[0].state, State::Draining);
1175 running.complete();
1176 pool.unregister(id).unwrap();
1177 assert!(matches!(pool.accept(leases[2]), Err(Error::Stale(_))));
1178 assert!(pool.candidates(KEY).is_empty());
1179 assert_eq!(pool.budget(KEY), (0, 0));
1180 }
1181
1182 #[test]
1183 fn a_reused_slot_makes_the_old_id_stale() {
1184 let pool = pool("pool-stale");
1185 let first = pool.register(KEY, 1).unwrap();
1186 let lease = pool.reserve(first).unwrap();
1187 pool.unregister(first).unwrap();
1188 for _ in 0..POOL_RESOURCE_LANE_CAPACITY - 1 {
1189 pool.unregister(pool.register(KEY, 1).unwrap()).unwrap();
1190 }
1191 let second = pool.register(KEY, 1).unwrap();
1192 assert_eq!(second.slot(), first.slot());
1193 assert_ne!(second.generation(), first.generation());
1194 assert!(matches!(pool.reserve(first), Err(Error::Stale(_))));
1195 assert!(matches!(pool.accept(lease), Err(Error::Stale(_))));
1196 assert!(!pool.is_current(lease));
1197 }
1198
1199 #[test]
1200 fn creation_claims_hold_the_fleet_under_max_live() {
1201 let pool = pool("pool-claims");
1202 let first = pool.claim_create(KEY, 2).unwrap();
1203 let second = pool.claim_create(KEY, 2).unwrap();
1204 assert!(matches!(
1205 pool.claim_create(KEY, 2),
1206 Err(Error::CreationBudget { max_live: 2, .. })
1207 ));
1208 assert_eq!(pool.budget(KEY), (0, 2));
1209 let id = pool.register(KEY, 1).unwrap();
1210 first.finish();
1211 assert_eq!(pool.budget(KEY), (1, 1));
1212 drop(second);
1213 assert_eq!(pool.budget(KEY), (1, 0));
1214 assert!(pool.claim_create(KEY, 2).is_ok());
1215 pool.unregister(id).unwrap();
1216 assert_eq!(pool.budget(KEY), (0, 0));
1217
1218 let threads = (0..8)
1220 .map(|_| {
1221 let pool = pool.clone();
1222 std::thread::spawn(move || pool.claim_create(KEY, 3).ok())
1223 })
1224 .collect::<Vec<_>>();
1225 let held = threads
1226 .into_iter()
1227 .map(|thread| thread.join().unwrap())
1228 .collect::<Vec<_>>();
1229 assert_eq!(held.iter().filter(|claim| claim.is_some()).count(), 3);
1230 }
1231
1232 #[test]
1233 fn a_waiter_is_woken_when_the_owner_completes() {
1234 let pool = pool("pool-wait");
1235 let id = pool.register(KEY, 1).unwrap();
1236 let execution = pool.accept(pool.reserve(id).unwrap()).unwrap();
1237 let since = pool.version(KEY).unwrap();
1238 let waiter = {
1239 let pool = pool.clone();
1240 std::thread::spawn(move || pool.wait_capacity(KEY, since))
1241 };
1242 std::thread::sleep(Duration::from_millis(30));
1243 execution.complete();
1244 assert!(waiter.join().unwrap().unwrap() > since);
1245 assert!(pool.reserve(id).is_ok());
1246 }
1247
1248 #[test]
1249 fn a_death_report_closes_that_incarnations_resources_and_claims() {
1250 let pool = pool("pool-dead");
1251 let id = pool.register(KEY, 2).unwrap();
1252 let lease = pool.reserve(id).unwrap();
1253 let _claim = pool.claim_create(KEY, 4).unwrap();
1254 assert_eq!(pool.budget(KEY), (1, 1));
1255
1256 pool.node_dead(NodeId::ZERO, Incarnation::new(2));
1257 assert!(pool.is_current(lease));
1258
1259 pool.node_dead(NodeId::ZERO, Incarnation::new(1));
1260 assert!(!pool.is_current(lease));
1261 assert!(matches!(pool.reserve(id), Err(Error::Stale(_))));
1262 assert!(pool.candidates(KEY).is_empty());
1263 assert_eq!(pool.budget(KEY), (0, 0));
1264 }
1266
1267 #[test]
1268 fn a_reset_starts_a_new_epoch_and_a_slot_can_be_exhausted() {
1269 let pool = pool("pool-epoch");
1270 let id = pool.register(KEY, 1).unwrap();
1271 let before = pool.epoch();
1272 pool.reset_all();
1273 assert!(pool.epoch() > before);
1274 assert!(matches!(pool.reserve(id), Err(Error::Stale(_))));
1275 assert!(pool.candidates(KEY).is_empty());
1276
1277 for _ in 0..POOL_RESOURCE_LANE_CAPACITY * 4 {
1279 pool.unregister(pool.register(KEY, 1).unwrap()).unwrap();
1280 }
1281 assert!(matches!(pool.register(KEY, 1), Err(Error::Full { .. })));
1282 pool.reset_all();
1283 assert!(pool.register(KEY, 1).is_ok());
1284 }
1285
1286 #[test]
1287 fn acquire_walks_reuse_create_and_wait() {
1288 let pool = pool("pool-acquire");
1289 let limits = Limits {
1290 max_live: 2,
1291 attempts: 2,
1292 };
1293 let Plan::Create(permit) = pool.acquire(KEY, &limits, &LocalFirst).unwrap() else {
1295 panic!("expected Create");
1296 };
1297 let id = pool.register(KEY, 1).unwrap();
1298 permit.finish();
1299 let Plan::LocalReuse(lease) = pool.acquire(KEY, &limits, &LocalFirst).unwrap() else {
1301 panic!("expected LocalReuse");
1302 };
1303 let execution = pool.accept(lease).unwrap();
1304 assert!(matches!(
1306 pool.acquire(KEY, &limits, &LocalFirst).unwrap(),
1307 Plan::Create(_)
1308 ));
1309 let _held = pool.claim_create(KEY, 2).unwrap();
1312 assert!(matches!(
1313 pool.acquire(KEY, &limits, &LocalFirst).unwrap(),
1314 Plan::Wait(_)
1315 ));
1316 assert!(matches!(
1317 pool.acquire(KEY, &limits, &LocalOnly).unwrap(),
1318 Plan::Wait(_)
1319 ));
1320 execution.complete();
1321 assert!(matches!(
1322 pool.acquire(KEY, &limits, &LocalOnly).unwrap(),
1323 Plan::LocalReuse(_)
1324 ));
1325 let _ = id;
1326 }
1327}