1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

#[cfg(test)]
use std::thread;

/// MonteFlake creates a configurable Snowflake ID generator where the epoch
/// and bitwidths may be adjusted to your liking from the defaults.
/// MonteFlake id generation will always increase in value from the time
/// its instantiated. No other guarantees are made. If a MonteFlake generator
/// is created after a clock moves back from other MonteFlake generated ids
/// conflicting ID values are entirely possible.
pub struct MonteFlake {
    id: u64,
    epoch: u64,
    bitwidths: (u64, u64, u64),
    seq: u64,
    start_ts: u64,
    start_instant: Instant,
    duration: u64,
}

trait AsMillis {
    fn as_millis(self) -> u64;
}

impl AsMillis for Duration {
    fn as_millis(self) -> u64 {
        self.as_secs()*1000 + (self.subsec_nanos() as u64)/1000000
    }
}

impl MonteFlake {
    /// Build a new monteflake with default options
    /// id: 0
    /// epoch: 2013-01-01T00:00:00Z in milliseconds since the unix epoch
    /// bitwidths (42 timestamp bits, 10 id bits, 12 sequence bits)
    pub fn new(id: u64) -> MonteFlake {
        let since_unix = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
        let ts = since_unix.as_millis();
        let instant = Instant::now();
        MonteFlake {
            id: id,
            seq: 0,
            epoch: 1356998400000,
            bitwidths: (42, 10, 12),
            start_ts: ts,
            start_instant: instant,
            duration: 0,
        }
    }

    /// Set the epoch of a MonteFlake generator
    pub fn epoch(mut self, epoch: u64) -> MonteFlake {
        self.epoch = epoch;
        self
    }

    /// Set the id of a MonteFlake generator
    pub fn id(mut self, id: u64) -> MonteFlake {
        self.id = id;
        self
    }

    /// Set the bitwidths of a MonteFlake generator
    pub fn bitwidths(mut self, ts_bits: u64, id_bits: u64) -> MonteFlake {
        assert!(ts_bits + id_bits < 64);
        self.bitwidths = (ts_bits, id_bits, 64-(ts_bits+id_bits));
        self
    }

    /// generate the next id
    /// internally this updates at least the current sequence value, possibly
    /// the timestamp value if enough time has elapsed to matter
    pub fn next(&mut self) -> u64 {
        let duration = self.start_instant.elapsed().as_millis();
        if duration != self.duration {
            self.seq = 0;
        }
        let ts = self.start_ts + duration;
        let id = self.encode(ts, self.id, self.seq);
        self.duration = duration;
        self.seq += 1;
        id
    }

    /// encode into an id the current time, id of the generator (machine, node, shard), and
    /// sequence value
    /// the current time (ts) is the number of milliseconds passed since the unix epoch
    pub fn encode(&self, ts: u64, id: u64, seq: u64) -> u64 {
        let ts0 = ts - self.epoch;
        let (_, id_shift, seq_shift) = self.bitwidths;
        let ts_mask = bitmask(id_shift+seq_shift);
        let id_mask = bitmask(seq_shift) ^ ts_mask;
        let seq_mask = (bitmask(0) ^ ts_mask) ^ id_mask;
        ((ts0 << (id_shift + seq_shift)) & ts_mask) | ((id << seq_shift) & id_mask) | (seq & seq_mask)
    }

    /// decode from an encoded id the time, id of the generator, and sequence id
    /// the current time (ts) is the number of milliseconds passed since the unix epoch
    pub fn decode(&self, flake_id: u64) -> (u64, u64, u64) {
        let (_, id_shift, seq_shift) = self.bitwidths;
        let ts_mask = bitmask(id_shift+seq_shift);
        let id_mask = bitmask(seq_shift) ^ ts_mask;
        let seq_mask = (bitmask(0) ^ ts_mask) ^ id_mask;
        let ts = (flake_id & ts_mask) >> (id_shift+seq_shift);
        let id = (flake_id & id_mask) >> seq_shift;
        let seq = flake_id & seq_mask;
        (ts + self.epoch, id, seq)
    }
}

fn bitmask(left_shift: u64) -> u64 {
    0xFFFFFFFFFFFFFFFF << left_shift
}

#[test]
fn test_bitmask() {
    assert_eq!(bitmask(4), 0xFFFFFFFFFFFFFFF0);
    assert_eq!(bitmask(7), 0xFFFFFFFFFFFFFF80);
}

#[test]
fn test_encode_decode() {
    let flake = MonteFlake::new(0);
    let vals = (13+flake.start_ts,24,81);
    let id = flake.encode(vals.0, vals.1, vals.2); 
    assert_eq!(flake.decode(id), vals);
}

#[test]
fn test_next() {
    let new_epoch = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis();
    let mut flake = MonteFlake::new(0).epoch(new_epoch);
    let id0 = flake.next();
    let id1 = flake.next();
    let (ts0, id0, seq0) = flake.decode(id0);
    let (ts1, id1, seq1) = flake.decode(id1);
    assert!((ts0-new_epoch) < 1);
    assert_eq!(id0, 0);
    assert_eq!(seq0, 0);
    assert!((ts1-new_epoch) < 1);
    assert_eq!(id1, 0);
    assert_eq!(seq1, 1);
    let mut flake1 = flake.id(100);
    let id2 = flake1.next();
    let (ts2, id2, seq2) = flake1.decode(id2);
    assert!((ts2-new_epoch) < 1);
    assert_eq!(id2, 100);
    assert_eq!(seq2, 2);
    thread::sleep(Duration::from_millis(10));
    let id3 = flake1.next();
    let (ts3, id3, seq3) = flake1.decode(id3);
    assert!((ts3-new_epoch) >= 10);
    assert_eq!(id3, 100);
    assert_eq!(seq3, 0);
}