iotlcm/internal/producer/imei/lease_io_test.go

package imei

import (
	"bufio"
	"context"
	"errors"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"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"
	"github.com/gogf/gf/v2/database/gredis"
	goredis "github.com/redis/go-redis/v9"
)

// This proxy withholds a real Redis request/reply without closing its TCP socket.
func readLeaseRESP(r *bufio.Reader) ([]byte, error) {
	line, err := r.ReadBytes('\n')
	if err != nil {
		return nil, err
	}
	switch line[0] {
	case '$':
		n, err := strconv.Atoi(strings.TrimSpace(string(line[1:])))
		if err != nil {
			return nil, err
		}
		if n >= 0 {
			payload := make([]byte, n+2)
			if _, err := io.ReadFull(r, payload); err != nil {
				return nil, err
			}
			line = append(line, payload...)
		}
	case '*', '%', '~', '>', '|':
		n, err := strconv.Atoi(strings.TrimSpace(string(line[1:])))
		if err != nil {
			return nil, err
		}
		if line[0] == '%' || line[0] == '|' {
			n *= 2
		}
		for i := 0; i < n; i++ {
			payload, err := readLeaseRESP(r)
			if err != nil {
				return nil, err
			}
			line = append(line, payload...)
		}
	}
	return line, nil
}

func TestProduceBoundsRealRedisRenewalIO(t *testing.T) {
	for _, mode := range []string{"standalone", "sentinel"} {
		t.Run(mode, func(t *testing.T) { testProduceBoundsRealRedisRenewalIO(t, mode) })
	}
}

func testProduceBoundsRealRedisRenewalIO(t *testing.T, mode string) {
	for _, callerCancel := range []bool{false, true} {
		t.Run(strconv.FormatBool(callerCancel), func(t *testing.T) {
			backend := testutil.Redis(t)
			real, held, release := leaseRedisProxy(t, backend, "redis.call('PEXPIRE'", true)
			// Release the gate even on a red-test failure so no worker leaks.
			defer release()
			opts := real.GetAdapter().Client().(*goredis.Client).Options()
			if opts.ReadTimeout != 0 || opts.WriteTimeout != 0 || opts.ContextTimeoutEnabled {
				t.Fatal("test must start with production-default unbounded socket I/O")
			}
			cause := errors.New("snapshot failed independently")
			ctx, cancel := context.WithCancelCause(context.Background())
			defer cancel(nil)
			scanner := &lockTestScanner{snapshot: func(workCtx context.Context) (int64, error) {
				select {
				case <-held:
					if callerCancel {
						cancel(cause)
						<-workCtx.Done()
						return 0, workCtx.Err()
					}
					return 0, cause
				case <-workCtx.Done():
					return 0, workCtx.Err()
				}
			}}
			store, err := queue.New(real, "imei:")
			if err != nil {
				t.Fatal(err)
			}
			config := Config{RunLockTTL: time.Second}
			if mode == "sentinel" {
				real, config.LeaseRedisConfig = sentinelLeaseRedis(t, real)
				store, err = queue.New(real, "imei:")
				if err != nil {
					t.Fatal(err)
				}
			}
			producer, err := New(scanner, store, config)
			if err != nil {
				t.Fatal(err)
			}
			result := make(chan error, 1)
			go func() { result <- producer.Produce(ctx, model.ImeiQueueStartMessage{RunID: "lock-test"}) }()
			select {
			case <-held:
			case <-time.After(2 * time.Second):
				t.Fatal("proxy did not hold renewal response")
			}
			select {
			case err := <-result:
				if !errors.Is(err, cause) {
					t.Fatalf("lost original cause: %v", err)
				}
			case <-time.After(600 * time.Millisecond):
				release()
				<-result
				t.Fatal("canceled renewal socket I/O blocked cleanup beyond its lease budget")
			}
			if producer.running.Load() != 0 {
				t.Fatal("running flag was not reset")
			}
			lock, err := backend.Do(context.Background(), "GET", producer.config.RunLockKey)
			if err != nil || !lock.IsNil() {
				t.Fatalf("lock not released: %v, %v", lock, err)
			}
			if opts.ReadTimeout != 0 || opts.WriteTimeout != 0 {
				t.Fatal("lease client mutated the shared production client")
			}
			// A fresh run must start without opening the stalled gate.
			producer.scanner = cardScanner{}
			if err := producer.Produce(context.Background(), model.ImeiQueueStartMessage{RunID: "next-run"}); err != nil {
				t.Fatalf("producer cannot restart: %v", err)
			}
		})
	}
}

func TestProduceFencesDelayedRealRedisWrites(t *testing.T) {
	for _, mode := range []string{"standalone", "sentinel"} {
		t.Run(mode, func(t *testing.T) { testProduceFencesDelayedRealRedisWrites(t, mode) })
	}
}

func testProduceFencesDelayedRealRedisWrites(t *testing.T, mode string) {
	for _, stage := range []string{"enqueue", "producer_done"} {
		t.Run(stage, func(t *testing.T) {
			backend := testutil.Redis(t)
			match := "local taskID=ARGV[i]"
			cards := []Card{{ID: 1, ICCID: "a", ChanKey: "ch", TempKey: "sim_yd_v5"}}
			if stage == "producer_done" {
				match = "'producerCompletedAt'"
				cards = nil
			}
			real, held, release := leaseRedisProxy(t, backend, match, false)
			defer release()
			workContexts := make(chan context.Context, 1)
			scanner := &lockTestScanner{cardScanner: cardScanner{cards: cards}}
			scanner.snapshot = func(ctx context.Context) (int64, error) { workContexts <- ctx; return int64(len(cards)), nil }
			store, err := queue.New(real, "imei:")
			if err != nil {
				t.Fatal(err)
			}
			config := Config{RunLockTTL: time.Second}
			if mode == "sentinel" {
				real, config.LeaseRedisConfig = sentinelLeaseRedis(t, real)
				store, err = queue.New(real, "imei:")
				if err != nil {
					t.Fatal(err)
				}
			}
			producer, err := New(scanner, store, config)
			if err != nil {
				t.Fatal(err)
			}
			result := startLockTestProducer(t, producer)
			select {
			case <-held:
			case <-time.After(2 * time.Second):
				t.Fatal("proxy did not hold old-owner EVAL")
			}
			workCtx := <-workContexts
			observer, err := queue.New(backend, "imei:")
			if err != nil {
				t.Fatal(err)
			}
			ok, err := observer.AcquireLockOrTakeover(context.Background(), producer.config.RunLockKey, "new-owner", "lock-test", time.Minute)
			if err != nil || !ok {
				t.Fatalf("takeover: %v, %v", ok, err)
			}
			select {
			case <-workCtx.Done():
			case <-time.After(2 * time.Second):
				t.Fatal("real renewal did not detect takeover")
			}
			if !errors.Is(context.Cause(workCtx), ErrAlreadyRunning) {
				t.Fatalf("wrong loss cause: %v", context.Cause(workCtx))
			}
			release()
			select {
			case err := <-result:
				if !errors.Is(err, ErrAlreadyRunning) {
					t.Fatalf("wrong result: %v", err)
				}
			case <-time.After(2 * time.Second):
				t.Fatal("producer did not finish")
			}
			meta, err := observer.GetRunMeta(context.Background(), "lock-test")
			if err != nil || meta.ActiveCount != 0 || meta.ProducerDone {
				t.Fatalf("stale owner changed run state: %+v, %v", meta, err)
			}
			channels, err := observer.RunChannels(context.Background(), "lock-test")
			if err != nil || len(channels) != 0 {
				t.Fatalf("stale owner indexed channel: %v, %v", channels, err)
			}
			ok, err = observer.RenewLock(context.Background(), producer.config.RunLockKey, "new-owner", time.Minute)
			if err != nil || !ok {
				t.Fatalf("old cleanup removed new lease: %v, %v", ok, err)
			}
			added, err := observer.EnqueueRunBatch(context.Background(), "lock-test", "ch", []queue.Task{{TaskID: "next", RunID: "lock-test", ChanKey: "ch", CardID: 2}}, time.Now())
			if err != nil || added != 1 {
				t.Fatalf("new work rejected: added=%d, %v", added, err)
			}
		})
	}
}

func leaseRedisProxy(t *testing.T, upstream *gredis.Redis, match string, holdReply bool) (*gredis.Redis, <-chan struct{}, func()) {
	t.Helper()
	raw := upstream.GetAdapter().Client().(*goredis.Client)
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	held, gate := make(chan struct{}), make(chan struct{})
	var heldOnce, releaseOnce sync.Once
	release := func() { releaseOnce.Do(func() { close(gate) }) }
	var connections sync.Map
	go func() {
		for {
			down, err := listener.Accept()
			if err != nil {
				return
			}
			connections.Store(down, struct{}{})
			go func() {
				defer down.Close()
				defer connections.Delete(down)
				up, err := net.Dial("tcp", raw.Options().Addr)
				if err != nil {
					return
				}
				defer up.Close()
				dr, ur := bufio.NewReader(down), bufio.NewReader(up)
				for {
					request, err := readLeaseRESP(dr)
					if err != nil {
						return
					}
					holding := strings.Contains(string(request), match)
					if holding && !holdReply {
						heldOnce.Do(func() { close(held) })
						<-gate
					}
					if _, err := up.Write(request); err != nil {
						return
					}
					response, err := readLeaseRESP(ur)
					if err != nil {
						return
					}
					if holding && holdReply {
						heldOnce.Do(func() { close(held) })
						<-gate
					}
					if _, err := down.Write(response); err != nil {
						return
					}
				}
			}()
		}
	}()
	t.Cleanup(func() {
		release()
		_ = listener.Close()
		connections.Range(func(key, _ any) bool { _ = key.(net.Conn).Close(); return true })
	})
	client, err := gredis.New(&gredis.Config{Address: listener.Addr().String(), DialTimeout: 100 * time.Millisecond, Protocol: 2})
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = client.Close(context.Background()) })
	if _, err := client.Do(context.Background(), "PING"); err != nil {
		t.Fatal(err)
	}
	opts := client.GetAdapter().Client().(*goredis.Client).Options()
	t.Logf("real driver options: readTimeout=%s writeTimeout=%s contextTimeoutEnabled=%v", opts.ReadTimeout, opts.WriteTimeout, opts.ContextTimeoutEnabled)
	return client, held, release
}

// A real disposable Sentinel discovers the proxy as its master; neither client
// configuration nor server state comes from the application's Redis settings.
func sentinelLeaseRedis(t *testing.T, master *gredis.Redis) (*gredis.Redis, *gredis.Config) {
	t.Helper()
	binary := os.Getenv("REDIS_SERVER_BINARY")
	if binary == "" {
		var err error
		binary, err = exec.LookPath("redis-server")
		if err != nil {
			t.Skip("Sentinel tests require redis-server")
		}
	}
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	address := listener.Addr().String()
	_, port, _ := net.SplitHostPort(address)
	_ = listener.Close()
	masterAddress := master.GetAdapter().Client().(*goredis.Client).Options().Addr
	host, masterPort, err := net.SplitHostPort(masterAddress)
	if err != nil {
		t.Fatal(err)
	}
	dir := t.TempDir()
	path := filepath.Join(dir, "sentinel.conf")
	contents := "bind 127.0.0.1\nport " + port + "\nsentinel monitor lease-test " + host + " " + masterPort + " 1\nsentinel down-after-milliseconds lease-test 60000\n"
	if err := os.WriteFile(path, []byte(contents), 0600); err != nil {
		t.Fatal(err)
	}
	command := exec.Command(binary, filepath.Base(path), "--sentinel")
	command.Dir = dir
	logPath := filepath.Join(dir, "sentinel.log")
	logFile, err := os.Create(logPath)
	if err != nil {
		t.Fatal(err)
	}
	defer logFile.Close()
	command.Stdout, command.Stderr = logFile, logFile
	if err := command.Start(); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = command.Process.Kill(); _ = command.Wait() })
	probe, err := gredis.New(&gredis.Config{Address: address, ReadTimeout: 100 * time.Millisecond, WriteTimeout: 100 * time.Millisecond, DialTimeout: 100 * time.Millisecond, Protocol: 2})
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = probe.Close(context.Background()) })
	for deadline := time.Now().Add(5 * time.Second); ; {
		if _, err = probe.Do(context.Background(), "SENTINEL", "get-master-addr-by-name", "lease-test"); err == nil {
			break
		}
		if time.Now().After(deadline) {
			log, _ := os.ReadFile(logPath)
			t.Fatalf("isolated Sentinel did not start: %v; %s", err, log)
		}
		time.Sleep(20 * time.Millisecond)
	}
	config := &gredis.Config{Address: address, MasterName: "lease-test", DialTimeout: 100 * time.Millisecond, Protocol: 2}
	client, err := gredis.New(config)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = client.Close(context.Background()) })
	return client, config
}

func TestLeaseClientBoundsSentinelDiscoveryIO(t *testing.T) {
	master := testutil.Redis(t)
	_, config := sentinelLeaseRedis(t, master)
	sentinel, err := gredis.New(&gredis.Config{Address: config.Address, ReadTimeout: time.Second, WriteTimeout: time.Second, Protocol: 2})
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = sentinel.Close(context.Background()) })
	proxy, held, release := leaseRedisProxy(t, sentinel, "get-master-addr-by-name", true)
	defer release()
	config.Address = proxy.GetAdapter().Client().(*goredis.Client).Options().Addr
	shared, err := gredis.New(config)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = shared.Close(context.Background()) })
	factory, err := newLeaseClientFactory(shared, config, 100*time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}
	client, closeClient, err := factory()
	if err != nil {
		t.Fatal(err)
	}
	defer closeClient()
	result := make(chan error, 1)
	go func() { _, err := client.Do(context.Background(), "PING"); result <- err }()
	select {
	case <-held:
	case <-time.After(2 * time.Second):
		t.Fatal("Sentinel discovery request was not held")
	}
	select {
	case err := <-result:
		if err == nil {
			t.Fatal("stalled Sentinel discovery unexpectedly succeeded")
		}
	case <-time.After(600 * time.Millisecond):
		release()
		<-result
		t.Fatal("Sentinel discovery has no socket I/O upper bound")
	}
}

func TestLeaseClientRejectsClusterAndUnknownAdapter(t *testing.T) {
	cluster, err := gredis.New(&gredis.Config{Address: "127.0.0.1:1", Cluster: true})
	if err != nil {
		t.Fatal(err)
	}
	defer cluster.Close(context.Background())
	store, err := queue.New(cluster, "imei:")
	if err != nil {
		t.Fatal(err)
	}
	if _, err := New(cardScanner{}, store, Config{}); err == nil || !strings.Contains(err.Error(), "Cluster") {
		t.Fatalf("Cluster must be rejected at initialization before any write: %v", err)
	}
	if _, err := newLeaseClientFactory(&leaseRedis{}, nil, time.Second); err == nil {
		t.Fatal("unknown adapter must not fall back to unbounded context-only I/O")
	}
}

func TestProduceBoundsReleaseIO(t *testing.T) {
	backend := testutil.Redis(t)
	real, held, release := leaseRedisProxy(t, backend, "redis.call('DEL', KEYS[1])", true)
	defer release()
	store, err := queue.New(real, "imei:")
	if err != nil {
		t.Fatal(err)
	}
	producer, err := New(cardScanner{}, store, Config{RunLockTTL: time.Second})
	if err != nil {
		t.Fatal(err)
	}
	result := startLockTestProducer(t, producer)
	select {
	case <-held:
	case <-time.After(2 * time.Second):
		t.Fatal("release response was not held")
	}
	select {
	case err := <-result:
		if err != nil {
			t.Fatal(err)
		}
	case <-time.After(600 * time.Millisecond):
		release()
		<-result
		t.Fatal("release socket I/O blocked cleanup")
	}
	if producer.running.Load() != 0 {
		t.Fatal("cleanup did not reset running")
	}
}

func TestLeaseClientBoundsSocketWriteIO(t *testing.T) {
	backend := testutil.Redis(t)
	upstream := backend.GetAdapter().Client().(*goredis.Client).Options().Addr
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	var stopReading atomic.Bool
	held, gate := make(chan struct{}), make(chan struct{})
	var releaseOnce sync.Once
	release := func() { releaseOnce.Do(func() { close(gate) }) }
	defer release()
	t.Cleanup(func() { _ = listener.Close() })
	go func() {
		down, err := listener.Accept()
		if err != nil {
			return
		}
		defer down.Close()
		up, err := net.Dial("tcp", upstream)
		if err != nil {
			return
		}
		defer up.Close()
		dr, ur := bufio.NewReader(down), bufio.NewReader(up)
		for {
			// Read only the array header once armed: the large EVAL payload
			// cannot enter the upstream or fit in TCP buffers, so Write blocks.
			header, err := dr.ReadBytes('\n')
			if err != nil {
				return
			}
			if stopReading.Load() {
				close(held)
				<-gate
				return
			}
			n, err := strconv.Atoi(strings.TrimSpace(string(header[1:])))
			if err != nil {
				return
			}
			for i := 0; i < n; i++ {
				part, err := readLeaseRESP(dr)
				if err != nil {
					return
				}
				header = append(header, part...)
			}
			if _, err := up.Write(header); err != nil {
				return
			}
			response, err := readLeaseRESP(ur)
			if err != nil {
				return
			}
			if _, err := down.Write(response); err != nil {
				return
			}
		}
	}()
	shared, err := gredis.New(&gredis.Config{Address: listener.Addr().String(), Protocol: 2})
	if err != nil {
		t.Fatal(err)
	}
	defer shared.Close(context.Background())
	factory, err := newLeaseClientFactory(shared, nil, 100*time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}
	client, closeClient, err := factory()
	if err != nil {
		t.Fatal(err)
	}
	defer closeClient()
	if _, err := client.Do(context.Background(), "PING"); err != nil {
		t.Fatal(err)
	}
	opts := client.(*leaseRedis).client.Options()
	if opts.ReadTimeout != 100*time.Millisecond || opts.WriteTimeout != 100*time.Millisecond || !opts.ContextTimeoutEnabled || opts.MaxRetries != 0 || opts.DialerRetries != 1 {
		t.Fatalf("lease client has unsafe timeout/retry options")
	}
	stopReading.Store(true)
	result := make(chan error, 1)
	go func() {
		_, err := client.Eval(context.Background(), "return 1 --"+strings.Repeat("x", 32<<20), 0, nil, nil)
		result <- err
	}()
	select {
	case <-held:
	case <-time.After(2 * time.Second):
		t.Fatal("write proxy did not stop reading")
	}
	select {
	case err := <-result:
		var timeout net.Error
		if !errors.As(err, &timeout) || !timeout.Timeout() || !strings.Contains(err.Error(), "write") {
			t.Fatalf("expected real socket write timeout, got %v", err)
		}
	case <-time.After(600 * time.Millisecond):
		release()
		<-result
		t.Fatal("socket write has no upper bound")
	}
}