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