tests: Support iterations in Heimdall simulator

polygon/heimdall/simulator/heimdall_simulator.go

@@ -3,6 +3,7 @@ package simulator
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -30,12 +31,13 @@ type HeimdallSimulator struct {
 	downloader         *sync.TorrentClient
 	chain              string
 
-	lastDownloadedBlockNumber uint64
+	iterations               []uint64 // list of final block numbers for an iteration
+	lastAvailableBlockNumber uint64
 }
 
 type IndexFnType func(context.Context, snaptype.FileInfo, uint32, string, *background.Progress, log.Lvl, log.Logger) error
 
-func NewHeimdall(ctx context.Context, chain string, snapshotLocation string, logger log.Logger) (HeimdallSimulator, error) {
+func NewHeimdall(ctx context.Context, chain string, snapshotLocation string, logger log.Logger, iterations []uint64) (HeimdallSimulator, error) {
 	cfg := snapcfg.KnownCfg(chain)
 	torrentDir := filepath.Join(snapshotLocation, "torrents", chain)
 
@@ -89,15 +91,25 @@ func NewHeimdall(ctx context.Context, chain string, snapshotLocation string, log
 		return HeimdallSimulator{}, err
 	}
 
+	var lastAvailableBlockNum uint64
+	if len(iterations) == 0 {
+		lastAvailableBlockNum = 0
+	} else {
+		lastAvailableBlockNum = iterations[0]
+		iterations = iterations[1:]
+	}
+
 	s := HeimdallSimulator{
-		ctx:                       ctx,
-		knownBorSnapshots:         knownBorSnapshots,
-		activeBorSnapshots:        activeBorSnapshots,
-		blockReader:               freezeblocks.NewBlockReader(nil, activeBorSnapshots),
-		logger:                    logger,
-		downloader:                downloader,
-		chain:                     chain,
-		lastDownloadedBlockNumber: 0,
+		ctx:                ctx,
+		knownBorSnapshots:  knownBorSnapshots,
+		activeBorSnapshots: activeBorSnapshots,
+		blockReader:        freezeblocks.NewBlockReader(nil, activeBorSnapshots),
+		logger:             logger,
+		downloader:         downloader,
+		chain:              chain,
+
+		iterations:               iterations,
+		lastAvailableBlockNumber: lastAvailableBlockNum,
 	}
 
 	go func() {
@@ -109,7 +121,7 @@ func NewHeimdall(ctx context.Context, chain string, snapshotLocation string, log
 }
 
 func (h *HeimdallSimulator) FetchLatestSpan(ctx context.Context) (*heimdall.Span, error) {
-	latestSpan := h.blockReader.LastFrozenSpanId()
+	latestSpan := uint64(heimdall.SpanIdAt(h.lastAvailableBlockNumber))
 
 	span, err := h.getSpan(h.ctx, latestSpan)
 	if err != nil {
@@ -120,6 +132,20 @@ func (h *HeimdallSimulator) FetchLatestSpan(ctx context.Context) (*heimdall.Span
 }
 
 func (h *HeimdallSimulator) FetchSpan(ctx context.Context, spanID uint64) (*heimdall.Span, error) {
+	// move to next iteration
+	if spanID == uint64(heimdall.SpanIdAt(h.lastAvailableBlockNumber)) {
+		if len(h.iterations) == 0 {
+			h.lastAvailableBlockNumber++
+		} else {
+			h.lastAvailableBlockNumber = h.iterations[0]
+			h.iterations = h.iterations[1:]
+		}
+	}
+
+	if spanID > uint64(heimdall.SpanIdAt(h.lastAvailableBlockNumber)) {
+		return nil, errors.New("span not found")
+	}
+
 	span, err := h.getSpan(h.ctx, spanID)
 	if err != nil {
 		return nil, err
@@ -138,12 +164,19 @@ func (h *HeimdallSimulator) FetchStateSyncEvents(ctx context.Context, fromId uin
 	defer view.Close()
 
 	for !maxTime && len(events) != limit {
-		if seg, ok := view.EventsSegment(h.lastDownloadedBlockNumber); ok {
+		if seg, ok := view.EventsSegment(h.lastAvailableBlockNumber); ok {
 			if err := h.downloadData(ctx, seg, snaptype.BorEvents, freezeblocks.BorEventsIdx); err != nil {
 				return nil, err
 			}
 		}
-		h.lastDownloadedBlockNumber += 500000
+
+		if len(h.iterations) == 0 {
+			// we increment by 500k because the events we need are not in this snapshot file
+			h.lastAvailableBlockNumber += 500000
+		} else {
+			h.lastAvailableBlockNumber = h.iterations[0]
+			h.iterations = h.iterations[1:]
+		}
 
 		events, maxTime, err = h.blockReader.EventsByIdFromSnapshot(fromId, to, limit)
 		if err != nil {

polygon/heimdall/simulator/simulator_test.go

@@ -52,7 +52,7 @@ func createFiles(dataDir, chain string) error {
 	return nil
 }
 
-func setup(t *testing.T, ctx context.Context) simulator.HeimdallSimulator {
+func setup(t *testing.T, ctx context.Context, iterations []uint64) simulator.HeimdallSimulator {
 	chain := "mumbai"
 	logger := log.New()
 	logger.SetHandler(log.StdoutHandler)
@@ -63,7 +63,7 @@ func setup(t *testing.T, ctx context.Context) simulator.HeimdallSimulator {
 		t.Fatal(err)
 	}
 
-	sim, err := simulator.NewHeimdall(ctx, chain, dataDir, logger)
+	sim, err := simulator.NewHeimdall(ctx, chain, dataDir, logger, iterations)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -75,7 +75,7 @@ func TestSimulatorEvents(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	sim := setup(t, ctx)
+	sim := setup(t, ctx, []uint64{1_000_000})
 
 	res, err := sim.FetchStateSyncEvents(ctx, 0, time.Now(), 100)
 	assert.NoError(t, err)
@@ -103,20 +103,40 @@ func TestSimulatorSpans(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	sim := setup(t, ctx)
+	sim := setup(t, ctx, []uint64{100_000, 205_055})
 
+	// should have the final span from first iteration
 	span, err := sim.FetchLatestSpan(ctx)
 	assert.NoError(t, err)
-	assert.Equal(t, heimdall.SpanId(78), span.Id)
-	assert.Equal(t, uint64(493056), span.StartBlock)
-	assert.Equal(t, uint64(499455), span.EndBlock)
+	assert.Equal(t, heimdall.SpanIdAt(100_000), span.Id)
+	assert.Equal(t, uint64(96_256), span.StartBlock)
+	assert.Equal(t, uint64(102_655), span.EndBlock)
 
-	span2, err := sim.FetchSpan(ctx, 78)
+	// get last span to move to next iteration
+	span2, err := sim.FetchSpan(ctx, uint64(heimdall.SpanIdAt(100_000)))
 	assert.NoError(t, err)
 	assert.Equal(t, span, span2)
 
-	span3, err := sim.FetchSpan(ctx, 77)
+	// check if we are in the next iteration
+	span3, err := sim.FetchLatestSpan(ctx)
 	assert.NoError(t, err)
-	assert.Equal(t, heimdall.SpanId(77), span3.Id)
-	assert.Equal(t, span.StartBlock-1, span3.EndBlock)
+	assert.Equal(t, heimdall.SpanIdAt(205_055), span3.Id)
+	assert.Equal(t, uint64(198_656), span3.StartBlock)
+	assert.Equal(t, uint64(205_055), span3.EndBlock)
+
+	// higher spans should not be available
+	_, err = sim.FetchSpan(ctx, uint64(heimdall.SpanIdAt(205_055)+1))
+	assert.Error(t, err, "span not found")
+
+	// move to next iteration (should be +1 block since we have no more iterations defined)
+	span4, err := sim.FetchSpan(ctx, uint64(heimdall.SpanIdAt(205_055)))
+	assert.NoError(t, err)
+	assert.Equal(t, span4, span3)
+
+	// check latest span (should be the same since we are
+	span5, err := sim.FetchLatestSpan(ctx)
+	assert.NoError(t, err)
+	assert.Equal(t, heimdall.SpanIdAt(205_056), span5.Id)
+	assert.Equal(t, uint64(205_056), span5.StartBlock)
+	assert.Equal(t, uint64(211_455), span5.EndBlock)
 }