Adding Statuscode Handler to Agent Health Extension

cmd/amazon-cloudwatch-agent/amazon-cloudwatch-agent.go

@@ -245,7 +245,6 @@ func runAgent(ctx context.Context,
 	c.OutputFilters = outputFilters
 	c.InputFilters = inputFilters
 	c.AllowUnusedFields = true
-
 	err = loadTomlConfigIntoAgent(c)
 	if err != nil {
 		return err

extension/agenthealth/config.go

@@ -12,6 +12,7 @@ import (
 type Config struct {
 	IsUsageDataEnabled bool              `mapstructure:"is_usage_data_enabled"`
 	Stats              agent.StatsConfig `mapstructure:"stats"`
+	StatusCodeOnly     *bool             `mapstructure:"is_status_code_only,omitempty"`
 }
 
 var _ component.Config = (*Config)(nil)

extension/agenthealth/extension.go

@@ -24,11 +24,20 @@ 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
+	}
+
+	statusCodeOnly := false
+	if ah.cfg.StatusCodeOnly != nil {
+		statusCodeOnly = *ah.cfg.StatusCodeOnly
 	}
+	statsRequestHandlers, statsResponseHandlers := stats.NewHandlers(ah.logger, ah.cfg.Stats, statusCodeOnly)
+	requestHandlers = append(requestHandlers, statsRequestHandlers...)
+	responseHandlers = append(responseHandlers, statsResponseHandlers...)
+
 	return requestHandlers, responseHandlers
 }
 

extension/agenthealth/extension_test.go

@@ -20,7 +20,7 @@ func TestExtension(t *testing.T) {
 	assert.NoError(t, extension.Start(ctx, componenttest.NewNopHost()))
 	requestHandlers, responseHandlers := extension.Handlers()
 	// user agent, client stats, stats
-	assert.Len(t, requestHandlers, 3)
+	assert.Len(t, requestHandlers, 4)
 	// client stats
 	assert.Len(t, responseHandlers, 1)
 	cfg.IsUsageDataEnabled = false

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

@@ -5,6 +5,7 @@ package agent
 
 import (
 	"encoding/json"
+	"log"
 	"strings"
 
 	"github.com/aws/amazon-cloudwatch-agent/internal/util/collections"
@@ -15,21 +16,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 +82,45 @@ func (s *Stats) Merge(other Stats) {
 	if other.EntityRejected != nil {
 		s.EntityRejected = other.EntityRejected
 	}
+	if other.StatusCodes != nil {
+		log.Println("Merging status codes from another source.")
+
+		if s.StatusCodes == nil {
+			log.Println("Initializing status codes map as it was nil.")
+			s.StatusCodes = make(map[string][5]int)
+		}
+
+		for key, value := range other.StatusCodes {
+			log.Printf("Processing key: %s with value: 200=%d, 400=%d, 408=%d, 413=%d, 429=%d", key, value[0], value[1], value[2], value[3], value[4])
+
+			if existing, ok := s.StatusCodes[key]; ok {
+				log.Printf(
+					"Key %s already exists. Existing: 200=%d, 400=%d, 408=%d, 413=%d, 429=%d. Merging with: 200=%d, 400=%d, 408=%d, 413=%d, 429=%d",
+					key, existing[0], existing[1], existing[2], existing[3], existing[4],
+					value[0], value[1], value[2], value[3], value[4],
+				)
+
+				//Merge the values for each status code
+				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
+				}
+
+				log.Printf(
+					"Updated key %s: 200=%d, 400=%d, 408=%d, 413=%d, 429=%d",
+					key, s.StatusCodes[key][0], s.StatusCodes[key][1], s.StatusCodes[key][2], s.StatusCodes[key][3], s.StatusCodes[key][4],
+				)
+			} else {
+				log.Printf("Key %s does not exist. Adding it with: 200=%d, 400=%d, 408=%d, 413=%d, 429=%d", key, value[0], value[1], value[2], value[3], value[4])
+				s.StatusCodes[key] = value
+			}
+		}
+		log.Println("Merging of status codes completed.")
+	}
+
 }
 
 func (s *Stats) Marshal() (string, error) {
@@ -118,3 +159,39 @@ type StatsConfig struct {
 	// UsageFlags are the usage flags to set on start up.
 	UsageFlags map[Flag]any `mapstructure:"usage_flags,omitempty"`
 }
+
+var StatusCodeOperations = []string{ // all the operations that are allowed
+	"DescribeInstances",
+	"DescribeTags",
+	"DescribeVolumes",
+	"DescribeContainerInstances",
+	"DescribeServices",
+	"DescribeTaskDefinition",
+	"ListServices",
+	"ListTasks",
+	"CreateLogGroup",
+	"CreateLogStream",
+}
+
+// NewStatusCodeOperationsFilter creates a new filter for allowed operations and status codes.
+func NewStatusCodeOperationsFilter() OperationsFilter {
+	allowed := collections.NewSet[string](StatusCodeOperations...)
+
+	return OperationsFilter{
+		operations: allowed,
+		allowAll:   allowed.Contains(AllowAllOperations),
+	}
+}
+
+func NewStatusCodeAndOtherOperationsFilter(operations []string) OperationsFilter {
+	allowed := collections.NewSet[string](StatusCodeOperations...)
+
+	for _, operation := range operations {
+		allowed.Add(operation)
+	}
+
+	return OperationsFilter{
+		operations: allowed,
+		allowAll:   allowed.Contains(AllowAllOperations),
+	}
+}

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

@@ -10,105 +10,72 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestMerge(t *testing.T) {
-	stats := &Stats{CpuPercent: aws.Float64(1.2)}
-	assert.EqualValues(t, 1.2, *stats.CpuPercent)
-	assert.Nil(t, stats.MemoryBytes)
+func TestMergeWithStatusCodes(t *testing.T) {
+	stats := &Stats{
+		StatusCodes: map[string][5]int{
+			"operation1": {1, 2, 3, 4, 5},
+		},
+	}
+
 	stats.Merge(Stats{
-		CpuPercent:  aws.Float64(1.3),
-		MemoryBytes: aws.Uint64(123),
+		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.EqualValues(t, 1.3, *stats.CpuPercent)
-	assert.EqualValues(t, 123, *stats.MemoryBytes)
+
+	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{
-		CpuPercent:                aws.Float64(1.5),
-		MemoryBytes:               aws.Uint64(133),
-		FileDescriptorCount:       aws.Int32(456),
-		ThreadCount:               aws.Int32(789),
-		LatencyMillis:             aws.Int64(1234),
-		PayloadBytes:              aws.Int(5678),
-		StatusCode:                aws.Int(200),
-		SharedConfigFallback:      aws.Int(1),
-		ImdsFallbackSucceed:       aws.Int(1),
-		AppSignals:                aws.Int(1),
-		EnhancedContainerInsights: aws.Int(1),
-		RunningInContainer:        aws.Int(0),
-		RegionType:                aws.String("RegionType"),
-		Mode:                      aws.String("Mode"),
+		StatusCodes: nil,
 	})
-	assert.EqualValues(t, 1.5, *stats.CpuPercent)
-	assert.EqualValues(t, 133, *stats.MemoryBytes)
-	assert.EqualValues(t, 456, *stats.FileDescriptorCount)
-	assert.EqualValues(t, 789, *stats.ThreadCount)
-	assert.EqualValues(t, 1234, *stats.LatencyMillis)
-	assert.EqualValues(t, 5678, *stats.PayloadBytes)
-	assert.EqualValues(t, 200, *stats.StatusCode)
-	assert.EqualValues(t, 1, *stats.ImdsFallbackSucceed)
-	assert.EqualValues(t, 1, *stats.SharedConfigFallback)
-	assert.EqualValues(t, 1, *stats.AppSignals)
-	assert.EqualValues(t, 1, *stats.EnhancedContainerInsights)
-	assert.EqualValues(t, 0, *stats.RunningInContainer)
-	assert.EqualValues(t, "RegionType", *stats.RegionType)
-	assert.EqualValues(t, "Mode", *stats.Mode)
+
+	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 TestMarshal(t *testing.T) {
+func TestMarshalWithStatusCodes(t *testing.T) {
 	testCases := map[string]struct {
 		stats *Stats
 		want  string
 	}{
-		"WithEmpty": {
-			stats: &Stats{},
-			want:  "",
-		},
-		"WithPartial": {
+		"WithEmptyStatusCodes": {
 			stats: &Stats{
-				CpuPercent:   aws.Float64(1.2),
-				MemoryBytes:  aws.Uint64(123),
-				ThreadCount:  aws.Int32(789),
-				PayloadBytes: aws.Int(5678),
+				StatusCodes: map[string][5]int{},
 			},
-			want: `"cpu":1.2,"mem":123,"th":789,"load":5678`,
+			want: "",
 		},
-		"WithFull": {
+		"WithStatusCodes": {
 			stats: &Stats{
-				CpuPercent:          aws.Float64(1.2),
-				MemoryBytes:         aws.Uint64(123),
-				FileDescriptorCount: aws.Int32(456),
-				ThreadCount:         aws.Int32(789),
-				LatencyMillis:       aws.Int64(1234),
-				PayloadBytes:        aws.Int(5678),
-				StatusCode:          aws.Int(200),
-				ImdsFallbackSucceed: aws.Int(1),
+				StatusCodes: map[string][5]int{
+					"operation1": {1, 2, 3, 4, 5},
+					"operation2": {0, 1, 2, 3, 4},
+				},
 			},
-			want: `"cpu":1.2,"mem":123,"fd":456,"th":789,"lat":1234,"load":5678,"code":200,"ifs":1`,
+			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.Equal(t, testCase.want, got)
+			assert.Contains(t, got, testCase.want)
 		})
 	}
 }
 
-func TestOperationFilter(t *testing.T) {
-	testCases := map[string]struct {
-		allowedOperations []string
-		testOperations    []string
-		want              []bool
-	}{
-		"WithNoneAllowed": {allowedOperations: nil, testOperations: []string{"nothing", "is", "allowed"}, want: []bool{false, false, false}},
-		"WithSomeAllowed": {allowedOperations: []string{"are"}, testOperations: []string{"some", "are", "allowed"}, want: []bool{false, true, false}},
-		"WithAllAllowed":  {allowedOperations: []string{"*"}, testOperations: []string{"all", "are", "allowed"}, want: []bool{true, true, true}},
-	}
-	for name, testCase := range testCases {
-		t.Run(name, func(t *testing.T) {
-			filter := NewOperationsFilter(testCase.allowedOperations...)
-			for index, testOperation := range testCase.testOperations {
-				assert.Equal(t, testCase.want[index], filter.IsAllowed(testOperation))
-			}
-		})
+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"])
 }

extension/agenthealth/handler/stats/handler.go

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

extension/agenthealth/handler/stats/handler_test.go

@@ -37,9 +37,12 @@ func TestStatsHandler(t *testing.T) {
 		ThreadCount:          aws.Int32(789),
 		LatencyMillis:        aws.Int64(1234),
 		PayloadBytes:         aws.Int(5678),
-		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 +62,9 @@ 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,"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))
-}
-
-func TestNewHandlers(t *testing.T) {
-	requestHandlers, responseHandlers := NewHandlers(zap.NewNop(), agent.StatsConfig{})
-	assert.Len(t, requestHandlers, 2)
-	assert.Len(t, responseHandlers, 1)
+	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))
 }