Adding Statuscode Handler to Agent Health Extension

extension/agenthealth/config.go

@@ -10,8 +10,9 @@ import (
 )
 
 type Config struct {
-	IsUsageDataEnabled bool              `mapstructure:"is_usage_data_enabled"`
-	Stats              agent.StatsConfig `mapstructure:"stats"`
+	IsUsageDataEnabled      bool              `mapstructure:"is_usage_data_enabled"`
+	Stats                   agent.StatsConfig `mapstructure:"stats"`
+	IsOnlyStatusCodeEnabled bool              `mapstructure:"is_status_code_enabled,omitempty"`
 }
 
 var _ component.Config = (*Config)(nil)

extension/agenthealth/extension.go

@@ -7,6 +7,7 @@ import (
 	"github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware"
 	"go.opentelemetry.io/collector/component"
 	"go.uber.org/zap"
+	"golang.org/x/exp/slices"
 
 	"github.com/aws/amazon-cloudwatch-agent/extension/agenthealth/handler/stats"
 	"github.com/aws/amazon-cloudwatch-agent/extension/agenthealth/handler/useragent"
@@ -24,11 +25,24 @@ var _ awsmiddleware.Extension = (*agentHealth)(nil)
 func (ah *agentHealth) Handlers() ([]awsmiddleware.RequestHandler, []awsmiddleware.ResponseHandler) {
 	var responseHandlers []awsmiddleware.ResponseHandler
 	requestHandlers := []awsmiddleware.RequestHandler{useragent.NewHandler(ah.cfg.IsUsageDataEnabled)}
-	if ah.cfg.IsUsageDataEnabled {
-		req, res := stats.NewHandlers(ah.logger, ah.cfg.Stats)
-		requestHandlers = append(requestHandlers, req...)
-		responseHandlers = append(responseHandlers, res...)
+
+	if !ah.cfg.IsUsageDataEnabled {
+		ah.logger.Debug("Usage data is disabled, skipping stats handlers")
+		return requestHandlers, responseHandlers
 	}
+
+	statusCodeEnabled := false
+	statusCodeEnabled = ah.cfg.IsOnlyStatusCodeEnabled
+
+	agentStatsEnabled :=
+		slices.Contains(ah.cfg.Stats.Operations, "PutMetricData") ||
+			slices.Contains(ah.cfg.Stats.Operations, "PutLogEvents") ||
+			slices.Contains(ah.cfg.Stats.Operations, "PutTraceSegment")
+
+	statsRequestHandlers, statsResponseHandlers := stats.NewHandlers(ah.logger, ah.cfg.Stats, statusCodeEnabled, agentStatsEnabled)
+	requestHandlers = append(requestHandlers, statsRequestHandlers...)
+	responseHandlers = append(responseHandlers, statsResponseHandlers...)
+
 	return requestHandlers, responseHandlers
 }
 

extension/agenthealth/extension_test.go

@@ -10,11 +10,13 @@ import (
 	"github.com/stretchr/testify/assert"
 	"go.opentelemetry.io/collector/component/componenttest"
 	"go.uber.org/zap"
+
+	"github.com/aws/amazon-cloudwatch-agent/extension/agenthealth/handler/stats/agent"
 )
 
 func TestExtension(t *testing.T) {
 	ctx := context.Background()
-	cfg := &Config{IsUsageDataEnabled: true}
+	cfg := &Config{IsUsageDataEnabled: true, Stats: agent.StatsConfig{Operations: []string{"PutLogEvents"}}}
 	extension := NewAgentHealth(zap.NewNop(), cfg)
 	assert.NotNil(t, extension)
 	assert.NoError(t, extension.Start(ctx, componenttest.NewNopHost()))

extension/agenthealth/handler/stats/agent/agent.go

@@ -15,21 +15,22 @@ const (
 )
 
 type Stats struct {
-	CpuPercent                *float64 `json:"cpu,omitempty"`
-	MemoryBytes               *uint64  `json:"mem,omitempty"`
-	FileDescriptorCount       *int32   `json:"fd,omitempty"`
-	ThreadCount               *int32   `json:"th,omitempty"`
-	LatencyMillis             *int64   `json:"lat,omitempty"`
-	PayloadBytes              *int     `json:"load,omitempty"`
-	StatusCode                *int     `json:"code,omitempty"`
-	SharedConfigFallback      *int     `json:"scfb,omitempty"`
-	ImdsFallbackSucceed       *int     `json:"ifs,omitempty"`
-	AppSignals                *int     `json:"as,omitempty"`
-	EnhancedContainerInsights *int     `json:"eci,omitempty"`
-	RunningInContainer        *int     `json:"ric,omitempty"`
-	RegionType                *string  `json:"rt,omitempty"`
-	Mode                      *string  `json:"m,omitempty"`
-	EntityRejected            *int     `json:"ent,omitempty"`
+	CpuPercent                *float64          `json:"cpu,omitempty"`
+	MemoryBytes               *uint64           `json:"mem,omitempty"`
+	FileDescriptorCount       *int32            `json:"fd,omitempty"`
+	ThreadCount               *int32            `json:"th,omitempty"`
+	LatencyMillis             *int64            `json:"lat,omitempty"`
+	PayloadBytes              *int              `json:"load,omitempty"`
+	StatusCode                *int              `json:"code,omitempty"`
+	SharedConfigFallback      *int              `json:"scfb,omitempty"`
+	ImdsFallbackSucceed       *int              `json:"ifs,omitempty"`
+	AppSignals                *int              `json:"as,omitempty"`
+	EnhancedContainerInsights *int              `json:"eci,omitempty"`
+	RunningInContainer        *int              `json:"ric,omitempty"`
+	RegionType                *string           `json:"rt,omitempty"`
+	Mode                      *string           `json:"m,omitempty"`
+	EntityRejected            *int              `json:"ent,omitempty"`
+	StatusCodes               map[string][5]int `json:"codes,omitempty"` //represents status codes 200,400,408,413,429,
 }
 
 // Merge the other Stats into the current. If the field is not nil,
@@ -80,6 +81,28 @@ func (s *Stats) Merge(other Stats) {
 	if other.EntityRejected != nil {
 		s.EntityRejected = other.EntityRejected
 	}
+	if other.StatusCodes == nil {
+		return
+	}
+
+	if s.StatusCodes == nil {
+		s.StatusCodes = make(map[string][5]int)
+	}
+
+	for key, value := range other.StatusCodes {
+		if existing, ok := s.StatusCodes[key]; ok {
+			s.StatusCodes[key] = [5]int{
+				existing[0] + value[0], // 200
+				existing[1] + value[1], // 400
+				existing[2] + value[2], // 408
+				existing[3] + value[3], // 413
+				existing[4] + value[4], // 429
+			}
+		} else {
+			s.StatusCodes[key] = value
+		}
+	}
+
 }
 
 func (s *Stats) Marshal() (string, error) {
@@ -104,6 +127,26 @@ func (of OperationsFilter) IsAllowed(operationName string) bool {
 	return of.allowAll || of.operations.Contains(operationName)
 }
 
+var StatusCodeOperations = []string{ // all the operations that are allowed
+	"DescribeInstances",
+	"DescribeTags",
+	"DescribeVolumes",
+	"DescribeContainerInstances",
+	"DescribeServices",
+	"DescribeTaskDefinition",
+	"ListServices",
+	"ListTasks",
+	"CreateLogGroup",
+	"CreateLogStream",
+}
+
+type StatsConfig struct {
+	// Operations are the allowed operation names to gather stats for.
+	Operations []string `mapstructure:"operations,omitempty"`
+	// UsageFlags are the usage flags to set on start up.
+	UsageFlags map[Flag]any `mapstructure:"usage_flags,omitempty"`
+}
+
 func NewOperationsFilter(operations ...string) OperationsFilter {
 	allowed := collections.NewSet[string](operations...)
 	return OperationsFilter{
@@ -112,9 +155,21 @@ func NewOperationsFilter(operations ...string) OperationsFilter {
 	}
 }
 
-type StatsConfig struct {
-	// Operations are the allowed operation names to gather stats for.
-	Operations []string `mapstructure:"operations,omitempty"`
-	// UsageFlags are the usage flags to set on start up.
-	UsageFlags map[Flag]any `mapstructure:"usage_flags,omitempty"`
+// NewStatusCodeOperationsFilter creates a new filter for allowed operations and status codes.
+func NewStatusCodeOperationsFilter() OperationsFilter {
+	allowed := make(map[string]struct{}, len(StatusCodeOperations))
+
+	return OperationsFilter{
+		operations: allowed,
+		allowAll:   false,
+	}
+}
+
+func NewStatusCodeAndOtherOperationsFilter(operations []string) OperationsFilter {
+	filter := NewStatusCodeOperationsFilter()
+	for _, operation := range operations {
+		filter.operations.Add(operation)
+	}
+
+	return filter
 }

extension/agenthealth/handler/stats/agent/agent_test.go

@@ -52,6 +52,76 @@ func TestMerge(t *testing.T) {
 	assert.EqualValues(t, "Mode", *stats.Mode)
 }
 
+func TestMergeWithStatusCodes(t *testing.T) {
+	stats := &Stats{
+		StatusCodes: map[string][5]int{
+			"operation1": {1, 2, 3, 4, 5},
+		},
+	}
+
+	stats.Merge(Stats{
+		StatusCodes: map[string][5]int{
+			"operation1": {2, 3, 4, 5, 6}, // Existing operation with new values
+			"operation2": {0, 1, 2, 3, 4}, // New operation
+		},
+	})
+
+	assert.Equal(t, [5]int{3, 5, 7, 9, 11}, stats.StatusCodes["operation1"]) // Values should sum
+	assert.Equal(t, [5]int{0, 1, 2, 3, 4}, stats.StatusCodes["operation2"])  // New operation added
+
+	stats.Merge(Stats{
+		StatusCodes: nil,
+	})
+
+	assert.Equal(t, [5]int{3, 5, 7, 9, 11}, stats.StatusCodes["operation1"])
+	assert.Equal(t, [5]int{0, 1, 2, 3, 4}, stats.StatusCodes["operation2"])
+}
+
+func TestMarshalWithStatusCodes(t *testing.T) {
+	testCases := map[string]struct {
+		stats *Stats
+		want  string
+	}{
+		"WithEmptyStatusCodes": {
+			stats: &Stats{
+				StatusCodes: map[string][5]int{},
+			},
+			want: "",
+		},
+		"WithStatusCodes": {
+			stats: &Stats{
+				StatusCodes: map[string][5]int{
+					"operation1": {1, 2, 3, 4, 5},
+					"operation2": {0, 1, 2, 3, 4},
+				},
+			},
+			want: `"codes":{"operation1":[1,2,3,4,5],"operation2":[0,1,2,3,4]}`,
+		},
+	}
+	for name, testCase := range testCases {
+		t.Run(name, func(t *testing.T) {
+			got, err := testCase.stats.Marshal()
+			assert.NoError(t, err)
+			assert.Contains(t, got, testCase.want)
+		})
+	}
+}
+
+func TestMergeFullWithStatusCodes(t *testing.T) {
+	stats := &Stats{
+		CpuPercent:  aws.Float64(1.0),
+		StatusCodes: map[string][5]int{"operation1": {1, 0, 0, 0, 0}},
+	}
+	stats.Merge(Stats{
+		CpuPercent:  aws.Float64(2.0),
+		StatusCodes: map[string][5]int{"operation1": {0, 1, 0, 0, 0}, "operation2": {1, 1, 1, 1, 1}},
+	})
+
+	assert.Equal(t, 2.0, *stats.CpuPercent)
+	assert.Equal(t, [5]int{1, 1, 0, 0, 0}, stats.StatusCodes["operation1"])
+	assert.Equal(t, [5]int{1, 1, 1, 1, 1}, stats.StatusCodes["operation2"])
+}
+
 func TestMarshal(t *testing.T) {
 	testCases := map[string]struct {
 		stats *Stats

extension/agenthealth/handler/stats/handler.go

@@ -21,12 +21,36 @@ const (
 	headerKeyAgentStats = "X-Amz-Agent-Stats"
 )
 
-func NewHandlers(logger *zap.Logger, cfg agent.StatsConfig) ([]awsmiddleware.RequestHandler, []awsmiddleware.ResponseHandler) {
-	filter := agent.NewOperationsFilter(cfg.Operations...)
-	clientStats := client.NewHandler(filter)
-	stats := newStatsHandler(logger, filter, []agent.StatsProvider{clientStats, provider.GetProcessStats(), provider.GetFlagsStats()})
+func NewHandlers(logger *zap.Logger, cfg agent.StatsConfig, statusCodeEnabled bool, agentStatsEnabled bool) ([]awsmiddleware.RequestHandler, []awsmiddleware.ResponseHandler) {
+	var requestHandlers []awsmiddleware.RequestHandler
+	var responseHandlers []awsmiddleware.ResponseHandler
+	var statsProviders []agent.StatsProvider
+
+	if !statusCodeEnabled && !agentStatsEnabled {
+		return nil, nil
+	}
+
+	statusCodeFilter := agent.NewStatusCodeOperationsFilter()
+	statusCodeStats := provider.GetStatusCodeStats(statusCodeFilter)
+	if statusCodeEnabled {
+		requestHandlers = append(requestHandlers, statusCodeStats)
+		responseHandlers = append(responseHandlers, statusCodeStats)
+		statsProviders = append(statsProviders, statusCodeStats)
+	}
+
+	if agentStatsEnabled {
+		clientStats := client.NewHandler(agent.NewOperationsFilter())
+		statsProviders = append(statsProviders, clientStats, provider.GetProcessStats(), provider.GetFlagsStats())
+		responseHandlers = append(responseHandlers, clientStats)
+		requestHandlers = append(requestHandlers, clientStats)
+
+	}
+	filter := agent.NewStatusCodeAndOtherOperationsFilter(cfg.Operations)
+	stats := newStatsHandler(logger, filter, statsProviders)
+	requestHandlers = append(requestHandlers, stats)
+
 	agent.UsageFlags().SetValues(cfg.UsageFlags)
-	return []awsmiddleware.RequestHandler{stats, clientStats}, []awsmiddleware.ResponseHandler{clientStats}
+	return requestHandlers, responseHandlers
 }
 
 type statsHandler struct {

extension/agenthealth/handler/stats/handler_test.go

@@ -40,6 +40,10 @@ func TestStatsHandler(t *testing.T) {
 		StatusCode:           aws.Int(200),
 		ImdsFallbackSucceed:  aws.Int(1),
 		SharedConfigFallback: aws.Int(1),
+		StatusCodes: map[string][5]int{
+			"pmd": {1, 0, 0, 0, 0},
+			"di":  {0, 1, 0, 0, 0},
+		},
 	}
 	handler := newStatsHandler(
 		zap.NewNop(),
@@ -59,15 +63,31 @@ func TestStatsHandler(t *testing.T) {
 	assert.Equal(t, "", req.Header.Get(headerKeyAgentStats))
 	handler.filter = agent.NewOperationsFilter(agent.AllowAllOperations)
 	handler.HandleRequest(ctx, req)
-	assert.Equal(t, `"cpu":1.2,"mem":123,"fd":456,"th":789,"lat":1234,"load":5678,"code":200,"scfb":1,"ifs":1`, req.Header.Get(headerKeyAgentStats))
+	assert.Equal(t, `"cpu":1.2,"mem":123,"fd":456,"th":789,"lat":1234,"load":5678,"code":200,"scfb":1,"ifs":1,"codes":{"di":[0,1,0,0,0],"pmd":[1,0,0,0,0]}`, req.Header.Get(headerKeyAgentStats))
 	stats.StatusCode = aws.Int(404)
 	stats.LatencyMillis = nil
 	handler.HandleRequest(ctx, req)
-	assert.Equal(t, `"cpu":1.2,"mem":123,"fd":456,"th":789,"load":5678,"code":404,"scfb":1,"ifs":1`, req.Header.Get(headerKeyAgentStats))
+	assert.Equal(t, `"cpu":1.2,"mem":123,"fd":456,"th":789,"load":5678,"code":404,"scfb":1,"ifs":1,"codes":{"di":[0,1,0,0,0],"pmd":[1,0,0,0,0]}`, req.Header.Get(headerKeyAgentStats))
 }
 
-func TestNewHandlers(t *testing.T) {
-	requestHandlers, responseHandlers := NewHandlers(zap.NewNop(), agent.StatsConfig{})
+func TestNewHandlersWithStatusCodeOnly(t *testing.T) {
+	requestHandlers, responseHandlers := NewHandlers(zap.NewNop(), agent.StatsConfig{}, true, false)
+	assert.Len(t, requestHandlers, 2)
+	assert.Len(t, responseHandlers, 1)
+}
+func TestNewHandlersWithAgentStatsOnly(t *testing.T) {
+	requestHandlers, responseHandlers := NewHandlers(zap.NewNop(), agent.StatsConfig{}, false, true)
 	assert.Len(t, requestHandlers, 2)
 	assert.Len(t, responseHandlers, 1)
 }
+
+func TestNewHandlersWithStatusCodeAndAgenthStats(t *testing.T) {
+	requestHandlers, responseHandlers := NewHandlers(zap.NewNop(), agent.StatsConfig{}, true, true)
+	assert.Len(t, requestHandlers, 3)
+	assert.Len(t, responseHandlers, 2)
+}
+func TestNewHandlersWithoutStatusCodeAndAgenthStats(t *testing.T) {
+	requestHandlers, responseHandlers := NewHandlers(zap.NewNop(), agent.StatsConfig{}, false, false)
+	assert.Len(t, requestHandlers, 0)
+	assert.Len(t, responseHandlers, 0)
+}

extension/agenthealth/handler/stats/provider/process_test.go

@@ -11,8 +11,6 @@ import (
 
 	"github.com/shirou/gopsutil/v3/process"
 	"github.com/stretchr/testify/assert"
-
-	"github.com/aws/amazon-cloudwatch-agent/extension/agenthealth/handler/stats/agent"
 )
 
 type mockProcessMetrics struct {
@@ -77,6 +75,6 @@ func TestProcessStats(t *testing.T) {
 	mock.mu.Unlock()
 	provider.refresh()
 	assert.Eventually(t, func() bool {
-		return provider.getStats() == agent.Stats{}
+		return len(provider.getStats().StatusCodes) == 0 // map isn't comparable, so we check the length
 	}, 5*time.Millisecond, time.Millisecond)
 }