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