Add CPU profile collection to diagnostics handler (#4394)

Add the optional CPU profile collection to the diagnostic action handler.
CPU profiles will be collected if the REQUEST_DIAGNOSTICS action has the
optional additional_metrics parameter list contains "CPU".

Author: michel-laterman

changelog/fragments/1710153529-Add-CPU-profile-collection-to-diagnostics-action-handler.yaml

@@ -0,0 +1,35 @@
+# Kind can be one of:
+# - breaking-change: a change to previously-documented behavior
+# - deprecation: functionality that is being removed in a later release
+# - bug-fix: fixes a problem in a previous version
+# - enhancement: extends functionality but does not break or fix existing behavior
+# - feature: new functionality
+# - known-issue: problems that we are aware of in a given version
+# - security: impacts on the security of a product or a user’s deployment.
+# - upgrade: important information for someone upgrading from a prior version
+# - other: does not fit into any of the other categories
+kind: enhancement
+
+# Change summary; a 80ish characters long description of the change.
+summary: Add CPU profile collection to the Fleet diagnostics action handler
+
+# Long description; in case the summary is not enough to describe the change
+# this field accommodate a description without length limits.
+# NOTE: This field will be rendered only for breaking-change and known-issue kinds at the moment.
+description: |
+  Add the optional CPU profile collection to the diagnostic action handler.
+  CPU profiles will be collected if the REQUEST_DIAGNOSTICS action has the
+  optional additional_metrics parameter list contains "CPU".
+
+# Affected component; a word indicating the component this changeset affects.
+component:
+
+# PR URL; optional; the PR number that added the changeset.
+# If not present is automatically filled by the tooling finding the PR where this changelog fragment has been added.
+# NOTE: the tooling supports backports, so it's able to fill the original PR number instead of the backport PR number.
+# Please provide it if you are adding a fragment for a different PR.
+pr: 4394
+
+# Issue URL; optional; the GitHub issue related to this changeset (either closes or is part of).
+# If not present is automatically filled by the tooling with the issue linked to the PR number.
+issue: 3491

internal/pkg/agent/application/actions/handlers/handler_action_diagnostics.go

@@ -31,6 +31,11 @@ import (
 // In either case the 1st action will succeed and the others will ack with an the error.
 var ErrRateLimit = fmt.Errorf("rate limit exceeded")
 
+// getCPUDiag is a wrapper around diagnostics.CreateCPUProfile so it can be replaced in unit-tests.
+var getCPUDiag = func(ctx context.Context, d time.Duration) ([]byte, error) {
+	return diagnostics.CreateCPUProfile(ctx, d)
+}
+
 // Uploader is the interface used to upload a diagnostics bundle to fleet-server.
 type Uploader interface {
 	UploadDiagnostics(context.Context, string, string, int64, io.Reader) (string, error)
@@ -122,7 +127,7 @@ func (h *Diagnostics) collectDiag(ctx context.Context, action *fleetapi.ActionDi
 	}
 
 	h.log.Debug("Gathering agent diagnostics.")
-	aDiag, err := h.runHooks(ctx)
+	aDiag, err := h.runHooks(ctx, action)
 	if err != nil {
 		action.Err = err
 		h.log.Errorw("diagnostics action handler failed to run diagnostics hooks",
@@ -134,7 +139,7 @@ func (h *Diagnostics) collectDiag(ctx context.Context, action *fleetapi.ActionDi
 	uDiag := h.diagUnits(ctx)
 
 	h.log.Debug("Gathering component diagnostics.")
-	cDiag := h.diagComponents(ctx)
+	cDiag := h.diagComponents(ctx, action)
 
 	var r io.Reader
 	// attempt to create the a temporary diagnostics file on disk in order to avoid loading a
@@ -185,9 +190,26 @@ func (h *Diagnostics) collectDiag(ctx context.Context, action *fleetapi.ActionDi
 }
 
 // runHooks runs the agent diagnostics hooks.
-func (h *Diagnostics) runHooks(ctx context.Context) ([]client.DiagnosticFileResult, error) {
+func (h *Diagnostics) runHooks(ctx context.Context, action *fleetapi.ActionDiagnostics) ([]client.DiagnosticFileResult, error) {
 	hooks := append(h.diagProvider.DiagnosticHooks(), diagnostics.GlobalHooks()...)
-	diags := make([]client.DiagnosticFileResult, 0, len(hooks))
+
+	// Currently CPU is the only additional metric we can collect.
+	// If this changes we would need to change how we scan AdditionalMetrics.
+	collectCPU := false
+	for _, metric := range action.AdditionalMetrics {
+		if metric == "CPU" {
+			h.log.Debug("Diagnostics will collect CPU profile.")
+			collectCPU = true
+			break
+		}
+	}
+
+	resultLen := len(hooks)
+	if collectCPU {
+		resultLen++
+	}
+	diags := make([]client.DiagnosticFileResult, 0, resultLen)
+
 	for _, hook := range hooks {
 		if ctx.Err() != nil {
 			return diags, ctx.Err()
@@ -205,6 +227,20 @@ func (h *Diagnostics) runHooks(ctx context.Context) ([]client.DiagnosticFileResu
 		elapsed := time.Since(startTime)
 		h.log.Debugw(fmt.Sprintf("Hook %s execution complete, took %s", hook.Name, elapsed.String()), "hook", hook.Name, "filename", hook.Filename, "elapsed", elapsed.String())
 	}
+	if collectCPU {
+		p, err := getCPUDiag(ctx, diagnostics.DiagCPUDuration)
+		if err != nil {
+			return diags, fmt.Errorf("unable to gather CPU profile: %w", err)
+		}
+		diags = append(diags, client.DiagnosticFileResult{
+			Name:        diagnostics.DiagCPUName,
+			Filename:    diagnostics.DiagCPUFilename,
+			Description: diagnostics.DiagCPUDescription,
+			ContentType: diagnostics.DiagCPUContentType,
+			Content:     p,
+			Generated:   time.Now().UTC(),
+		})
+	}
 	return diags, nil
 }
 
@@ -246,14 +282,20 @@ func (h *Diagnostics) diagUnits(ctx context.Context) []client.DiagnosticUnitResu
 }
 
 // diagUnits gathers diagnostics from components.
-func (h *Diagnostics) diagComponents(ctx context.Context) []client.DiagnosticComponentResult {
+func (h *Diagnostics) diagComponents(ctx context.Context, action *fleetapi.ActionDiagnostics) []client.DiagnosticComponentResult {
 	cDiag := make([]client.DiagnosticComponentResult, 0)
 	h.log.Debug("Performing component diagnostics")
 	startTime := time.Now()
 	defer func() {
 		h.log.Debugf("Component diagnostics complete. Took: %s", time.Since(startTime))
 	}()
-	rr, err := h.diagProvider.PerformComponentDiagnostics(ctx, []cproto.AdditionalDiagnosticRequest{})
+	additionalMetrics := []cproto.AdditionalDiagnosticRequest{}
+	for _, metric := range action.AdditionalMetrics {
+		if metric == "CPU" {
+			additionalMetrics = append(additionalMetrics, cproto.AdditionalDiagnosticRequest_CPU)
+		}
+	}
+	rr, err := h.diagProvider.PerformComponentDiagnostics(ctx, additionalMetrics)
 	if err != nil {
 		h.log.Errorf("Error fetching component-level diagnostics: %w", err)
 	}

internal/pkg/agent/application/actions/handlers/handler_action_diagnostics_test.go

@@ -18,6 +18,7 @@ import (
 
 	"github.com/elastic/elastic-agent-client/v7/pkg/client"
 	"github.com/elastic/elastic-agent-client/v7/pkg/proto"
+	"github.com/elastic/elastic-agent/pkg/control/v2/cproto"
 
 	"github.com/elastic/elastic-agent/internal/pkg/agent/application/actions/handlers/mocks"
 	"github.com/elastic/elastic-agent/internal/pkg/agent/application/paths"
@@ -350,3 +351,54 @@ func TestDiagnosticHandlerContexteExpiredErrorWithLogs(t *testing.T) {
 		1)
 	// we could assert the logs for the hooks, but those will be the same as the happy path, so for brevity we won't
 }
+
+func TestDiagnosticHandlerWithCPUProfile(t *testing.T) {
+	tempAgentRoot := t.TempDir()
+	paths.SetTop(tempAgentRoot)
+	err := os.MkdirAll(path.Join(tempAgentRoot, "data"), 0755)
+	require.NoError(t, err)
+
+	// make a flag to check if a CPU profile is collected.
+	cpuCalled := false
+	getCPUDiag = func(_ context.Context, _ time.Duration) ([]byte, error) {
+		cpuCalled = true
+		return []byte(`hello, world!`), nil
+	}
+
+	mockDiagProvider := mocks.NewDiagnosticsProvider(t)
+	mockUploader := mocks.NewUploader(t)
+	testLogger, _ := logger.NewTesting("diagnostic-handler-test")
+	handler := NewDiagnostics(testLogger, mockDiagProvider, defaultRateLimit, mockUploader)
+
+	mockDiagProvider.EXPECT().DiagnosticHooks().Return([]diagnostics.Hook{hook1})
+	mockDiagProvider.EXPECT().PerformDiagnostics(mock.Anything, mock.Anything).Return([]runtime.ComponentUnitDiagnostic{mockUnitDiagnostic})
+
+	// only match if CPU metrics are requested.
+	mockDiagProvider.EXPECT().PerformComponentDiagnostics(mock.Anything, mock.MatchedBy(func(additionalMetrics []cproto.AdditionalDiagnosticRequest) bool {
+		for _, metric := range additionalMetrics {
+			if metric == cproto.AdditionalDiagnosticRequest_CPU {
+				return true
+			}
+		}
+		return false
+	})).Return([]runtime.ComponentDiagnostic{}, nil)
+
+	mockAcker := mocks.NewAcker(t)
+	mockAcker.EXPECT().Ack(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, a fleetapi.Action) error {
+		require.IsType(t, new(fleetapi.ActionDiagnostics), a)
+		assert.NoError(t, a.(*fleetapi.ActionDiagnostics).Err)
+		return nil
+	})
+	mockAcker.EXPECT().Commit(mock.Anything).Return(nil)
+
+	mockUploader.EXPECT().UploadDiagnostics(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return("upload-id", nil)
+
+	diagAction := &fleetapi.ActionDiagnostics{
+		AdditionalMetrics: []string{"CPU"},
+	}
+	handler.collectDiag(context.Background(), diagAction, mockAcker)
+
+	// Check that CPU profile was collected and passed to PerformComponentDiagnostics
+	assert.True(t, cpuCalled, "CPU profile collector was not called.")
+	mockDiagProvider.AssertExpectations(t)
+}

internal/pkg/diagnostics/diagnostics.go

@@ -37,6 +37,15 @@ const (
 	agentName = "elastic-agent"
 )
 
+// DiagCPU* are contstants to describe the CPU profile that is collected when the --cpu-profile flag is used with the diagnostics command, or the diagnostics action contains "CPU" in the additional_metrics list.
+const (
+	DiagCPUName        = "cpuprofile"
+	DiagCPUFilename    = "cpu.pprof"
+	DiagCPUDescription = "CPU profile"
+	DiagCPUContentType = "application/octet-stream"
+	DiagCPUDuration    = 30 * time.Second
+)
+
 // Hook is a hook that gets used when diagnostic information is requested from the Elastic Agent.
 type Hook struct {
 	Name        string

internal/pkg/fleetapi/action.go

@@ -434,10 +434,11 @@ func (a *ActionCancel) AckEvent() AckEvent {
 
 // ActionDiagnostics is a request to gather and upload a diagnostics bundle.
 type ActionDiagnostics struct {
-	ActionID   string `json:"action_id"`
-	ActionType string `json:"type"`
-	UploadID   string `json:"-"`
-	Err        error  `json:"-"`
+	ActionID          string   `json:"action_id"`
+	ActionType        string   `json:"type"`
+	AdditionalMetrics []string `json:"additional_metrics"`
+	UploadID          string   `json:"-"`
+	Err               error    `json:"-"`
 }
 
 // ID returns the ID of the action.

internal/pkg/fleetapi/action_test.go

@@ -152,6 +152,29 @@ func TestActionsUnmarshalJSON(t *testing.T) {
 		assert.Equal(t, "http://example.com", action.SourceURI)
 		assert.Equal(t, 1, action.Retry)
 	})
+	t.Run("ActionDiagnostics with no additional metrics", func(t *testing.T) {
+		p := []byte(`[{"id":"testid","type":"REQUEST_DIAGNOSTICS","data":{}}]`)
+		a := &Actions{}
+		err := a.UnmarshalJSON(p)
+		require.Nil(t, err)
+		action, ok := (*a)[0].(*ActionDiagnostics)
+		require.True(t, ok, "unable to cast action to specific type")
+		assert.Equal(t, "testid", action.ActionID)
+		assert.Equal(t, ActionTypeDiagnostics, action.ActionType)
+		assert.Empty(t, action.AdditionalMetrics)
+	})
+	t.Run("ActionDiagnostics with additional CPU metrics", func(t *testing.T) {
+		p := []byte(`[{"id":"testid","type":"REQUEST_DIAGNOSTICS","data":{"additional_metrics":["CPU"]}}]`)
+		a := &Actions{}
+		err := a.UnmarshalJSON(p)
+		require.Nil(t, err)
+		action, ok := (*a)[0].(*ActionDiagnostics)
+		require.True(t, ok, "unable to cast action to specific type")
+		assert.Equal(t, "testid", action.ActionID)
+		assert.Equal(t, ActionTypeDiagnostics, action.ActionType)
+		require.Len(t, action.AdditionalMetrics, 1)
+		assert.Equal(t, "CPU", action.AdditionalMetrics[0])
+	})
 }
 
 func TestActionUnenrollMarshalMap(t *testing.T) {

pkg/control/v2/server/server.go

@@ -205,17 +205,17 @@ func (s *Server) DiagnosticAgent(ctx context.Context, req *cproto.DiagnosticAgen
 	for _, metric := range req.AdditionalMetrics {
 		switch metric {
 		case cproto.AdditionalDiagnosticRequest_CPU:
-			duration := time.Second * 30
+			duration := diagnostics.DiagCPUDuration
 			s.logger.Infof("Collecting CPU metrics, waiting for %s", duration)
 			cpuResults, err := diagnostics.CreateCPUProfile(ctx, duration)
 			if err != nil {
 				return nil, fmt.Errorf("error gathering CPU profile: %w", err)
 			}
 			res = append(res, &cproto.DiagnosticFileResult{
-				Name:        "cpuprofile",
-				Filename:    "cpu.pprof",
-				Description: "CPU profile",
-				ContentType: "application/octet-stream",
+				Name:        diagnostics.DiagCPUName,
+				Filename:    diagnostics.DiagCPUFilename,
+				Description: diagnostics.DiagCPUDescription,
+				ContentType: diagnostics.DiagCPUContentType,
 				Content:     cpuResults,
 				Generated:   timestamppb.New(time.Now().UTC()),
 			})