1use std::collections::HashMap;
4use std::marker::PhantomData;
5use std::sync::atomic::{AtomicU8, AtomicU64, Ordering};
6use std::sync::{Arc, LazyLock, Mutex, MutexGuard, Weak};
7
8#[cfg(unix)]
9use orbit_core::shm::{ShmRegion, ring_segment_name};
10use orbit_core::{Fleet, NetId64, NodeId};
11
12use crate::layout::LockLayout;
13use crate::{Error, LockAcquire, LockKey, LockLease, LockOwner, LockTransition, Result};
14
15pub const LOCK_STATE_CAPACITY: usize =
16 orbit_core::compile::usize_from_env(option_env!("ORBIT_LOCK_STATE_CAPACITY"), 256);
17pub const LOCK_STATE_PAYLOAD_MAX: usize =
18 orbit_core::compile::usize_from_env(option_env!("ORBIT_LOCK_STATE_PAYLOAD_CAPACITY"), 960);
19
20const LOCK_COUNTER_MAX: u64 = 0xFF_FFFF_FFFF;
21const SLOT_EMPTY: u8 = 0;
22const SLOT_OCCUPIED: u8 = 1;
23const SLOT_TOMBSTONE: u8 = 2;
24const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
25const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
26
27#[cfg(unix)]
28const STATE_MAGIC: u32 = 0x4F_4C_53_54; #[cfg(unix)]
30const STATE_VERSION: u16 = 2;
31
32#[derive(Clone)]
33pub(crate) struct LockStateStore<L: LockLayout> {
34 backend: StateBacking,
35 state_kind: u8,
36 node_id: NodeId,
37 _layout: PhantomData<L>
38}
39
40#[derive(Clone)]
41enum StateBacking {
42 InMemory(Arc<Mutex<MemoryTable>>),
43 #[cfg(unix)]
44 Shm(Arc<ShmStateStore>)
45}
46
47impl<L: LockLayout> LockStateStore<L> {
48 pub(crate) fn new(fleet: &Arc<Fleet>) -> Self {
49 let backend = if fleet.is_shm() {
50 #[cfg(unix)]
51 {
52 StateBacking::Shm(Arc::new(ShmStateStore::new(&ring_segment_name(
53 fleet.name(),
54 L::STATE_KIND
55 ))))
56 }
57 #[cfg(not(unix))]
58 unreachable!("non-Unix fleets cannot use POSIX SHM")
59 } else {
60 StateBacking::InMemory(memory_table(fleet, L::STATE_KIND))
61 };
62 Self {
63 backend,
64 state_kind: L::STATE_KIND,
65 node_id: fleet.node_id(),
66 _layout: PhantomData
67 }
68 }
69
70 pub(crate) fn acquire(
71 &self,
72 key: &LockKey,
73 owner: &LockOwner,
74 now_ms: u64,
75 expires_at_ms: u64,
76 publish: impl FnOnce(&LockTransition) -> Result<()>
77 ) -> Result<LockAcquire> {
78 validate_entry(key, owner)?;
79 self.with_table(|header, slots| {
80 let hash = lock_hash(key);
81 let mut insertion = None;
82 for offset in 0..slots.len() {
83 let index = probe_index(hash, offset, slots.len());
84 let slot = &mut slots[index];
85 match slot.state() {
86 SLOT_EMPTY => {
87 insertion.get_or_insert(index);
88 break;
89 }
90 SLOT_TOMBSTONE => {
91 insertion.get_or_insert(index);
92 }
93 SLOT_OCCUPIED if slot.expires_at_ms <= now_ms => {
94 slot.tombstone();
95 insertion.get_or_insert(index);
96 }
97 SLOT_OCCUPIED if slot.matches(hash, key) => {
98 return Ok(LockAcquire::Occupied(slot.lease()));
99 }
100 _ => {}
101 }
102 }
103
104 let index = insertion.ok_or(Error::StateFull { capacity: slots.len() })?;
105 let lease = LockLease {
106 lock_id: mint_lock_id(header, self.state_kind, self.node_id)?,
107 key: key.clone(),
108 owner: owner.clone(),
109 acquired_at_ms: now_ms,
110 expires_at_ms,
111 state_revision: mint_revision(header)?
112 };
113 publish(&LockTransition::Acquired(lease.clone()))?;
114 slots[index].write(hash, &lease);
115 Ok(LockAcquire::Acquired(lease))
116 })
117 }
118
119 pub(crate) fn renew_id(
120 &self,
121 lease: &LockLease,
122 now_ms: u64,
123 requested_expires_at_ms: u64,
124 publish: impl FnOnce(&LockTransition) -> Result<()>
125 ) -> Result<Option<LockLease>> {
126 validate_entry(&lease.key, &lease.owner)?;
127 self.with_table(|header, slots| {
128 let Some(slot) = live_slot(slots, &lease.key, now_ms) else {
129 return Ok(None);
130 };
131 if slot.lock_id != lease.lock_id.raw() {
132 return Ok(None);
133 }
134 renew_slot(header, slot, requested_expires_at_ms, publish)
135 })
136 }
137
138 pub(crate) fn renew_owner(
139 &self,
140 key: &LockKey,
141 owner: &LockOwner,
142 now_ms: u64,
143 requested_expires_at_ms: u64,
144 publish: impl FnOnce(&LockTransition) -> Result<()>
145 ) -> Result<Option<LockLease>> {
146 validate_entry(key, owner)?;
147 self.with_table(|header, slots| {
148 let Some(slot) = live_slot(slots, key, now_ms) else {
149 return Ok(None);
150 };
151 if slot.owner() != owner.as_bytes() {
152 return Ok(None);
153 }
154 renew_slot(header, slot, requested_expires_at_ms, publish)
155 })
156 }
157
158 pub(crate) fn release_id(
159 &self,
160 lease: &LockLease,
161 now_ms: u64,
162 publish: impl FnOnce(&LockTransition) -> Result<()>
163 ) -> Result<bool> {
164 validate_entry(&lease.key, &lease.owner)?;
165 self.with_table(|header, slots| {
166 let Some(slot) = live_slot(slots, &lease.key, now_ms) else {
167 return Ok(false);
168 };
169 if slot.lock_id != lease.lock_id.raw() {
170 return Ok(false);
171 }
172 release_slot(header, slot, publish)
173 })
174 }
175
176 pub(crate) fn release_owner(
177 &self,
178 key: &LockKey,
179 owner: &LockOwner,
180 now_ms: u64,
181 publish: impl FnOnce(&LockTransition) -> Result<()>
182 ) -> Result<bool> {
183 validate_entry(key, owner)?;
184 self.with_table(|header, slots| {
185 let Some(slot) = live_slot(slots, key, now_ms) else {
186 return Ok(false);
187 };
188 if slot.owner() != owner.as_bytes() {
189 return Ok(false);
190 }
191 release_slot(header, slot, publish)
192 })
193 }
194
195 pub(crate) fn force_release(
196 &self,
197 key: &LockKey,
198 now_ms: u64,
199 publish: impl FnOnce(&LockTransition) -> Result<()>
200 ) -> Result<bool> {
201 validate_key(key)?;
202 self.with_table(|header, slots| {
203 let Some(index) = find_slot(slots, lock_hash(key), key) else {
204 return Ok(false);
205 };
206 let slot = &mut slots[index];
207 if slot.expires_at_ms <= now_ms {
208 slot.tombstone();
209 return Ok(false);
210 }
211 release_slot(header, slot, publish)
212 })
213 }
214
215 pub(crate) fn current(
216 &self,
217 key: &LockKey,
218 now_ms: u64
219 ) -> Result<Option<LockLease>> {
220 validate_key(key)?;
221 self.with_table(|_header, slots| Ok(live_slot(slots, key, now_ms).map(|slot| slot.lease())))
222 }
223
224 pub(crate) fn reset(&self) -> Result<()> {
225 self.with_table(|header, slots| {
226 for slot in slots {
227 slot.clear();
228 }
229 header.next_lock.store(0, Ordering::Relaxed);
230 header.next_revision.store(0, Ordering::Relaxed);
231 Ok(())
232 })
233 }
234
235 #[cfg(unix)]
236 pub(crate) fn unlink(&self) -> Result<()> {
237 match &self.backend {
238 StateBacking::InMemory(_) => self.reset(),
239 StateBacking::Shm(store) => store.open()?.unlink().map_err(Error::Io)
240 }
241 }
242
243 fn with_table<T>(
244 &self,
245 operation: impl FnOnce(&LockStateHeader, &mut [LockStateSlot]) -> Result<T>
246 ) -> Result<T> {
247 match &self.backend {
248 StateBacking::InMemory(table) => {
249 let mut table = lock_unpoisoned(table);
250 let MemoryTable { header, slots, .. } = &mut *table;
251 operation(header, slots)
252 }
253 #[cfg(unix)]
254 StateBacking::Shm(store) => store.open()?.with_table(operation).map_err(Error::Io)?
255 }
256 }
257}
258
259fn renew_slot(
260 header: &LockStateHeader,
261 slot: &mut LockStateSlot,
262 requested_expires_at_ms: u64,
263 publish: impl FnOnce(&LockTransition) -> Result<()>
264) -> Result<Option<LockLease>> {
265 let mut lease = slot.lease();
266 lease.expires_at_ms = lease.expires_at_ms.max(requested_expires_at_ms);
267 lease.state_revision = mint_revision(header)?;
268 publish(&LockTransition::Renewed(lease.clone()))?;
269 let hash = slot.key_hash;
270 slot.write(hash, &lease);
271 Ok(Some(lease))
272}
273
274fn release_slot(
275 header: &LockStateHeader,
276 slot: &mut LockStateSlot,
277 publish: impl FnOnce(&LockTransition) -> Result<()>
278) -> Result<bool> {
279 let lease = slot.lease();
280 let state_revision = mint_revision(header)?;
281 publish(&LockTransition::Released {
282 lock_id: lease.lock_id,
283 key: lease.key,
284 owner: lease.owner,
285 state_revision
286 })?;
287 slot.tombstone();
288 Ok(true)
289}
290
291fn live_slot<'a>(
292 slots: &'a mut [LockStateSlot],
293 key: &LockKey,
294 now_ms: u64
295) -> Option<&'a mut LockStateSlot> {
296 let hash = lock_hash(key);
297 let index = find_slot(slots, hash, key)?;
298 let slot = &mut slots[index];
299 if slot.expires_at_ms <= now_ms {
300 slot.tombstone();
301 None
302 } else {
303 Some(slot)
304 }
305}
306
307fn find_slot(
308 slots: &[LockStateSlot],
309 hash: u64,
310 key: &LockKey
311) -> Option<usize> {
312 for offset in 0..slots.len() {
313 let index = probe_index(hash, offset, slots.len());
314 let slot = &slots[index];
315 match slot.state() {
316 SLOT_EMPTY => return None,
317 SLOT_OCCUPIED if slot.matches(hash, key) => return Some(index),
318 _ => {}
319 }
320 }
321 None
322}
323
324fn validate_key(key: &LockKey) -> Result<()> {
325 if key.namespace().is_empty() {
326 return Err(Error::NamespaceEmpty);
327 }
328 if key.label().is_empty() {
329 return Err(Error::KeyEmpty);
330 }
331 if key.namespace().len() > u16::MAX as usize
332 || key.label().len() > u16::MAX as usize
333 || key.namespace().len().saturating_add(key.label().len()) > LOCK_STATE_PAYLOAD_MAX
334 {
335 return Err(Error::EntryTooLarge {
336 namespace_len: key.namespace().len(),
337 key_len: key.label().len(),
338 owner_len: 0,
339 max_payload: LOCK_STATE_PAYLOAD_MAX
340 });
341 }
342 Ok(())
343}
344
345fn validate_entry(
346 key: &LockKey,
347 owner: &LockOwner
348) -> Result<()> {
349 validate_key(key)?;
350 if owner.as_bytes().is_empty() {
351 return Err(Error::OwnerEmpty);
352 }
353 if key.namespace().len() > u16::MAX as usize
354 || key.label().len() > u16::MAX as usize
355 || owner.as_bytes().len() > u16::MAX as usize
356 || key
357 .namespace()
358 .len()
359 .saturating_add(key.label().len())
360 .saturating_add(owner.as_bytes().len())
361 > LOCK_STATE_PAYLOAD_MAX
362 {
363 return Err(Error::EntryTooLarge {
364 namespace_len: key.namespace().len(),
365 key_len: key.label().len(),
366 owner_len: owner.as_bytes().len(),
367 max_payload: LOCK_STATE_PAYLOAD_MAX
368 });
369 }
370 Ok(())
371}
372
373fn mint_lock_id(
374 header: &LockStateHeader,
375 state_kind: u8,
376 node_id: NodeId
377) -> Result<NetId64> {
378 let current = header.next_lock.load(Ordering::Relaxed);
379 let counter = current.checked_add(1).ok_or(Error::IdExhausted)?;
380 if counter > LOCK_COUNTER_MAX {
381 return Err(Error::IdExhausted);
382 }
383 header.next_lock.store(counter, Ordering::Relaxed);
384 Ok(NetId64::make(state_kind, node_id.get(), counter))
385}
386
387fn mint_revision(header: &LockStateHeader) -> Result<u64> {
388 let current = header.next_revision.load(Ordering::Relaxed);
389 let revision = current.checked_add(1).ok_or(Error::IdExhausted)?;
390 header.next_revision.store(revision, Ordering::Relaxed);
391 Ok(revision)
392}
393
394fn lock_hash(key: &LockKey) -> u64 {
395 let mut hash = FNV_OFFSET_BASIS;
396 for byte in key.namespace().iter().chain(key.label()) {
397 hash ^= u64::from(*byte);
398 hash = hash.wrapping_mul(FNV_PRIME);
399 }
400 hash
401}
402
403fn probe_index(
404 hash: u64,
405 offset: usize,
406 capacity: usize
407) -> usize {
408 debug_assert!(capacity.is_power_of_two());
409 (hash as usize).wrapping_add(offset) & (capacity - 1)
410}
411
412fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
413 mutex.lock().unwrap_or_else(|error| error.into_inner())
414}
415
416struct MemoryTable {
417 _fleet: Arc<Fleet>,
420 header: LockStateHeader,
421 slots: Vec<LockStateSlot>
422}
423
424impl MemoryTable {
425 fn new(fleet: Arc<Fleet>) -> Self {
426 Self {
427 _fleet: fleet,
428 header: LockStateHeader::new(),
429 slots: (0..LOCK_STATE_CAPACITY).map(|_| LockStateSlot::empty()).collect()
430 }
431 }
432}
433
434type MemoryRegistry = HashMap<(usize, u8), Weak<Mutex<MemoryTable>>>;
435static MEMORY_TABLES: LazyLock<Mutex<MemoryRegistry>> =
436 LazyLock::new(|| Mutex::new(HashMap::new()));
437
438fn memory_table(
439 fleet: &Arc<Fleet>,
440 state_kind: u8
441) -> Arc<Mutex<MemoryTable>> {
442 let fleet_identity = Arc::as_ptr(fleet) as usize;
443 let mut tables = lock_unpoisoned(&MEMORY_TABLES);
444 tables.retain(|_, table| table.strong_count() > 0);
445 if let Some(table) = tables.get(&(fleet_identity, state_kind)).and_then(Weak::upgrade) {
446 return table;
447 }
448 let table = Arc::new(Mutex::new(MemoryTable::new(Arc::clone(fleet))));
449 tables.insert((fleet_identity, state_kind), Arc::downgrade(&table));
450 table
451}
452
453#[repr(C, align(64))]
454struct LockStateHeader {
455 #[cfg(unix)]
456 magic: u32,
457 #[cfg(not(unix))]
458 _magic: u32,
459 #[cfg(unix)]
460 version: u16,
461 #[cfg(not(unix))]
462 _version: u16,
463 header_size: u16,
464 capacity: u32,
465 slot_size: u32,
466 next_lock: AtomicU64,
467 next_revision: AtomicU64,
468 _reserved: [u8; 32]
469}
470
471impl LockStateHeader {
472 fn new() -> Self {
473 Self {
474 #[cfg(unix)]
475 magic: STATE_MAGIC,
476 #[cfg(not(unix))]
477 _magic: 0,
478 #[cfg(unix)]
479 version: STATE_VERSION,
480 #[cfg(not(unix))]
481 _version: 0,
482 header_size: std::mem::size_of::<Self>() as u16,
483 capacity: LOCK_STATE_CAPACITY as u32,
484 slot_size: std::mem::size_of::<LockStateSlot>() as u32,
485 next_lock: AtomicU64::new(0),
486 next_revision: AtomicU64::new(0),
487 _reserved: [0; 32]
488 }
489 }
490}
491
492#[repr(C, align(64))]
493struct LockStateSlot {
494 state: AtomicU8,
495 _reserved: u8,
496 namespace_len: u16,
497 key_len: u16,
498 owner_len: u16,
499 key_hash: u64,
500 lock_id: u64,
501 state_revision: u64,
502 acquired_at_ms: u64,
503 expires_at_ms: u64,
504 payload: [u8; LOCK_STATE_PAYLOAD_MAX]
505}
506
507impl LockStateSlot {
508 fn empty() -> Self {
509 Self {
510 state: AtomicU8::new(SLOT_EMPTY),
511 _reserved: 0,
512 namespace_len: 0,
513 key_len: 0,
514 owner_len: 0,
515 key_hash: 0,
516 lock_id: 0,
517 state_revision: 0,
518 acquired_at_ms: 0,
519 expires_at_ms: 0,
520 payload: [0; LOCK_STATE_PAYLOAD_MAX]
521 }
522 }
523
524 fn state(&self) -> u8 {
525 self.state.load(Ordering::Acquire)
526 }
527
528 fn namespace_bytes(&self) -> &[u8] {
529 &self.payload[..usize::from(self.namespace_len)]
530 }
531
532 fn key_bytes(&self) -> &[u8] {
533 let start = usize::from(self.namespace_len);
534 let end = start + usize::from(self.key_len);
535 &self.payload[start..end]
536 }
537
538 fn owner(&self) -> &[u8] {
539 let start = usize::from(self.namespace_len) + usize::from(self.key_len);
540 let end = start + usize::from(self.owner_len);
541 &self.payload[start..end]
542 }
543
544 fn matches(
545 &self,
546 hash: u64,
547 key: &LockKey
548 ) -> bool {
549 self.state() == SLOT_OCCUPIED
550 && self.key_hash == hash
551 && self.namespace_bytes() == key.namespace()
552 && self.key_bytes() == key.label()
553 }
554
555 fn lease(&self) -> LockLease {
556 LockLease {
557 lock_id: NetId64::from_raw(self.lock_id),
558 key: LockKey::from_parts(
559 bytes::Bytes::copy_from_slice(self.namespace_bytes()),
560 bytes::Bytes::copy_from_slice(self.key_bytes())
561 ),
562 owner: LockOwner::from(bytes::Bytes::copy_from_slice(self.owner())),
563 acquired_at_ms: self.acquired_at_ms,
564 expires_at_ms: self.expires_at_ms,
565 state_revision: self.state_revision
566 }
567 }
568
569 fn write(
570 &mut self,
571 hash: u64,
572 lease: &LockLease
573 ) {
574 self.state.store(SLOT_TOMBSTONE, Ordering::Release);
575 self.namespace_len = lease.key.namespace().len() as u16;
576 self.key_len = lease.key.label().len() as u16;
577 self.owner_len = lease.owner.as_bytes().len() as u16;
578 self.key_hash = hash;
579 self.lock_id = lease.lock_id.raw();
580 self.state_revision = lease.state_revision;
581 self.acquired_at_ms = lease.acquired_at_ms;
582 self.expires_at_ms = lease.expires_at_ms;
583 self.payload.fill(0);
584 let namespace_end = lease.key.namespace().len();
585 self.payload[..namespace_end].copy_from_slice(lease.key.namespace());
586 let key_end = namespace_end + lease.key.label().len();
587 self.payload[namespace_end..key_end].copy_from_slice(lease.key.label());
588 let owner_start = key_end;
589 let owner_end = owner_start + lease.owner.as_bytes().len();
590 self.payload[owner_start..owner_end].copy_from_slice(lease.owner.as_bytes());
591 self.state.store(SLOT_OCCUPIED, Ordering::Release);
592 }
593
594 fn tombstone(&self) {
595 self.state.store(SLOT_TOMBSTONE, Ordering::Release);
596 }
597
598 fn clear(&self) {
599 self.state.store(SLOT_EMPTY, Ordering::Release);
600 }
601}
602
603#[cfg(unix)]
604struct ShmStateStore {
605 name: String,
606 opened: Mutex<Option<Arc<ShmLockTable>>>
607}
608
609#[cfg(unix)]
610impl ShmStateStore {
611 fn new(name: &str) -> Self {
612 Self { name: name.to_owned(), opened: Mutex::new(None) }
613 }
614
615 fn open(&self) -> Result<Arc<ShmLockTable>> {
616 let mut opened = lock_unpoisoned(&self.opened);
617 if let Some(table) = opened.as_ref() {
618 return Ok(table.clone());
619 }
620 let table = Arc::new(ShmLockTable::open_or_create(&self.name)?);
621 *opened = Some(table.clone());
622 Ok(table)
623 }
624}
625
626#[cfg(unix)]
627struct ShmLockTable {
628 region: ShmRegion,
629 local_lock: Mutex<()>
630}
631
632#[cfg(unix)]
633impl ShmLockTable {
634 fn open_or_create(name: &str) -> Result<Self> {
635 use std::{io, ptr};
636
637 let (region, _initialization_lock) =
638 ShmRegion::open_or_create_locked(name, shm_segment_size())?;
639 if region.created() {
640 unsafe {
641 ptr::write(region.as_ptr().cast::<LockStateHeader>(), LockStateHeader::new());
642 let slots = region.as_ptr().add(std::mem::size_of::<LockStateHeader>());
643 ptr::write_bytes(
644 slots,
645 0,
646 LOCK_STATE_CAPACITY * std::mem::size_of::<LockStateSlot>()
647 );
648 }
649 } else {
650 let header = unsafe { &*region.as_ptr().cast::<LockStateHeader>() };
651 if header.magic != STATE_MAGIC {
652 return Err(Error::Io(io::Error::new(
653 io::ErrorKind::InvalidData,
654 format!(
655 "SHM segment {name} has wrong magic 0x{:08X} (expected 0x{STATE_MAGIC:08X})",
656 header.magic
657 )
658 )));
659 }
660 if header.version != STATE_VERSION
661 || usize::from(header.header_size) != std::mem::size_of::<LockStateHeader>()
662 || header.capacity as usize != LOCK_STATE_CAPACITY
663 || header.slot_size as usize != std::mem::size_of::<LockStateSlot>()
664 {
665 return Err(Error::Io(io::Error::new(
666 io::ErrorKind::InvalidData,
667 format!("SHM segment {name} has an incompatible lock-state layout")
668 )));
669 }
670 }
671
672 Ok(Self { region, local_lock: Mutex::new(()) })
673 }
674
675 fn with_table<T>(
676 &self,
677 operation: impl FnOnce(&LockStateHeader, &mut [LockStateSlot]) -> Result<T>
678 ) -> std::io::Result<Result<T>> {
679 let _local = lock_unpoisoned(&self.local_lock);
680 let _process = self.region.lock_exclusive()?;
681 let header = unsafe { &*self.region.as_ptr().cast::<LockStateHeader>() };
682 let slots = unsafe {
683 std::slice::from_raw_parts_mut(
684 self.region
685 .as_ptr()
686 .add(std::mem::size_of::<LockStateHeader>())
687 .cast::<LockStateSlot>(),
688 LOCK_STATE_CAPACITY
689 )
690 };
691 Ok(operation(header, slots))
692 }
693
694 fn unlink(&self) -> std::io::Result<()> {
695 self.region.unlink()
696 }
697}
698
699#[cfg(unix)]
700fn shm_segment_size() -> usize {
701 std::mem::size_of::<LockStateHeader>()
702 + LOCK_STATE_CAPACITY * std::mem::size_of::<LockStateSlot>()
703}
704
705const _: () = assert!(LOCK_STATE_CAPACITY.is_power_of_two());
706const _: () = assert!(LOCK_STATE_CAPACITY <= u32::MAX as usize);
707const _: () = assert!(LOCK_STATE_PAYLOAD_MAX > 0);
708const _: () = assert!(std::mem::size_of::<LockStateHeader>() == 64);
709const _: () = assert!(std::mem::size_of::<LockStateSlot>() <= u32::MAX as usize);