Add context to functions in setup.go

Signed-off-by: Takeshi Arabiki <takeshi.arabiki@datachain.jp>

Author: abicky

module/prover.go

@@ -118,8 +118,8 @@ func (pr *Prover) SetupHeadersForUpdate(ctx context.Context, counterparty core.F
 }
 
 func (pr *Prover) SetupHeadersForUpdateByLatestHeight(ctx context.Context, clientStateLatestHeight exported.Height, latestFinalizedHeader *Header) ([]core.Header, error) {
-	queryVerifiableNeighboringEpochHeader := func(height uint64, limitHeight uint64) (core.Header, error) {
-		ethHeaders, err := queryFinalizedHeader(context.TODO(), pr.chain.Header, height, limitHeight)
+	queryVerifiableNeighboringEpochHeader := func(ctx context.Context, height uint64, limitHeight uint64) (core.Header, error) {
+		ethHeaders, err := queryFinalizedHeader(ctx, pr.chain.Header, height, limitHeight)
 		if err != nil {
 			return nil, err
 		}
@@ -134,6 +134,7 @@ func (pr *Prover) SetupHeadersForUpdateByLatestHeight(ctx context.Context, clien
 		return nil, err
 	}
 	return setupHeadersForUpdate(
+		context.TODO(),
 		queryVerifiableNeighboringEpochHeader,
 		pr.chain.Header,
 		clientStateLatestHeight,

module/setup.go

@@ -11,9 +11,10 @@ import (
 	"github.com/hyperledger-labs/yui-relayer/log"
 )
 
-type queryVerifiableNeighboringEpochHeaderFn = func(uint64, uint64) (core.Header, error)
+type queryVerifiableNeighboringEpochHeaderFn = func(context.Context, uint64, uint64) (core.Header, error)
 
 func setupHeadersForUpdate(
+	ctx context.Context,
 	queryVerifiableNeighboringEpochHeader queryVerifiableNeighboringEpochHeaderFn,
 	getHeader getHeaderFn,
 	clientStateLatestHeight exported.Height,
@@ -39,7 +40,7 @@ func setupHeadersForUpdate(
 
 	// Append insufficient epoch blocks
 	for epochHeight := firstUnsavedEpoch; epochHeight < latestFinalizedHeight; epochHeight += constant.BlocksPerEpoch {
-		verifiableEpoch, err := setupNeighboringEpochHeader(getHeader, queryVerifiableNeighboringEpochHeader, epochHeight, trustedEpochHeight, latestHeight)
+		verifiableEpoch, err := setupNeighboringEpochHeader(ctx, getHeader, queryVerifiableNeighboringEpochHeader, epochHeight, trustedEpochHeight, latestHeight)
 		if err != nil {
 			return nil, err
 		}
@@ -55,21 +56,22 @@ func setupHeadersForUpdate(
 }
 
 func setupNeighboringEpochHeader(
+	ctx context.Context,
 	getHeader getHeaderFn,
 	queryVerifiableHeader queryVerifiableNeighboringEpochHeaderFn,
 	epochHeight uint64,
 	trustedEpochHeight uint64,
 	latestHeight exported.Height,
 ) (core.Header, error) {
 	// neighboring epoch needs block before checkpoint
-	currentValidatorSet, currentTurnLength, err := queryValidatorSetAndTurnLength(context.TODO(), getHeader, epochHeight)
+	currentValidatorSet, currentTurnLength, err := queryValidatorSetAndTurnLength(ctx, getHeader, epochHeight)
 	if err != nil {
 		return nil, fmt.Errorf("setupNeighboringEpochHeader: failed to get current validator set: epochHeight=%d : %+v", epochHeight, err)
 	}
 	// ex) trusted(prevSaved = 200), epochHeight = 400 must be finalized by min(610,latest)
 	nextCheckpoint := currentValidatorSet.Checkpoint(currentTurnLength) + (epochHeight + constant.BlocksPerEpoch)
 	limit := minUint64(nextCheckpoint-1, latestHeight.GetRevisionHeight())
-	return queryVerifiableHeader(epochHeight, limit)
+	return queryVerifiableHeader(ctx, epochHeight, limit)
 }
 
 func withTrustedHeight(targetHeaders []core.Header, clientStateLatestHeight exported.Height) []core.Header {

module/setup_test.go

@@ -2,14 +2,15 @@ package module
 
 import (
 	"context"
+	"math/big"
+	"testing"
+
 	clienttypes "github.com/cosmos/ibc-go/v8/modules/core/02-client/types"
 	"github.com/datachainlab/ibc-parlia-relay/module/constant"
 	types2 "github.com/ethereum/go-ethereum/core/types"
 	"github.com/hyperledger-labs/yui-relayer/core"
 	"github.com/hyperledger-labs/yui-relayer/log"
 	"github.com/stretchr/testify/suite"
-	"math/big"
-	"testing"
 )
 
 type SetupTestSuite struct {
@@ -38,7 +39,7 @@ func (ts *SetupTestSuite) TestSuccess_setupHeadersForUpdate_neighboringEpoch() {
 			CurrentValidators:  [][]byte{{1}},
 			PreviousValidators: [][]byte{{1}},
 		}
-		neighborFn := func(height uint64, _ uint64) (core.Header, error) {
+		neighborFn := func(_ context.Context, height uint64, _ uint64) (core.Header, error) {
 			h, e := newETHHeader(&types2.Header{
 				Number: big.NewInt(int64(height)),
 			})
@@ -53,7 +54,7 @@ func (ts *SetupTestSuite) TestSuccess_setupHeadersForUpdate_neighboringEpoch() {
 			}, nil
 		}
 
-		targets, err := setupHeadersForUpdate(neighborFn, headerFn, clientStateLatestHeight, latestFinalizedHeader, clienttypes.NewHeight(0, 100000))
+		targets, err := setupHeadersForUpdate(context.Background(), neighborFn, headerFn, clientStateLatestHeight, latestFinalizedHeader, clienttypes.NewHeight(0, 100000))
 		ts.Require().NoError(err)
 		ts.Require().Len(targets, expected)
 		for i, h := range targets {
@@ -101,7 +102,7 @@ func (ts *SetupTestSuite) TestSuccess_setupHeadersForUpdate_allEmpty() {
 		latestFinalizedHeader := &Header{
 			Headers: []*ETHHeader{target},
 		}
-		neighboringEpochFn := func(height uint64, _ uint64) (core.Header, error) {
+		neighboringEpochFn := func(_ context.Context, height uint64, _ uint64) (core.Header, error) {
 			// insufficient vote attestation
 			return nil, nil
 		}
@@ -111,7 +112,7 @@ func (ts *SetupTestSuite) TestSuccess_setupHeadersForUpdate_allEmpty() {
 				Extra:  epochHeader().Extra,
 			}, nil
 		}
-		targets, err := setupHeadersForUpdate(neighboringEpochFn, headerFn, clientStateLatestHeight, latestFinalizedHeader, clienttypes.NewHeight(0, 1000000))
+		targets, err := setupHeadersForUpdate(context.Background(), neighboringEpochFn, headerFn, clientStateLatestHeight, latestFinalizedHeader, clienttypes.NewHeight(0, 1000000))
 		ts.Require().NoError(err)
 		ts.Require().Len(targets, expected)
 	}
@@ -144,7 +145,7 @@ func (ts *SetupTestSuite) TestSuccess_setupNeighboringEpochHeader() {
 	epochHeight := uint64(400)
 	trustedEpochHeight := uint64(200)
 
-	neighboringEpochFn := func(height uint64, limit uint64) (core.Header, error) {
+	neighboringEpochFn := func(_ context.Context, height uint64, limit uint64) (core.Header, error) {
 		target, err := newETHHeader(&types2.Header{
 			Number: big.NewInt(int64(limit)),
 		})
@@ -156,7 +157,7 @@ func (ts *SetupTestSuite) TestSuccess_setupNeighboringEpochHeader() {
 	headerFn := func(_ context.Context, height uint64) (*types2.Header, error) {
 		return headerByHeight(int64(height)), nil
 	}
-	hs, err := setupNeighboringEpochHeader(headerFn, neighboringEpochFn, epochHeight, trustedEpochHeight, clienttypes.NewHeight(0, 10000))
+	hs, err := setupNeighboringEpochHeader(context.Background(), headerFn, neighboringEpochFn, epochHeight, trustedEpochHeight, clienttypes.NewHeight(0, 10000))
 	ts.Require().NoError(err)
 	target, err := hs.(*Header).Target()
 	ts.Require().NoError(err)