pkg/capabilities: support replacing registered capabilities after shutdown

Author: jmank88

pkg/capabilities/capabilities.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	p2ptypes "github.com/smartcontractkit/libocr/ragep2p/types"
+	"google.golang.org/grpc/connectivity"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
 
@@ -209,6 +210,7 @@ type Validatable interface {
 // or extension in the future.
 type BaseCapability interface {
 	Info(ctx context.Context) (CapabilityInfo, error)
+	GetState() connectivity.State
 }
 
 type TriggerRegistrationRequest struct {
@@ -397,6 +399,11 @@ type CapabilityInfo struct {
 	SpendTypes []CapabilitySpendType
 }
 
+// GetState is included to implement BaseCapability.
+func (c CapabilityInfo) GetState() connectivity.State {
+	return connectivity.Idle
+}
+
 // Parse out the version from the ID.
 func (c CapabilityInfo) Version() string {
 	return c.ID[strings.Index(c.ID, "@")+1:]

pkg/capabilities/registry/base.go

@@ -6,8 +6,10 @@ import (
 	"fmt"
 	"strings"
 	"sync"
+	"sync/atomic"
 
 	"github.com/Masterminds/semver/v3"
+	"google.golang.org/grpc/connectivity"
 
 	"github.com/smartcontractkit/chainlink-common/pkg/capabilities"
 	"github.com/smartcontractkit/chainlink-common/pkg/logger"
@@ -18,8 +20,13 @@ var (
 	ErrCapabilityAlreadyExists = errors.New("capability already exists")
 )
 
+type atomicBaseCapability interface {
+	capabilities.BaseCapability
+	Update(capabilities.BaseCapability) error
+}
+
 type baseRegistry struct {
-	m    map[string]capabilities.BaseCapability
+	m    map[string]atomicBaseCapability
 	lggr logger.Logger
 	mu   sync.RWMutex
 }
@@ -28,7 +35,7 @@ var _ core.CapabilitiesRegistryBase = (*baseRegistry)(nil)
 
 func NewBaseRegistry(lggr logger.Logger) core.CapabilitiesRegistryBase {
 	return &baseRegistry{
-		m:    map[string]capabilities.BaseCapability{},
+		m:    map[string]atomicBaseCapability{},
 		lggr: logger.Named(lggr, "registries.basic"),
 	}
 }
@@ -142,46 +149,231 @@ func (r *baseRegistry) Add(ctx context.Context, c capabilities.BaseCapability) e
 		return err
 	}
 
-	switch info.CapabilityType {
-	case capabilities.CapabilityTypeTrigger:
-		_, ok := c.(capabilities.TriggerCapability)
-		if !ok {
-			return errors.New("trigger capability does not satisfy TriggerCapability interface")
+	id := info.ID
+	bc, ok := r.m[id]
+	if ok {
+		if bc.GetState() != connectivity.Shutdown {
+			return fmt.Errorf("%w: id %s found in registry", ErrCapabilityAlreadyExists, id)
 		}
-	case capabilities.CapabilityTypeAction, capabilities.CapabilityTypeConsensus, capabilities.CapabilityTypeTarget:
-		_, ok := c.(capabilities.ExecutableCapability)
-		if !ok {
-			return errors.New("action does not satisfy ExecutableCapability interface")
+		if err := bc.Update(c); err != nil {
+			return fmt.Errorf("failed to update capability %s: %w", id, err)
 		}
-	case capabilities.CapabilityTypeCombined:
-		_, ok := c.(capabilities.ExecutableAndTriggerCapability)
-		if !ok {
-			return errors.New("target capability does not satisfy ExecutableAndTriggerCapability interface")
+	} else {
+		var ac atomicBaseCapability
+		switch info.CapabilityType {
+		case capabilities.CapabilityTypeTrigger:
+			ac = &atomicTriggerCapability{}
+		case capabilities.CapabilityTypeAction, capabilities.CapabilityTypeConsensus, capabilities.CapabilityTypeTarget:
+			ac = &atomicExecuteCapability{}
+		case capabilities.CapabilityTypeCombined:
+			ac = &atomicExecuteAndTriggerCapability{}
+		default:
+			return fmt.Errorf("unknown capability type: %s", info.CapabilityType)
 		}
-	default:
-		return fmt.Errorf("unknown capability type: %s", info.CapabilityType)
-	}
-
-	id := info.ID
-	_, ok := r.m[id]
-	if ok {
-		return fmt.Errorf("%w: id %s found in registry", ErrCapabilityAlreadyExists, id)
+		if err := ac.Update(c); err != nil {
+			return err
+		}
+		r.m[id] = ac
 	}
-
-	r.m[id] = c
 	r.lggr.Infow("capability added", "id", id, "type", info.CapabilityType, "description", info.Description, "version", info.Version())
 	return nil
 }
 
 func (r *baseRegistry) Remove(_ context.Context, id string) error {
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	_, ok := r.m[id]
+	ac, ok := r.m[id]
 	if !ok {
 		return fmt.Errorf("unable to remove, capability not found: %s", id)
 	}
-
-	delete(r.m, id)
+	if err := ac.Update(nil); err != nil {
+		return fmt.Errorf("failed to remove capability %s: %w", id, err)
+	}
 	r.lggr.Infow("capability removed", "id", id)
 	return nil
 }
+
+var _ capabilities.TriggerCapability = &atomicTriggerCapability{}
+
+type atomicTriggerCapability struct {
+	atomic.Pointer[capabilities.TriggerCapability]
+}
+
+func (a *atomicTriggerCapability) Update(c capabilities.BaseCapability) error {
+	if c == nil {
+		a.Store(nil)
+		return nil
+	}
+	tc, ok := c.(capabilities.TriggerCapability)
+	if !ok {
+		return errors.New("trigger capability does not satisfy TriggerCapability interface")
+	}
+	a.Store(&tc)
+	return nil
+}
+
+func (a *atomicTriggerCapability) Info(ctx context.Context) (capabilities.CapabilityInfo, error) {
+	c := a.Load()
+	if c == nil {
+		return capabilities.CapabilityInfo{}, errors.New("capability unavailable")
+	}
+	return (*c).Info(ctx)
+}
+
+func (a *atomicTriggerCapability) GetState() connectivity.State {
+	c := a.Load()
+	if c == nil {
+		return connectivity.Shutdown
+	}
+	return (*c).GetState()
+}
+
+func (a *atomicTriggerCapability) RegisterTrigger(ctx context.Context, request capabilities.TriggerRegistrationRequest) (<-chan capabilities.TriggerResponse, error) {
+	c := a.Load()
+	if c == nil {
+		return nil, errors.New("capability unavailable")
+	}
+	return (*c).RegisterTrigger(ctx, request)
+}
+
+func (a *atomicTriggerCapability) UnregisterTrigger(ctx context.Context, request capabilities.TriggerRegistrationRequest) error {
+	c := a.Load()
+	if c == nil {
+		return errors.New("capability unavailable")
+	}
+	return (*c).UnregisterTrigger(ctx, request)
+}
+
+var _ capabilities.ExecutableCapability = &atomicExecuteCapability{}
+
+type atomicExecuteCapability struct {
+	atomic.Pointer[capabilities.ExecutableCapability]
+}
+
+func (a *atomicExecuteCapability) Update(c capabilities.BaseCapability) error {
+	if c == nil {
+		a.Store(nil)
+		return nil
+	}
+	tc, ok := c.(capabilities.ExecutableCapability)
+	if !ok {
+		return errors.New("action does not satisfy ExecutableCapability interface")
+	}
+	a.Store(&tc)
+	return nil
+}
+
+func (a *atomicExecuteCapability) Info(ctx context.Context) (capabilities.CapabilityInfo, error) {
+	c := a.Load()
+	if c == nil {
+		return capabilities.CapabilityInfo{}, errors.New("capability unavailable")
+	}
+	return (*c).Info(ctx)
+}
+
+func (a *atomicExecuteCapability) GetState() connectivity.State {
+	c := a.Load()
+	if c == nil {
+		return connectivity.Shutdown
+	}
+	return (*c).GetState()
+}
+
+func (a *atomicExecuteCapability) RegisterToWorkflow(ctx context.Context, request capabilities.RegisterToWorkflowRequest) error {
+	c := a.Load()
+	if c == nil {
+		return errors.New("capability unavailable")
+	}
+	return (*c).RegisterToWorkflow(ctx, request)
+}
+
+func (a *atomicExecuteCapability) UnregisterFromWorkflow(ctx context.Context, request capabilities.UnregisterFromWorkflowRequest) error {
+	c := a.Load()
+	if c == nil {
+		return errors.New("capability unavailable")
+	}
+	return (*c).UnregisterFromWorkflow(ctx, request)
+}
+
+func (a *atomicExecuteCapability) Execute(ctx context.Context, request capabilities.CapabilityRequest) (capabilities.CapabilityResponse, error) {
+	c := a.Load()
+	if c == nil {
+		return capabilities.CapabilityResponse{}, errors.New("capability unavailable")
+	}
+	return (*c).Execute(ctx, request)
+}
+
+var _ capabilities.ExecutableAndTriggerCapability = &atomicExecuteAndTriggerCapability{}
+
+type atomicExecuteAndTriggerCapability struct {
+	atomic.Pointer[capabilities.ExecutableAndTriggerCapability]
+}
+
+func (a *atomicExecuteAndTriggerCapability) Update(c capabilities.BaseCapability) error {
+	if c == nil {
+		a.Store(nil)
+		return nil
+	}
+	tc, ok := c.(capabilities.ExecutableAndTriggerCapability)
+	if !ok {
+		return errors.New("target capability does not satisfy ExecutableAndTriggerCapability interface")
+	}
+	a.Store(&tc)
+	return nil
+}
+
+func (a *atomicExecuteAndTriggerCapability) Info(ctx context.Context) (capabilities.CapabilityInfo, error) {
+	c := a.Load()
+	if c == nil {
+		return capabilities.CapabilityInfo{}, errors.New("capability unavailable")
+	}
+	return (*c).Info(ctx)
+}
+
+func (a *atomicExecuteAndTriggerCapability) GetState() connectivity.State {
+	c := a.Load()
+	if c == nil {
+		return connectivity.Shutdown
+	}
+	return (*c).GetState()
+}
+
+func (a *atomicExecuteAndTriggerCapability) RegisterTrigger(ctx context.Context, request capabilities.TriggerRegistrationRequest) (<-chan capabilities.TriggerResponse, error) {
+	c := a.Load()
+	if c == nil {
+		return nil, errors.New("capability unavailable")
+	}
+	return (*c).RegisterTrigger(ctx, request)
+}
+
+func (a *atomicExecuteAndTriggerCapability) UnregisterTrigger(ctx context.Context, request capabilities.TriggerRegistrationRequest) error {
+	c := a.Load()
+	if c == nil {
+		return errors.New("capability unavailable")
+	}
+	return (*c).UnregisterTrigger(ctx, request)
+}
+
+func (a *atomicExecuteAndTriggerCapability) RegisterToWorkflow(ctx context.Context, request capabilities.RegisterToWorkflowRequest) error {
+	c := a.Load()
+	if c == nil {
+		return errors.New("capability unavailable")
+	}
+	return (*c).RegisterToWorkflow(ctx, request)
+}
+
+func (a *atomicExecuteAndTriggerCapability) UnregisterFromWorkflow(ctx context.Context, request capabilities.UnregisterFromWorkflowRequest) error {
+	c := a.Load()
+	if c == nil {
+		return errors.New("capability unavailable")
+	}
+	return (*c).UnregisterFromWorkflow(ctx, request)
+}
+
+func (a *atomicExecuteAndTriggerCapability) Execute(ctx context.Context, request capabilities.CapabilityRequest) (capabilities.CapabilityResponse, error) {
+	c := a.Load()
+	if c == nil {
+		return capabilities.CapabilityResponse{}, errors.New("capability unavailable")
+	}
+	return (*c).Execute(ctx, request)
+}

pkg/loop/internal/core/services/capability/capabilities.go

@@ -8,6 +8,7 @@ import (
 	"sync"
 
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
 	"google.golang.org/protobuf/types/known/emptypb"
 
 	"github.com/smartcontractkit/chainlink-common/pkg/capabilities"
@@ -135,14 +136,19 @@ func InfoToReply(info capabilities.CapabilityInfo) *capabilitiespb.CapabilityInf
 }
 
 type baseCapabilityClient struct {
+	c    *grpc.ClientConn
 	grpc capabilitiespb.BaseCapabilityClient
 	*net.BrokerExt
 }
 
 var _ capabilities.BaseCapability = (*baseCapabilityClient)(nil)
 
 func newBaseCapabilityClient(brokerExt *net.BrokerExt, conn *grpc.ClientConn) *baseCapabilityClient {
-	return &baseCapabilityClient{grpc: capabilitiespb.NewBaseCapabilityClient(conn), BrokerExt: brokerExt}
+	return &baseCapabilityClient{c: conn, grpc: capabilitiespb.NewBaseCapabilityClient(conn), BrokerExt: brokerExt}
+}
+
+func (c *baseCapabilityClient) GetState() connectivity.State {
+	return c.c.GetState()
 }
 
 func (c *baseCapabilityClient) Info(ctx context.Context) (capabilities.CapabilityInfo, error) {

pkg/loop/internal/core/services/capability/capabilities_registry_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
 
 	p2ptypes "github.com/smartcontractkit/libocr/ragep2p/types"
 
@@ -29,6 +30,10 @@ type mockBaseCapability struct {
 	info capabilities.CapabilityInfo
 }
 
+func (f *mockBaseCapability) GetState() connectivity.State {
+	return connectivity.Idle
+}
+
 func (f *mockBaseCapability) Info(ctx context.Context) (capabilities.CapabilityInfo, error) {
 	return f.info, nil
 }

pkg/loop/internal/test/test.go

@@ -3,6 +3,8 @@ package test
 import (
 	"context"
 
+	"google.golang.org/grpc/connectivity"
+
 	"github.com/smartcontractkit/chainlink-common/pkg/capabilities"
 	"github.com/smartcontractkit/chainlink-common/pkg/services"
 
@@ -33,6 +35,10 @@ var _ capabilities.BaseCapability = (*baseCapability)(nil)
 type baseCapability struct {
 }
 
+func (e baseCapability) GetState() connectivity.State {
+	return connectivity.Idle
+}
+
 func (e baseCapability) Info(ctx context.Context) (capabilities.CapabilityInfo, error) {
 	return CapabilityInfo, nil
 }