Add agenthealth extension with user agent handler. (#919)

cfg/envconfig/envconfig.go

@@ -3,6 +3,12 @@
 
 package envconfig
 
+import (
+	"os"
+	"strconv"
+	"sync"
+)
+
 const (
 	//the following are the names of environment variables
 	HTTP_PROXY         = "HTTP_PROXY"
@@ -15,3 +21,22 @@ const (
 	CWAGENT_USAGE_DATA = "CWAGENT_USAGE_DATA"
 	IMDS_NUMBER_RETRY  = "IMDS_NUMBER_RETRY"
 )
+
+var (
+	usageDataEnabled bool
+	onceUsageData    sync.Once
+)
+
+// getUsageDataEnabled returns true for true or invalid
+// examples of invalid are not set env var, "", "invalid"
+func getUsageDataEnabled() bool {
+	ok, err := strconv.ParseBool(os.Getenv(CWAGENT_USAGE_DATA))
+	return ok || err != nil
+}
+
+func IsUsageDataEnabled() bool {
+	onceUsageData.Do(func() {
+		usageDataEnabled = getUsageDataEnabled()
+	})
+	return usageDataEnabled
+}

cfg/envconfig/envconfig_test.go

@@ -0,0 +1,23 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT
+
+package envconfig
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIsUsageDataEnabled(t *testing.T) {
+	assert.True(t, getUsageDataEnabled())
+
+	t.Setenv(CWAGENT_USAGE_DATA, "TRUE")
+	assert.True(t, getUsageDataEnabled())
+
+	t.Setenv(CWAGENT_USAGE_DATA, "INVALID")
+	assert.True(t, getUsageDataEnabled())
+
+	t.Setenv(CWAGENT_USAGE_DATA, "FALSE")
+	assert.False(t, getUsageDataEnabled())
+}

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

@@ -34,7 +34,8 @@ import (
 	configaws "github.com/aws/amazon-cloudwatch-agent/cfg/aws"
 	"github.com/aws/amazon-cloudwatch-agent/cfg/envconfig"
 	"github.com/aws/amazon-cloudwatch-agent/cmd/amazon-cloudwatch-agent/internal"
-	"github.com/aws/amazon-cloudwatch-agent/handlers/agentinfo"
+	"github.com/aws/amazon-cloudwatch-agent/extension/agenthealth/handler/useragent"
+	"github.com/aws/amazon-cloudwatch-agent/internal/version"
 	"github.com/aws/amazon-cloudwatch-agent/logs"
 	_ "github.com/aws/amazon-cloudwatch-agent/plugins"
 	"github.com/aws/amazon-cloudwatch-agent/profiler"
@@ -258,7 +259,7 @@ func runAgent(ctx context.Context,
 
 	logger.SetupLogging(logConfig)
 
-	log.Printf("I! Starting AmazonCloudWatchAgent %s\n", agentinfo.FullVersion())
+	log.Printf("I! Starting AmazonCloudWatchAgent %s\n", version.Full())
 	// Need to set SDK log level before plugins get loaded.
 	// Some aws.Config objects get created early and live forever which means
 	// we cannot change the sdk log level without restarting the Agent.
@@ -304,7 +305,7 @@ func runAgent(ctx context.Context,
 		// So just start Telegraf.
 		_, err = os.Stat(*fOtelConfig)
 		if errors.Is(err, os.ErrNotExist) {
-			agentinfo.SetComponents(&otelcol.Config{}, c)
+			useragent.Get().SetComponents(&otelcol.Config{}, c)
 			return ag.Run(ctx)
 		}
 	}
@@ -328,7 +329,7 @@ func runAgent(ctx context.Context,
 		return err
 	}
 
-	agentinfo.SetComponents(cfg, c)
+	useragent.Get().SetComponents(cfg, c)
 
 	params := getCollectorParams(factories, provider)
 
@@ -351,7 +352,7 @@ func getCollectorParams(factories otelcol.Factories, provider otelcol.ConfigProv
 		BuildInfo: component.BuildInfo{
 			Command:     "CWAgent",
 			Description: "CloudWatch Agent",
-			Version:     agentinfo.Version(),
+			Version:     version.Number(),
 		},
 	}
 	return params
@@ -453,7 +454,7 @@ func main() {
 	if len(args) > 0 {
 		switch args[0] {
 		case "version":
-			fmt.Println(agentinfo.FullVersion())
+			fmt.Println(version.Full())
 			return
 		case "config":
 			config.PrintSampleConfig(
@@ -492,7 +493,7 @@ func main() {
 		}
 		return
 	case *fVersion:
-		fmt.Println(agentinfo.FullVersion())
+		fmt.Println(version.Full())
 		return
 	case *fSampleConfig:
 		config.PrintSampleConfig(
@@ -637,7 +638,7 @@ func validateAgentFinalConfigAndPlugins(c *config.Config) error {
 
 	if *fSchemaTest {
 		//up to this point, the given config file must be valid
-		fmt.Println(agentinfo.FullVersion())
+		fmt.Println(version.Full())
 		fmt.Printf("The given config: %v is valid\n", *fTomlConfig)
 		os.Exit(0)
 	}

extension/agenthealth/config.go

@@ -0,0 +1,12 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT
+
+package agenthealth
+
+import "go.opentelemetry.io/collector/component"
+
+type Config struct {
+	IsUsageDataEnabled bool `mapstructure:"is_usage_data_enabled"`
+}
+
+var _ component.Config = (*Config)(nil)

extension/agenthealth/config_test.go

@@ -0,0 +1,41 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT
+
+package agenthealth
+
+import (
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/collector/component"
+	"go.opentelemetry.io/collector/confmap/confmaptest"
+)
+
+func TestLoadConfig(t *testing.T) {
+	testCases := []struct {
+		id   component.ID
+		want component.Config
+	}{
+		{
+			id:   component.NewID(TypeStr),
+			want: NewFactory().CreateDefaultConfig(),
+		},
+		{
+			id:   component.NewIDWithName(TypeStr, "1"),
+			want: &Config{IsUsageDataEnabled: false},
+		},
+	}
+	for _, testCase := range testCases {
+		conf, err := confmaptest.LoadConf(filepath.Join("testdata", "config.yaml"))
+		require.NoError(t, err)
+		cfg := NewFactory().CreateDefaultConfig()
+		sub, err := conf.Sub(testCase.id.String())
+		require.NoError(t, err)
+		require.NoError(t, component.UnmarshalConfig(sub, cfg))
+
+		assert.NoError(t, component.ValidateConfig(cfg))
+		assert.Equal(t, testCase.want, cfg)
+	}
+}

extension/agenthealth/extension.go

@@ -0,0 +1,31 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT
+
+package agenthealth
+
+import (
+	"github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware"
+	"go.opentelemetry.io/collector/component"
+	"go.uber.org/zap"
+
+	"github.com/aws/amazon-cloudwatch-agent/extension/agenthealth/handler/useragent"
+)
+
+type agentHealth struct {
+	logger *zap.Logger
+	cfg    *Config
+	component.StartFunc
+	component.ShutdownFunc
+}
+
+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)}
+	return requestHandlers, responseHandlers
+}
+
+func newAgentHealth(logger *zap.Logger, cfg *Config) (*agentHealth, error) {
+	return &agentHealth{logger: logger, cfg: cfg}, nil
+}

extension/agenthealth/extension_test.go

@@ -0,0 +1,30 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT
+
+package agenthealth
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.opentelemetry.io/collector/component/componenttest"
+	"go.uber.org/zap"
+)
+
+func TestExtension(t *testing.T) {
+	ctx := context.Background()
+	extension, err := newAgentHealth(zap.NewNop(), &Config{IsUsageDataEnabled: true})
+	assert.NoError(t, err)
+	assert.NotNil(t, extension)
+	assert.NoError(t, extension.Start(ctx, componenttest.NewNopHost()))
+	requests, responses := extension.Handlers()
+	assert.Len(t, requests, 1)
+	assert.Len(t, responses, 0)
+	extension.cfg.IsUsageDataEnabled = false
+	extension.Handlers()
+	requests, responses = extension.Handlers()
+	assert.Len(t, requests, 1)
+	assert.Len(t, responses, 0)
+	assert.NoError(t, extension.Shutdown(ctx))
+}

extension/agenthealth/factory.go

@@ -0,0 +1,34 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT
+
+package agenthealth
+
+import (
+	"context"
+
+	"go.opentelemetry.io/collector/component"
+	"go.opentelemetry.io/collector/extension"
+)
+
+const (
+	TypeStr = "agenthealth"
+)
+
+func NewFactory() extension.Factory {
+	return extension.NewFactory(
+		TypeStr,
+		createDefaultConfig,
+		createExtension,
+		component.StabilityLevelAlpha,
+	)
+}
+
+func createDefaultConfig() component.Config {
+	return &Config{
+		IsUsageDataEnabled: true,
+	}
+}
+
+func createExtension(_ context.Context, settings extension.CreateSettings, cfg component.Config) (extension.Extension, error) {
+	return newAgentHealth(settings.Logger, cfg.(*Config))
+}

extension/agenthealth/factory_test.go

@@ -0,0 +1,26 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT
+
+package agenthealth
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.opentelemetry.io/collector/component/componenttest"
+	"go.opentelemetry.io/collector/extension/extensiontest"
+)
+
+func TestCreateDefaultConfig(t *testing.T) {
+	cfg := NewFactory().CreateDefaultConfig()
+	assert.Equal(t, &Config{IsUsageDataEnabled: true}, cfg)
+	assert.NoError(t, componenttest.CheckConfigStruct(cfg))
+}
+
+func TestCreateExtension(t *testing.T) {
+	cfg := &Config{}
+	got, err := NewFactory().CreateExtension(context.Background(), extensiontest.NewNopCreateSettings(), cfg)
+	assert.NoError(t, err)
+	assert.NotNil(t, got)
+}

extension/agenthealth/handler/useragent/handler.go

@@ -0,0 +1,62 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT
+
+package useragent
+
+import (
+	"context"
+	"net/http"
+
+	"github.com/amazon-contributing/opentelemetry-collector-contrib/extension/awsmiddleware"
+	"go.uber.org/atomic"
+)
+
+type userAgentHandler struct {
+	userAgent          UserAgent
+	isUsageDataEnabled bool
+	header             *atomic.String
+}
+
+var _ awsmiddleware.RequestHandler = (*userAgentHandler)(nil)
+
+func (uah *userAgentHandler) ID() string {
+	return handlerID
+}
+
+func (uah *userAgentHandler) Position() awsmiddleware.HandlerPosition {
+	return awsmiddleware.After
+}
+
+// HandleRequest prepends the User-Agent header with the CloudWatch Agent's
+// user agent string.
+func (uah *userAgentHandler) HandleRequest(_ context.Context, r *http.Request) {
+	newHeader := uah.Header()
+	current := r.Header.Get(headerKeyUserAgent)
+	if current != "" {
+		newHeader += separator + current
+	}
+	r.Header.Set(headerKeyUserAgent, newHeader)
+}
+
+func (uah *userAgentHandler) Header() string {
+	return uah.header.Load()
+}
+
+func (uah *userAgentHandler) refreshHeader() {
+	uah.header.Store(uah.userAgent.Header(uah.isUsageDataEnabled))
+}
+
+func newHandler(userAgent UserAgent, isUsageDataEnabled bool) *userAgentHandler {
+	handler := &userAgentHandler{
+		userAgent:          userAgent,
+		header:             &atomic.String{},
+		isUsageDataEnabled: isUsageDataEnabled,
+	}
+	handler.refreshHeader()
+	userAgent.Listen(handler.refreshHeader)
+	return handler
+}
+
+func NewHandler(isUsageDataEnabled bool) awsmiddleware.RequestHandler {
+	return newHandler(Get(), isUsageDataEnabled)
+}