fix: use redis lua to claim ranges safely across instances

Author: vietddude

internal/worker/catchup.go

@@ -156,15 +156,8 @@ func (cw *CatchupWorker) processCatchupBlocksParallel() error {
 		return nil
 	}
 
-	// Process multiple ranges in parallel
+	// Claim multiple ranges in parallel so multiple instances can split work.
 	var wg sync.WaitGroup
-	rangeChan := make(chan blockstore.CatchupRange, len(cw.blockRanges))
-
-	// Fill channel with ranges
-	for _, r := range cw.blockRanges {
-		rangeChan <- r
-	}
-	close(rangeChan)
 
 	// Start parallel workers
 	for i := 0; i < CATCHUP_WORKERS; i++ {
@@ -174,7 +167,11 @@ func (cw *CatchupWorker) processCatchupBlocksParallel() error {
 			defer cw.recoverPanic(fmt.Sprintf("catchup range worker %d", workerID))
 			cw.logger.Debug("Starting catchup worker", "worker_id", workerID)
 
-			for r := range rangeChan {
+			for {
+				r, ok := cw.claimNextRange()
+				if !ok {
+					return
+				}
 				if err := cw.processRange(r, workerID); err != nil {
 					cw.logger.Error("Failed to process range",
 						"worker_id", workerID,
@@ -194,6 +191,27 @@ func (cw *CatchupWorker) processCatchupBlocksParallel() error {
 	return nil
 }
 
+func (cw *CatchupWorker) claimNextRange() (blockstore.CatchupRange, bool) {
+	claimed, err := cw.catchupStore.GetNextRange(cw.ctx, cw.chain.GetNetworkInternalCode())
+	if err != nil {
+		cw.logger.Warn("Failed to claim catchup range",
+			"chain", cw.chain.GetName(),
+			"error", err,
+		)
+		return blockstore.CatchupRange{}, false
+	}
+	if claimed == nil {
+		return blockstore.CatchupRange{}, false
+	}
+
+	cw.logger.Info("Claimed catchup range",
+		"chain", cw.chain.GetName(),
+		"range", fmt.Sprintf("%d-%d", claimed.Start, claimed.End),
+		"current", claimed.Current,
+	)
+	return *claimed, true
+}
+
 func (cw *CatchupWorker) processRange(r blockstore.CatchupRange, workerID int) error {
 	batchCount := 0
 	startTime := time.Now()

internal/worker/catchup_test.go

@@ -2,15 +2,21 @@ package worker
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"log/slog"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/fystack/multichain-indexer/internal/status"
+	"github.com/fystack/multichain-indexer/internal/indexer"
 	"github.com/fystack/multichain-indexer/pkg/common/config"
 	"github.com/fystack/multichain-indexer/pkg/common/enum"
+	"github.com/fystack/multichain-indexer/pkg/infra"
 	"github.com/fystack/multichain-indexer/pkg/store/blockstore"
+	"github.com/fystack/multichain-indexer/pkg/store/catchupstore"
+	"github.com/redis/go-redis/v9"
 	"github.com/stretchr/testify/require"
 )
 
@@ -82,3 +88,152 @@ func TestCatchupWorkerPollsForNewRangesInsteadOfExiting(t *testing.T) {
 		t.Fatal("catchup worker did not exit after context cancellation")
 	}
 }
+
+func TestCatchupWorkersClaimDistinctRangesAcrossInstances(t *testing.T) {
+	t.Parallel()
+
+	client, cleanup := setupWorkerTestRedis(t)
+	defer cleanup()
+
+	statusRegistry := status.NewRegistry()
+	statusRegistry.RegisterChain("aptos", "aptos_testnet", config.ChainConfig{
+		NetworkId:    "aptos_testnet",
+		InternalCode: "APTOS_TESTNET",
+		Type:         enum.NetworkTypeApt,
+	})
+
+	store := catchupstore.New(&workerTestRedisClient{client: client})
+	ctx1, cancel1 := context.WithCancel(context.Background())
+	ctx2, cancel2 := context.WithCancel(context.Background())
+	defer cancel1()
+	defer cancel2()
+
+	var mu sync.Mutex
+	seen := make([]string, 0, 2)
+	indexerStub := &stubIndexer{
+		name:         "aptos",
+		internalCode: "APTOS_TESTNET",
+		networkType:  enum.NetworkTypeApt,
+		getBlocksFunc: func(_ context.Context, from, to uint64, _ bool) ([]indexer.BlockResult, error) {
+			mu.Lock()
+			seen = append(seen, fmt.Sprintf("%d-%d", from, to))
+			mu.Unlock()
+			return nil, nil
+		},
+	}
+
+	require.NoError(t, store.SaveRanges(context.Background(), "APTOS_TESTNET", []blockstore.CatchupRange{
+		{Start: 1, End: 10, Current: 0},
+		{Start: 11, End: 20, Current: 10},
+	}))
+
+	newWorker := func(ctx context.Context) *CatchupWorker {
+		return &CatchupWorker{
+			BaseWorker: &BaseWorker{
+				ctx:    ctx,
+				cancel: func() {},
+				logger: slog.New(slog.NewTextHandler(io.Discard, nil)),
+				config: config.ChainConfig{
+					PollInterval: time.Millisecond,
+					Throttle:     config.Throttle{BatchSize: 20},
+				},
+				chain:          indexerStub,
+				blockStore:     &stubBlockStore{},
+				catchupStore:   store,
+				statusRegistry: statusRegistry,
+			},
+			blockRanges: []blockstore.CatchupRange{},
+			workerPool:  make(chan struct{}, CATCHUP_WORKERS),
+		}
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		w := newWorker(ctx1)
+		w.runCatchup()
+	}()
+	go func() {
+		defer wg.Done()
+		w := newWorker(ctx2)
+		w.runCatchup()
+	}()
+
+	time.Sleep(300 * time.Millisecond)
+	cancel1()
+	cancel2()
+
+	wg.Wait()
+	mu.Lock()
+	defer mu.Unlock()
+	require.ElementsMatch(t, []string{"1-10", "11-20"}, seen)
+}
+
+type workerTestRedisClient struct {
+	client *redis.Client
+}
+
+var _ infra.RedisClient = (*workerTestRedisClient)(nil)
+
+func (r *workerTestRedisClient) GetClient() *redis.Client {
+	return r.client
+}
+
+func (r *workerTestRedisClient) Set(key string, value any, expiration time.Duration) error {
+	return r.client.Set(context.Background(), key, value, expiration).Err()
+}
+
+func (r *workerTestRedisClient) Get(key string) (string, error) {
+	return r.client.Get(context.Background(), key).Result()
+}
+
+func (r *workerTestRedisClient) Del(keys ...string) error {
+	return r.client.Del(context.Background(), keys...).Err()
+}
+
+func (r *workerTestRedisClient) ZAdd(key string, members ...redis.Z) error {
+	return r.client.ZAdd(context.Background(), key, members...).Err()
+}
+
+func (r *workerTestRedisClient) ZRem(key string, members ...interface{}) error {
+	return r.client.ZRem(context.Background(), key, members...).Err()
+}
+
+func (r *workerTestRedisClient) ZRange(key string, start, stop int64) ([]string, error) {
+	return r.client.ZRange(context.Background(), key, start, stop).Result()
+}
+
+func (r *workerTestRedisClient) ZRangeWithScores(key string, start, stop int64) ([]redis.Z, error) {
+	return r.client.ZRangeWithScores(context.Background(), key, start, stop).Result()
+}
+
+func (r *workerTestRedisClient) ZRevRangeWithScores(key string, start, stop int64) ([]redis.Z, error) {
+	return r.client.ZRevRangeWithScores(context.Background(), key, start, stop).Result()
+}
+
+func (r *workerTestRedisClient) Close() error {
+	return r.client.Close()
+}
+
+func setupWorkerTestRedis(t *testing.T) (*redis.Client, func()) {
+	t.Helper()
+
+	client := redis.NewClient(&redis.Options{
+		Addr: "localhost:6379",
+		DB:   13,
+	})
+
+	ctx := context.Background()
+	if _, err := client.Ping(ctx).Result(); err != nil {
+		t.Skip("Redis not available")
+	}
+	require.NoError(t, client.FlushDB(ctx).Err())
+
+	cleanup := func() {
+		_ = client.FlushDB(ctx).Err()
+		_ = client.Close()
+	}
+
+	return client, cleanup
+}

internal/worker/regular_test.go

@@ -401,6 +401,14 @@ func (s *stubCatchupStore) GetProgress(_ context.Context, _ string) ([]blockstor
 	return append([]blockstore.CatchupRange(nil), s.catchupProgress...), nil
 }
 
+func (s *stubCatchupStore) GetNextRange(_ context.Context, _ string) (*blockstore.CatchupRange, error) {
+	if len(s.catchupProgress) == 0 {
+		return nil, nil
+	}
+	rng := s.catchupProgress[0]
+	return &rng, nil
+}
+
 func (s *stubCatchupStore) DeleteRange(_ context.Context, _ string, start, end uint64) error {
 	if s.deleteCatchupErr != nil {
 		return s.deleteCatchupErr