rbatis 4.7.2

The Rust SQL Toolkit and ORM Library. An async, pure Rust SQL crate featuring compile-time Dynamic SQL
Documentation 
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

use std::hint::spin_loop;
use std::sync::atomic::{AtomicI64, AtomicU16, Ordering};
use std::sync::LazyLock;
use std::time::{SystemTime, UNIX_EPOCH};
use parking_lot::ReentrantMutex;

/// ID generator trait for generating unique task IDs
pub trait IdGenerator: Send + Sync + std::fmt::Debug {
    /// Generate a unique ID (i64)
    fn generate(&self) -> i64;
}

///Snowflakes algorithm
#[derive(Debug)]
pub struct Snowflake {
    pub epoch: SystemTime,
    pub last_timestamp: AtomicI64, // Atomic to replace mutable i64
    pub machine_id: i32,
    pub node_id: i32,
    pub mode: i32, // mode [0=fast_generate,1=realtime_generate]
    pub idx: AtomicU16, // Atomic to replace mutable u16
    lock: ReentrantMutex<()>,
}

impl Clone for Snowflake {
    fn clone(&self) -> Self {
        Self {
            epoch: self.epoch,
            machine_id: self.machine_id,
            node_id: self.node_id,
            last_timestamp: AtomicI64::new(self.last_timestamp.load(Ordering::SeqCst)),
            idx: AtomicU16::new(self.idx.load(Ordering::SeqCst)),
            lock: ReentrantMutex::new(()),
            mode: self.mode,
        }
    }
}

impl Snowflake {
    pub fn new(machine_id: i32, node_id: i32, mode: i32) -> Snowflake {
        Self::with_epoch(machine_id, node_id, mode, UNIX_EPOCH)
    }

    /// from epoch for example:
    ///
    /// ```rust
    /// use std::time::UNIX_EPOCH;
    /// use rbatis::id_generator::Snowflake;
    /// let snowflake = Snowflake::with_epoch(1,1,0,UNIX_EPOCH);
    /// ```
    pub fn with_epoch(machine_id: i32, node_id: i32, mode: i32, epoch: SystemTime) -> Snowflake {
        let last_time_millis = Self::get_time_millis(epoch);
        Snowflake {
            epoch,
            last_timestamp: AtomicI64::new(last_time_millis),
            machine_id,
            node_id,
            mode,
            idx: AtomicU16::new(0),
            lock: ReentrantMutex::new(()),
        }
    }

    /// from last_timestamp for example:
    ///
    /// ```rust
    /// use std::time::UNIX_EPOCH;
    /// use rbatis::id_generator::Snowflake;
    /// let snowflake = Snowflake::with_last_timestamp(1,1,0, 1726417159);
    /// ```
    pub fn with_last_timestamp(machine_id: i32, node_id: i32, mode: i32, last_timestamp: i64) -> Snowflake {
        Snowflake {
            epoch: UNIX_EPOCH,
            last_timestamp: AtomicI64::new(last_timestamp),
            machine_id,
            node_id,
            mode,
            idx: AtomicU16::new(0),
            lock: ReentrantMutex::new(()),
        }
    }

    pub fn default() -> Snowflake {
        Snowflake::new(1, 1, 0)
    }

    #[inline]
    pub fn generate_id(&self) -> i64 {
        let g = self.lock.lock();
        let idx = self.idx.fetch_add(1, Ordering::SeqCst) % 4096;
        if idx == 0 {
            // use timestamp
            if self.mode == 1 {
                let mut now_millis = Self::get_time_millis(self.epoch);
                let last_time = self.last_timestamp.load(Ordering::SeqCst);
                if now_millis == last_time {
                    now_millis = Self::biding_time_conditions(last_time, self.epoch);
                }
                self.last_timestamp.store(now_millis, Ordering::SeqCst);
            } else {
                //mode = 0,this is fast mode
                self.last_timestamp.fetch_add(1, Ordering::SeqCst);
            }
        }
        let last_time = self.last_timestamp.load(Ordering::SeqCst);
        drop(g);
        last_time << 22
            | ((self.machine_id << 17) as i64)
            | ((self.node_id << 12) as i64)
            | (idx as i64)
    }


    #[inline(always)]
    pub fn get_time_millis(epoch: SystemTime) -> i64 {
        SystemTime::now()
            .duration_since(epoch)
            .expect("Time went backward")
            .as_millis() as i64
    }

    #[inline(always)]
    fn biding_time_conditions(last_time_millis: i64, epoch: SystemTime) -> i64 {
        let mut latest_time_millis: i64;
        loop {
            latest_time_millis = Self::get_time_millis(epoch);
            if latest_time_millis > last_time_millis {
                return latest_time_millis;
            }
            spin_loop();
        }
    }
}

/// Note:  if you have multiple machines/server, please modify the machine ID and node ID
pub static SNOWFLAKE: LazyLock<Snowflake> = LazyLock::new(|| {
    Snowflake::new(1, 1, 0)
});

///gen new snowflake_id
pub fn new_snowflake_id() -> i64 {
    SNOWFLAKE.generate_id()
}

impl IdGenerator for Snowflake {
    fn generate(&self) -> i64 {
        self.generate_id()
    }
}