chore: 并入 iotlcm-worker-svr 历史到 worker-svr/（原提交 96f0f4e）

2026-11-02 13:28:03 +08:00
parent 6de2e503cc 96f0f4eacd
commit 2923c78755
422 changed files with 80993 additions and 0 deletions
@@ -0,0 +1,424 @@
+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
+`