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
use chrono::Utc; use parking_lot::Mutex; use std::sync::Arc; pub struct Snowflake { epoch: i64, worker_id: i64, datacenter_id: i64, sequence: i64, time: Arc<Mutex<i64>>, } impl Snowflake { pub fn default() -> Snowflake { Snowflake { epoch: 1_564_790_400_000, worker_id: 1, datacenter_id: 1, sequence: 0, time: Arc::new(Mutex::new(0)), } } pub fn new(epoch: i64, worker_id: i64, datacenter_id: i64) -> Snowflake { Snowflake { epoch, worker_id, datacenter_id, sequence: 0, time: Arc::new(Mutex::new(0)), } } pub fn epoch(&mut self, epoch: i64) -> &mut Self { self.epoch = epoch; self } pub fn worker_id(&mut self, worker_id: i64) -> &mut Self { self.worker_id = worker_id; self } pub fn datacenter_id(&mut self, datacenter_id: i64) -> &mut Self { self.datacenter_id = datacenter_id; self } pub fn generate(&mut self) -> i64 { let mut last_timestamp = self.time.lock(); let mut timestamp = self.get_time(); if timestamp == *last_timestamp { self.sequence = (self.sequence + 1) & (-1 ^ (-1 << 12)); if self.sequence == 0 && timestamp <= *last_timestamp { timestamp = self.get_time(); } } else { self.sequence = 0; } *last_timestamp = timestamp; (timestamp << 22) | (self.worker_id << 17) | (self.datacenter_id << 12) | self.sequence } fn get_time(&self) -> i64 { Utc::now().timestamp_millis() - self.epoch } }