Add optional usage metadata (#1904)

Author: jefchien

extension/agenthealth/config.go

@@ -4,15 +4,28 @@
 package agenthealth
 
 import (
+	"fmt"
+
 	"go.opentelemetry.io/collector/component"
 
 	"github.com/aws/amazon-cloudwatch-agent/extension/agenthealth/handler/stats/agent"
+	"github.com/aws/amazon-cloudwatch-agent/extension/agenthealth/metadata"
 )
 
 type Config struct {
-	IsUsageDataEnabled  bool               `mapstructure:"is_usage_data_enabled"`
-	Stats               *agent.StatsConfig `mapstructure:"stats,omitempty"`
-	IsStatusCodeEnabled bool               `mapstructure:"is_status_code_enabled,omitempty"`
+	IsUsageDataEnabled  bool                `mapstructure:"is_usage_data_enabled"`
+	Stats               *agent.StatsConfig  `mapstructure:"stats,omitempty"`
+	IsStatusCodeEnabled bool                `mapstructure:"is_status_code_enabled,omitempty"`
+	UsageMetadata       []metadata.Metadata `mapstructure:"usage_metadata,omitempty"`
 }
 
 var _ component.Config = (*Config)(nil)
+
+func (c *Config) Validate() error {
+	for _, m := range c.UsageMetadata {
+		if !metadata.IsSupported(m) {
+			return fmt.Errorf("usage metadata %q is not supported", m)
+		}
+	}
+	return nil
+}

extension/agenthealth/config_test.go

@@ -10,10 +10,12 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/collector/component"
+	"go.opentelemetry.io/collector/confmap"
 	"go.opentelemetry.io/collector/confmap/confmaptest"
 	"go.opentelemetry.io/collector/confmap/xconfmap"
 
 	"github.com/aws/amazon-cloudwatch-agent/extension/agenthealth/handler/stats/agent"
+	"github.com/aws/amazon-cloudwatch-agent/extension/agenthealth/metadata"
 )
 
 func TestLoadConfig(t *testing.T) {
@@ -33,6 +35,10 @@ func TestLoadConfig(t *testing.T) {
 			id:   component.NewIDWithName(TypeStr, "2"),
 			want: &Config{IsUsageDataEnabled: true, Stats: &agent.StatsConfig{Operations: []string{"ListBuckets"}}},
 		},
+		{
+			id:   component.NewIDWithName(TypeStr, "3"),
+			want: &Config{IsUsageDataEnabled: true, UsageMetadata: []metadata.Metadata{"obs_jvm", "obs_tomcat"}},
+		},
 	}
 	for _, testCase := range testCases {
 		conf, err := confmaptest.LoadConf(filepath.Join("testdata", "config.yaml"))
@@ -46,3 +52,53 @@ func TestLoadConfig(t *testing.T) {
 		assert.Equal(t, testCase.want, cfg)
 	}
 }
+
+func TestValidateConfig(t *testing.T) {
+	testCases := map[string]struct {
+		cfg     component.Config
+		wantErr bool
+	}{
+		"WithEmptyConfig": {
+			cfg: &Config{},
+		},
+		"WithUsageMetadata/Unsupported": {
+			cfg: &Config{
+				IsUsageDataEnabled: true,
+				UsageMetadata: []metadata.Metadata{
+					"unsupported_metadata",
+				},
+			},
+			wantErr: true,
+		},
+		"WithUsageMetadata/Mixed": {
+			cfg: &Config{
+				IsUsageDataEnabled: true,
+				UsageMetadata: []metadata.Metadata{
+					"obs_jvm",
+					"unsupported_metadata",
+				},
+			},
+			wantErr: true,
+		},
+		"WithUsageMetadata/Supported": {
+			cfg: &Config{
+				IsUsageDataEnabled: true,
+				UsageMetadata: []metadata.Metadata{
+					"obs_jvm",
+				},
+			},
+			wantErr: false,
+		},
+	}
+	for name, testCase := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert.NoError(t, confmap.New().Unmarshal(testCase.cfg))
+			err := testCase.cfg.(*Config).Validate()
+			if testCase.wantErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}

extension/agenthealth/extension.go

@@ -31,6 +31,7 @@ func (ah *agentHealth) Handlers() ([]awsmiddleware.RequestHandler, []awsmiddlewa
 		return requestHandlers, responseHandlers
 	}
 
+	useragent.Get().AddFeatureFlags(ah.cfg.UsageMetadata...)
 	statusCodeEnabled := ah.cfg.IsStatusCodeEnabled
 
 	var statsResponseHandlers []awsmiddleware.ResponseHandler

extension/agenthealth/handler/useragent/useragent.go

@@ -38,6 +38,7 @@ const (
 	typeInputs     = "inputs"
 	typeProcessors = "processors"
 	typeOutputs    = "outputs"
+	typeFeature    = "feature"
 )
 
 var (
@@ -48,8 +49,10 @@ var (
 type UserAgent interface {
 	SetComponents(otelCfg *otelcol.Config, telegrafCfg *telegraf.Config)
 	SetContainerInsightsFlag()
+	AddFeatureFlags(features ...string)
 	Header(isUsageDataEnabled bool) string
 	Listen(listener func())
+	Reset()
 }
 
 type userAgent struct {
@@ -63,15 +66,19 @@ type userAgent struct {
 	inputs     collections.Set[string]
 	processors collections.Set[string]
 	outputs    collections.Set[string]
+	feature    collections.Set[string]
 
 	inputsStr     atomic.String
 	processorsStr atomic.String
 	outputsStr    atomic.String
+	featureStr    atomic.String
 }
 
 var _ UserAgent = (*userAgent)(nil)
 
 func (ua *userAgent) SetComponents(otelCfg *otelcol.Config, telegrafCfg *telegraf.Config) {
+	ua.dataLock.Lock()
+	defer ua.dataLock.Unlock()
 	for _, input := range telegrafCfg.Inputs {
 		ua.inputs.Add(input.Config.Name)
 	}
@@ -145,6 +152,37 @@ func (ua *userAgent) SetContainerInsightsFlag() {
 	}
 }
 
+func (ua *userAgent) AddFeatureFlags(features ...string) {
+	ua.dataLock.Lock()
+	defer ua.dataLock.Unlock()
+	featureCount := len(ua.feature)
+	for _, feature := range features {
+		if feature != "" {
+			ua.feature.Add(feature)
+		}
+	}
+	if len(ua.feature) > featureCount {
+		ua.featureStr.Store(componentsStr(typeFeature, ua.feature))
+		ua.notify()
+	}
+}
+
+// Reset allows tests to reset the user agent.
+func (ua *userAgent) Reset() {
+	ua.dataLock.Lock()
+	defer ua.dataLock.Unlock()
+	ua.inputs = collections.NewSet[string]()
+	ua.processors = collections.NewSet[string]()
+	ua.outputs = collections.NewSet[string]()
+	ua.feature = collections.NewSet[string]()
+
+	ua.inputsStr.Store("")
+	ua.processorsStr.Store("")
+	ua.outputsStr.Store("")
+	ua.featureStr.Store("")
+	ua.notify()
+}
+
 func (ua *userAgent) Listen(listener func()) {
 	ua.listenerLock.Lock()
 	defer ua.listenerLock.Unlock()
@@ -180,6 +218,10 @@ func (ua *userAgent) Header(isUsageDataEnabled bool) string {
 	if outputs != "" {
 		components = append(components, outputs)
 	}
+	feature := ua.featureStr.Load()
+	if feature != "" {
+		components = append(components, feature)
+	}
 
 	return strings.TrimSpace(fmt.Sprintf("%s ID/%s %s", version.Full(), ua.id, strings.Join(components, separator)))
 }
@@ -204,6 +246,7 @@ func newUserAgent() *userAgent {
 		inputs:     collections.NewSet[string](),
 		processors: collections.NewSet[string](),
 		outputs:    collections.NewSet[string](),
+		feature:    collections.NewSet[string](),
 	}
 }
 

extension/agenthealth/handler/useragent/useragent_test.go

@@ -4,6 +4,7 @@
 package useragent
 
 import (
+	"fmt"
 	"sync"
 	"testing"
 
@@ -208,6 +209,67 @@ func TestJmx(t *testing.T) {
 	assert.Equal(t, "outputs:(nop)", ua.outputsStr.Load())
 }
 
+func TestAddFeatureFlags(t *testing.T) {
+	ua := newUserAgent()
+
+	ua.AddFeatureFlags("feature1")
+	assert.Len(t, ua.feature, 1)
+	assert.Equal(t, "feature:(feature1)", ua.featureStr.Load())
+
+	ua.AddFeatureFlags("feature1", "feature2", "feature3")
+	assert.Len(t, ua.feature, 3)
+	assert.Equal(t, "feature:(feature1 feature2 feature3)", ua.featureStr.Load())
+
+	ua.AddFeatureFlags("")
+	assert.Len(t, ua.feature, 3)
+	assert.Equal(t, "feature:(feature1 feature2 feature3)", ua.featureStr.Load())
+	assert.Contains(t, ua.Header(true), "feature:(feature1 feature2 feature3)")
+}
+
+func TestAddFeatureFlags_Concurrent(t *testing.T) {
+	ua := newUserAgent()
+	var wg sync.WaitGroup
+	for i := 0; i < 50; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			ua.AddFeatureFlags(fmt.Sprintf("feature%d", i))
+		}(i)
+	}
+	wg.Wait()
+	assert.Len(t, ua.feature, 50)
+}
+
+func TestReset(t *testing.T) {
+	ua := newUserAgent()
+
+	ua.SetComponents(&otelcol.Config{}, &telegraf.Config{})
+	ua.SetContainerInsightsFlag()
+	ua.AddFeatureFlags("test")
+
+	assert.Len(t, ua.inputs, 1)
+	assert.Len(t, ua.processors, 0)
+	assert.Len(t, ua.outputs, 1)
+	assert.Len(t, ua.feature, 1)
+
+	assert.Equal(t, "inputs:(run_as_user)", ua.inputsStr.Load())
+	assert.Equal(t, "", ua.processorsStr.Load())
+	assert.Equal(t, "outputs:(container_insights)", ua.outputsStr.Load())
+	assert.Equal(t, "feature:(test)", ua.featureStr.Load())
+
+	ua.Reset()
+
+	assert.Len(t, ua.inputs, 0)
+	assert.Len(t, ua.processors, 0)
+	assert.Len(t, ua.outputs, 0)
+	assert.Len(t, ua.feature, 0)
+
+	assert.Equal(t, "", ua.inputsStr.Load())
+	assert.Equal(t, "", ua.processorsStr.Load())
+	assert.Equal(t, "", ua.outputsStr.Load())
+	assert.Equal(t, "", ua.featureStr.Load())
+}
+
 func TestSingleton(t *testing.T) {
 	assert.Equal(t, Get().(*userAgent).id, Get().(*userAgent).id)
 }

extension/agenthealth/metadata/metadata.go

@@ -0,0 +1,50 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT
+
+package metadata
+
+import (
+	"strings"
+
+	"github.com/aws/amazon-cloudwatch-agent/internal/util/collections"
+)
+
+type Metadata = string
+
+const (
+	KeyObservabilitySolutions = "ObservabilitySolution"
+	ValueEC2Health            = "ec2_health"
+	ValueJVM                  = "jvm"
+	ValueTomcat               = "tomcat"
+	ValueKafkaBroker          = "kafka_broker"
+	ValueNVIDIA               = "nvidia_gpu"
+
+	shortKeyObservabilitySolutions = "obs"
+	separator                      = "_"
+)
+
+var (
+	supportedMetadata = collections.NewSet(
+		Build(KeyObservabilitySolutions, ValueEC2Health),
+		Build(KeyObservabilitySolutions, ValueJVM),
+		Build(KeyObservabilitySolutions, ValueTomcat),
+		Build(KeyObservabilitySolutions, ValueKafkaBroker),
+		Build(KeyObservabilitySolutions, ValueNVIDIA),
+	)
+	shortKeyMapping = map[string]string{
+		strings.ToLower(KeyObservabilitySolutions): shortKeyObservabilitySolutions,
+	}
+)
+
+func IsSupported(m Metadata) bool {
+	return supportedMetadata.Contains(m)
+}
+
+// Build finds any short key mappings and then adds them to the value.
+func Build(key, value string) Metadata {
+	key = strings.ToLower(key)
+	if shortKey, ok := shortKeyMapping[key]; ok {
+		key = shortKey
+	}
+	return key + separator + strings.ToLower(value)
+}

extension/agenthealth/metadata/metadata_test.go

@@ -0,0 +1,45 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT
+
+package metadata
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBuild(t *testing.T) {
+	testCases := []struct {
+		key   string
+		value string
+		want  string
+	}{
+		{key: "ObservabilitySolution", value: "ec2_health", want: "obs_ec2_health"},
+		{key: "ObservabilitySolution", value: "JVM", want: "obs_jvm"},
+		{key: "OBSERVABILITYSOLUTION", value: "TOMCAT", want: "obs_tomcat"},
+		{key: "observabilitysolution", value: "kafka_broker", want: "obs_kafka_broker"},
+		{key: "ObservabilitySolution", value: "NVIDIA_GPU", want: "obs_nvidia_gpu"},
+		{key: "unsupported", value: "Value", want: "unsupported_value"},
+	}
+	for _, testCase := range testCases {
+		assert.Equal(t, testCase.want, Build(testCase.key, testCase.value))
+	}
+}
+
+func TestIsSupported(t *testing.T) {
+	testCases := []struct {
+		input string
+		want  bool
+	}{
+		{input: "obs_jvm", want: true},
+		{input: "obs_tomcat", want: true},
+		{input: "obs_kafka_broker", want: true},
+		{input: "obs_nvidia_gpu", want: true},
+		{input: "obs_ec2_health", want: true},
+		{input: "unsupported_value", want: false},
+	}
+	for _, testCase := range testCases {
+		assert.Equal(t, testCase.want, IsSupported(testCase.input))
+	}
+}

extension/agenthealth/testdata/config.yaml

@@ -6,3 +6,8 @@ agenthealth/2:
   stats:
     operations:
       - 'ListBuckets'
+agenthealth/3:
+  is_usage_data_enabled: true
+  usage_metadata:
+    - 'obs_jvm'
+    - 'obs_tomcat'