ferroid/generator/snowflake/atomic.rs
1use crate::{IdGenStatus, Result, SnowflakeGenerator, SnowflakeId, TimeSource};
2use core::{
3 cmp,
4 marker::PhantomData,
5 sync::atomic::{AtomicU64, Ordering},
6};
7#[cfg(feature = "tracing")]
8use tracing::instrument;
9
10/// A lock-free Snowflake ID generator suitable for multi-threaded environments.
11///
12/// This generator stores the Snowflake state in an [`AtomicU64`], allowing safe
13/// shared use across threads.
14///
15/// ## Features
16///
17/// - ✅ Thread-safe
18/// - ❌ Safely implement any [`SnowflakeId`] layout
19///
20/// ## Caveats
21/// This implementation uses an `AtomicU64` internally, so it only supports ID
22/// layouts where the underlying type is `u64`. You cannot use layouts with
23/// larger or smaller representations (i.e., `ID::Ty` must be `u64`).
24///
25/// ## Recommended When
26/// - You're in a multi-threaded environment
27/// - Fair access is sacrificed for higher throughput
28///
29/// ## See Also
30/// - [`BasicSnowflakeGenerator`]
31/// - [`LockSnowflakeGenerator`]
32///
33/// [`BasicSnowflakeGenerator`]: crate::BasicSnowflakeGenerator
34/// [`LockSnowflakeGenerator`]: crate::LockSnowflakeGenerator
35pub struct AtomicSnowflakeGenerator<ID, T>
36where
37 ID: SnowflakeId<Ty = u64>,
38 T: TimeSource<ID::Ty>,
39{
40 state: AtomicU64,
41 machine_id: u64,
42 time: T,
43 _id: PhantomData<ID>,
44}
45
46impl<ID, T> AtomicSnowflakeGenerator<ID, T>
47where
48 ID: SnowflakeId<Ty = u64>,
49 T: TimeSource<ID::Ty>,
50{
51 /// Creates a new [`AtomicSnowflakeGenerator`] initialized with the current
52 /// time and a given machine ID.
53 ///
54 /// This constructor sets the initial timestamp and sequence to zero, and
55 /// uses the provided `clock` to fetch the current time during ID
56 /// generation. It is the recommended way to create a new atomic generator
57 /// for typical use cases.
58 ///
59 /// # Parameters
60 ///
61 /// - `machine_id`: A unique identifier for the node or instance generating
62 /// IDs. This value will be encoded into every generated ID.
63 /// - `clock`: A [`TimeSource`] implementation (e.g., [`MonotonicClock`])
64 /// that determines how timestamps are generated.
65 ///
66 /// # Returns
67 ///
68 /// A new [`AtomicSnowflakeGenerator`] ready to produce unique, time-ordered
69 /// IDs.
70 ///
71 /// # Example
72 ///
73 /// ```
74 /// #[cfg(all(feature = "std", feature = "alloc", feature = "snowflake"))]
75 /// {
76 /// use ferroid::{AtomicSnowflakeGenerator, IdGenStatus, SnowflakeTwitterId, TWITTER_EPOCH, MonotonicClock};
77 ///
78 /// let generator = AtomicSnowflakeGenerator::new(0, MonotonicClock::with_epoch(TWITTER_EPOCH));
79 ///
80 /// let id: SnowflakeTwitterId = loop {
81 /// match generator.next_id() {
82 /// IdGenStatus::Ready { id } => break id,
83 /// IdGenStatus::Pending { .. } => core::hint::spin_loop(),
84 /// }
85 /// };
86 /// }
87 /// ```
88 ///
89 /// [`TimeSource`]: crate::TimeSource
90 /// [`MonotonicClock`]: crate::MonotonicClock
91 pub fn new(machine_id: ID::Ty, clock: T) -> Self {
92 Self::from_components(ID::ZERO, machine_id, ID::ZERO, clock)
93 }
94
95 /// Creates a new ID generator from explicit component values.
96 ///
97 /// This constructor is primarily useful for advanced use cases such as
98 /// restoring state from persistent storage or controlling the starting
99 /// point of the generator manually.
100 ///
101 /// # Parameters
102 /// - `timestamp`: The initial timestamp component (usually in milliseconds)
103 /// - `machine_id`: The machine or worker identifier
104 /// - `sequence`: The initial sequence number
105 /// - `clock`: A [`TimeSource`] implementation used to fetch the current
106 /// time
107 ///
108 /// # Returns
109 /// A new generator instance preloaded with the given state.
110 ///
111 /// # ⚠️ Note
112 /// In typical use cases, you should prefer [`Self::new`] to let the
113 /// generator initialize itself from the current time.
114 pub fn from_components(
115 timestamp: ID::Ty,
116 machine_id: ID::Ty,
117 sequence: ID::Ty,
118 clock: T,
119 ) -> Self {
120 let initial = ID::from_components(timestamp, machine_id, sequence);
121 Self {
122 state: AtomicU64::new(initial.to_raw()),
123 machine_id,
124 time: clock,
125 _id: PhantomData,
126 }
127 }
128
129 /// Attempts to generate the next available ID.
130 ///
131 /// Returns a new, time-ordered, unique ID if generation succeeds. If the
132 /// generator is temporarily exhausted (e.g., the sequence is full and the
133 /// clock has not advanced, or CAS fails), it returns
134 /// [`IdGenStatus::Pending`].
135 ///
136 /// # Panics
137 /// This method currently has no fallible code paths, but may panic if an
138 /// internal error occurs in future implementations. For explicitly fallible
139 /// behavior, use [`Self::try_next_id`] instead.
140 ///
141 /// # Example
142 /// ```
143 /// #[cfg(all(feature = "std", feature = "alloc", feature = "snowflake"))]
144 /// {
145 /// use ferroid::{AtomicSnowflakeGenerator, IdGenStatus, SnowflakeTwitterId, TWITTER_EPOCH, MonotonicClock};
146 ///
147 /// let generator = AtomicSnowflakeGenerator::new(0, MonotonicClock::with_epoch(TWITTER_EPOCH));
148 ///
149 /// let id: SnowflakeTwitterId = loop {
150 /// match generator.next_id() {
151 /// IdGenStatus::Ready { id } => break id,
152 /// IdGenStatus::Pending { .. } => std::thread::yield_now(),
153 /// }
154 /// };
155 /// }
156 /// ```
157 pub fn next_id(&self) -> IdGenStatus<ID> {
158 self.try_next_id().unwrap()
159 }
160
161 /// A fallible version of [`Self::next_id`] that returns a [`Result`].
162 ///
163 /// This method attempts to generate the next ID based on the current time
164 /// and internal state. If successful, it returns [`IdGenStatus::Ready`]
165 /// with a newly generated ID. If the generator is temporarily exhausted or
166 /// CAS fails, it returns [`IdGenStatus::Pending`]. If an internal failure
167 /// occurs (e.g., a time source or lock error), it returns an error.
168 ///
169 /// # Returns
170 /// - `Ok(IdGenStatus::Ready { id })`: A new ID is available
171 /// - `Ok(IdGenStatus::Pending { yield_for })`: The time to wait (in
172 /// milliseconds) before trying again
173 /// - `Err(_)`: infallible for this generator
174 ///
175 /// # Errors
176 /// - This method currently does not return any errors and always returns
177 /// `Ok`. It is marked as fallible to allow for future extensibility
178 ///
179 /// # Example
180 /// ```
181 /// #[cfg(all(feature = "std", feature = "alloc", feature = "snowflake"))]
182 /// {
183 /// use ferroid::{AtomicSnowflakeGenerator, ToU64, IdGenStatus, SnowflakeTwitterId, TWITTER_EPOCH, MonotonicClock};
184 ///
185 /// let generator = AtomicSnowflakeGenerator::new(0, MonotonicClock::with_epoch(TWITTER_EPOCH));
186 ///
187 /// // Attempt to generate a new ID
188 /// let id: SnowflakeTwitterId = loop {
189 /// match generator.try_next_id() {
190 /// Ok(IdGenStatus::Ready { id }) => break id,
191 /// Ok(IdGenStatus::Pending { yield_for }) => {
192 /// std::thread::sleep(core::time::Duration::from_millis(yield_for.to_u64()));
193 /// }
194 /// Err(_) => unreachable!(),
195 /// }
196 /// };
197 /// }
198 /// ```
199 #[cfg_attr(feature = "tracing", instrument(level = "trace", skip(self)))]
200 pub fn try_next_id(&self) -> Result<IdGenStatus<ID>> {
201 let now = self.time.current_millis();
202
203 let current_raw = self.state.load(Ordering::Relaxed);
204 let current_id = ID::from_raw(current_raw);
205
206 let current_ts = current_id.timestamp();
207 let seq = current_id.sequence();
208
209 let (next_ts, next_seq) = match now.cmp(¤t_ts) {
210 cmp::Ordering::Less => {
211 let yield_for = current_ts - now;
212 debug_assert!(yield_for >= ID::ZERO);
213 return Ok(IdGenStatus::Pending { yield_for });
214 }
215 cmp::Ordering::Greater => (now, ID::ZERO),
216 cmp::Ordering::Equal => {
217 if seq < ID::max_sequence() {
218 (current_ts, seq + ID::ONE)
219 } else {
220 return Ok(IdGenStatus::Pending { yield_for: ID::ONE });
221 }
222 }
223 };
224
225 let next_id = ID::from_components(next_ts, self.machine_id, next_seq);
226 let next_raw = next_id.to_raw();
227
228 if self
229 .state
230 .compare_exchange(current_raw, next_raw, Ordering::Relaxed, Ordering::Relaxed)
231 .is_ok()
232 {
233 Ok(IdGenStatus::Ready { id: next_id })
234 } else {
235 // CAS failed - another thread won the race. Yield 0 to retry
236 // immediately.
237 Ok(IdGenStatus::Pending {
238 yield_for: ID::ZERO,
239 })
240 }
241 }
242}
243
244impl<ID, T> SnowflakeGenerator<ID, T> for AtomicSnowflakeGenerator<ID, T>
245where
246 ID: SnowflakeId<Ty = u64>,
247 T: TimeSource<ID::Ty>,
248{
249 type Err = core::convert::Infallible;
250
251 fn new(machine_id: ID::Ty, clock: T) -> Self {
252 Self::new(machine_id, clock)
253 }
254
255 fn next_id(&self) -> IdGenStatus<ID> {
256 self.next_id()
257 }
258
259 fn try_next_id(&self) -> Result<IdGenStatus<ID>, Self::Err> {
260 self.try_next_id()
261 }
262}