[LM] Fix context passing across plugin and bridge interfaces (#1081)

- We were using context.Background() in most places in the bridge interfaces - This PR passes the context down from the plugin where possible
smartcontractkit · Jul 2, 2024 · 1b1964b · 1b1964b
1 parent a6245df
commit 1b1964b
Show file tree

Hide file tree

Showing 10 changed files with 59 additions and 23 deletions.
diff --git a/core/services/ocr2/plugins/liquiditymanager/bridge/arb/l1_to_l2.go b/core/services/ocr2/plugins/liquiditymanager/bridge/arb/l1_to_l2.go
@@ -56,6 +56,7 @@ type l1ToL2Bridge struct {
 }
 
 func NewL1ToL2Bridge(
+	ctx context.Context,
 	lggr logger.Logger,
 	localSelector,
 	remoteSelector models.NetworkSelector,
@@ -95,8 +96,6 @@ func NewL1ToL2Bridge(
 		remoteChain.Name,
 		"",
 	)
-	// FIXME Makram please pass the valid context
-	ctx := context.Background()
 	err = l1LogPoller.RegisterFilter(ctx, logpoller.Filter{
 		Addresses: []common.Address{l1LiquidityManagerAddress},
 		Name:      l1FilterName,

diff --git a/core/services/ocr2/plugins/liquiditymanager/bridge/arb/l2_to_l1.go b/core/services/ocr2/plugins/liquiditymanager/bridge/arb/l2_to_l1.go
@@ -54,6 +54,7 @@ type l2ToL1Bridge struct {
 }
 
 func NewL2ToL1Bridge(
+	ctx context.Context,
 	lggr logger.Logger,
 	localSelector,
 	remoteSelector models.NetworkSelector,
@@ -81,8 +82,6 @@ func NewL2ToL1Bridge(
 		remoteChain.Name,
 		"",
 	)
-	// FIXME Makram fix the context plax
-	ctx := context.Background()
 	err := l2LogPoller.RegisterFilter(
 		ctx,
 		logpoller.Filter{

diff --git a/core/services/ocr2/plugins/liquiditymanager/bridge/bridge.go b/core/services/ocr2/plugins/liquiditymanager/bridge/bridge.go
@@ -60,7 +60,8 @@ type Bridge interface {
 
 //go:generate mockery --name Factory --output ./mocks --filename bridge_factory_mock.go --case=underscore
 type Factory interface {
-	NewBridge(source, dest models.NetworkSelector) (Bridge, error)
+	NewBridge(ctx context.Context, source, dest models.NetworkSelector) (Bridge, error)
+	GetBridge(source, dest models.NetworkSelector) (Bridge, error)
 }
 
 type Opt func(c *factory)
@@ -106,20 +107,20 @@ func WithEvmDep(
 	}
 }
 
-func (f *factory) NewBridge(source, dest models.NetworkSelector) (Bridge, error) {
+func (f *factory) NewBridge(ctx context.Context, source, dest models.NetworkSelector) (Bridge, error) {
 	if source == dest {
 		return nil, fmt.Errorf("no bridge between the same network and itself: %d", source)
 	}
 
 	bridge, err := f.GetBridge(source, dest)
 	if errors.Is(err, ErrBridgeNotFound) {
 		f.lggr.Infow("Bridge not found, initializing new bridge", "source", source, "dest", dest)
-		return f.initBridge(source, dest)
+		return f.initBridge(ctx, source, dest)
 	}
 	return bridge, err
 }
 
-func (f *factory) initBridge(source, dest models.NetworkSelector) (Bridge, error) {
+func (f *factory) initBridge(ctx context.Context, source, dest models.NetworkSelector) (Bridge, error) {
 	f.lggr.Debugw("Initializing bridge", "source", source, "dest", dest)
 
 	var bridge Bridge
@@ -155,6 +156,7 @@ func (f *factory) initBridge(source, dest models.NetworkSelector) (Bridge, error
 			"l2BridgeAdapter", l2BridgeAdapter,
 		)
 		bridge, err = arb.NewL2ToL1Bridge(
+			ctx,
 			f.lggr,
 			source,
 			dest,
@@ -198,6 +200,7 @@ func (f *factory) initBridge(source, dest models.NetworkSelector) (Bridge, error
 			"l2BridgeAdapter", l2BridgeAdapter,
 		)
 		bridge, err = opstack.NewL2ToL1Bridge(
+			ctx,
 			f.lggr,
 			source,
 			dest,
@@ -245,6 +248,7 @@ func (f *factory) initBridge(source, dest models.NetworkSelector) (Bridge, error
 				"l1BridgeAdapter", l1BridgeAdapter,
 			)
 			bridge, err = arb.NewL1ToL2Bridge(
+				ctx,
 				f.lggr,
 				source,
 				dest,
@@ -267,6 +271,7 @@ func (f *factory) initBridge(source, dest models.NetworkSelector) (Bridge, error
 				"l1BridgeAdapter", l1BridgeAdapter,
 			)
 			bridge, err = opstack.NewL1ToL2Bridge(
+				ctx,
 				f.lggr,
 				source,
 				dest,
@@ -310,6 +315,7 @@ func (f *factory) initBridge(source, dest models.NetworkSelector) (Bridge, error
 			return nil, fmt.Errorf("bridge adapter not found for dest selector %d in deps for selector %d", dest, source)
 		}
 		bridge, err = testonlybridge.New(
+			ctx,
 			source,
 			dest,
 			sourceDeps.liquidityManagerAddress,

diff --git a/core/services/ocr2/plugins/liquiditymanager/bridge/mocks/bridge_factory_mock.go b/core/services/ocr2/plugins/liquiditymanager/bridge/mocks/bridge_factory_mock.go
diff --git a/core/services/ocr2/plugins/liquiditymanager/bridge/opstack/l1_to_l2.go b/core/services/ocr2/plugins/liquiditymanager/bridge/opstack/l1_to_l2.go
@@ -47,6 +47,7 @@ type l1ToL2Bridge struct {
 }
 
 func NewL1ToL2Bridge(
+	ctx context.Context,
 	lggr logger.Logger,
 	localSelector,
 	remoteSelector models.NetworkSelector,
@@ -77,8 +78,6 @@ func NewL1ToL2Bridge(
 		"",
 	)
 
-	// TODO: FIXME pass valid context
-	ctx := context.Background()
 	err := l1LogPoller.RegisterFilter(ctx, logpoller.Filter{
 		Addresses: []common.Address{l1LiquidityManagerAddress}, // emits LiquidityTransferred
 		Name:      l1FilterName,

diff --git a/core/services/ocr2/plugins/liquiditymanager/bridge/opstack/l2_to_l1.go b/core/services/ocr2/plugins/liquiditymanager/bridge/opstack/l2_to_l1.go
@@ -41,6 +41,7 @@ type l2ToL1Bridge struct {
 }
 
 func NewL2ToL1Bridge(
+	ctx context.Context,
 	lggr logger.Logger,
 	localSelector,
 	remoteSelector models.NetworkSelector,
@@ -68,8 +69,6 @@ func NewL2ToL1Bridge(
 		remoteChain.Name,
 		"",
 	)
-	// TODO (ogtownsend): pass context from above
-	ctx := context.Background()
 	err := l2LogPoller.RegisterFilter(
 		ctx,
 		logpoller.Filter{

diff --git a/core/services/ocr2/plugins/liquiditymanager/bridge/testonlybridge/l1_to_l1.go b/core/services/ocr2/plugins/liquiditymanager/bridge/testonlybridge/l1_to_l1.go
@@ -51,13 +51,13 @@ type testBridge struct {
 }
 
 func New(
+	ctx context.Context,
 	sourceSelector, destSelector models.NetworkSelector,
 	sourceLiquidityManagerAddress, destLiquidityManagerAddress, sourceAdapter, destAdapter models.Address,
 	sourceLogPoller, destLogPoller logpoller.LogPoller,
 	sourceClient, destClient client.Client,
 	lggr logger.Logger,
 ) (*testBridge, error) {
-	ctx := context.Background()
 	err := sourceLogPoller.RegisterFilter(
 		ctx,
 		logpoller.Filter{

diff --git a/core/services/ocr2/plugins/liquiditymanager/bridge/testonlybridge/l1_to_l1_test.go b/core/services/ocr2/plugins/liquiditymanager/bridge/testonlybridge/l1_to_l1_test.go
@@ -728,7 +728,7 @@ func TestNew(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			tt.expect(t, tt.args)
 			defer tt.assert(t, tt.args)
-			got, err := New(tt.args.sourceSelector, tt.args.destSelector, tt.args.sourceLiquidityManagerAddress, tt.args.destLiquidityManagerAddress, tt.args.sourceAdapter, tt.args.destAdapter, tt.args.sourceLogPoller, tt.args.destLogPoller, tt.args.sourceClient, tt.args.destClient, tt.args.lggr)
+			got, err := New(testutils.Context(t), tt.args.sourceSelector, tt.args.destSelector, tt.args.sourceLiquidityManagerAddress, tt.args.destLiquidityManagerAddress, tt.args.sourceAdapter, tt.args.destAdapter, tt.args.sourceLogPoller, tt.args.destLogPoller, tt.args.sourceClient, tt.args.destClient, tt.args.lggr)
 			if tt.wantErr {
 				require.Error(t, err)
 			} else {

diff --git a/core/services/ocr2/plugins/liquiditymanager/plugin.go b/core/services/ocr2/plugins/liquiditymanager/plugin.go
@@ -507,7 +507,7 @@ func (p *Plugin) loadPendingTransfers(ctx context.Context, lggr logger.Logger) (
 	}
 	for _, edge := range edges {
 		logger := lggr.With("sourceNetwork", edge.Source, "sourceChainID", edge.Source.ChainID(), "destNetwork", edge.Dest, "destChainID", edge.Dest.ChainID())
-		bridge, err := p.bridgeFactory.NewBridge(edge.Source, edge.Dest)
+		bridge, err := p.bridgeFactory.NewBridge(ctx, edge.Source, edge.Dest)
 		if err != nil {
 			return nil, fmt.Errorf("init bridge: %w", err)
 		}
@@ -607,7 +607,7 @@ func (p *Plugin) computeResolvedTransfersQuorum(observations []models.Observatio
 			}
 			medianizedNativeFee := rebalcalc.BigIntSortedMiddle(bridgeFees)
 			medianizedDateUnix := rebalcalc.BigIntSortedMiddle(datesUnix)
-			bridge, err := p.bridgeFactory.NewBridge(k.From, k.To)
+			bridge, err := p.bridgeFactory.GetBridge(k.From, k.To)
 			if err != nil {
 				return nil, fmt.Errorf("init bridge: %w", err)
 			}
@@ -650,7 +650,7 @@ func (p *Plugin) resolveProposedTransfers(ctx context.Context, lggr logger.Logge
 
 	resolvedTransfers := make([]models.Transfer, 0, len(outcome.ProposedTransfers))
 	for _, proposedTransfer := range outcome.ProposedTransfers {
-		bridge, err := p.bridgeFactory.NewBridge(proposedTransfer.From, proposedTransfer.To)
+		bridge, err := p.bridgeFactory.NewBridge(ctx, proposedTransfer.From, proposedTransfer.To)
 		if err != nil {
 			return nil, fmt.Errorf("init bridge: %w", err)
 		}

diff --git a/core/services/ocr2/plugins/liquiditymanager/plugin_test.go b/core/services/ocr2/plugins/liquiditymanager/plugin_test.go
@@ -289,7 +289,7 @@ func TestPlugin_Observation(t *testing.T) {
 			for sourceDest, bridgeFn := range tc.bridges {
 				br, err2 := bridgeFn(t)
 				p.bridgeFactory.
-					On("NewBridge", sourceDest[0], sourceDest[1]).
+					On("NewBridge", ctx, sourceDest[0], sourceDest[1]).
 					Return(br, err2)
 			}
 
@@ -660,7 +660,7 @@ func TestPlugin_Outcome(t *testing.T) {
 			for sourceDest, bridgeFn := range tc.bridges {
 				br, err := bridgeFn(t)
 				p.bridgeFactory.
-					On("NewBridge", sourceDest[0], sourceDest[1]).
+					On("GetBridge", sourceDest[0], sourceDest[1]).
 					Return(br, err)
 			}
 
@@ -1135,7 +1135,8 @@ func TestPlugin_E2EWithMocks(t *testing.T) {
 					for _, edge := range edges {
 						br, ok := n.bridges[[2]models.NetworkSelector{edge.Source, edge.Dest}]
 						require.True(t, ok, "the test case is wrong, bridge is not defined %d->%d", edge.Source, edge.Dest)
-						n.bridgeFactory.On("NewBridge", edge.Source, edge.Dest).Return(br, nil).Maybe()
+						n.bridgeFactory.On("NewBridge", mock.Anything /* cancelContext */, edge.Source, edge.Dest).Return(br, nil).Maybe()
+						n.bridgeFactory.On("GetBridge", edge.Source, edge.Dest).Return(br, nil).Maybe()
 
 						pendingTransfers := make([]models.PendingTransfer, 0)
 						for _, tr := range round.pendingTransfersPerNode[i] {