fix: make rpc_timeout and max_bundle_workers configurable on the agent

cmd/spire-agent/cli/run/run.go

@@ -129,6 +129,8 @@ type workloadAPIRateLimitConfig struct {
 type experimentalConfig struct {
 	SyncInterval             string `hcl:"sync_interval"`
 	JWTSVIDCacheHitTimeout   string `hcl:"jwt_svid_cache_hit_timeout"`
+	RPCTimeout               string `hcl:"rpc_timeout"`
+	MaxBundleWorkers         int    `hcl:"max_bundle_workers"`
 	NamedPipeName            string `hcl:"named_pipe_name"`
 	AdminNamedPipeName       string `hcl:"admin_named_pipe_name"`
 	UseSyncAuthorizedEntries *bool  `hcl:"use_sync_authorized_entries"`
@@ -488,6 +490,26 @@ func NewAgentConfig(c *Config, logOptions []log.Option, allowUnknownConfig bool)
 		logger.Warn("The use of 'jwt_svid_cache_hit_timeout' is experimental")
 	}
 
+	if c.Agent.Experimental.RPCTimeout != "" {
+		timeout, err := time.ParseDuration(c.Agent.Experimental.RPCTimeout)
+		if err != nil {
+			return nil, fmt.Errorf("could not parse rpc_timeout: %w", err)
+		}
+		if timeout <= 0 {
+			return nil, fmt.Errorf("rpc_timeout (%s) must be greater than 0", timeout)
+		}
+		client.SetRPCTimeout(timeout)
+		logger.Warn("The use of 'rpc_timeout' is experimental")
+	}
+
+	if c.Agent.Experimental.MaxBundleWorkers != 0 {
+		if c.Agent.Experimental.MaxBundleWorkers < 1 {
+			return nil, fmt.Errorf("max_bundle_workers (%d) must be greater than 0", c.Agent.Experimental.MaxBundleWorkers)
+		}
+		client.SetMaxBundleWorkers(c.Agent.Experimental.MaxBundleWorkers)
+		logger.Warn("The use of 'max_bundle_workers' is experimental")
+	}
+
 	ac.UseSyncAuthorizedEntries = true
 	if c.Agent.Experimental.UseSyncAuthorizedEntries != nil {
 		ac.Log.WithFields(logrus.Fields{

cmd/spire-agent/cli/run/run_test.go

@@ -1113,6 +1113,89 @@
 				require.Nil(t, ac)
 			},
 		},
+		{
+			msg: "rpc_timeout sets the client timeout and logs warning",
+			input: func(c *Config) {
+				c.Agent.Experimental.RPCTimeout = "10s"
+			},
+			logOptions: func(t *testing.T) []log.Option {
+				return []log.Option{
+					func(logger *log.Logger) error {
+						logger.SetOutput(io.Discard)
+						hook := test.NewLocal(logger.Logger)
+						t.Cleanup(func() {
+							spiretest.AssertLogsContainEntries(t, hook.AllEntries(), []spiretest.LogEntry{
+								{
+									Level:   logrus.WarnLevel,
+									Message: "The use of 'rpc_timeout' is experimental",
+								},
+							})
+						})
+						return nil
+					},
+				}
+			},
+			require.NotNil(t, ac)
+			assert.Equal(t, 10*time.Second, client.RPCTimeout)
+		},
+		{
+			msg:                "rpc_timeout returns an error if <= 0",
+			expectError:        true,
+			requireErrorPrefix: "rpc_timeout (0s) must be greater than 0",
+			input: func(c *Config) {
+				c.Agent.Experimental.RPCTimeout = "0s"
+			},
+			test: func(t *testing.T, ac *agent.Config) {
+				require.Nil(t, ac)
+			},
+		},
+		{
+			msg:                "rpc_timeout returns an error if invalid duration",
+			expectError:        true,
+			requireErrorPrefix: "could not parse rpc_timeout:",
+			input: func(c *Config) {
+				c.Agent.Experimental.RPCTimeout = "invalid"
+			},
+			test: func(t *testing.T, ac *agent.Config) {
+				require.Nil(t, ac)
+			},
+		},
+		{
+			msg: "max_bundle_workers sets the worker count and logs warning",
+			input: func(c *Config) {
+				c.Agent.Experimental.MaxBundleWorkers = 5
+			},
+			logOptions: func(t *testing.T) []log.Option {
+				return []log.Option{
+					func(logger *log.Logger) error {
+						logger.SetOutput(io.Discard)
+						hook := test.NewLocal(logger.Logger)
+						t.Cleanup(func() {
+							spiretest.AssertLogsContainEntries(t, hook.AllEntries(), []spiretest.LogEntry{
+								{
+									Level:   logrus.WarnLevel,
+									Message: "The use of 'max_bundle_workers' is experimental",
+								},
+							})
+						})
+						return nil
+					},
+				}
+			},
+			require.NotNil(t, ac)
+			assert.Equal(t, 5, client.MaxBundleWorkers)
+		},
+		{
+			msg:                "max_bundle_workers returns an error if < 1",
+			expectError:        true,
+			requireErrorPrefix: "max_bundle_workers (-1) must be greater than 0",
+			input: func(c *Config) {
+				c.Agent.Experimental.MaxBundleWorkers = -1
+			},
+			test: func(t *testing.T, ac *agent.Config) {
+				require.Nil(t, ac)
+			},
+		},
 		{
 			msg: "ratelimit defaults to zero (disabled)",
 			input: func(c *Config) {

pkg/agent/client/client.go

@@ -29,15 +29,16 @@ import (
 )
 
 const (
-	rpcTimeout = 30 * time.Second
-
-	// maxBundleWorkers is the maximum number of worker goroutines to use when fetching bundles.
+	rpcTimeout       = 30 * time.Second
 	maxBundleWorkers = 10
 )
 
 var (
 	ErrUnableToGetStream = errors.New("unable to get a stream")
 
+	RPCTimeout       = rpcTimeout
+	MaxBundleWorkers = maxBundleWorkers
+
 	entryOutputMask = &types.EntryMask{
 		SpiffeId:             true,
 		Selectors:            true,
@@ -57,6 +58,14 @@ var (
 	RPCTimeoutWithCacheHit = rpcTimeout
 )
 
+func SetRPCTimeout(d time.Duration) {
+	RPCTimeout = d
+}
+
+func SetMaxBundleWorkers(n int) {
+	MaxBundleWorkers = n
+}
+
 func SetJWTSVIDCacheHitTimeout(d time.Duration) {
 	RPCTimeoutWithCacheHit = d
 }
@@ -146,7 +155,7 @@ func (c *client) FetchUpdates(ctx context.Context) (*Update, error) {
 	c.c.RotMtx.RLock()
 	defer c.c.RotMtx.RUnlock()
 
-	ctx, cancel := context.WithTimeout(ctx, rpcTimeout)
+	ctx, cancel := context.WithTimeout(ctx, RPCTimeout)
 	defer cancel()
 
 	protoEntries, err := c.fetchEntries(ctx)
@@ -212,7 +221,7 @@ func (c *client) SyncUpdates(ctx context.Context, cachedEntries map[string]*comm
 	c.c.RotMtx.RLock()
 	defer c.c.RotMtx.RUnlock()
 
-	ctx, cancel := context.WithTimeout(ctx, rpcTimeout)
+	ctx, cancel := context.WithTimeout(ctx, RPCTimeout)
 	defer cancel()
 
 	entriesStats, err := c.syncEntries(ctx, cachedEntries)
@@ -254,7 +263,7 @@ func (c *client) SyncUpdates(ctx context.Context, cachedEntries map[string]*comm
 }
 
 func (c *client) RenewSVID(ctx context.Context, csr []byte) (*X509SVID, error) {
-	ctx, cancel := context.WithTimeout(ctx, rpcTimeout)
+	ctx, cancel := context.WithTimeout(ctx, RPCTimeout)
 	defer cancel()
 
 	agentClient, connection, err := c.newAgentClient()
@@ -288,7 +297,7 @@ func (c *client) PostStatus(ctx context.Context, agentVersion string) error {
 	c.c.RotMtx.RLock()
 	defer c.c.RotMtx.RUnlock()
 
-	ctx, cancel := context.WithTimeout(ctx, rpcTimeout)
+	ctx, cancel := context.WithTimeout(ctx, RPCTimeout)
 	defer cancel()
 
 	agentClient, connection, err := c.newAgentClient()
@@ -313,7 +322,7 @@ func (c *client) NewX509SVIDs(ctx context.Context, csrs map[string][]byte) (map[
 	c.c.RotMtx.RLock()
 	defer c.c.RotMtx.RUnlock()
 
-	ctx, cancel := context.WithTimeout(ctx, rpcTimeout)
+	ctx, cancel := context.WithTimeout(ctx, RPCTimeout)
 	defer cancel()
 
 	svids := make(map[string]*X509SVID)
@@ -354,7 +363,7 @@ func (c *client) NewJWTSVID(ctx context.Context, entryID string, audience []stri
 	c.c.RotMtx.RLock()
 	defer c.c.RotMtx.RUnlock()
 
-	timeout := rpcTimeout
+	timeout := RPCTimeout
 	if hasCacheHit {
 		timeout = RPCTimeoutWithCacheHit
 	}
@@ -679,13 +688,13 @@ func (c *client) fetchBundles(ctx context.Context, federatedBundles []string) ([
 
 // fetchFederatedBundlesConcurrently fetches federated bundles concurrently.
 // This is done to improve sync times when there are many federations. This should ensure that the
-// sync does not exceed rpcTimeout.
+// sync does not exceed RPCTimeout.
 func (c *client) fetchFederatedBundlesConcurrently(ctx context.Context, bundleClient bundlev1.BundleClient, trustDomains []string, bundles []*types.Bundle) ([]*types.Bundle, error) {
 	jobCh := make(chan string)
 	resultCh := make(chan fetchBundleResult, len(trustDomains))
 	// Start a set of worker goroutines.
 	wg := sync.WaitGroup{}
-	for range min(maxBundleWorkers, len(trustDomains)) {
+	for range min(MaxBundleWorkers, len(trustDomains)) {
 		wg.Go(func() {
 			for trustDomain := range jobCh {
 				bundle, err := c.fetchFederatedBundle(ctx, bundleClient, trustDomain)