forked from baron/baron-sso
80 lines
1.8 KiB
Go
80 lines
1.8 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"os"
|
|
"time"
|
|
|
|
"github.com/go-redis/redis/v8"
|
|
)
|
|
|
|
const (
|
|
worksmobileRelayLeaderLockKey = "baron:worksmobile:relay:leader"
|
|
worksmobileRelayLeaderLockTTL = 30 * time.Second
|
|
)
|
|
|
|
const worksmobileRelayLeaderRenewScript = `
|
|
if redis.call("GET", KEYS[1]) == ARGV[1] then
|
|
return redis.call("EXPIRE", KEYS[1], ARGV[2])
|
|
end
|
|
return 0
|
|
`
|
|
|
|
type WorksmobileRedisRelayLeaderLock struct {
|
|
client *redis.Client
|
|
key string
|
|
ttl time.Duration
|
|
ownerID string
|
|
}
|
|
|
|
func NewWorksmobileRedisRelayLeaderLock(redisService *RedisService) *WorksmobileRedisRelayLeaderLock {
|
|
if redisService == nil || redisService.Client == nil {
|
|
return nil
|
|
}
|
|
return &WorksmobileRedisRelayLeaderLock{
|
|
client: redisService.Client,
|
|
key: worksmobileRelayLeaderLockKey,
|
|
ttl: worksmobileRelayLeaderLockTTL,
|
|
ownerID: newWorksmobileRelayLeaderOwnerID(),
|
|
}
|
|
}
|
|
|
|
func (l *WorksmobileRedisRelayLeaderLock) EnsureLeadership(ctx context.Context) (bool, error) {
|
|
if l == nil || l.client == nil {
|
|
return true, nil
|
|
}
|
|
|
|
acquired, err := l.client.SetNX(ctx, l.key, l.ownerID, l.ttl).Result()
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
if acquired {
|
|
return true, nil
|
|
}
|
|
|
|
ttlSeconds := int64(l.ttl / time.Second)
|
|
if ttlSeconds <= 0 {
|
|
ttlSeconds = 30
|
|
}
|
|
result, err := l.client.Eval(ctx, worksmobileRelayLeaderRenewScript, []string{l.key}, l.ownerID, ttlSeconds).Int()
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
return result == 1, nil
|
|
}
|
|
|
|
func newWorksmobileRelayLeaderOwnerID() string {
|
|
hostname, _ := os.Hostname()
|
|
if hostname == "" {
|
|
hostname = "unknown-host"
|
|
}
|
|
randomBytes := make([]byte, 8)
|
|
if _, err := rand.Read(randomBytes); err != nil {
|
|
return fmt.Sprintf("%s:%d:%d", hostname, os.Getpid(), time.Now().UnixNano())
|
|
}
|
|
return fmt.Sprintf("%s:%d:%s", hostname, os.Getpid(), hex.EncodeToString(randomBytes))
|
|
}
|