feat(trace): Trigger API Key refresh on 403 (#45674)

### What does this PR do?
The agent now supports API Key rotation. Instead of dropping payloads when a 403 is received from the backend, we will instead trigger an API Key (secrets) refresh, and mark the payload as retriable.

This PR also refactors the API Key management within `pkg/trace/writer/sender.go`, making reads and writes safe for concurrent use.

### Motivation
See above

### Describe how you validated your changes
Added unit tests

### Additional Notes


Co-authored-by: inigo.lopezdeheredia <inigo.lopezdeheredia@datadoghq.com>

Author: ichinaski

comp/trace/agent/impl/agent.go

@@ -137,9 +137,17 @@ func NewAgent(deps dependencies) (traceagent.Component, error) {
 
 	prepGoRuntime(tracecfg)
 
+	tracecfg.SecretsRefreshFn = func() (string, error) {
+		if deps.Secrets == nil {
+			log.Error("Secrets component not available, cannot trigger refresh")
+			return "", errors.New("secrets component not available")
+		}
+		return deps.Secrets.Refresh(true)
+	}
+
 	c.Agent = pkgagent.NewAgent(
 		ctx,
-		c.config.Object(),
+		tracecfg,
 		c.telemetryCollector,
 		statsdCl,
 		deps.Compressor,

comp/trace/config/setup.go

@@ -548,6 +548,10 @@ func applyDatadogConfig(c *config.AgentConfig, core corecompcfg.Component) error
 		// Default of 4 was chosen through experimentation, but may not be the optimal value.
 		c.MaxSenderRetries = 4
 	}
+	if core.IsSet("secret_refresh_on_api_key_failure_interval") {
+		// Use the global secret refresh interval for throttling API key refresh at the sender level
+		c.APIKeyRefreshThrottleInterval = time.Duration(core.GetInt("secret_refresh_on_api_key_failure_interval")) * time.Minute
+	}
 	if core.IsConfigured("apm_config.sync_flushing") {
 		c.SynchronousFlushing = core.GetBool("apm_config.sync_flushing")
 	}

pkg/trace/config/config.go

@@ -416,6 +416,8 @@ type AgentConfig struct {
 	// case, the sender will drop failed payloads when it is unable to enqueue
 	// them for another retry.
 	MaxSenderRetries int
+	// APIKeyRefreshThrottleInterval is the minimum time between API key refresh attempts.
+	APIKeyRefreshThrottleInterval time.Duration
 	// HTTP Transport used in writer connections. If nil, default transport values will be used.
 	HTTPTransportFunc func() *http.Transport `json:"-"`
 	// ClientStatsFlushInterval specifies the frequency at which the client stats aggregator will flush its buffer.
@@ -555,6 +557,11 @@ type AgentConfig struct {
 
 	// APMMode specifies whether using "edge" APM mode. May support other modes in the future. If unset, it has no impact.
 	APMMode string
+
+	// SecretsRefreshFn is called when a 403 response is received to trigger
+	// API key refresh from the secrets backend. It blocks until the refresh
+	// completes and returns a message and any error encountered.
+	SecretsRefreshFn func() (string, error) `json:"-"`
 }
 
 // RemoteClient client is used to APM Sampling Updates from a remote source.
@@ -610,11 +617,12 @@ func New() *AgentConfig {
 		PipeSecurityDescriptor: "D:AI(A;;GA;;;WD)",
 		GUIPort:                "5002",
 
-		StatsWriter:              new(WriterConfig),
-		TraceWriter:              new(WriterConfig),
-		ConnectionResetInterval:  0, // disabled
-		MaxSenderRetries:         4,
-		ClientStatsFlushInterval: 2 * time.Second, // bucket duration (2s)
+		StatsWriter:                   new(WriterConfig),
+		TraceWriter:                   new(WriterConfig),
+		ConnectionResetInterval:       0, // disabled
+		MaxSenderRetries:              4,
+		APIKeyRefreshThrottleInterval: 2 * time.Minute,
+		ClientStatsFlushInterval:      2 * time.Second, // bucket duration (2s)
 
 		StatsdHost:    "localhost",
 		StatsdPort:    8125,

pkg/trace/writer/sender.go

@@ -45,18 +45,24 @@ func newSenders(cfg *config.AgentConfig, r eventRecorder, path string, climit, q
 			log.Criticalf("Invalid host endpoint: %q", endpoint.Host)
 			os.Exit(1)
 		}
-		senders[i] = newSender(&senderConfig{
+		scfg := &senderConfig{
 			client:         cfg.NewHTTPClient(),
 			maxConns:       int(maxConns),
 			maxQueued:      qsize,
 			maxRetries:     cfg.MaxSenderRetries,
 			url:            url,
-			apiKey:         endpoint.APIKey,
 			recorder:       r,
 			userAgent:      fmt.Sprintf("Datadog Trace Agent/%s/%s", cfg.AgentVersion, cfg.GitCommit),
 			isMRF:          endpoint.IsMRF,
 			MRFFailoverAPM: cfg.MRFFailoverAPM,
-		}, statsd)
+		}
+		apiKeyManager := &apiKeyManager{
+			apiKey:           endpoint.APIKey,
+			refreshFn:        cfg.SecretsRefreshFn,
+			throttleInterval: cfg.APIKeyRefreshThrottleInterval,
+		}
+
+		senders[i] = newSender(scfg, apiKeyManager, statsd)
 	}
 	return senders
 }
@@ -140,8 +146,6 @@ type senderConfig struct {
 	client *config.ResetClient
 	// url specifies the URL to send requests too.
 	url *url.URL
-	// apiKey specifies the Datadog API key to use.
-	apiKey string
 	// maxConns specifies the maximum number of allowed concurrent ougoing
 	// connections.
 	maxConns int
@@ -162,10 +166,62 @@ type senderConfig struct {
 	MRFFailoverAPM func() bool
 }
 
+// apiKeyManager handles API Key access for concurrent use
+type apiKeyManager struct {
+	sync.RWMutex
+	// apiKey specifies the Datadog API key to use.
+	apiKey string
+
+	// refreshFn triggers blocking API key refresh from the secrets backend
+	refreshFn func() (string, error)
+
+	// throttleInterval specifies minimum time between refresh calls
+	throttleInterval time.Duration
+
+	// lastRefresh tracks when the last refresh occurred
+	lastRefresh time.Time
+}
+
+func (m *apiKeyManager) Get() string {
+	m.RLock()
+	defer m.RUnlock()
+	return m.apiKey
+}
+
+func (m *apiKeyManager) Update(newKey string) {
+	m.Lock()
+	defer m.Unlock()
+	m.apiKey = newKey
+}
+
+func (m *apiKeyManager) refresh() {
+	if m.refreshFn == nil || m.throttleInterval == 0 {
+		return
+	}
+
+	m.Lock()
+	if time.Since(m.lastRefresh) < m.throttleInterval {
+		m.Unlock()
+		log.Debugf("API Key refresh throttled, last refresh was %v ago", time.Since(m.lastRefresh))
+		return
+	}
+
+	// Update the last refresh time before calling refresh to prevent concurrent calls
+	m.lastRefresh = time.Now()
+	m.Unlock()
+
+	if result, err := m.refreshFn(); err != nil {
+		log.Debugf("API Key refresh failed: %v", err)
+	} else if result != "" {
+		log.Infof("API Key refresh completed: %s", result)
+	}
+}
+
 // sender is responsible for sending payloads to a given URL. It uses a size-limited
 // retry queue with a backoff mechanism in case of retriable errors.
 type sender struct {
-	cfg *senderConfig
+	cfg           *senderConfig
+	apiKeyManager *apiKeyManager
 
 	queue      chan *payload // payload queue
 	inflight   *atomic.Int32 // inflight payloads
@@ -178,14 +234,15 @@ type sender struct {
 }
 
 // newSender returns a new sender based on the given config cfg.
-func newSender(cfg *senderConfig, statsd statsd.ClientInterface) *sender {
+func newSender(cfg *senderConfig, apiKeyManager *apiKeyManager, statsd statsd.ClientInterface) *sender {
 	s := sender{
-		cfg:        cfg,
-		queue:      make(chan *payload, cfg.maxQueued),
-		inflight:   atomic.NewInt32(0),
-		maxRetries: int32(cfg.maxRetries),
-		statsd:     statsd,
-		enabled:    true,
+		cfg:           cfg,
+		apiKeyManager: apiKeyManager,
+		queue:         make(chan *payload, cfg.maxQueued),
+		inflight:      atomic.NewInt32(0),
+		maxRetries:    int32(cfg.maxRetries),
+		statsd:        statsd,
+		enabled:       true,
 	}
 	for i := 0; i < cfg.maxConns; i++ {
 		go s.loop()
@@ -388,7 +445,7 @@ const (
 )
 
 func (s *sender) do(req *http.Request) error {
-	req.Header.Set(headerAPIKey, s.cfg.apiKey)
+	req.Header.Set(headerAPIKey, s.apiKeyManager.Get())
 	req.Header.Set(headerUserAgent, s.cfg.userAgent)
 	resp, err := s.cfg.client.Do(req)
 	if err != nil {
@@ -404,6 +461,11 @@ func (s *sender) do(req *http.Request) error {
 	}
 	resp.Body.Close()
 
+	if resp.StatusCode == http.StatusForbidden {
+		log.Debugf("API Key invalid (403), triggering secret refresh")
+		s.apiKeyManager.refresh()
+	}
+
 	if isRetriable(resp.StatusCode) {
 		return &retriableError{
 			fmt.Errorf("server responded with %q", resp.Status),
@@ -419,7 +481,8 @@ func (s *sender) do(req *http.Request) error {
 
 // isRetriable reports whether the give HTTP status code should be retried.
 func isRetriable(code int) bool {
-	if code == http.StatusRequestTimeout || code == http.StatusTooManyRequests {
+	// TODO: Double check what response codes are expected from the backend when API Key is invalid
+	if code == http.StatusRequestTimeout || code == http.StatusTooManyRequests || code == http.StatusForbidden {
 		return true
 	}
 	// 5xx errors can be retried