snowflaked/
lib.rs

1//! A crate for working with snowflake ids.
2//!
3//! Most notably this provides [`Snowflake`] for working with custom snowflake ids and
4//! [`Generator`] creating new snowflake ids.
5//!
6//! # Snowflake structure
7//!
8//! A snowflake id is a 64-bit integer generated using the current timestamp in milliseconds, a
9//! constant instance and a sequence number.
10//!
11//! ```text
12//! 0               1               2                 3               4                 5               6               7
13//! 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 8 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7
14//! +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
15//! |                                       Timestamp                                       |     Instance      |       Sequence      |
16//! +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
17//! ```
18//!
19//! - The [`Snowflake`] implementation for `u64` uses 42 bits for the timestamp, 10 bits for the
20//! instance and 12 bits for the sequence.
21//! - The [`Snowflake`] implementation for `i64` uses 41 bits for the timestamp, 10 bits for the
22//! instance and 12 bits for the sequence.
23//!
24//! ## Timestamp overflow
25//!
26//! Since the timestamp range is limited it is possible for the timestamp to overflow and wrap
27//! around after a specific date. For the by-default configured UNIX epoch these dates are:
28//! - For `i64`: Sep 07 2039 15:47:35 (`2039-09-07T15:47:35Z`)
29//! - For `u64`: May 15 2109 07:35:11 (`2109-05-15T07:35:11Z`)
30//!
31//! If overflowing after these dates is not acceptable for you [`Builder::epoch`] allows
32//! configuring a custom epoch.
33//!
34//! # Custom snowflake ids
35//!
36//! Custom snowflake ids can be created with the [`Snowflake`] trait.
37//!
38//! ## Example
39//! ```
40//! use snowflaked::Snowflake;
41//!
42//! struct UserId(u64);
43//!
44//! impl Snowflake for UserId {
45//!     fn from_parts(timestamp: u64, instance: u64, sequence: u64) -> Self {
46//!         Self(u64::from_parts(timestamp, instance, sequence))
47//!     }
48//!
49//!     fn timestamp(&self) -> u64 {
50//!         self.0.timestamp()
51//!     }
52//!
53//!     fn instance(&self) -> u64 {
54//!         self.0.instance()
55//!     }
56//!
57//!     fn sequence(&self) -> u64 {
58//!         self.0.sequence()
59//!     }
60//! }
61//! ```
62//!
63//! # Generating snowflake ids
64//!
65//! [`Generator`] can be used to generate unique snowflake ids. Additionally [`sync::Generator`]
66//! can be used when working with multiple threads (requires the `sync` feature).
67//!
68//! ## Example
69//! ```
70//! use snowflaked::Generator;
71//!
72//! let mut generator = Generator::new(0);
73//! let id: u64 = generator.generate();
74//! ```
75//!
76//! [`Generator::generate`] can also generate custom snowflake ids:
77//! ```
78//! # use snowflaked::Snowflake;
79//! #
80//! # pub struct UserId(u64);
81//! #
82//! # impl Snowflake for UserId {
83//! #    fn from_parts(timestamp: u64, instance: u64, sequence: u64) -> Self {
84//! #       Self(u64::from_parts(timestamp, instance, sequence))
85//! #    }
86//! #
87//! #    fn timestamp(&self) -> u64 {
88//! #        self.0.timestamp()
89//! #    }
90//! #
91//! #    fn instance(&self) -> u64 {
92//! #        self.0.instance()
93//! #    }
94//! #
95//! #    fn sequence(&self) -> u64 {
96//! #        self.0.sequence()
97//! #    }
98//! # }
99//!
100//! use snowflaked::Generator;
101//!
102//! let mut generator = Generator::new(0);
103//! let id: UserId = generator.generate();
104//! ```
105//!
106//! For more details on [`sync::Generator`] see the [`sync`] module.
107//!
108//! # Feature flags
109//! `sync`: Enables the [`sync`] module.
110#![cfg_attr(docsrs, feature(doc_cfg))]
111
112mod builder;
113mod loom;
114mod time;
115
116#[cfg(feature = "sync")]
117#[cfg_attr(docsrs, doc(cfg(feature = "sync")))]
118pub mod sync;
119
120pub use builder::Builder;
121
122use std::time::{SystemTime, UNIX_EPOCH};
123
124const BITMASK_INSTANCE: u64 = 0x3FF000;
125const BITMASK_SEQUENCE: u64 = 0xFFF;
126
127const INSTANCE_MAX: u16 = 2_u16.pow(10) - 1;
128
129/// A type that can be used as a snowflake id.
130pub trait Snowflake {
131    /// Creates a new value from the snowflake parts.
132    fn from_parts(timestamp: u64, instance: u64, sequence: u64) -> Self;
133
134    /// Returns the timestamp component of the snowflake.
135    fn timestamp(&self) -> u64;
136
137    /// Returns the instance component of the snowflake.
138    fn instance(&self) -> u64;
139
140    /// Returns the sequence component of the snowflake.
141    fn sequence(&self) -> u64;
142}
143
144impl Snowflake for u64 {
145    fn from_parts(timestamp: u64, instance: u64, sequence: u64) -> Self {
146        let timestamp = timestamp << 22;
147        let instance = (instance << 12) & BITMASK_INSTANCE;
148        timestamp | instance | sequence
149    }
150
151    #[inline]
152    fn timestamp(&self) -> u64 {
153        self >> 22
154    }
155
156    #[inline]
157    fn instance(&self) -> u64 {
158        (self & BITMASK_INSTANCE) >> 12
159    }
160
161    #[inline]
162    fn sequence(&self) -> u64 {
163        self & BITMASK_SEQUENCE
164    }
165}
166
167impl Snowflake for i64 {
168    fn from_parts(timestamp: u64, instance: u64, sequence: u64) -> Self {
169        // Don't set the sign bit.
170        let timestamp = (timestamp << 22) & (i64::MAX as u64);
171        let instance = (instance << 12) & BITMASK_INSTANCE;
172        (timestamp | instance | sequence) as i64
173    }
174
175    #[inline]
176    fn timestamp(&self) -> u64 {
177        *self as u64 >> 22
178    }
179
180    #[inline]
181    fn instance(&self) -> u64 {
182        (*self as u64 & BITMASK_INSTANCE) >> 12
183    }
184
185    #[inline]
186    fn sequence(&self) -> u64 {
187        *self as u64 & BITMASK_SEQUENCE
188    }
189}
190
191/// A generator for new unique [`Snowflake`] ids.
192///
193/// # Examples
194///
195/// ```
196/// # use snowflaked::Generator;
197/// #
198/// let mut generator = Generator::new(0);
199///
200/// let id: u64 = generator.generate();
201/// ```
202#[derive(Debug)]
203pub struct Generator {
204    components: Components,
205    epoch: SystemTime,
206}
207
208impl Generator {
209    /// Creates a new `Generator` using the given `instance`.
210    ///
211    /// # Panics
212    ///
213    /// Panics if `instance` exceeds the maximum value (2^10 - 1).
214    ///
215    /// # Examples
216    ///
217    /// ```
218    /// # use snowflaked::Generator;
219    /// #
220    /// let mut generator = Generator::new(123);
221    ///
222    /// assert_eq!(generator.instance(), 123);
223    /// ```
224    ///
225    /// Providing an invalid `instance` will panic.
226    ///
227    /// ```should_panic
228    /// # use snowflaked::Generator;
229    /// #
230    /// let mut generator = Generator::new(10_000);
231    /// ```
232    #[inline]
233    pub const fn new(instance: u16) -> Self {
234        match Self::new_checked(instance) {
235            Some(this) => this,
236            None => const_panic_new(),
237        }
238    }
239
240    /// Creates a new `Generator` using the given `instance`. Returns `None` if the instance
241    /// exceeds the maximum instance value (2^10 - 1).
242    ///
243    /// # Examples
244    ///
245    /// ```
246    /// # use snowflaked::Generator;
247    /// #
248    /// let mut generator = Generator::new_checked(123).unwrap();
249    ///
250    /// assert_eq!(generator.instance(), 123);
251    /// ```
252    #[inline]
253    pub const fn new_checked(instance: u16) -> Option<Self> {
254        if instance > INSTANCE_MAX {
255            None
256        } else {
257            Some(Self::new_unchecked(instance))
258        }
259    }
260
261    /// Creates a new `Generator` using the given `instance` without checking if instance exceeds
262    /// the maximum value (2^10 -1).
263    ///
264    /// Note: If `instance` exceeds the maximum instance value the `Generator` will create
265    /// incorrect snowflakes.
266    ///
267    /// # Examples
268    ///
269    /// ```
270    /// # use snowflaked::Generator;
271    /// #
272    /// let mut generator = Generator::new_unchecked(123);
273    ///
274    /// assert_eq!(generator.instance(), 123);
275    /// ```
276    #[inline]
277    pub const fn new_unchecked(instance: u16) -> Self {
278        Self {
279            components: Components::new(instance as u64),
280            epoch: UNIX_EPOCH,
281        }
282    }
283
284    /// Creates a new `Builder` used to configure a `Generator`.
285    ///
286    /// # Examples
287    ///
288    /// ```
289    /// # use snowflaked::Generator;
290    /// use std::time::SystemTime;
291    ///
292    /// let epoch = SystemTime::now();
293    /// let generator: Generator = Generator::builder().instance(123).epoch(epoch).build();
294    ///
295    /// assert_eq!(generator.instance(), 123);
296    /// assert_eq!(generator.epoch(), epoch);
297    /// ```
298    #[inline]
299    pub const fn builder() -> Builder {
300        Builder::new()
301    }
302
303    /// Returns the configured instace component of this `Generator`.
304    ///
305    /// # Examples
306    ///
307    /// ```
308    /// # use snowflaked::Generator;
309    /// #
310    /// let mut generator = Generator::new(123);
311    ///
312    /// assert_eq!(generator.instance(), 123);
313    /// ```
314    #[inline]
315    pub fn instance(&self) -> u16 {
316        self.components.instance() as u16
317    }
318
319    /// Returns the configured epoch of this `Generator`. By default this is [`UNIX_EPOCH`].
320    ///
321    /// # Examples
322    ///
323    /// ```
324    /// # use snowflaked::Generator;
325    /// use std::time::UNIX_EPOCH;
326    ///
327    /// let generator = Generator::new(123);
328    /// assert_eq!(generator.epoch(), UNIX_EPOCH);
329    /// ```
330    #[inline]
331    pub fn epoch(&self) -> SystemTime {
332        self.epoch
333    }
334
335    /// Generate a new unique snowflake id.
336    ///
337    /// # Examples
338    ///
339    /// ```
340    /// # use snowflaked::Generator;
341    /// #
342    /// let mut generator = Generator::new(0);
343    ///
344    /// let id1: u64 = generator.generate();
345    /// let id2: i64 = generator.generate();
346    /// ```
347    pub fn generate<T>(&mut self) -> T
348    where
349        T: Snowflake,
350    {
351        use std::cmp::Ordering;
352
353        let mut now = self.epoch.elapsed().unwrap().as_millis() as u64;
354        let instance = self.components.instance();
355        match now.cmp(&self.components.timestamp()) {
356            Ordering::Less => {
357                panic!("Clock has moved backwards! This is not supported.");
358            }
359            Ordering::Greater => {
360                self.components.reset_sequence();
361                self.components.set_timestamp(now);
362                T::from_parts(now, instance, 0)
363            }
364            Ordering::Equal => {
365                let sequence = self.components.take_sequence();
366                if sequence == 0 {
367                    now = self.wait_until_next_millisecond(now);
368                }
369                self.components.set_timestamp(now);
370                T::from_parts(now, instance, sequence)
371            }
372        }
373    }
374
375    fn wait_until_next_millisecond(&mut self, last_millisecond: u64) -> u64 {
376        loop {
377            let now = self.epoch.elapsed().unwrap().as_millis() as u64;
378            if now > last_millisecond {
379                return now;
380            }
381            std::hint::spin_loop();
382        }
383    }
384}
385
386impl From<Builder> for Generator {
387    fn from(builder: Builder) -> Self {
388        Self {
389            components: Components::new(builder.instance as u64),
390            epoch: builder.epoch,
391        }
392    }
393}
394
395/// A single `u64` representing the complete current state of a generator.
396#[derive(Copy, Clone, Debug, PartialEq, Eq)]
397#[repr(transparent)]
398pub(crate) struct Components(u64);
399
400impl Components {
401    #[inline]
402    pub(crate) const fn new(instance: u64) -> Self {
403        // Fill in the given instance, and set the sequence at '1'
404        Self((instance << 12) & BITMASK_INSTANCE | 1)
405    }
406
407    #[inline]
408    pub(crate) fn sequence(&self) -> u64 {
409        self.0 & BITMASK_SEQUENCE
410    }
411
412    #[inline]
413    pub(crate) fn instance(&self) -> u64 {
414        (self.0 & BITMASK_INSTANCE) >> 12
415    }
416
417    #[inline]
418    pub(crate) fn timestamp(&self) -> u64 {
419        self.0 >> 22
420    }
421
422    pub(crate) fn set_sequence(&mut self, seq: u64) {
423        let timestamp = self.timestamp() << 22;
424        let instance = (self.instance() << 12) & BITMASK_INSTANCE;
425        *self = Self(timestamp + instance + seq)
426    }
427
428    pub(crate) fn reset_sequence(&mut self) {
429        self.set_sequence(1)
430    }
431
432    pub(crate) fn set_timestamp(&mut self, ts: u64) {
433        let timestamp = ts << 22;
434        let instance = (self.instance() << 12) & BITMASK_INSTANCE;
435        *self = Self(timestamp + instance + self.sequence())
436    }
437
438    #[inline]
439    pub(crate) fn take_sequence(&mut self) -> u64 {
440        match self.sequence() {
441            4095 => {
442                self.set_sequence(0);
443                4095
444            }
445            n => {
446                self.set_sequence(n + 1);
447                n
448            }
449        }
450    }
451
452    #[inline]
453    pub(crate) const fn from_bits(bits: u64) -> Self {
454        Self(bits)
455    }
456
457    #[inline]
458    pub(crate) const fn to_bits(self) -> u64 {
459        self.0
460    }
461}
462
463/// Emits a panic with the appropriate message for providing an invalid instance id to
464/// a generator.
465#[inline(never)]
466#[cold]
467pub(crate) const fn const_panic_new() -> ! {
468    panic!("invalid instance provided for snowflake generator");
469}
470
471#[cfg(all(test, not(loom)))]
472mod tests {
473    use std::time::UNIX_EPOCH;
474
475    use super::Generator;
476    use crate::{Components, Snowflake};
477
478    #[test]
479    fn test_components() {
480        let mut components = Components::new(0);
481        assert_eq!(components.sequence(), 1);
482        assert_eq!(components.timestamp(), 0);
483
484        components.set_sequence(1024);
485        assert_eq!(components.sequence(), 1024);
486        assert_eq!(components.timestamp(), 0);
487
488        components.set_timestamp(1024);
489        assert_eq!(components.sequence(), 1024);
490        assert_eq!(components.timestamp(), 1024);
491
492        let now = UNIX_EPOCH.elapsed().unwrap().as_millis() as u64;
493        components.set_timestamp(now);
494        assert_eq!(components.sequence(), 1024);
495        assert_eq!(components.timestamp(), now);
496    }
497
498    #[test]
499    fn test_generate_ordered() {
500        const INSTANCE: u64 = 0;
501
502        let mut last_id = None;
503        let mut generator = Generator::new_unchecked(INSTANCE as u16);
504
505        for _ in 0..10_000 {
506            let id: u64 = generator.generate();
507            assert_eq!(id.instance(), INSTANCE);
508            assert!(
509                last_id < Some(id),
510                "expected {:?} to be less than {:?}",
511                last_id,
512                Some(id)
513            );
514            last_id = Some(id);
515        }
516    }
517
518    #[test]
519    fn test_generate_no_duplicates() {
520        let mut generator = Generator::new_unchecked(0);
521        let mut ids: Vec<u64> = Vec::with_capacity(10_000);
522
523        for _ in 0..ids.capacity() {
524            ids.push(generator.generate());
525        }
526
527        for (index, id) in ids.iter().enumerate() {
528            for (index2, id2) in ids.iter().enumerate() {
529                if index != index2 && id == id2 {
530                    panic!(
531                        "Found duplicate id {} (SEQ {}, INS {}, TS {}) at index {} and {}",
532                        id,
533                        id.sequence(),
534                        id.instance(),
535                        id.timestamp(),
536                        index,
537                        index2
538                    );
539                }
540            }
541        }
542    }
543
544    #[test]
545    fn test_snowflake_u64() {
546        let id = 442252698964721669_u64;
547        assert_eq!(id.timestamp(), 105441260091);
548        assert_eq!(id.instance(), 0);
549        assert_eq!(id.sequence(), 5);
550
551        let id = 862026798833074178_u64;
552        assert_eq!(id.timestamp(), 205523204525);
553        assert_eq!(id.instance(), 256);
554        assert_eq!(id.sequence(), 2);
555    }
556
557    #[test]
558    fn test_snowflake_i64() {
559        let id = 442252698964721669_i64;
560        assert_eq!(id.timestamp(), 105441260091);
561        assert_eq!(id.instance(), 0);
562        assert_eq!(id.sequence(), 5);
563
564        let id = 862026798833074178_i64;
565        assert_eq!(id.timestamp(), 205523204525);
566        assert_eq!(id.instance(), 256);
567        assert_eq!(id.sequence(), 2);
568    }
569
570    #[test]
571    fn i64_no_negative() {
572        let timestamp = 2297684976571;
573        let instance = 0;
574        let sequence = 0;
575
576        let id = i64::from_parts(timestamp, instance, sequence);
577        assert!(id >= 0);
578    }
579
580    #[test]
581    fn i64_reflexive() {
582        let timestamp = 1684917000190;
583        let instance = 0;
584        let sequence = 2056;
585
586        let id = i64::from_parts(timestamp, instance, sequence);
587
588        assert_eq!(id.timestamp(), timestamp);
589        assert_eq!(id.instance(), instance);
590        assert_eq!(id.sequence(), sequence);
591    }
592
593    #[test]
594    fn u64_reflexive() {
595        let timestamp = 1684917075097;
596        let instance = 0;
597        let sequence = 1086;
598
599        let id = u64::from_parts(timestamp, instance, sequence);
600
601        assert_eq!(id.timestamp(), timestamp);
602        assert_eq!(id.instance(), instance);
603        assert_eq!(id.sequence(), sequence);
604    }
605
606    #[test]
607    fn u64_i64_equivalent() {
608        let timestamp = 1684917177537;
609        let instance = 0;
610        let sequence = 3060;
611
612        let id0 = u64::from_parts(timestamp, instance, sequence);
613        let id1 = i64::from_parts(timestamp, instance, sequence);
614
615        assert_eq!(id0.timestamp(), id1.timestamp());
616        assert_eq!(id0.instance(), id1.instance());
617        assert_eq!(id0.sequence(), id1.sequence());
618    }
619}