CCIP - re-sync chain accessor on gRPC client refresh

Author: ogtownsend

pkg/loop/internal/net/broker.go

@@ -83,6 +83,16 @@ func (b *BrokerExt) NewClientConn(name string, newClient newClientFn) *clientCon
 	}
 }
 
+// NewClientConnWithCallback return a new *clientConn backed by this *BrokerExt with a refresh callback.
+func (b *BrokerExt) NewClientConnWithCallback(name string, newClient newClientFn, refreshCallback RefreshCallback) *clientConn {
+	return &clientConn{
+		BrokerExt:       b.WithName(name),
+		newClient:       newClient,
+		name:            name,
+		refreshCallback: refreshCallback,
+	}
+}
+
 func (b *BrokerExt) StopCtx() (context.Context, context.CancelFunc) {
 	return utils.ContextFromChan(b.StopCh)
 }

pkg/loop/internal/net/client.go

@@ -35,6 +35,10 @@ var _ grpc.ClientConnInterface = (*clientConn)(nil)
 // newClientFn returns a new client connection id to dial, and a set of Resource dependencies to close.
 type newClientFn func(context.Context) (id uint32, deps Resources, err error)
 
+// RefreshCallback is called after a successful refresh, but before the connection is made available.
+// This allows clients to perform state management operations before the refreshed client is used again.
+type RefreshCallback func(ctx context.Context) error
+
 // clientConn is a [grpc.ClientConnInterface] backed by a [*grpc.ClientConn] which can be recreated and swapped out
 // via the provided [newClientFn].
 // New instances should be created via BrokerExt.NewClientConn.
@@ -43,9 +47,10 @@ type clientConn struct {
 	newClient newClientFn
 	name      string
 
-	mu   sync.RWMutex
-	deps Resources
-	cc   *grpc.ClientConn
+	mu              sync.RWMutex
+	deps            Resources
+	cc              *grpc.ClientConn
+	refreshCallback RefreshCallback
 }
 
 func (c *clientConn) Invoke(ctx context.Context, method string, args interface{}, reply interface{}, opts ...grpc.CallOption) error {
@@ -93,6 +98,14 @@ func (c *clientConn) NewStream(ctx context.Context, desc *grpc.StreamDesc, metho
 	return nil, context.Cause(ctx)
 }
 
+// SetRefreshCallback registers a callback to be invoked after a successful refresh,
+// but before the connection is made available for use.
+func (c *clientConn) SetRefreshCallback(callback RefreshCallback) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.refreshCallback = callback
+}
+
 // refresh replaces c.cc with a new (different from orig) *grpc.ClientConn, and returns it as well.
 // It will block until a new connection is successfully dialed, or return nil if the context expires.
 func (c *clientConn) refresh(ctx context.Context, orig *grpc.ClientConn) *grpc.ClientConn {
@@ -150,5 +163,12 @@ func (c *clientConn) refresh(ctx context.Context, orig *grpc.ClientConn) *grpc.C
 		}
 	}
 
+	if c.refreshCallback != nil {
+		if err := c.refreshCallback(ctx); err != nil {
+			// Log error but make the refreshed connection available anyway.
+			c.Logger.Errorw("Refresh callback failed", "err", err)
+		}
+	}
+
 	return c.cc
 }

pkg/loop/internal/relayer/pluginprovider/ext/ccipocr3/ccip_provider.go

@@ -1,6 +1,8 @@
 package ccipocr3
 
 import (
+	"context"
+
 	"google.golang.org/grpc"
 
 	"github.com/smartcontractkit/libocr/offchainreporting2plus/ocr3types"
@@ -70,6 +72,13 @@ func (p *CCIPProviderClient) Codec() ccipocr3.Codec {
 	}
 }
 
+func (p *CCIPProviderClient) RefreshChainAccessor(ctx context.Context) error {
+	if refresher, ok := p.chainAccessor.(*chainAccessorClient); ok {
+		return refresher.Refresh(ctx)
+	}
+	return nil
+}
+
 // Server implementation
 type CCIPProviderServer struct{}
 

pkg/loop/internal/relayer/pluginprovider/ext/ccipocr3/chainaccessor.go

@@ -2,6 +2,7 @@ package ccipocr3
 
 import (
 	"context"
+	"sync"
 	"time"
 
 	"google.golang.org/grpc"
@@ -19,12 +20,17 @@ var _ ccipocr3.ChainAccessor = (*chainAccessorClient)(nil)
 type chainAccessorClient struct {
 	*net.BrokerExt
 	grpc ccipocr3pb.ChainAccessorClient
+
+	// Local persistence for refresh functionality
+	mu              sync.RWMutex
+	syncedContracts map[string]ccipocr3.UnknownAddress // contractName -> contractAddress
 }
 
 func NewChainAccessorClient(broker *net.BrokerExt, cc grpc.ClientConnInterface) ccipocr3.ChainAccessor {
 	return &chainAccessorClient{
-		BrokerExt: broker,
-		grpc:      ccipocr3pb.NewChainAccessorClient(cc),
+		BrokerExt:       broker,
+		grpc:            ccipocr3pb.NewChainAccessorClient(cc),
+		syncedContracts: make(map[string]ccipocr3.UnknownAddress),
 	}
 }
 
@@ -82,9 +88,52 @@ func (c *chainAccessorClient) Sync(ctx context.Context, contractName string, con
 		ContractName:    contractName,
 		ContractAddress: contractAddress,
 	})
+
+	if err == nil {
+		// Persist the synced contract locally for client refresh
+		c.mu.Lock()
+		c.syncedContracts[contractName] = contractAddress
+		c.mu.Unlock()
+		c.Logger.Debugw("Persisted synced contract", "contractName", contractName, "contractAddress", contractAddress)
+	}
+
 	return err
 }
 
+// Refresh re-syncs all previously synced contracts after a connection refresh.
+func (c *chainAccessorClient) Refresh(ctx context.Context) error {
+	c.mu.RLock()
+	contracts := make(map[string]ccipocr3.UnknownAddress)
+	for name, addr := range c.syncedContracts {
+		contracts[name] = addr
+	}
+	c.mu.RUnlock()
+
+	if len(contracts) == 0 {
+		c.Logger.Debug("No previously synced contracts to refresh")
+		return nil
+	}
+
+	c.Logger.Infow("Refreshing synced contracts", "count", len(contracts))
+
+	// Re-sync all previously synced contracts
+	for contractName, contractAddress := range contracts {
+		if err := c.Sync(ctx, contractName, contractAddress); err != nil {
+			c.Logger.Errorw("Failed to refresh contract sync",
+				"contractName", contractName,
+				"contractAddress", contractAddress,
+				"err", err)
+			// Continue refreshing other contracts even if one fails
+		} else {
+			c.Logger.Debugw("Successfully refreshed contract sync",
+				"contractName", contractName,
+				"contractAddress", contractAddress)
+		}
+	}
+
+	return nil
+}
+
 // DestinationAccessor methods
 func (c *chainAccessorClient) CommitReportsGTETimestamp(
 	ctx context.Context,

pkg/loop/internal/relayer/relayer.go

@@ -289,23 +289,38 @@ func (r *relayerClient) NewLLOProvider(ctx context.Context, rargs types.RelayArg
 }
 
 func (r *relayerClient) NewCCIPProvider(ctx context.Context, cargs types.CCIPProviderArgs) (types.CCIPProvider, error) {
-	cc := r.NewClientConn("CCIPProvider", func(ctx context.Context) (uint32, net.Resources, error) {
-		reply, err := r.relayer.NewCCIPProvider(ctx, &pb.NewCCIPProviderRequest{
-			CcipProviderArgs: &pb.CCIPProviderArgs{
-				ExternalJobID:        cargs.ExternalJobID[:],
-				ContractReaderConfig: cargs.ContractReaderConfig,
-				ChainWriterConfig:    cargs.ChainWriterConfig,
-				OffRampAddress:       cargs.OffRampAddress,
-				PluginType:           cargs.PluginType,
-			},
-		})
-		if err != nil {
-			return 0, nil, err
-		}
-		return reply.CcipProviderID, nil, nil
-	})
+	var provider types.CCIPProvider
+
+	// Create client connection with refresh callback
+	clientConn := r.NewClientConnWithCallback("CCIPProvider",
+		func(ctx context.Context) (uint32, net.Resources, error) {
+			reply, err := r.relayer.NewCCIPProvider(ctx, &pb.NewCCIPProviderRequest{
+				CcipProviderArgs: &pb.CCIPProviderArgs{
+					ExternalJobID:        cargs.ExternalJobID[:],
+					ContractReaderConfig: cargs.ContractReaderConfig,
+					ChainWriterConfig:    cargs.ChainWriterConfig,
+					OffRampAddress:       cargs.OffRampAddress,
+					PluginType:           cargs.PluginType,
+				},
+			})
+			if err != nil {
+				return 0, nil, err
+			}
+			return reply.CcipProviderID, nil, nil
+		},
+		func(ctx context.Context) error {
+			r.Logger.Info("CCIP Provider connection refreshed - refreshing chain accessor contracts")
+			if ccipProvider, ok := provider.(*ccipocr3.CCIPProviderClient); ok {
+				return ccipProvider.RefreshChainAccessor(ctx)
+			}
+			return nil
+		},
+	)
+
+	// Create the CCIP provider client
+	provider = ccipocr3.NewCCIPProviderClient(r.WithName(cargs.ExternalJobID.String()).WithName("CCIPProviderClient"), clientConn)
 
-	return ccipocr3.NewCCIPProviderClient(r.WithName(cargs.ExternalJobID.String()).WithName("CCIPProviderClient"), cc), nil
+	return provider, nil
 }
 
 func (r *relayerClient) LatestHead(ctx context.Context) (types.Head, error) {