iotlcm/internal/producer/imei/lease_client.go

package imei

import (
	"context"
	"strings"
	"time"

	"github.com/gogf/gf/v2/container/gvar"
	"github.com/gogf/gf/v2/database/gredis"
	"github.com/gogf/gf/v2/errors/gerror"
	goredis "github.com/redis/go-redis/v9"
	"github.com/redis/go-redis/v9/maintnotifications"

	queue "git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/queue/zset"
)

type leaseClientFactory func() (queue.RedisClient, func() error, error)

// A custom client must explicitly supply the same bounded-I/O contract. There
// is no fallback to context cancellation on an unknown Redis implementation.
type boundedLeaseProvider interface {
	newBoundedLeaseClient(time.Duration) (queue.RedisClient, func() error, error)
}

func newLeaseClientFactory(client queue.RedisClient, config *gredis.Config, budget time.Duration) (leaseClientFactory, error) {
	if budget <= 0 {
		return nil, gerror.New("IMEI lease I/O budget must be positive")
	}
	if provider, ok := client.(boundedLeaseProvider); ok {
		return func() (queue.RedisClient, func() error, error) { return provider.newBoundedLeaseClient(budget) }, nil
	}
	getter, ok := client.(interface{ GetAdapter() gredis.Adapter })
	if !ok || getter.GetAdapter() == nil {
		return nil, gerror.New("IMEI producer requires a bounded GoFrame Redis adapter")
	}
	raw, ok := getter.GetAdapter().Client().(*goredis.Client)
	if !ok {
		return nil, gerror.New("IMEI producer does not support Redis Cluster or an unknown Redis adapter")
	}
	opts := *raw.Options()
	if opts.Addr == "FailoverClient" {
		if config == nil || config.MasterName == "" || config.Cluster || config.SlaveOnly {
			return nil, gerror.New("IMEI Sentinel lease client requires its master configuration")
		}
		// Rebuild Sentinel discovery too: sharing the old dialer would retain
		// its unbounded Sentinel socket reads even if master I/O were bounded.
		cfg := *config
		return func() (queue.RedisClient, func() error, error) {
			client := goredis.NewFailoverClient(&goredis.FailoverOptions{
				MasterName: cfg.MasterName, SentinelAddrs: sentinelAddresses(cfg.Address),
				SentinelUsername: cfg.SentinelUser, SentinelPassword: cfg.SentinelPass,
				Username: cfg.User, Password: cfg.Pass, DB: cfg.Db, Protocol: cfg.Protocol,
				TLSConfig:   opts.TLSConfig,
				DialTimeout: budget, ReadTimeout: budget, WriteTimeout: budget, PoolTimeout: budget,
				ContextTimeoutEnabled: true, MaxRetries: -1, DialerRetries: 1,
				PoolSize: 1, MaxActiveConns: 1,
			})
			return &leaseRedis{client}, client.Close, nil
		}, nil
	}
	// Own a separate pool so a timeout/close never closes a consumer's socket.
	opts.DialTimeout, opts.ReadTimeout, opts.WriteTimeout, opts.PoolTimeout = budget, budget, budget, budget
	opts.ContextTimeoutEnabled = true
	// Maintenance handoffs can relax socket deadlines; leases must never do so.
	opts.MaintNotificationsConfig = &maintnotifications.Config{Mode: maintnotifications.ModeDisabled}
	opts.MaxRetries, opts.DialerRetries = -1, 1
	opts.PoolSize, opts.MaxActiveConns, opts.MinIdleConns, opts.MaxIdleConns = 1, 1, 0, 1
	opts.Dialer = goredis.NewDialer(&opts)
	return func() (queue.RedisClient, func() error, error) {
		client := goredis.NewClient(&opts)
		return &leaseRedis{client}, client.Close, nil
	}, nil
}

func sentinelAddresses(address string) []string {
	var addresses []string
	for _, addr := range strings.Split(address, ",") {
		if addr = strings.TrimSpace(addr); addr != "" {
			addresses = append(addresses, addr)
		}
	}
	return addresses
}

type leaseRedis struct{ client *goredis.Client }

func (r *leaseRedis) Do(ctx context.Context, command string, args ...any) (*gvar.Var, error) {
	result, err := r.client.Do(ctx, append([]any{command}, args...)...).Result()
	return gvar.New(result), err
}

func (r *leaseRedis) Eval(ctx context.Context, script string, _ int64, keys []string, args []any) (*gvar.Var, error) {
	result, err := r.client.Eval(ctx, script, keys, args...).Result()
	return gvar.New(result), err
}