Added a shared secrets counter to both secret call paths

core/services/workflows/v2/capability_executor.go

@@ -38,10 +38,12 @@ type ExecutionHelper struct {
 	TimeProvider
 	SecretsFetcher
 
-	chainAllowed limits.GateLimiter
-	callLimiters map[capCall]limits.BoundLimiter[int]
-	mu           sync.Mutex
-	callCounts   map[limits.Limiter[int]]int
+	chainAllowed         limits.GateLimiter
+	callLimiters         map[capCall]limits.BoundLimiter[int]
+	concurrencyLimiters  map[capCall]limits.ResourcePoolLimiter[int]
+	mu                   sync.Mutex
+	callCounts           map[limits.Limiter[int]]int
+	sharedSecretsCounter *secretsCallCounter
 
 	executionProfile *executionProfileCollector
 
@@ -89,6 +91,11 @@ func (c *ExecutionHelper) initLimiters(limiters *EngineLimiters) {
 
 		{"http-actions", "SendRequest"}:      limiters.HTTPActionCalls,
 		{"confidential-http", "SendRequest"}: limiters.ConfidentialHTTPCalls,
+
+		{"vault", "vault.secrets.get"}: limiters.SecretsCalls,
+	}
+	c.concurrencyLimiters = map[capCall]limits.ResourcePoolLimiter[int]{
+		{"vault", "vault.secrets.get"}: limiters.SecretsConcurrency,
 	}
 }
 
@@ -121,20 +128,37 @@ func (c *ExecutionHelper) CallCapability(ctx context.Context, request *sdkpb.Cap
 
 	limiter, ok := c.callLimiters[capCall{name: capName, method: request.Method}]
 	if ok {
-		c.mu.Lock()
-		if c.callCounts == nil {
-			c.callCounts = make(map[limits.Limiter[int]]int)
-		}
-		cnt := c.callCounts[limiter] + 1
-		if err := limiter.Check(ctx, cnt); err != nil {
+		if c.sharedSecretsCounter != nil && capName == "vault" {
+			if err := c.sharedSecretsCounter.acquire(ctx, limiter); err != nil {
+				return nil, caperrors.NewPublicUserError(
+					fmt.Errorf("capability call limit exceeded for %s.%s: %w", capName, request.Method, err),
+					caperrors.LimitExceeded,
+				)
+			}
+		} else {
+			c.mu.Lock()
+			if c.callCounts == nil {
+				c.callCounts = make(map[limits.Limiter[int]]int)
+			}
+			cnt := c.callCounts[limiter] + 1
+			if err := limiter.Check(ctx, cnt); err != nil {
+				c.mu.Unlock()
+				return nil, caperrors.NewPublicUserError(
+					fmt.Errorf("capability call limit exceeded for %s.%s: %w", capName, request.Method, err),
+					caperrors.LimitExceeded,
+				)
+			}
+			c.callCounts[limiter] = cnt
 			c.mu.Unlock()
-			return nil, caperrors.NewPublicUserError(
-				fmt.Errorf("capability call limit exceeded for %s.%s: %w", capName, request.Method, err),
-				caperrors.LimitExceeded,
-			)
 		}
-		c.callCounts[limiter] = cnt
-		c.mu.Unlock()
+	}
+
+	if concurrencyLimiter, ok := c.concurrencyLimiters[capCall{name: capName, method: request.Method}]; ok {
+		freeConcurrency, err := concurrencyLimiter.Wait(ctx, 1)
+		if err != nil {
+			return nil, err
+		}
+		defer freeConcurrency()
 	}
 
 	free, err := c.capCallsSemaphore.Wait(ctx, 1)

core/services/workflows/v2/capability_executor_test.go

@@ -1,12 +1,15 @@
 package v2
 
 import (
+	"context"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	caperrors "github.com/smartcontractkit/chainlink-common/pkg/capabilities/errors"
+	"github.com/smartcontractkit/chainlink-common/pkg/contexts"
 	"github.com/smartcontractkit/chainlink-common/pkg/settings"
 	"github.com/smartcontractkit/chainlink-common/pkg/settings/cresettings"
 	"github.com/smartcontractkit/chainlink-common/pkg/settings/limits"
@@ -75,6 +78,132 @@ func TestExecutionHelper_ConfidentialHTTPPerWorkflowLimit(t *testing.T) {
 	require.Equal(t, caperrors.LimitExceeded, capErr.Code())
 }
 
+func TestExecutionHelper_VaultSecretsGetPerWorkflowLimit(t *testing.T) {
+	t.Parallel()
+
+	lggr := logger.TestLogger(t)
+	lf := limits.Factory{Logger: lggr}
+
+	cfgFn := func(w *cresettings.Workflows) {
+		w.Secrets.CallLimit = settings.Int(1)
+	}
+
+	limiters, err := NewLimiters(lf, cfgFn)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = limiters.Close() })
+
+	exec := &ExecutionHelper{}
+	exec.initLimiters(limiters)
+	exec.sharedSecretsCounter = &secretsCallCounter{called: 1}
+
+	req := &sdk.CapabilityRequest{
+		Id:         "vault",
+		Method:     "vault.secrets.get",
+		CallbackId: 1,
+	}
+
+	_, err = exec.CallCapability(t.Context(), req)
+	require.Error(t, err, "expected CallCapability to fail when per-workflow secrets call limit is exceeded")
+	var capErr caperrors.Error
+	require.ErrorAs(t, err, &capErr, "expected per-workflow call limit exceedance to be classified as capability user error")
+	require.Equal(t, caperrors.OriginUser, capErr.Origin())
+	require.Equal(t, caperrors.LimitExceeded, capErr.Code())
+}
+
+func TestExecutionHelper_VaultSecretsGetSharedCallCounter(t *testing.T) {
+	t.Parallel()
+
+	lggr := logger.TestLogger(t)
+	lf := limits.Factory{Logger: lggr}
+
+	cfgFn := func(w *cresettings.Workflows) {
+		w.Secrets.CallLimit = settings.Int(3)
+		w.SecretsConcurrencyLimit = settings.Int(100)
+	}
+
+	limiters, err := NewLimiters(lf, cfgFn)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = limiters.Close() })
+
+	exec := &ExecutionHelper{Engine: &Engine{}}
+	exec.initLimiters(limiters)
+	exec.capCallsSemaphore = limits.GlobalResourcePoolLimiter(0)
+	exec.sharedSecretsCounter = &secretsCallCounter{}
+
+	// Simulate 2 prior calls via the GetSecrets path (secretsFetcher increments
+	// the same shared counter).
+	exec.sharedSecretsCounter.called = 2
+	require.Equal(t, 2, exec.sharedSecretsCounter.called)
+
+	req := &sdk.CapabilityRequest{
+		Id:         "vault",
+		Method:     "vault.secrets.get",
+		CallbackId: 1,
+	}
+
+	// 3rd call: limit check should pass (counter 2->3, limit 3), then fail at
+	// the zero-capacity semaphore. The counter must be incremented to 3.
+	callCtx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond)
+	defer cancel()
+	_, err = exec.CallCapability(callCtx, req)
+	require.Error(t, err, "expected call to fail at zero-capacity semaphore")
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	require.Equal(t, 3, exec.sharedSecretsCounter.called, "shared counter should be incremented to 3 after a non-limit failure")
+
+	// 4th call: should be rejected at the limit check (counter 3 > limit 3).
+	_, err = exec.CallCapability(t.Context(), req)
+	require.Error(t, err, "expected CallCapability to fail when shared counter exceeds limit")
+	var capErr caperrors.Error
+	require.ErrorAs(t, err, &capErr)
+	require.Equal(t, caperrors.LimitExceeded, capErr.Code())
+	require.Equal(t, 3, exec.sharedSecretsCounter.called, "shared counter must not increment on a failed limit check")
+}
+
+func TestExecutionHelper_VaultSecretsGetConcurrencyLimit(t *testing.T) {
+	t.Parallel()
+
+	lggr := logger.TestLogger(t)
+	lf := limits.Factory{Logger: lggr}
+
+	cfgFn := func(w *cresettings.Workflows) {
+		w.Secrets.CallLimit = settings.Int(100)
+		w.SecretsConcurrencyLimit = settings.Int(1)
+	}
+
+	limiters, err := NewLimiters(lf, cfgFn)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = limiters.Close() })
+
+	exec := &ExecutionHelper{}
+	exec.initLimiters(limiters)
+
+	capCallValue := capCall{name: "vault", method: "vault.secrets.get"}
+	concurrencyLimiter, ok := exec.concurrencyLimiters[capCallValue]
+	require.True(t, ok, "expected vault.secrets.get concurrency limiter to be configured")
+
+	ctx := contexts.WithCRE(t.Context(), contexts.CRE{
+		Owner:    "1111111111111111111111111111111111111111",
+		Workflow: "22222222222222222222222222222222222222222222222222222222222222222",
+	})
+
+	free, err := concurrencyLimiter.Wait(ctx, 1)
+	require.NoError(t, err)
+	defer free()
+
+	req := &sdk.CapabilityRequest{
+		Id:         "vault",
+		Method:     "vault.secrets.get",
+		CallbackId: 1,
+	}
+
+	callCtx, cancel := context.WithTimeout(ctx, 50*time.Millisecond)
+	defer cancel()
+
+	_, err = exec.CallCapability(callCtx, req)
+	require.Error(t, err, "expected CallCapability to fail when secrets concurrency limit is exhausted")
+	assert.ErrorIs(t, err, context.DeadlineExceeded)
+}
+
 func TestUserMetricTypeSuffix(t *testing.T) {
 	t.Parallel()
 

core/services/workflows/v2/engine.go

@@ -1008,6 +1008,9 @@ func (e *Engine) startExecution(ctx context.Context, wrappedTriggerEvent enqueue
 	}
 
 	execHelper.initLimiters(e.cfg.LocalLimiters)
+	if sf, ok := execHelper.SecretsFetcher.(*secretsFetcher); ok {
+		execHelper.sharedSecretsCounter = sf.callCounter
+	}
 	e.metrics.With(platform.KeyTriggerID, wrappedTriggerEvent.triggerCapID).RecordTriggerPayloadBytes(ctx, int64(proto.Size(triggerEvent.Payload)))
 	var result *sdkpb.ExecutionResult
 	result, execErr = e.cfg.Module.Execute(execCtx, &sdkpb.ExecuteRequest{

core/services/workflows/v2/secrets.go

@@ -69,13 +69,27 @@ type secretsFetcher struct {
 }
 
 // secretsCallCounter holds the mutable call-count state shared between a
-// secretsFetcher and any clones created via WithEncryptionKeyFetcher. It is
-// referenced by pointer so the struct can be copied without copying the lock.
+// secretsFetcher and the ExecutionHelper so that both the GetSecrets and
+// CallCapability code paths increment the same counter. It is referenced by
+// pointer so the struct can be copied without copying the lock.
 type secretsCallCounter struct {
 	mu     sync.Mutex
 	called int
 }
 
+// acquire increments the counter and checks it against limiter. The counter
+// is only incremented when the check passes, a failed check leaves the counter unchanged.
+func (c *secretsCallCounter) acquire(ctx context.Context, limiter limits.BoundLimiter[int]) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	next := c.called + 1
+	if err := limiter.Check(ctx, next); err != nil {
+		return err
+	}
+	c.called = next
+	return nil
+}
+
 func NewSecretsFetcher(
 	metrics *monitoring.WorkflowsMetricLabeler,
 	capRegistry core.CapabilitiesRegistry,
@@ -163,14 +177,9 @@ func (s *secretsFetcher) GetSecrets(ctx context.Context, request *sdkpb.GetSecre
 	}
 	vaultRequestID := vault.BuildWorkflowGetSecretsRequestID(metadata)
 	s.lggr.Debugw("get secrets request received", "vaultRequestID", vaultRequestID, "metadata", metadata)
-	s.callCounter.mu.Lock()
-	secretsCalled := s.callCounter.called + 1
-	if err := s.secretsCallsLimit.Check(ctx, secretsCalled); err != nil {
-		s.callCounter.mu.Unlock()
-		return nil, err
+	if acqErr := s.callCounter.acquire(ctx, s.secretsCallsLimit); acqErr != nil {
+		return nil, acqErr
 	}
-	s.callCounter.called = secretsCalled
-	s.callCounter.mu.Unlock()
 	start := time.Now()
 	resp, err := func() ([]*sdkpb.SecretResponse, error) {
 		free, err := s.semaphore.Wait(ctx, 1)