iotlcm/worker-svr/internal/queue/zset/store.go

package zset

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/gogf/gf/v2/container/gvar"
)

// RedisClient is the small Redis surface required by Store. Keeping this
// interface here makes the claim/ack protocol testable without a live Redis.
type RedisClient interface {
	Do(ctx context.Context, command string, args ...any) (*gvar.Var, error)
	Eval(ctx context.Context, script string, numKeys int64, keys []string, args []any) (*gvar.Var, error)
}

const defaultPrefix = "flow:"

// Task is the payload stored as a ZSet member. The member itself is retained
// so a task can be recovered after a consumer process exits.
type Task struct {
	TaskID    string `json:"taskId"`
	RunID     string `json:"runId"`
	Operation string `json:"operation"`
	// QueryMonth is the run's period identifier, not necessarily a month.
	// Monthly usage sync stores yyyyMM; daily usage sync stores yyyyMMdd.
	// The whole queue, run, archive and replay path only ever copies this
	// string, so a period change needs no schema change. Do not parse it as a
	// month without checking the operation first.
	QueryMonth string `json:"queryMonth"`
	ChanKey    string `json:"chanKey"`
	TempKey    string `json:"tempKey"`
	Operator   int    `json:"operator,omitempty"`
	CardID     int64  `json:"cardId"`
	ICCID      string `json:"iccid"`
	MSISDN     string `json:"msisdn"`
	IMSI       string `json:"imsi"`
	// PackageCode is the card's current package at production time. Daily
	// usage writes it to the TDengine child table where the diagnosis page
	// reads it back. It is additive: older members simply decode as empty.
	PackageCode string `json:"packageCode,omitempty"`
	// NextExecuteAt is the ZSet score and is intentionally excluded from the
	// member payload so retrying the same producer run remains idempotent.
	NextExecuteAt int64 `json:"-"`
}

// ClaimedTask is returned after a successful Lua lease claim.
type ClaimedTask struct {
	Task       Task
	Member     string
	LeaseToken string
	LeaseUntil int64
	// Round is populated by ClaimRun. Legacy channel-scoped claims leave it
	// zero and continue using the original lease keys.
	Round int
}

// StableTaskID is the logical id used for idempotency across producer retries.
func StableTaskID(runID, operation string, cardID int64) string {
	return fmt.Sprintf("%s:%s:%d", runID, operation, cardID)
}

// Store owns the Redis key and lease protocol for one task namespace.
type Store struct {
	redis          RedisClient
	prefix         string
	owner          string
	archiveEnabled bool
}

// RedisClient exposes the narrow Redis surface for adjacent full-sync
// repositories. Callers must continue to use Store key builders for namespacing.
func (s *Store) RedisClient() RedisClient { return s.redis }

// Prefix returns the configured Redis namespace prefix.
func (s *Store) Prefix() string { return s.prefix }

// EnableArchiveIndex makes terminal runs discoverable by the optional JSON
// archive worker. It is opt-in so the legacy queue path remains unchanged.
func (s *Store) EnableArchiveIndex(enabled bool) { s.archiveEnabled = enabled }

func New(redis RedisClient, prefix string) (*Store, error) {
	if redis == nil {
		return nil, fmt.Errorf("redis client is nil")
	}
	if strings.TrimSpace(prefix) == "" {
		prefix = defaultPrefix
	}
	return &Store{redis: redis, prefix: prefix, owner: newOwner()}, nil
}

func (s *Store) Owner() string { return s.owner }

func (s *Store) Key(chanKey string) (string, error) {
	if strings.TrimSpace(chanKey) == "" {
		return "", fmt.Errorf("chanKey is empty")
	}
	return s.prefix + chanKey, nil
}

func (s *Store) leaseKey(chanKey string) (string, error) {
	if strings.TrimSpace(chanKey) == "" {
		return "", fmt.Errorf("chanKey is empty")
	}
	return s.prefix + "lease:" + chanKey, nil
}

func (s *Store) blockedKey(chanKey string) (string, error) {
	if strings.TrimSpace(chanKey) == "" {
		return "", fmt.Errorf("chanKey is empty")
	}
	return s.prefix + "blocked:" + chanKey, nil
}

func (s *Store) ChannelsKey() string { return s.prefix + "channels" }

func (s *Store) RegisterChannel(ctx context.Context, chanKey string) error {
	if strings.TrimSpace(chanKey) == "" {
		return fmt.Errorf("chanKey is empty")
	}
	if _, err := s.redis.Do(ctx, "SADD", s.ChannelsKey(), chanKey); err != nil {
		return fmt.Errorf("register channel: %w", err)
	}
	return nil
}

func (s *Store) ActiveChannels(ctx context.Context) ([]string, error) {
	result, err := s.redis.Do(ctx, "SMEMBERS", s.ChannelsKey())
	if err != nil {
		return nil, fmt.Errorf("list active channels: %w", err)
	}
	return result.Strings(), nil
}

// Enqueue adds or replaces a task at the requested execution time. A
// deterministic TaskID lets a producer retry its write without creating a
// second logical task.
func (s *Store) Enqueue(ctx context.Context, task Task, nextExecuteAt time.Time) error {
	_, err := s.EnqueueBatch(ctx, task.ChanKey, []Task{task}, nextExecuteAt)
	return err
}

// EnqueueBatch registers one channel and writes all of its tasks atomically.
// It returns zero when the channel has already been stopped for this run.
func (s *Store) EnqueueBatch(ctx context.Context, chanKey string, tasks []Task, nextExecuteAt time.Time) (int, error) {
	if len(tasks) == 0 {
		return 0, nil
	}
	key, err := s.Key(chanKey)
	if err != nil {
		return 0, err
	}
	blockedKey, err := s.blockedKey(chanKey)
	if err != nil {
		return 0, err
	}
	if nextExecuteAt.IsZero() {
		nextExecuteAt = time.Now()
	}
	runID := tasks[0].RunID
	args := make([]any, 0, len(tasks)+3)
	args = append(args, chanKey, runID, nextExecuteAt.UnixMilli())
	for index := range tasks {
		task := tasks[index]
		if task.TaskID == "" {
			return 0, fmt.Errorf("taskId is empty")
		}
		if task.ChanKey != chanKey {
			return 0, fmt.Errorf("task chanKey %q does not match batch chanKey %q", task.ChanKey, chanKey)
		}
		if task.RunID != runID {
			return 0, fmt.Errorf("task runId %q does not match batch runId %q", task.RunID, runID)
		}
		task.NextExecuteAt = nextExecuteAt.UnixMilli()
		member, marshalErr := json.Marshal(task)
		if marshalErr != nil {
			return 0, fmt.Errorf("marshal task: %w", marshalErr)
		}
		args = append(args, string(member))
	}
	result, err := s.redis.Eval(ctx, enqueueBatchScript, 3, []string{key, s.ChannelsKey(), blockedKey}, args)
	if err != nil {
		return 0, fmt.Errorf("enqueue task batch: %w", err)
	}
	return result.Int(), nil
}

// Claim atomically leases due tasks. It intentionally does not remove the
// members: a process crash leaves them recoverable after the lease expires.
func (s *Store) Claim(ctx context.Context, chanKey string, now time.Time, leaseUntil time.Time, limit int) ([]ClaimedTask, error) {
	if limit <= 0 {
		return nil, nil
	}
	key, err := s.Key(chanKey)
	if err != nil {
		return nil, err
	}
	leaseKey, err := s.leaseKey(chanKey)
	if err != nil {
		return nil, err
	}
	if now.IsZero() {
		now = time.Now()
	}
	if leaseUntil.IsZero() || !leaseUntil.After(now) {
		return nil, fmt.Errorf("leaseUntil must be after now")
	}
	result, err := s.redis.Eval(ctx, claimScript, 2, []string{key, leaseKey}, []any{now.UnixMilli(), leaseUntil.UnixMilli(), limit, s.owner})
	if err != nil {
		return nil, fmt.Errorf("claim tasks: %w", err)
	}
	values := result.Vars()
	if len(values)%2 != 0 {
		return nil, fmt.Errorf("claim returned odd number of values: %d", len(values))
	}
	tasks := make([]ClaimedTask, 0, len(values)/2)
	for i := 0; i < len(values); i += 2 {
		member := values[i].String()
		var task Task
		if err := json.Unmarshal([]byte(member), &task); err != nil {
			return nil, fmt.Errorf("decode claimed task: %w", err)
		}
		tasks = append(tasks, ClaimedTask{
			Task:       task,
			Member:     member,
			LeaseToken: values[i+1].String(),
			LeaseUntil: leaseUntil.UnixMilli(),
		})
	}
	return tasks, nil
}

// Renew extends a lease only when the caller still owns it.
func (s *Store) Renew(ctx context.Context, task ClaimedTask, leaseUntil time.Time) (bool, error) {
	count, err := s.RenewBatch(ctx, []ClaimedTask{task}, leaseUntil)
	return count == 1, err
}

// RenewBatch extends every lease in one Redis round trip. The Lua script
// validates the complete batch before changing any score, so callers never
// continue with only part of a claimed batch.
func (s *Store) RenewBatch(ctx context.Context, tasks []ClaimedTask, leaseUntil time.Time) (int, error) {
	if len(tasks) == 0 {
		return 0, nil
	}
	if leaseUntil.IsZero() {
		return 0, fmt.Errorf("leaseUntil is zero")
	}
	chanKey := tasks[0].Task.ChanKey
	key, err := s.Key(chanKey)
	if err != nil {
		return 0, err
	}
	leaseKey, err := s.leaseKey(chanKey)
	if err != nil {
		return 0, err
	}
	args := make([]any, 0, len(tasks)*2+1)
	args = append(args, leaseUntil.UnixMilli())
	for _, task := range tasks {
		if task.Task.ChanKey != chanKey {
			return 0, fmt.Errorf("renew batch contains multiple channels")
		}
		args = append(args, task.Member, task.LeaseToken)
	}
	result, err := s.redis.Eval(ctx, renewBatchScript, 2, []string{key, leaseKey}, args)
	if err != nil {
		return 0, fmt.Errorf("renew task batch lease: %w", err)
	}
	return result.Int(), nil
}

// Ack removes a task only when the lease token still matches. A stale worker
// cannot delete a task reclaimed by another worker.
func (s *Store) Ack(ctx context.Context, task ClaimedTask) (bool, error) {
	count, err := s.AckBatch(ctx, []ClaimedTask{task})
	return count == 1, err
}

// AckBatch removes every task whose lease token is still owned by this
// consumer. A partial count means at least one lease expired or was reclaimed.
func (s *Store) AckBatch(ctx context.Context, tasks []ClaimedTask) (int, error) {
	if len(tasks) == 0 {
		return 0, nil
	}
	chanKey := tasks[0].Task.ChanKey
	key, err := s.Key(chanKey)
	if err != nil {
		return 0, err
	}
	leaseKey, err := s.leaseKey(chanKey)
	if err != nil {
		return 0, err
	}
	args := make([]any, 0, len(tasks)*2)
	for _, task := range tasks {
		if task.Task.ChanKey != chanKey {
			return 0, fmt.Errorf("ack batch contains multiple channels")
		}
		args = append(args, task.Member, task.LeaseToken)
	}
	result, err := s.redis.Eval(ctx, ackBatchScript, 2, []string{key, leaseKey}, args)
	if err != nil {
		return 0, fmt.Errorf("ack task batch: %w", err)
	}
	return result.Int(), nil
}

// BlockRun stops one channel for the current run and removes its queued and
// leased tasks. Producer batches for the same run are rejected atomically.
func (s *Store) BlockRun(ctx context.Context, chanKey, runID string) (int64, error) {
	if strings.TrimSpace(runID) == "" {
		return 0, fmt.Errorf("runId is empty")
	}
	key, err := s.Key(chanKey)
	if err != nil {
		return 0, err
	}
	leaseKey, err := s.leaseKey(chanKey)
	if err != nil {
		return 0, err
	}
	blockedKey, err := s.blockedKey(chanKey)
	if err != nil {
		return 0, err
	}
	result, err := s.redis.Eval(ctx, blockRunScript, 3, []string{key, leaseKey, blockedKey}, []any{runID})
	if err != nil {
		return 0, fmt.Errorf("block channel run: %w", err)
	}
	return result.Int64(), nil
}

func (s *Store) Size(ctx context.Context, chanKey string) (int64, error) {
	key, err := s.Key(chanKey)
	if err != nil {
		return 0, err
	}
	result, err := s.redis.Do(ctx, "ZCARD", key)
	if err != nil {
		return 0, fmt.Errorf("get task count: %w", err)
	}
	return result.Int64(), nil
}

func newOwner() string {
	var bytes [12]byte
	if _, err := rand.Read(bytes[:]); err == nil {
		return hex.EncodeToString(bytes[:])
	}
	return fmt.Sprintf("consumer-%d", time.Now().UnixNano())
}

const claimScript = `
local members = redis.call('ZRANGEBYSCORE', KEYS[1], '-inf', ARGV[1], 'LIMIT', 0, ARGV[3])
local result = {}
local serverTime = redis.call('TIME')
for index, member in ipairs(members) do
  local token = ARGV[4] .. ':' .. serverTime[1] .. ':' .. serverTime[2] .. ':' .. index
  redis.call('ZADD', KEYS[1], ARGV[2], member)
  redis.call('HSET', KEYS[2], member, token)
  table.insert(result, member)
  table.insert(result, token)
end
return result
`

const enqueueBatchScript = `
if redis.call('GET', KEYS[3]) == ARGV[2] then return 0 end
redis.call('SADD', KEYS[2], ARGV[1])
for index = 4, #ARGV do
  redis.call('ZADD', KEYS[1], ARGV[3], ARGV[index])
end
return #ARGV - 3
`

const renewBatchScript = `
for index = 2, #ARGV, 2 do
  local member = ARGV[index]
  local token = ARGV[index + 1]
  if redis.call('HGET', KEYS[2], member) ~= token then return 0 end
  if redis.call('ZSCORE', KEYS[1], member) == false then return 0 end
end
for index = 2, #ARGV, 2 do
  redis.call('ZADD', KEYS[1], ARGV[1], ARGV[index])
end
return (#ARGV - 1) / 2
`

const ackBatchScript = `
local removed = 0
for index = 1, #ARGV, 2 do
  local member = ARGV[index]
  local token = ARGV[index + 1]
  if redis.call('HGET', KEYS[2], member) == token then
    redis.call('ZREM', KEYS[1], member)
    redis.call('HDEL', KEYS[2], member)
    removed = removed + 1
  end
end
return removed
`

const blockRunScript = `
redis.call('SET', KEYS[3], ARGV[1])
local pending = 0
local members = redis.call('ZRANGE', KEYS[1], 0, -1)
for _, member in ipairs(members) do
  -- Task members are JSON produced by this package.  Match the complete
  -- runId field so blocking run A cannot delete run B's work on the same
  -- legacy channel queue.
  if string.find(member, '"runId":"' .. ARGV[1] .. '"', 1, true) then
    redis.call('ZREM', KEYS[1], member)
    redis.call('HDEL', KEYS[2], member)
    pending = pending + 1
  end
end
return pending
`