snowflake_rust/
lib.rs

1use std::{thread, time};
2use std::net::UdpSocket;
3use std::sync::{Arc, Mutex};
4
5use chrono::Utc;
6
7pub struct Snowflake {
8    epoch: i64,
9    worker_id: i64,
10    sequence: i64,
11    time: Arc<Mutex<i64>>,
12    sequence_mask: i64
13}
14
15impl Snowflake {
16    pub fn kubernetes() -> Snowflake {
17        let ip = get_ip().unwrap();
18        let ip_split: Vec<&str> = ip.split(".").collect();
19        let ip_low = ip_split[2].to_string().parse::<i64>().unwrap() << 8 | ip_split[3].to_string().parse::<i64>().unwrap();
20        Snowflake {
21            epoch: 1575129600000,
22            worker_id: ip_low,
23            sequence: 0,
24            time: Arc::new(Mutex::new(0)),
25            sequence_mask: -1 ^ -1 << 12,
26        }
27    }
28
29    pub fn new(worker_id: i64) -> Snowflake {
30        Snowflake {
31            epoch: 1575129600000,
32            worker_id,
33            sequence: 0,
34            time: Arc::new(Mutex::new(0)),
35            sequence_mask: -1 ^ (-1 << 12),
36        }
37    }
38
39    pub fn generate(&mut self) -> Option<i64> {
40        let mut last_timestamp = self.time.lock().unwrap();
41        let mut timestamp = self.get_time();
42        if timestamp < *last_timestamp {
43            if *last_timestamp - timestamp > 150 {
44                return None
45            } else {
46                thread::sleep(time::Duration::from_millis((*last_timestamp - timestamp + 1) as u64));
47                timestamp = self.get_time();
48            }
49        } else if timestamp == *last_timestamp {
50            self.sequence = (self.sequence + 1) & self.sequence_mask
51        } else {
52            self.sequence = 0
53        }
54
55        *last_timestamp = timestamp;
56        Option::from((timestamp << 28) | (self.worker_id << 12) | self.sequence)
57    }
58
59    fn get_time(&self) -> i64 {
60        Utc::now().timestamp_millis() - self.epoch
61    }
62}
63
64fn get_ip() -> Option<String> {
65    let socket = match UdpSocket::bind("0.0.0.0:0") {
66        Ok(s) => s,
67        Err(_) => return None,
68    };
69
70    match socket.connect("8.8.8.8:80") {
71        Ok(()) => (),
72        Err(_) => return None,
73    };
74
75    return match socket.local_addr() {
76        Ok(addr) => Some(addr.ip().to_string()),
77        Err(_) => None,
78    };
79}
snowflake_rust/lib.rs

snowflake_rust/
lib.rs
