MiFaZhan
172 lines
6.6 KiB
Go
172 lines
6.6 KiB
Go
package zset
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"time"
|
|
)
|
|
|
|
// AcquireLock obtains a Redis lease only when the key is absent.
|
|
func (s *Store) AcquireLock(ctx context.Context, key, token string, ttl time.Duration) (bool, error) {
|
|
if key == "" || token == "" {
|
|
return false, fmt.Errorf("lock key and token are required")
|
|
}
|
|
if ttl <= 0 {
|
|
return false, fmt.Errorf("lock ttl must be positive")
|
|
}
|
|
result, err := s.redis.Do(ctx, "SET", key, token, "NX", "PX", ttl.Milliseconds())
|
|
if err != nil {
|
|
return false, fmt.Errorf("acquire lock: %w", err)
|
|
}
|
|
return result.String() == "OK", nil
|
|
}
|
|
|
|
// RenewLock extends the lease only when the token still owns the key.
|
|
func (s *Store) RenewLock(ctx context.Context, key, token string, ttl time.Duration) (bool, error) {
|
|
if key == "" || token == "" {
|
|
return false, fmt.Errorf("lock key and token are required")
|
|
}
|
|
if ttl <= 0 {
|
|
return false, fmt.Errorf("lock ttl must be positive")
|
|
}
|
|
result, err := s.redis.Eval(ctx, renewLockScript, 1, []string{key}, []any{token, ttl.Milliseconds()})
|
|
if err != nil {
|
|
return false, fmt.Errorf("renew lock: %w", err)
|
|
}
|
|
return result.Int() == 1, nil
|
|
}
|
|
|
|
// ReleaseLock removes the key only when the token still owns it.
|
|
func (s *Store) ReleaseLock(ctx context.Context, key, token string) (bool, error) {
|
|
if key == "" || token == "" {
|
|
return false, fmt.Errorf("lock key and token are required")
|
|
}
|
|
result, err := s.redis.Eval(ctx, releaseLockScript, 1, []string{key}, []any{token})
|
|
if err != nil {
|
|
return false, fmt.Errorf("release lock: %w", err)
|
|
}
|
|
return result.Int() == 1, nil
|
|
}
|
|
|
|
// AcquireLockOrTakeover acquires key for token, or takes it over when the
|
|
// existing lock cannot belong to a live owner of different work.
|
|
//
|
|
// This exists because a producer that dies (for example on a Redis EOF) leaves
|
|
// its lock behind for the remainder of the TTL. A redelivered start message for
|
|
// that same run would otherwise fail to acquire the lock, be reported as
|
|
// "already running", and be acked and dropped by the consumer - the request
|
|
// would disappear without ever reaching the idempotent rescan path that
|
|
// hasPendingTasks deliberately allows.
|
|
//
|
|
// Takeover rules, in order:
|
|
//
|
|
// 1. No lock: acquire it.
|
|
// 2. Structured lock for the same runID: take it over. Owner tokens are
|
|
// regenerated per process, so only the runID can identify "the same work".
|
|
// 3. Structured lock for a different runID: refuse. Another run is live.
|
|
// 4. Bare (legacy) token: take it over only when it is provably stale (see
|
|
// below). A bare value can only have been written by a build that predates
|
|
// structured locks. Refusing outright would strand the lock for the whole
|
|
// TTL and drop the request, which is the exact failure this method exists to
|
|
// prevent. Taking it over unconditionally would be unsafe: during a rolling
|
|
// deploy an older process still writes and renews a bare token, and stealing
|
|
// a live lock lets two producers run at once.
|
|
//
|
|
// The staleness test for a legacy lock: a live holder renews every TTL/3 (see
|
|
// Producer.renewLock), so its remaining TTL never falls below 2*TTL/3. A
|
|
// remaining TTL at or below TTL/2 therefore proves the holder stopped renewing.
|
|
// A lock with no expiry at all is treated as stale, because every writer in this
|
|
// codebase sets one.
|
|
//
|
|
// runID may be empty, in which case the call degrades to AcquireLock.
|
|
func (s *Store) AcquireLockOrTakeover(ctx context.Context, key, token, runID string, ttl time.Duration) (bool, error) {
|
|
if key == "" || token == "" {
|
|
return false, fmt.Errorf("lock key and token are required")
|
|
}
|
|
if ttl <= 0 {
|
|
return false, fmt.Errorf("lock ttl must be positive")
|
|
}
|
|
if runID == "" {
|
|
return s.AcquireLock(ctx, key, token, ttl)
|
|
}
|
|
value, err := marshalLockValue(lockValue{Owner: token, RunID: runID})
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
result, err := s.redis.Eval(ctx, acquireOrTakeoverLockScript, 1, []string{key}, []any{value, ttl.Milliseconds()})
|
|
if err != nil {
|
|
return false, fmt.Errorf("acquire or take over lock: %w", err)
|
|
}
|
|
return result.Int() == 1, nil
|
|
}
|
|
|
|
// lockValue is the structured lock payload. Older locks hold a bare owner
|
|
// token; the scripts accept both so a rolling deploy is safe.
|
|
type lockValue struct {
|
|
Owner string `json:"owner"`
|
|
RunID string `json:"runId,omitempty"`
|
|
}
|
|
|
|
func marshalLockValue(value lockValue) (string, error) {
|
|
encoded, err := json.Marshal(value)
|
|
if err != nil {
|
|
return "", fmt.Errorf("encode lock value: %w", err)
|
|
}
|
|
return string(encoded), nil
|
|
}
|
|
|
|
// renewLockScript extends the lease when the token still owns the key. The
|
|
// stored value may be a bare legacy token or a structured lockValue.
|
|
const renewLockScript = `
|
|
local current = redis.call('GET', KEYS[1])
|
|
if not current then return 0 end
|
|
local ok, decoded = pcall(cjson.decode, current)
|
|
local owner = current
|
|
if ok and type(decoded) == 'table' and decoded['owner'] then owner = decoded['owner'] end
|
|
if owner ~= ARGV[1] then return 0 end
|
|
redis.call('PEXPIRE', KEYS[1], ARGV[2])
|
|
return 1
|
|
`
|
|
|
|
// releaseLockScript deletes the key only when the token still owns it.
|
|
const releaseLockScript = `
|
|
local current = redis.call('GET', KEYS[1])
|
|
if not current then return 0 end
|
|
local ok, decoded = pcall(cjson.decode, current)
|
|
local owner = current
|
|
if ok and type(decoded) == 'table' and decoded['owner'] then owner = decoded['owner'] end
|
|
if owner ~= ARGV[1] then return 0 end
|
|
redis.call('DEL', KEYS[1])
|
|
return 1
|
|
`
|
|
|
|
// acquireOrTakeoverLockScript sets the key when free, or replaces it when the
|
|
// existing lock either belongs to the same run or is a legacy bare token that
|
|
// cannot belong to a live structured-lock holder. Returns 1 on success.
|
|
const acquireOrTakeoverLockScript = `
|
|
local current = redis.call('GET', KEYS[1])
|
|
if not current then
|
|
redis.call('SET', KEYS[1], ARGV[1], 'PX', ARGV[2])
|
|
return 1
|
|
end
|
|
local okNext, wanted = pcall(cjson.decode, ARGV[1])
|
|
if not okNext or type(wanted) ~= 'table' or wanted['runId'] == nil then return 0 end
|
|
local ok, decoded = pcall(cjson.decode, current)
|
|
if not ok or type(decoded) ~= 'table' then
|
|
-- Legacy bare token. A live holder renews every TTL/3, so its remaining TTL
|
|
-- never drops below 2*TTL/3; at or below TTL/2 the holder has stopped
|
|
-- renewing and the lock is stale. A missing expiry is stale too, because
|
|
-- every writer sets one. Refuse a still-live legacy lock so a rolling deploy
|
|
-- cannot end up with two producers.
|
|
local pttl = redis.call('PTTL', KEYS[1])
|
|
if pttl > 0 and pttl > math.floor(tonumber(ARGV[2]) / 2) then return 0 end
|
|
redis.call('SET', KEYS[1], ARGV[1], 'PX', ARGV[2])
|
|
return 1
|
|
end
|
|
if decoded['runId'] == nil then return 0 end
|
|
if decoded['runId'] ~= wanted['runId'] then return 0 end
|
|
redis.call('SET', KEYS[1], ARGV[1], 'PX', ARGV[2])
|
|
return 1
|
|
`
|