新增部分模块
This commit is contained in:
zyj
32
cmd/main.go
32
cmd/main.go
@@ -1,5 +1,35 @@
|
||||
package main
|
||||
|
||||
func main() {
|
||||
import (
|
||||
"fmt"
|
||||
"go-user-service/config"
|
||||
"go-user-service/libs"
|
||||
userpb "go-user-service/proto/gen"
|
||||
"go-user-service/service"
|
||||
"log"
|
||||
"net"
|
||||
"strconv"
|
||||
|
||||
"google.golang.org/grpc"
|
||||
)
|
||||
|
||||
func main() {
|
||||
appConfig, err := config.LoadConfig()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
fmt.Println("读取配置失败")
|
||||
}
|
||||
libs.Connect()
|
||||
lis, err1 := net.Listen("tcp", ":"+strconv.Itoa(appConfig.Port))
|
||||
if err1 != nil {
|
||||
log.Fatalf("failed to listen: %v", err1)
|
||||
}
|
||||
grpcServer := grpc.NewServer()
|
||||
log.Println("gRPC用户服务启动成功，监听端口: ", appConfig.Port)
|
||||
userService := &service.UserService{}
|
||||
userpb.RegisterUserServiceServer(grpcServer, userService)
|
||||
// 启动服务
|
||||
if err := grpcServer.Serve(lis); err != nil {
|
||||
log.Fatalf("failed to serve: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
19
go.mod
19
go.mod
@@ -5,11 +5,24 @@ go 1.24.3
|
||||
require (
|
||||
google.golang.org/grpc v1.75.0
|
||||
google.golang.org/protobuf v1.36.8
|
||||
gopkg.in/yaml.v3 v3.0.1
|
||||
gorm.io/driver/postgres v1.6.0
|
||||
gorm.io/gorm v1.30.2
|
||||
)
|
||||
|
||||
require (
|
||||
golang.org/x/net v0.41.0 // indirect
|
||||
golang.org/x/sys v0.33.0 // indirect
|
||||
golang.org/x/text v0.26.0 // indirect
|
||||
github.com/jackc/pgpassfile v1.0.0 // indirect
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
|
||||
github.com/jackc/pgx/v5 v5.6.0 // indirect
|
||||
github.com/jackc/puddle/v2 v2.2.2 // indirect
|
||||
github.com/jinzhu/inflection v1.0.0 // indirect
|
||||
github.com/jinzhu/now v1.1.5 // indirect
|
||||
github.com/kr/text v0.2.0 // indirect
|
||||
github.com/rogpeppe/go-internal v1.14.1 // indirect
|
||||
golang.org/x/crypto v0.41.0 // indirect
|
||||
golang.org/x/net v0.42.0 // indirect
|
||||
golang.org/x/sync v0.16.0 // indirect
|
||||
golang.org/x/sys v0.35.0 // indirect
|
||||
golang.org/x/text v0.28.0 // indirect
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20250707201910-8d1bb00bc6a7 // indirect
|
||||
)
|
||||
|
||||
49
go.sum
49
go.sum
@@ -1,3 +1,7 @@
|
||||
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
|
||||
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
|
||||
github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
|
||||
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
|
||||
@@ -8,6 +12,31 @@ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
|
||||
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
|
||||
github.com/jackc/pgx/v5 v5.6.0 h1:SWJzexBzPL5jb0GEsrPMLIsi/3jOo7RHlzTjcAeDrPY=
|
||||
github.com/jackc/pgx/v5 v5.6.0/go.mod h1:DNZ/vlrUnhWCoFGxHAG8U2ljioxukquj7utPDgtQdTw=
|
||||
github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
|
||||
github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
|
||||
github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
|
||||
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
|
||||
github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
|
||||
github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
|
||||
github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0=
|
||||
github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk=
|
||||
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
|
||||
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
|
||||
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
|
||||
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
|
||||
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
|
||||
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
|
||||
github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
|
||||
github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
|
||||
go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA=
|
||||
go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A=
|
||||
go.opentelemetry.io/otel v1.37.0 h1:9zhNfelUvx0KBfu/gb+ZgeAfAgtWrfHJZcAqFC228wQ=
|
||||
@@ -20,12 +49,22 @@ go.opentelemetry.io/otel/sdk/metric v1.37.0 h1:90lI228XrB9jCMuSdA0673aubgRobVZFh
|
||||
go.opentelemetry.io/otel/sdk/metric v1.37.0/go.mod h1:cNen4ZWfiD37l5NhS+Keb5RXVWZWpRE+9WyVCpbo5ps=
|
||||
go.opentelemetry.io/otel/trace v1.37.0 h1:HLdcFNbRQBE2imdSEgm/kwqmQj1Or1l/7bW6mxVK7z4=
|
||||
go.opentelemetry.io/otel/trace v1.37.0/go.mod h1:TlgrlQ+PtQO5XFerSPUYG0JSgGyryXewPGyayAWSBS0=
|
||||
golang.org/x/crypto v0.39.0 h1:SHs+kF4LP+f+p14esP5jAoDpHU8Gu/v9lFRK6IT5imM=
|
||||
golang.org/x/crypto v0.39.0/go.mod h1:L+Xg3Wf6HoL4Bn4238Z6ft6KfEpN0tJGo53AAPC632U=
|
||||
golang.org/x/crypto v0.41.0 h1:WKYxWedPGCTVVl5+WHSSrOBT0O8lx32+zxmHxijgXp4=
|
||||
golang.org/x/crypto v0.41.0/go.mod h1:pO5AFd7FA68rFak7rOAGVuygIISepHftHnr8dr6+sUc=
|
||||
golang.org/x/net v0.41.0 h1:vBTly1HeNPEn3wtREYfy4GZ/NECgw2Cnl+nK6Nz3uvw=
|
||||
golang.org/x/net v0.41.0/go.mod h1:B/K4NNqkfmg07DQYrbwvSluqCJOOXwUjeb/5lOisjbA=
|
||||
golang.org/x/net v0.42.0/go.mod h1:FF1RA5d3u7nAYA4z2TkclSCKh68eSXtiFwcWQpPXdt8=
|
||||
golang.org/x/sync v0.15.0 h1:KWH3jNZsfyT6xfAfKiz6MRNmd46ByHDYaZ7KSkCtdW8=
|
||||
golang.org/x/sync v0.15.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
|
||||
golang.org/x/sync v0.16.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
|
||||
golang.org/x/sys v0.33.0 h1:q3i8TbbEz+JRD9ywIRlyRAQbM0qF7hu24q3teo2hbuw=
|
||||
golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
|
||||
golang.org/x/sys v0.35.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
|
||||
golang.org/x/text v0.26.0 h1:P42AVeLghgTYr4+xUnTRKDMqpar+PtX7KWuNQL21L8M=
|
||||
golang.org/x/text v0.26.0/go.mod h1:QK15LZJUUQVJxhz7wXgxSy/CJaTFjd0G+YLonydOVQA=
|
||||
golang.org/x/text v0.28.0/go.mod h1:U8nCwOR8jO/marOQ0QbDiOngZVEBB7MAiitBuMjXiNU=
|
||||
gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
|
||||
gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20250707201910-8d1bb00bc6a7 h1:pFyd6EwwL2TqFf8emdthzeX+gZE1ElRq3iM8pui4KBY=
|
||||
@@ -34,3 +73,13 @@ google.golang.org/grpc v1.75.0 h1:+TW+dqTd2Biwe6KKfhE5JpiYIBWq865PhKGSXiivqt4=
|
||||
google.golang.org/grpc v1.75.0/go.mod h1:JtPAzKiq4v1xcAB2hydNlWI2RnF85XXcV0mhKXr2ecQ=
|
||||
google.golang.org/protobuf v1.36.8 h1:xHScyCOEuuwZEc6UtSOvPbAT4zRh0xcNRYekJwfqyMc=
|
||||
google.golang.org/protobuf v1.36.8/go.mod h1:fuxRtAxBytpl4zzqUh6/eyUujkJdNiuEkXntxiD/uRU=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
|
||||
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
gorm.io/driver/postgres v1.6.0 h1:2dxzU8xJ+ivvqTRph34QX+WrRaJlmfyPqXmoGVjMBa4=
|
||||
gorm.io/driver/postgres v1.6.0/go.mod h1:vUw0mrGgrTK+uPHEhAdV4sfFELrByKVGnaVRkXDhtWo=
|
||||
gorm.io/gorm v1.30.2 h1:f7bevlVoVe4Byu3pmbWPVHnPsLoWaMjEb7/clyr9Ivs=
|
||||
gorm.io/gorm v1.30.2/go.mod h1:8Z33v652h4//uMA76KjeDH8mJXPm1QNCYrMeatR0DOE=
|
||||
|
||||
58
libs/database.go
Normal file
58
libs/database.go
Normal file
@@ -0,0 +1,58 @@
|
||||
package libs
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"go-user-service/config"
|
||||
"log"
|
||||
"time"
|
||||
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
var DB *gorm.DB
|
||||
|
||||
func Connect() (*gorm.DB, error) {
|
||||
appConfig, dberr := config.LoadConfig()
|
||||
if dberr != nil {
|
||||
fmt.Println("读取配置失败")
|
||||
}
|
||||
|
||||
dsn := fmt.Sprintf("host=%s user=%s password=%s dbname=%s port=%d sslmode=disable TimeZone=Asia/Shanghai",
|
||||
appConfig.Database.Host,
|
||||
appConfig.Database.Username,
|
||||
appConfig.Database.Password,
|
||||
appConfig.Database.Database,
|
||||
appConfig.Database.Port)
|
||||
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to connect to database: %w", err)
|
||||
}
|
||||
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get sql.DB: %w", err)
|
||||
}
|
||||
sqlDB.SetMaxIdleConns(10)
|
||||
sqlDB.SetMaxOpenConns(100)
|
||||
sqlDB.SetConnMaxLifetime(time.Hour)
|
||||
|
||||
log.Println("Database connection established")
|
||||
DB = db
|
||||
return db, nil
|
||||
}
|
||||
|
||||
func GetDB() *gorm.DB {
|
||||
if DB == nil {
|
||||
panic("database connection is not initialized")
|
||||
}
|
||||
return DB
|
||||
}
|
||||
|
||||
func Close() error {
|
||||
sqlDB, err := DB.DB()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return sqlDB.Close()
|
||||
}
|
||||
21
models/user.go
Normal file
21
models/user.go
Normal file
@@ -0,0 +1,21 @@
|
||||
package models
|
||||
|
||||
import "time"
|
||||
|
||||
type User struct {
|
||||
ID int8 `gorm:"id;primary_key"`
|
||||
Nickname string `gorm:"nickname"`
|
||||
Username string `gorm:"username"`
|
||||
Password string `gorm:"password"`
|
||||
Email string `gorm:"email"`
|
||||
Phone string `gorm:"phone"`
|
||||
Status int `gorm:"status"`
|
||||
Salt string `gorm:"salt"`
|
||||
CreatedAt time.Time `gorm:"created_at;type:timestamptz"`
|
||||
UpdatedAt time.Time `gorm:"updated_at;type:timestamptz"`
|
||||
}
|
||||
|
||||
// 实现 TableName 方法指定表名
|
||||
func (User) TableName() string {
|
||||
return "d_user"
|
||||
}
|
||||
@@ -1,7 +1,7 @@
|
||||
syntax = "proto3";
|
||||
|
||||
package user;
|
||||
option go_package = "./proto;userpb";
|
||||
option go_package = "./proto/gen;userpb";
|
||||
|
||||
// 登录请求
|
||||
message LoginRequest {
|
||||
|
||||
@@ -1,4 +1,39 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"go-user-service/libs"
|
||||
"go-user-service/models"
|
||||
userpb "go-user-service/proto/gen"
|
||||
"go-user-service/utils"
|
||||
"sync"
|
||||
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/status"
|
||||
)
|
||||
|
||||
type UserService struct {
|
||||
userpb.UnimplementedUserServiceServer
|
||||
tokens map[string]string // token -> userID
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
func (u *UserService) Login(ctx context.Context, req *userpb.LoginRequest) (*userpb.LoginResponse, error) {
|
||||
u.mu.RLock()
|
||||
defer u.mu.RUnlock()
|
||||
// 查询是否有这个用户名
|
||||
var user models.User
|
||||
db := libs.GetDB()
|
||||
db.Model(&models.User{}).Where("username = ?", req.Username).First(&user)
|
||||
if user.Username != "" {
|
||||
// 验证密码
|
||||
var crypto utils.Crypto
|
||||
hashPassword, _ := crypto.PasswordEncryption(user.Password + user.Salt)
|
||||
|
||||
if user.Password == hashPassword {
|
||||
|
||||
}
|
||||
return nil, status.Errorf(codes.Unauthenticated, "密码错误")
|
||||
}
|
||||
return nil, status.Errorf(codes.NotFound, "用户不存在")
|
||||
}
|
||||
|
||||
56
utils/crypto.go
Normal file
56
utils/crypto.go
Normal file
@@ -0,0 +1,56 @@
|
||||
package utils
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"log"
|
||||
"time"
|
||||
|
||||
"golang.org/x/crypto/bcrypt"
|
||||
)
|
||||
|
||||
type Crypto struct{}
|
||||
|
||||
func (*Crypto) PasswordEncryption(password string) (string, error) {
|
||||
cost := 10
|
||||
hashedBytes, err := bcrypt.GenerateFromPassword([]byte(password), cost)
|
||||
if err != nil {
|
||||
log.Fatal("加密失败：", err)
|
||||
return "", err
|
||||
}
|
||||
|
||||
hashedPassword := string(hashedBytes)
|
||||
|
||||
fmt.Println("加密后的哈希值:", hashedPassword)
|
||||
return hashedPassword, nil
|
||||
}
|
||||
|
||||
func (*Crypto) GenerateStateToken(tokenLength int) (string, error) {
|
||||
// 1. 生成密码学安全的随机字节
|
||||
randomBytes := make([]byte, tokenLength)
|
||||
_, err := rand.Read(randomBytes)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("生成随机数失败: %w", err)
|
||||
}
|
||||
|
||||
// 2. (可选) 加入时间戳作为额外熵源，进一步降低碰撞和重放风险
|
||||
// 注意：如果你需要验证 token 的新鲜度（如设置有效期），时间戳信息需要单独存储，因为哈希本身不可逆。
|
||||
currentTime := time.Now().UnixNano()
|
||||
timeBytes := make([]byte, 8)
|
||||
// 简单地将时间戳转为字节序列，这里使用大端序
|
||||
for i := 0; i < 8; i++ {
|
||||
timeBytes[i] = byte(currentTime >> (56 - i*8))
|
||||
}
|
||||
|
||||
// 将随机字节和时间戳字节组合
|
||||
dataToHash := append(randomBytes, timeBytes...)
|
||||
|
||||
// 3. 对组合后的数据进行 SHA-256 哈希
|
||||
hash := sha256.Sum256(dataToHash)
|
||||
|
||||
// 4. 将哈希结果转换为字符串（这里使用 Base64 编码）
|
||||
token := base64.URLEncoding.EncodeToString(hash[:])
|
||||
return token, nil
|
||||
}
|
||||
86
utils/generate_id.go
Normal file
86
utils/generate_id.go
Normal file
@@ -0,0 +1,86 @@
|
||||
package utils
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// 定义常量
|
||||
const (
|
||||
twepoch = int64(1577836800000) // 起始时间戳 (2020-01-01 00:00:00 UTC 的毫秒数)
|
||||
workerIDBits = 5 // 工作节点 ID 的位数
|
||||
datacenterIDBits = 5 // 数据中心 ID 的位数
|
||||
sequenceBits = 12 // 序列号的位数
|
||||
|
||||
// 计算最大值
|
||||
maxWorkerID = -1 ^ (-1 << workerIDBits) // 最大工作节点 ID
|
||||
maxDatacenterID = -1 ^ (-1 << datacenterIDBits) // 最大数据中心 ID
|
||||
maxSequence = -1 ^ (-1 << sequenceBits) // 最大序列号
|
||||
|
||||
// 定义移位偏移量
|
||||
workerIDShift = sequenceBits
|
||||
datacenterIDShift = sequenceBits + workerIDBits
|
||||
timestampShift = sequenceBits + workerIDBits + datacenterIDBits
|
||||
)
|
||||
|
||||
// Snowflake 结构体
|
||||
type Snowflake struct {
|
||||
mu sync.Mutex
|
||||
lastTimestamp int64
|
||||
workerID int64
|
||||
datacenterID int64
|
||||
sequence int64
|
||||
}
|
||||
|
||||
// NewSnowflake 初始化 Snowflake 实例
|
||||
func NewSnowflake(workerID, datacenterID int64) (*Snowflake, error) {
|
||||
if workerID < 0 || workerID > maxWorkerID {
|
||||
return nil, fmt.Errorf("worker ID must be between 0 and %d", maxWorkerID)
|
||||
}
|
||||
if datacenterID < 0 || datacenterID > maxDatacenterID {
|
||||
return nil, fmt.Errorf("datacenter ID must be between 0 and %d", maxDatacenterID)
|
||||
}
|
||||
return &Snowflake{
|
||||
lastTimestamp: -1,
|
||||
workerID: workerID,
|
||||
datacenterID: datacenterID,
|
||||
sequence: 0,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// NextID 生成唯一 ID
|
||||
func (s *Snowflake) NextID() int64 {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
timestamp := time.Now().UnixNano() / 1e6 // 当前毫秒时间戳
|
||||
|
||||
// 如果当前时间小于上次生成 ID 的时间，说明系统时钟可能回退，返回错误或等待
|
||||
if timestamp < s.lastTimestamp {
|
||||
panic(fmt.Sprintf("clock moved backwards. refusing to generate id for %d milliseconds", s.lastTimestamp-timestamp))
|
||||
}
|
||||
|
||||
// 如果是同一毫秒内生成的，则递增序列号
|
||||
if timestamp == s.lastTimestamp {
|
||||
s.sequence = (s.sequence + 1) & maxSequence
|
||||
// 如果序列号超过最大值，则等待下一毫秒
|
||||
if s.sequence == 0 {
|
||||
for timestamp <= s.lastTimestamp {
|
||||
timestamp = time.Now().UnixNano() / 1e6
|
||||
}
|
||||
}
|
||||
} else {
|
||||
s.sequence = 0
|
||||
}
|
||||
|
||||
s.lastTimestamp = timestamp
|
||||
|
||||
// 组合各部分生成最终 ID
|
||||
id := ((timestamp - twepoch) << timestampShift) |
|
||||
(s.datacenterID << datacenterIDShift) |
|
||||
(s.workerID << workerIDShift) |
|
||||
s.sequence
|
||||
|
||||
return id
|
||||
}
|
||||
Reference in New Issue
Block a user