typedflake/
state.rs

1//! Atomic state management and lazy state pooling.
2//!
3//! Each (worker_id, process_id) instance maintains its own [`State`] with packed atomic storage
4//! for timestamp and sequence. States are cache-line aligned to prevent false sharing.
5//!
6//! The [`StatePool`] provides lazy, on-demand state allocation using DashMap for lock-free
7//! concurrent access. Multiple generators for the same (worker_id, process_id) pair share
8//! the same underlying state.
9
10use crate::config::Config;
11use dashmap::DashMap;
12use std::sync::Arc;
13use std::sync::atomic::AtomicU64;
14
15/// Atomic state for each (worker_id, process_id) instance
16#[derive(Debug)]
17#[repr(align(64))] // Cache-line alignment to prevent false sharing across cores
18pub struct State {
19    /// Packed state: upper bits = timestamp, lower bits = sequence
20    pub packed: AtomicU64,
21}
22
23impl State {
24    pub fn new(packed_state: u64) -> Self {
25        Self {
26            packed: AtomicU64::new(packed_state),
27        }
28    }
29
30    /// Pack timestamp and sequence into single u64
31    #[inline(always)]
32    pub fn pack_state(&self, timestamp: u64, sequence: u64, config: Config) -> u64 {
33        (timestamp << config.layout().sequence()) | (sequence & config.layout().sequence_max())
34    }
35
36    /// Unpack timestamp and sequence from single u64
37    #[inline(always)]
38    pub fn unpack_state(&self, packed: u64, config: Config) -> (u64, u64) {
39        let timestamp = packed >> config.layout().sequence();
40        let sequence = packed & config.layout().sequence_max();
41        (timestamp, sequence)
42    }
43}
44
45impl Default for State {
46    fn default() -> Self {
47        Self {
48            packed: AtomicU64::new(0),
49        }
50    }
51}
52
53/// Lazy state pool with DashMap for on-demand allocation
54#[derive(Debug)]
55pub struct StatePool {
56    states: DashMap<u32, Arc<State>>,
57    config: Config,
58}
59
60impl StatePool {
61    /// Create with empty DashMap - states allocated on-demand
62    pub fn new(config: Config) -> Self {
63        Self {
64            states: DashMap::new(),
65            config,
66        }
67    }
68
69    /// Pack (worker_id, process_id) into u32 key using bit layout
70    #[inline(always)]
71    fn pack_key(&self, worker_id: u64, process_id: u64) -> u32 {
72        // Use max values from config for bounds safety
73        let masked_worker = worker_id & self.config.layout().worker_max();
74        let masked_process = process_id & self.config.layout().process_max();
75
76        // Mathematical mapping: key = worker_id * 2^process_bits + process_id
77        // This uses the same shift logic as ID composition but for key packing
78        let key = (masked_worker << self.config.layout().process()) | masked_process;
79        key as u32
80    }
81
82    /// Get state for (worker_id, process_id) - lazy initialization on first access
83    #[inline]
84    pub fn get_state(&self, worker_id: u64, process_id: u64) -> Arc<State> {
85        let key = self.pack_key(worker_id, process_id);
86
87        self.states
88            .entry(key)
89            .or_insert_with(|| Arc::new(State::default()))
90            .clone()
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97    use crate::Config;
98    use std::sync::Arc;
99
100    #[test]
101    fn pool_starts_empty() {
102        let config = Config::default();
103        let state_pool = StatePool::new(config);
104
105        // Should start with no states (lazy initialization)
106        assert_eq!(state_pool.states.len(), 0);
107    }
108
109    #[test]
110    fn pool_lazy_initialization() {
111        let config = Config::default();
112        let state_pool = StatePool::new(config);
113
114        // Initially empty
115        assert_eq!(state_pool.states.len(), 0);
116
117        // Access state (0, 0)
118        let _state1 = state_pool.get_state(0, 0);
119        assert_eq!(state_pool.states.len(), 1);
120
121        // Access different state (1, 0)
122        let _state2 = state_pool.get_state(1, 0);
123        assert_eq!(state_pool.states.len(), 2);
124
125        // Access another different state (0, 1)
126        let _state3 = state_pool.get_state(0, 1);
127        assert_eq!(state_pool.states.len(), 3);
128
129        // Re-accessing same state doesn't create new entry
130        let _state1_again = state_pool.get_state(0, 0);
131        assert_eq!(state_pool.states.len(), 3);
132    }
133
134    #[test]
135    fn pool_state_identity() {
136        let config = Config::default();
137        let state_pool = StatePool::new(config);
138
139        // Test state identity
140        let state1 = state_pool.get_state(0, 0);
141        let state2 = state_pool.get_state(1, 0);
142        let state3 = state_pool.get_state(0, 1);
143
144        // Different worker/process should give different states
145        assert!(!Arc::ptr_eq(&state1, &state2));
146        assert!(!Arc::ptr_eq(&state1, &state3));
147        assert!(!Arc::ptr_eq(&state2, &state3));
148
149        // Same worker/process should give same state
150        let state1_again = state_pool.get_state(0, 0);
151        assert!(Arc::ptr_eq(&state1, &state1_again));
152    }
153
154    #[test]
155    fn pool_key_packing() {
156        let config = Config::default();
157        let state_pool = StatePool::new(config);
158
159        // Test that key packing is deterministic
160        let key1 = state_pool.pack_key(5, 3);
161        let key2 = state_pool.pack_key(5, 3);
162        assert_eq!(key1, key2);
163
164        // Different (worker, process) should give different keys
165        let key3 = state_pool.pack_key(5, 4);
166        let key4 = state_pool.pack_key(6, 3);
167        assert_ne!(key1, key3);
168        assert_ne!(key1, key4);
169        assert_ne!(key3, key4);
170    }
171}
typedflake/state.rs

typedflake/
state.rs
