Add timeout to orgResolver (#2297)

* Add timeout to orgResolver

* update config name

* cleanup

Author: silaslenihan

pkg/services/orgresolver/linking.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"time"
 
 	"go.opentelemetry.io/otel/metric"
 	"google.golang.org/grpc"
@@ -17,6 +18,9 @@ import (
 	linkingclient "github.com/smartcontractkit/chainlink-protos/linking-service/go/v1"
 )
 
+// defaultRequestTimeout bounds Get(); callers sit in capability hot paths, so a hung connection must not propagate.
+const defaultRequestTimeout = 2 * time.Second
+
 // JWTGenerator interface for JWT token creation
 type JWTGenerator interface {
 	CreateJWTForRequest(req any) (string, error)
@@ -35,20 +39,23 @@ type Config struct {
 	WorkflowRegistryChainSelector uint64
 	JWTGenerator                  JWTGenerator
 
+	RequestTimeout time.Duration // bounds each Get() call; zero means defaultRequestTimeout
+
 	Client linkingclient.LinkingServiceClient // optional
 	Meter  metric.Meter                       // optional
 }
 
 // orgResolver makes direct calls to the linking service to resolve organization IDs from workflow owners.
-// This simplified implementation makes a network call for each Get() request.
+// This simplified implementation makes a network call for each Get() request, bounded by requestTimeout.
 type orgResolver struct {
 	workflowRegistryAddress       string
 	workflowRegistryChainSelector uint64
 
-	client       linkingclient.LinkingServiceClient
-	conn         *grpc.ClientConn // nil if client was injected
-	logger       log.SugaredLogger
-	jwtGenerator JWTGenerator
+	client         linkingclient.LinkingServiceClient
+	conn           *grpc.ClientConn // nil if client was injected
+	logger         log.SugaredLogger
+	jwtGenerator   JWTGenerator
+	requestTimeout time.Duration
 
 	passCount metric.Int64Counter
 	failCount metric.Int64Counter
@@ -72,11 +79,17 @@ func NewOrgResolverWithClient(cfg Config, client linkingclient.LinkingServiceCli
 }
 
 func (cfg *Config) New(logger log.Logger) (*orgResolver, error) {
+	requestTimeout := cfg.RequestTimeout
+	if requestTimeout <= 0 {
+		requestTimeout = defaultRequestTimeout
+	}
+
 	resolver := &orgResolver{
 		workflowRegistryAddress:       cfg.WorkflowRegistryAddress,
 		workflowRegistryChainSelector: cfg.WorkflowRegistryChainSelector,
 		logger:                        log.Sugared(logger).Named("OrgResolver"),
 		jwtGenerator:                  cfg.JWTGenerator,
+		requestTimeout:                requestTimeout,
 	}
 
 	if cfg.Client != nil {
@@ -135,6 +148,9 @@ func (o *orgResolver) addJWTAuth(ctx context.Context, req any) (context.Context,
 }
 
 func (o *orgResolver) Get(ctx context.Context, owner string) (string, error) {
+	ctx, cancel := context.WithTimeout(ctx, o.requestTimeout)
+	defer cancel()
+
 	req := &linkingclient.GetOrganizationFromWorkflowOwnerRequest{
 		WorkflowOwner:           owner,
 		WorkflowRegistryAddress: o.workflowRegistryAddress,

pkg/services/orgresolver/linking_test.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"net"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
@@ -26,6 +27,14 @@ func (m *mockLinkingClient) GetOrganizationFromWorkflowOwner(ctx context.Context
 	}, nil
 }
 
+// mockHangingLinkingClient blocks until the call context is done, simulating a blackholed connection
+type mockHangingLinkingClient struct{}
+
+func (m *mockHangingLinkingClient) GetOrganizationFromWorkflowOwner(ctx context.Context, req *linkingclient.GetOrganizationFromWorkflowOwnerRequest, opts ...grpc.CallOption) (*linkingclient.GetOrganizationFromWorkflowOwnerResponse, error) {
+	<-ctx.Done()
+	return nil, ctx.Err()
+}
+
 // mockJWTGenerator implements the JWTGenerator interface for testing
 type mockJWTGenerator struct {
 	token string
@@ -89,6 +98,39 @@ func TestOrgResolver_Get(t *testing.T) {
 	require.Equal(t, "org-"+workflowOwner, orgID)
 }
 
+func TestOrgResolver_Get_TimesOutWhenServiceHangs(t *testing.T) {
+	cfg := Config{
+		WorkflowRegistryAddress:       "0x1234567890abcdef",
+		WorkflowRegistryChainSelector: 1,
+		RequestTimeout:                50 * time.Millisecond,
+		Client:                        &mockHangingLinkingClient{},
+	}
+
+	resolver, err := cfg.New(logger.Test(t))
+	require.NoError(t, err)
+
+	start := time.Now()
+	_, err = resolver.Get(t.Context(), "0xabcdef1234567890")
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	require.Less(t, time.Since(start), 5*time.Second)
+}
+
+func TestOrgResolver_Get_DefaultTimeoutApplied(t *testing.T) {
+	cfg := Config{
+		WorkflowRegistryAddress:       "0x1234567890abcdef",
+		WorkflowRegistryChainSelector: 1,
+		Client:                        &mockLinkingClient{},
+	}
+
+	resolver, err := cfg.New(logger.Test(t))
+	require.NoError(t, err)
+	require.Equal(t, defaultRequestTimeout, resolver.requestTimeout)
+
+	orgID, err := resolver.Get(t.Context(), "0xabcdef1234567890")
+	require.NoError(t, err)
+	require.Equal(t, "org-0xabcdef1234567890", orgID)
+}
+
 func TestOrgResolver_NewOrgResolver_RequiresClientOrURL(t *testing.T) {
 	cfg := Config{
 		URL:                           "", // Empty URL should cause error