iotlcm/internal/queue/zset/lock.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
`