github.com/godaddy-x/freego@v1.0.156/utils/snowflake/id_worker.go (about) 1 // Package snowflake provides a very simple Twitter snowflake generator and parser. 2 package snowflake 3 4 import ( 5 "encoding/base64" 6 "encoding/binary" 7 "errors" 8 "fmt" 9 "strconv" 10 "sync" 11 "time" 12 _ "unsafe" 13 ) 14 15 var ( 16 // Epoch is set to the twitter snowflake epoch of Nov 04 2010 01:42:54 UTC 17 // You may customize this to set a different epoch for your application. 18 Epoch int64 = 1288834974657 19 20 // Number of bits to use for Node 21 // Remember, you have a total 22 bits to share between Node/Step 22 NodeBits uint8 = 10 23 24 // Number of bits to use for Step 25 // Remember, you have a total 22 bits to share between Node/Step 26 StepBits uint8 = 12 27 28 nodeMax int64 = -1 ^ (-1 << NodeBits) 29 nodeMask int64 = nodeMax << StepBits 30 stepMask int64 = -1 ^ (-1 << StepBits) 31 timeShift uint8 = NodeBits + StepBits 32 nodeShift uint8 = StepBits 33 ) 34 35 const encodeBase32Map = "ybndrfg8ejkmcpqxot1uwisza345h769" 36 37 var decodeBase32Map [256]byte 38 39 const encodeBase58Map = "123456789abcdefghijkmnopqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ" 40 41 var decodeBase58Map [256]byte 42 43 // A JSONSyntaxError is returned from UnmarshalJSON if an invalid ID is provided. 44 type JSONSyntaxError struct{ original []byte } 45 46 func (j JSONSyntaxError) Error() string { 47 return fmt.Sprintf("invalid snowflake ID %q", string(j.original)) 48 } 49 50 // Create a map for decoding Base58. This speeds up the process tremendously. 51 func init() { 52 53 for i := 0; i < len(encodeBase58Map); i++ { 54 decodeBase58Map[i] = 0xFF 55 } 56 57 for i := 0; i < len(encodeBase58Map); i++ { 58 decodeBase58Map[encodeBase58Map[i]] = byte(i) 59 } 60 61 for i := 0; i < len(encodeBase32Map); i++ { 62 decodeBase32Map[i] = 0xFF 63 } 64 65 for i := 0; i < len(encodeBase32Map); i++ { 66 decodeBase32Map[encodeBase32Map[i]] = byte(i) 67 } 68 } 69 70 // ErrInvalidBase58 is returned by ParseBase58 when given an invalid []byte 71 var ErrInvalidBase58 = errors.New("invalid base58") 72 73 // ErrInvalidBase32 is returned by ParseBase32 when given an invalid []byte 74 var ErrInvalidBase32 = errors.New("invalid base32") 75 76 // A Node struct holds the basic information needed for a snowflake generator 77 // node 78 type Node struct { 79 mu sync.Mutex 80 time int64 81 node int64 82 step int64 83 } 84 85 // An ID is a custom type used for a snowflake ID. This is used so we can 86 // attach methods onto the ID. 87 type ID int64 88 89 // NewNode returns a new snowflake node that can be used to generate snowflake 90 // IDs 91 func NewNode(node int64) (*Node, error) { 92 93 // re-calc in case custom NodeBits or StepBits were set 94 nodeMax = -1 ^ (-1 << NodeBits) 95 nodeMask = nodeMax << StepBits 96 stepMask = -1 ^ (-1 << StepBits) 97 timeShift = NodeBits + StepBits 98 nodeShift = StepBits 99 100 if node < 0 || node > nodeMax { 101 return nil, errors.New("Node number must be between 0 and " + strconv.FormatInt(nodeMax, 10)) 102 } 103 104 return &Node{ 105 time: 0, 106 node: node, 107 step: 0, 108 }, nil 109 } 110 111 //go:linkname now time.now 112 func now() (sec int64, nsec int32, mono int64) 113 114 func (n *Node) GetNow() int64 { 115 s, m, _ := now() 116 return (s*1e9 + int64(m)) / 1e6 117 //return time.Now().UnixNano() / 1000000 118 } 119 120 func (n *Node) GetValidNow() int64 { 121 now := n.GetNow() 122 if n.time > now { 123 time.Sleep(time.Duration((n.time-now)+1) * time.Millisecond) 124 } 125 return n.GetNow() 126 } 127 128 // Generate creates and returns a unique snowflake ID 129 func (n *Node) Generate() ID { 130 131 n.mu.Lock() 132 133 now := n.GetNow() 134 135 if n.time > now { 136 now = n.GetValidNow() 137 } 138 139 if n.time == now { 140 n.step = (n.step + 1) & stepMask 141 142 if n.step == 0 { 143 for now <= n.time { 144 now = n.GetNow() 145 } 146 } 147 } else { 148 n.step = 0 149 } 150 151 n.time = now 152 153 r := ID((now-Epoch)<<timeShift | 154 (n.node << nodeShift) | 155 (n.step), 156 ) 157 158 n.mu.Unlock() 159 return r 160 } 161 162 // Int64 returns an int64 of the snowflake ID 163 func (f ID) Int64() int64 { 164 return int64(f) 165 } 166 167 // String returns a string of the snowflake ID 168 func (f ID) String() string { 169 return strconv.FormatInt(int64(f), 10) 170 } 171 172 // Base2 returns a string base2 of the snowflake ID 173 func (f ID) Base2() string { 174 return strconv.FormatInt(int64(f), 2) 175 } 176 177 // Base36 returns a base36 string of the snowflake ID 178 func (f ID) Base36() string { 179 return strconv.FormatInt(int64(f), 36) 180 } 181 182 // Base32 uses the z-base-32 character set but encodes and decodes similar 183 // to base58, allowing it to create an even smaller result string. 184 // NOTE: There are many different base32 implementations so becareful when 185 // doing any interoperation interop with other packages. 186 func (f ID) Base32() string { 187 188 if f < 32 { 189 return string(encodeBase32Map[f]) 190 } 191 192 b := make([]byte, 0, 12) 193 for f >= 32 { 194 b = append(b, encodeBase32Map[f%32]) 195 f /= 32 196 } 197 b = append(b, encodeBase32Map[f]) 198 199 for x, y := 0, len(b)-1; x < y; x, y = x+1, y-1 { 200 b[x], b[y] = b[y], b[x] 201 } 202 203 return string(b) 204 } 205 206 // ParseBase32 parses a base32 []byte into a snowflake ID 207 // NOTE: There are many different base32 implementations so becareful when 208 // doing any interoperation interop with other packages. 209 func ParseBase32(b []byte) (ID, error) { 210 211 var id int64 212 213 for i := range b { 214 if decodeBase32Map[b[i]] == 0xFF { 215 return -1, ErrInvalidBase32 216 } 217 id = id*32 + int64(decodeBase32Map[b[i]]) 218 } 219 220 return ID(id), nil 221 } 222 223 // Base58 returns a base58 string of the snowflake ID 224 func (f ID) Base58() string { 225 226 if f < 58 { 227 return string(encodeBase58Map[f]) 228 } 229 230 b := make([]byte, 0, 11) 231 for f >= 58 { 232 b = append(b, encodeBase58Map[f%58]) 233 f /= 58 234 } 235 b = append(b, encodeBase58Map[f]) 236 237 for x, y := 0, len(b)-1; x < y; x, y = x+1, y-1 { 238 b[x], b[y] = b[y], b[x] 239 } 240 241 return string(b) 242 } 243 244 // ParseBase58 parses a base58 []byte into a snowflake ID 245 func ParseBase58(b []byte) (ID, error) { 246 247 var id int64 248 249 for i := range b { 250 if decodeBase58Map[b[i]] == 0xFF { 251 return -1, ErrInvalidBase58 252 } 253 id = id*58 + int64(decodeBase58Map[b[i]]) 254 } 255 256 return ID(id), nil 257 } 258 259 // Base64 returns a base64 string of the snowflake ID 260 func (f ID) Base64() string { 261 return base64.StdEncoding.EncodeToString(f.Bytes()) 262 } 263 264 // Bytes returns a byte slice of the snowflake ID 265 func (f ID) Bytes() []byte { 266 return []byte(f.String()) 267 } 268 269 // IntBytes returns an array of bytes of the snowflake ID, encoded as a 270 // big endian integer. 271 func (f ID) IntBytes() [8]byte { 272 var b [8]byte 273 binary.BigEndian.PutUint64(b[:], uint64(f)) 274 return b 275 } 276 277 // Time returns an int64 unix timestamp of the snowflake ID time 278 func (f ID) Time() int64 { 279 return (int64(f) >> timeShift) + Epoch 280 } 281 282 // Node returns an int64 of the snowflake ID node number 283 func (f ID) Node() int64 { 284 return int64(f) & nodeMask >> nodeShift 285 } 286 287 // Step returns an int64 of the snowflake step (or sequence) number 288 func (f ID) Step() int64 { 289 return int64(f) & stepMask 290 } 291 292 // MarshalJSON returns a json byte array string of the snowflake ID. 293 func (f ID) MarshalJSON() ([]byte, error) { 294 buff := make([]byte, 0, 22) 295 buff = append(buff, '"') 296 buff = strconv.AppendInt(buff, int64(f), 10) 297 buff = append(buff, '"') 298 return buff, nil 299 } 300 301 // UnmarshalJSON converts a json byte array of a snowflake ID into an ID type. 302 func (f *ID) UnmarshalJSON(b []byte) error { 303 if len(b) < 3 || b[0] != '"' || b[len(b)-1] != '"' { 304 return JSONSyntaxError{b} 305 } 306 307 i, err := strconv.ParseInt(string(b[1:len(b)-1]), 10, 64) 308 if err != nil { 309 return err 310 } 311 312 *f = ID(i) 313 return nil 314 }