MiFaZhan
65 lines
2.5 KiB
Go
65 lines
2.5 KiB
Go
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
|
|
`
|