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

worker-svr/internal/queue/zset/producer_lease.go

@@ -0,0 +1,64 @@
+package zset
+
+import (
+	"context"
+	"strings"
+
+	"github.com/gogf/gf/v2/container/gvar"
+	"github.com/gogf/gf/v2/errors/gerror"
+)
+
+var ErrProducerLeaseLost = gerror.New("producer lease lost")
+
+// WithIMEIProducerLease returns a private producer view; the shared store and
+// other producer/consumer protocols are unchanged. The IMEI caller must reject
+// Redis Cluster before using this view: the global lock and run keys cross slots.
+func (s *Store) WithIMEIProducerLease(key, owner, runID string) (*Store, error) {
+	if key == "" || owner == "" || runID == "" {
+		return nil, gerror.New("producer lease key, owner and runId are required")
+	}
+	view := *s
+	view.redis = &imeiProducerRedis{RedisClient: s.redis, key: key, owner: owner, runID: runID}
+	return &view, nil
+}
+
+type imeiProducerRedis struct {
+	RedisClient
+	key, owner, runID string
+}
+
+func (r *imeiProducerRedis) Eval(ctx context.Context, script string, numKeys int64, keys []string, args []any) (*gvar.Var, error) {
+	leaseKeys := append(append([]string(nil), keys...), r.key)
+	leaseArgs := append(append([]any(nil), args...), r.owner, r.runID)
+	result, err := r.RedisClient.Eval(ctx, imeiProducerLeaseGuard+script, numKeys+1, leaseKeys, leaseArgs)
+	if err != nil && strings.Contains(err.Error(), "IMEI_PRODUCER_LEASE_LOST") {
+		return nil, gerror.Wrap(ErrProducerLeaseLost, "reject stale IMEI producer write")
+	}
+	return result, err
+}
+
+func (r *imeiProducerRedis) Do(ctx context.Context, command string, args ...any) (*gvar.Var, error) {
+	switch strings.ToUpper(command) {
+	case "HGETALL", "SMEMBERS": // The producer's metadata/capacity reads are not mutations.
+		return r.RedisClient.Do(ctx, command, args...)
+	case "SADD", "HSET": // CreateRun also writes the index and metadata extensions.
+		commandArgs := append([]any{command}, args...)
+		return r.Eval(ctx, "return redis.call(unpack(ARGV))", 0, nil, commandArgs)
+	default:
+		return nil, gerror.Newf("unsupported command in IMEI producer lease view: %s", command)
+	}
+}
+
+// Remove guard-only arguments before executing the original queue script, so
+// task-pair loops and key indexing retain their original semantics.
+const imeiProducerLeaseGuard = `
+local leaseRunID = table.remove(ARGV)
+local leaseOwner = table.remove(ARGV)
+local leaseKey = table.remove(KEYS)
+local current = redis.call('GET', leaseKey)
+if not current then return redis.error_reply('IMEI_PRODUCER_LEASE_LOST') end
+local ok, lease = pcall(cjson.decode, current)
+if not ok or type(lease) ~= 'table' or lease['owner'] ~= leaseOwner or lease['runId'] ~= leaseRunID then
+ return redis.error_reply('IMEI_PRODUCER_LEASE_LOST')
+end
+`