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
use std::{thread, time}; use std::net::UdpSocket; use std::sync::{Arc, Mutex}; use chrono::Utc; pub struct Snowflake { epoch: i64, worker_id: i64, sequence: i64, time: Arc<Mutex<i64>>, sequence_mask: i64 } impl Snowflake { pub fn kubernetes() -> Snowflake { let ip = get_ip().unwrap(); let ip_split: Vec<&str> = ip.split(".").collect(); let ip_low = ip_split[2].to_string().parse::<i64>().unwrap() << 8 | ip_split[3].to_string().parse::<i64>().unwrap(); Snowflake { epoch: 1575129600000, worker_id: ip_low, sequence: 0, time: Arc::new(Mutex::new(0)), sequence_mask: -1 ^ -1 << 12, } } pub fn new(worker_id: i64) -> Snowflake { Snowflake { epoch: 1575129600000, worker_id, sequence: 0, time: Arc::new(Mutex::new(0)), sequence_mask: -1 ^ (-1 << 12), } } pub fn generate(&mut self) -> Option<i64> { let mut last_timestamp = self.time.lock().unwrap(); let mut timestamp = self.get_time(); if timestamp < *last_timestamp { if *last_timestamp - timestamp > 150 { return None } else { thread::sleep(time::Duration::from_millis((*last_timestamp - timestamp + 1) as u64)); timestamp = self.get_time(); } } else if timestamp == *last_timestamp { self.sequence = (self.sequence + 1) & self.sequence_mask } else { self.sequence = 0 } *last_timestamp = timestamp; Option::from((timestamp << 28) | (self.worker_id << 12) | self.sequence) } fn get_time(&self) -> i64 { Utc::now().timestamp_millis() - self.epoch } } fn get_ip() -> Option<String> { let socket = match UdpSocket::bind("0.0.0.0:0") { Ok(s) => s, Err(_) => return None, }; match socket.connect("8.8.8.8:80") { Ok(()) => (), Err(_) => return None, }; return match socket.local_addr() { Ok(addr) => Some(addr.ip().to_string()), Err(_) => None, }; }