MiFaZhan
406 lines
13 KiB
Go
406 lines
13 KiB
Go
package imei
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gogf/gf/v2/container/gvar"
|
|
|
|
"git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/model"
|
|
queue "git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/queue/zset"
|
|
"git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/testutil"
|
|
)
|
|
|
|
type lockTestRedis struct {
|
|
queue.RedisClient
|
|
do func(context.Context, string, []any) (*gvar.Var, error)
|
|
eval func(context.Context, string, int64, []string, []any) (*gvar.Var, error)
|
|
}
|
|
|
|
func (r *lockTestRedis) newBoundedLeaseClient(budget time.Duration) (queue.RedisClient, func() error, error) {
|
|
factory, err := newLeaseClientFactory(r.RedisClient, nil, budget)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
client, closeClient, err := factory()
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
return &lockTestRedis{RedisClient: client,
|
|
do: func(ctx context.Context, command string, args []any) (*gvar.Var, error) {
|
|
if r.do != nil {
|
|
return r.do(ctx, command, args)
|
|
}
|
|
return client.Do(ctx, command, args...)
|
|
},
|
|
eval: func(ctx context.Context, script string, numKeys int64, keys []string, args []any) (*gvar.Var, error) {
|
|
if r.eval != nil {
|
|
return r.eval(ctx, script, numKeys, keys, args)
|
|
}
|
|
return client.Eval(ctx, script, numKeys, keys, args)
|
|
},
|
|
}, closeClient, nil
|
|
}
|
|
|
|
func (r *lockTestRedis) Do(ctx context.Context, command string, args ...any) (*gvar.Var, error) {
|
|
if r.do != nil {
|
|
return r.do(ctx, command, args)
|
|
}
|
|
return r.RedisClient.Do(ctx, command, args...)
|
|
}
|
|
|
|
func (r *lockTestRedis) Eval(ctx context.Context, script string, numKeys int64, keys []string, args []any) (*gvar.Var, error) {
|
|
if r.eval != nil {
|
|
return r.eval(ctx, script, numKeys, keys, args)
|
|
}
|
|
return r.RedisClient.Eval(ctx, script, numKeys, keys, args)
|
|
}
|
|
|
|
type lockTestScanner struct {
|
|
cardScanner
|
|
scanCalls atomic.Int32
|
|
snapshot func(context.Context) (int64, error)
|
|
scan func(context.Context, int64, int64, int) ([]Card, error)
|
|
}
|
|
|
|
func (s *lockTestScanner) SnapshotMaxID(ctx context.Context) (int64, error) {
|
|
if s.snapshot != nil {
|
|
return s.snapshot(ctx)
|
|
}
|
|
return s.cardScanner.SnapshotMaxID(ctx)
|
|
}
|
|
|
|
func (s *lockTestScanner) Scan(ctx context.Context, last, max int64, limit int) ([]Card, error) {
|
|
s.scanCalls.Add(1)
|
|
if s.scan != nil {
|
|
return s.scan(ctx, last, max, limit)
|
|
}
|
|
return s.cardScanner.Scan(ctx, last, max, limit)
|
|
}
|
|
|
|
func startLockTestProducer(t *testing.T, producer *Producer) <-chan error {
|
|
t.Helper()
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
result := make(chan error, 1)
|
|
done := make(chan struct{})
|
|
go func() {
|
|
defer close(done)
|
|
result <- producer.Produce(ctx, model.ImeiQueueStartMessage{RunID: "lock-test"})
|
|
}()
|
|
t.Cleanup(func() {
|
|
cancel()
|
|
select {
|
|
case <-done:
|
|
case <-time.After(3 * time.Second):
|
|
t.Error("producer goroutine did not stop after cleanup cancellation")
|
|
}
|
|
})
|
|
return result
|
|
}
|
|
|
|
func TestProduceStopsOnLockLossWhileWaitingForCapacity(t *testing.T) {
|
|
backendErr := errors.New("renewal backend unavailable")
|
|
for _, test := range []struct {
|
|
name string
|
|
err error
|
|
}{
|
|
{name: "taken_over", err: ErrAlreadyRunning},
|
|
{name: "redis_error", err: backendErr},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
var (
|
|
client = &lockTestRedis{RedisClient: testutil.Redis(t)}
|
|
waiting = make(chan struct{})
|
|
failRenew = make(chan struct{})
|
|
waitingOnce sync.Once
|
|
capacityRead atomic.Bool
|
|
scanner = &lockTestScanner{cardScanner: cardScanner{cards: []Card{
|
|
{ID: 1, ICCID: "a", ChanKey: "ch", TempKey: "sim_yd_v5"},
|
|
{ID: 2, ICCID: "b", ChanKey: "ch", TempKey: "sim_yd_v5"},
|
|
{ID: 3, ICCID: "c", ChanKey: "ch", TempKey: "sim_yd_v5"},
|
|
}}}
|
|
)
|
|
store, err := queue.New(client, "imei:")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
producer, err := New(scanner, store, Config{
|
|
PageSize: 2, GlobalHighWaterMark: 1,
|
|
RunLockTTL: 300 * time.Millisecond, PollInterval: time.Hour,
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
client.do = func(ctx context.Context, command string, args []any) (*gvar.Var, error) {
|
|
result, err := client.RedisClient.Do(ctx, command, args...)
|
|
if err == nil && command == "SMEMBERS" && args[0] == store.ActiveRunsKey() {
|
|
capacityRead.Store(true)
|
|
}
|
|
if err == nil && command == "HGETALL" && capacityRead.Load() && result.MapStrStr()["activeCount"] == "2" {
|
|
waitingOnce.Do(func() { close(waiting) })
|
|
}
|
|
return result, err
|
|
}
|
|
client.eval = func(ctx context.Context, script string, numKeys int64, keys []string, args []any) (*gvar.Var, error) {
|
|
if strings.Contains(script, "redis.call('PEXPIRE'") {
|
|
select {
|
|
case <-failRenew:
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
}
|
|
if test.err == backendErr {
|
|
return nil, backendErr
|
|
}
|
|
}
|
|
return client.RedisClient.Eval(ctx, script, numKeys, keys, args)
|
|
}
|
|
result := startLockTestProducer(t, producer)
|
|
select {
|
|
case <-waiting:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("producer did not reach backpressure")
|
|
}
|
|
// The long poll cannot finish naturally; renewal failure must wake it.
|
|
select {
|
|
case err := <-result:
|
|
t.Fatalf("producer exited before lock loss: %v", err)
|
|
case <-time.After(20 * time.Millisecond):
|
|
}
|
|
if test.err == ErrAlreadyRunning {
|
|
ok, err := store.AcquireLockOrTakeover(context.Background(), producer.config.RunLockKey, "new-owner", "lock-test", time.Minute)
|
|
if err != nil || !ok {
|
|
t.Fatalf("take over lock: acquired=%v err=%v", ok, err)
|
|
}
|
|
}
|
|
close(failRenew)
|
|
select {
|
|
case err := <-result:
|
|
if !errors.Is(err, test.err) {
|
|
t.Fatalf("producer error = %v, want renewal cause %v", err, test.err)
|
|
}
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("lock loss did not interrupt backpressure before the next capacity poll")
|
|
}
|
|
if calls := scanner.scanCalls.Load(); calls != 1 {
|
|
t.Fatalf("scan calls = %d, want only the first page", calls)
|
|
}
|
|
meta, err := store.GetRunMeta(context.Background(), "lock-test")
|
|
if err != nil || meta.ProducerDone || meta.ActiveCount != 2 {
|
|
t.Fatalf("run must remain incomplete without additional tasks: meta=%+v err=%v", meta, err)
|
|
}
|
|
if producer.running.Load() != 0 {
|
|
t.Fatal("producer running flag was not reset")
|
|
}
|
|
if test.err == ErrAlreadyRunning {
|
|
ok, err := store.RenewLock(context.Background(), producer.config.RunLockKey, "new-owner", time.Minute)
|
|
if err != nil || !ok {
|
|
t.Fatalf("old producer released the new owner's lock: renewed=%v err=%v", ok, err)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestProduceCancelsLockedOperationsWithRenewalCause(t *testing.T) {
|
|
for _, stage := range []string{"run_init", "snapshot", "scan", "enqueue", "producer_done"} {
|
|
t.Run(stage, func(t *testing.T) {
|
|
var (
|
|
client = &lockTestRedis{RedisClient: testutil.Redis(t)}
|
|
entered = make(chan context.Context, 1)
|
|
failRenew = make(chan struct{})
|
|
cause = errors.New("renewal failed during " + stage)
|
|
scanner = &lockTestScanner{cardScanner: cardScanner{cards: []Card{
|
|
{ID: 1, ICCID: "a", ChanKey: "ch", TempKey: "sim_yd_v5"},
|
|
}}}
|
|
)
|
|
block := func(ctx context.Context) error {
|
|
entered <- ctx
|
|
<-ctx.Done()
|
|
return ctx.Err()
|
|
}
|
|
if stage == "snapshot" {
|
|
scanner.snapshot = func(ctx context.Context) (int64, error) { return 0, block(ctx) }
|
|
}
|
|
if stage == "scan" {
|
|
scanner.scan = func(ctx context.Context, _, _ int64, _ int) ([]Card, error) { return nil, block(ctx) }
|
|
}
|
|
store, err := queue.New(client, "imei:")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
producer, err := New(scanner, store, Config{RunLockTTL: 300 * time.Millisecond})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
client.eval = func(ctx context.Context, script string, numKeys int64, keys []string, args []any) (*gvar.Var, error) {
|
|
if strings.Contains(script, "redis.call('PEXPIRE'") {
|
|
select {
|
|
case <-failRenew:
|
|
return nil, cause
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
}
|
|
}
|
|
if stage == "run_init" && strings.Contains(script, "'createdAt', ARGV[12]") ||
|
|
stage == "enqueue" && strings.Contains(script, "local taskID=ARGV[i]") ||
|
|
stage == "producer_done" && strings.Contains(script, "'producerCompletedAt'") {
|
|
return nil, block(ctx)
|
|
}
|
|
return client.RedisClient.Eval(ctx, script, numKeys, keys, args)
|
|
}
|
|
result := startLockTestProducer(t, producer)
|
|
var workCtx context.Context
|
|
select {
|
|
case workCtx = <-entered:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("producer did not enter " + stage)
|
|
}
|
|
close(failRenew)
|
|
select {
|
|
case err := <-result:
|
|
if !errors.Is(err, cause) || !errors.Is(context.Cause(workCtx), cause) {
|
|
t.Fatalf("renewal cause lost: result=%v context cause=%v", err, context.Cause(workCtx))
|
|
}
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("lock loss did not interrupt " + stage)
|
|
}
|
|
if stage == "run_init" {
|
|
if scanner.scanCalls.Load() != 0 {
|
|
t.Fatal("producer scanned after run initialization failed")
|
|
}
|
|
return
|
|
}
|
|
meta, err := store.GetRunMeta(context.Background(), "lock-test")
|
|
wantActive := int64(0)
|
|
if stage == "producer_done" {
|
|
wantActive = 1
|
|
}
|
|
if err != nil || meta.ProducerDone || meta.ActiveCount != wantActive {
|
|
t.Fatalf("unexpected work after lock loss: meta=%+v err=%v", meta, err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestProduceJoinsRenewalBeforeReleasingLock(t *testing.T) {
|
|
scanErr := errors.New("snapshot failed")
|
|
for _, test := range []struct {
|
|
name string
|
|
err error
|
|
}{
|
|
{name: "success"},
|
|
{name: "scanner_error", err: scanErr},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
var (
|
|
client = &lockTestRedis{RedisClient: testutil.Redis(t)}
|
|
renewEntered = make(chan struct{})
|
|
renewCanceled = make(chan struct{})
|
|
allowRenewReturn = make(chan struct{})
|
|
resumeSnapshot = make(chan struct{})
|
|
released = make(chan struct{}, 1)
|
|
allowReturnOnce sync.Once
|
|
scanner = &lockTestScanner{}
|
|
)
|
|
scanner.snapshot = func(ctx context.Context) (int64, error) {
|
|
select {
|
|
case <-resumeSnapshot:
|
|
return 0, test.err
|
|
case <-ctx.Done():
|
|
return 0, ctx.Err()
|
|
}
|
|
}
|
|
store, err := queue.New(client, "imei:")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
producer, err := New(scanner, store, Config{RunLockTTL: 300 * time.Millisecond})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
client.eval = func(ctx context.Context, script string, numKeys int64, keys []string, args []any) (*gvar.Var, error) {
|
|
if strings.Contains(script, "redis.call('PEXPIRE'") {
|
|
close(renewEntered)
|
|
<-ctx.Done()
|
|
close(renewCanceled)
|
|
<-allowRenewReturn
|
|
return nil, ctx.Err()
|
|
}
|
|
if len(keys) == 1 && keys[0] == producer.config.RunLockKey && len(args) == 1 {
|
|
released <- struct{}{}
|
|
}
|
|
return client.RedisClient.Eval(ctx, script, numKeys, keys, args)
|
|
}
|
|
result := startLockTestProducer(t, producer)
|
|
t.Cleanup(func() { allowReturnOnce.Do(func() { close(allowRenewReturn) }) })
|
|
select {
|
|
case <-renewEntered:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("renewal did not start")
|
|
}
|
|
close(resumeSnapshot)
|
|
select {
|
|
case <-renewCanceled:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("producer did not cancel renewal on exit")
|
|
}
|
|
// Keep the in-flight renewal alive to verify cleanup actually joins it.
|
|
select {
|
|
case err := <-result:
|
|
t.Fatalf("producer returned before renewal exited: %v", err)
|
|
case <-released:
|
|
t.Fatal("producer released the lock before renewal exited")
|
|
case <-time.After(20 * time.Millisecond):
|
|
}
|
|
allowReturnOnce.Do(func() { close(allowRenewReturn) })
|
|
select {
|
|
case err := <-result:
|
|
if !errors.Is(err, test.err) {
|
|
t.Fatalf("cleanup changed the production error: got %v, want %v", err, test.err)
|
|
}
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("producer did not finish after renewal exited")
|
|
}
|
|
select {
|
|
case <-released:
|
|
default:
|
|
t.Fatal("producer did not release the lock after joining renewal")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestProducePreservesCallerCancellationCause(t *testing.T) {
|
|
producer, store := newProducer(t, nil)
|
|
ctx, cancel := context.WithCancelCause(context.Background())
|
|
defer cancel(nil)
|
|
cause := errors.New("caller stopped the run")
|
|
producer.scanner = &lockTestScanner{snapshot: func(ctx context.Context) (int64, error) {
|
|
cancel(cause)
|
|
<-ctx.Done()
|
|
return 0, ctx.Err()
|
|
}}
|
|
if err := producer.Produce(ctx, model.ImeiQueueStartMessage{RunID: "caller-canceled"}); !errors.Is(err, cause) {
|
|
t.Fatalf("caller cancellation cause lost: %v", err)
|
|
}
|
|
meta, err := store.GetRunMeta(context.Background(), "caller-canceled")
|
|
if err != nil || meta.ProducerDone || producer.running.Load() != 0 {
|
|
t.Fatalf("canceled producer did not clean up: meta=%+v err=%v running=%d", meta, err, producer.running.Load())
|
|
}
|
|
}
|
|
|
|
func TestWaitForCapacityHonorsCancellationWithoutWatermark(t *testing.T) {
|
|
producer := &Producer{}
|
|
ctx, cancel := context.WithCancelCause(context.Background())
|
|
cause := errors.New("lock lost")
|
|
cancel(cause)
|
|
if err := producer.waitForCapacity(ctx); !errors.Is(err, cause) {
|
|
t.Fatalf("disabled watermark ignored cancellation: %v", err)
|
|
}
|
|
}
|