add SDK-specific feature tracking (#2682)

Author: lucix-aws

.changelog/85f058b6d0bb40eea8548e539473d404.json

@@ -0,0 +1,9 @@
+{
+    "id": "a048e0bc-5088-4e76-a29c-9524e4ee34eb",
+    "type": "feature",
+    "collapse": true,
+    "description": "Add framework for tracking specific features in user-agent string.",
+    "modules": [
+        "."
+    ]
+}

.changelog/a048e0bc50884e76a29c9524e4ee34eb.json

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"runtime"
+	"sort"
 	"strings"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -30,6 +31,7 @@ const (
 	FrameworkMetadata
 	AdditionalMetadata
 	ApplicationIdentifier
+	FeatureMetadata2
 )
 
 func (k SDKAgentKeyType) string() string {
@@ -50,6 +52,8 @@ func (k SDKAgentKeyType) string() string {
 		return "lib"
 	case ApplicationIdentifier:
 		return "app"
+	case FeatureMetadata2:
+		return "m"
 	case AdditionalMetadata:
 		fallthrough
 	default:
@@ -64,9 +68,29 @@ var validChars = map[rune]bool{
 	'-': true, '.': true, '^': true, '_': true, '`': true, '|': true, '~': true,
 }
 
+// UserAgentFeature enumerates tracked SDK features.
+type UserAgentFeature string
+
+// Enumerates UserAgentFeature.
+const (
+	UserAgentFeatureResourceModel          UserAgentFeature = "A" // n/a (we don't generate separate resource types)
+	UserAgentFeatureWaiter                                  = "B"
+	UserAgentFeaturePaginator                               = "C"
+	UserAgentFeatureRetryModeLegacy                         = "D" // n/a (equivalent to standard)
+	UserAgentFeatureRetryModeStandard                       = "E"
+	UserAgentFeatureRetryModeAdaptive                       = "F"
+	UserAgentFeatureS3Transfer                              = "G"
+	UserAgentFeatureS3CryptoV1N                             = "H" // n/a (crypto client is external)
+	UserAgentFeatureS3CryptoV2                              = "I" // n/a
+	UserAgentFeatureS3ExpressBucket                         = "J"
+	UserAgentFeatureS3AccessGrants                          = "K" // not yet implemented
+	UserAgentFeatureGZIPRequestCompression                  = "L"
+)
+
 // RequestUserAgent is a build middleware that set the User-Agent for the request.
 type RequestUserAgent struct {
 	sdkAgent, userAgent *smithyhttp.UserAgentBuilder
+	features            map[UserAgentFeature]struct{}
 }
 
 // NewRequestUserAgent returns a new requestUserAgent which will set the User-Agent and X-Amz-User-Agent for the
@@ -87,6 +111,7 @@ func NewRequestUserAgent() *RequestUserAgent {
 	r := &RequestUserAgent{
 		sdkAgent:  sdkAgent,
 		userAgent: userAgent,
+		features:  map[UserAgentFeature]struct{}{},
 	}
 
 	addSDKMetadata(r)
@@ -191,6 +216,12 @@ func (u *RequestUserAgent) AddUserAgentKeyValue(key, value string) {
 	u.userAgent.AddKeyValue(strings.Map(rules, key), strings.Map(rules, value))
 }
 
+// AddUserAgentFeature adds the feature ID to the tracking list to be emitted
+// in the final User-Agent string.
+func (u *RequestUserAgent) AddUserAgentFeature(feature UserAgentFeature) {
+	u.features[feature] = struct{}{}
+}
+
 // AddSDKAgentKey adds the component identified by name to the User-Agent string.
 func (u *RequestUserAgent) AddSDKAgentKey(keyType SDKAgentKeyType, key string) {
 	// TODO: should target sdkAgent
@@ -227,6 +258,9 @@ func (u *RequestUserAgent) HandleBuild(ctx context.Context, in middleware.BuildI
 func (u *RequestUserAgent) addHTTPUserAgent(request *smithyhttp.Request) {
 	const userAgent = "User-Agent"
 	updateHTTPHeader(request, userAgent, u.userAgent.Build())
+	if len(u.features) > 0 {
+		updateHTTPHeader(request, userAgent, buildFeatureMetrics(u.features))
+	}
 }
 
 func (u *RequestUserAgent) addHTTPSDKAgent(request *smithyhttp.Request) {
@@ -259,3 +293,13 @@ func rules(r rune) rune {
 		return '-'
 	}
 }
+
+func buildFeatureMetrics(features map[UserAgentFeature]struct{}) string {
+	fs := make([]string, 0, len(features))
+	for f := range features {
+		fs = append(fs, string(f))
+	}
+
+	sort.Strings(fs)
+	return fmt.Sprintf("%s/%s", FeatureMetadata2.string(), strings.Join(fs, ","))
+}

aws/middleware/user_agent.go

@@ -231,6 +231,71 @@ func TestAddUserAgentKeyValue(t *testing.T) {
 	}
 }
 
+func TestAddUserAgentFeature(t *testing.T) {
+	restoreEnv := clearEnv()
+	defer restoreEnv()
+
+	cases := map[string]struct {
+		Features []UserAgentFeature
+		Expect   string
+	}{
+		"none": {
+			Features: []UserAgentFeature{},
+			Expect:   expectedAgent,
+		},
+		"one": {
+			Features: []UserAgentFeature{
+				UserAgentFeatureWaiter,
+			},
+			Expect: "m/B " + expectedAgent,
+		},
+		"two": {
+			Features: []UserAgentFeature{
+				UserAgentFeatureRetryModeAdaptive, // ensure stable order, and idempotent
+				UserAgentFeatureRetryModeAdaptive,
+				UserAgentFeatureWaiter,
+			},
+			Expect: "m/B,F " + expectedAgent,
+		},
+	}
+
+	for name, c := range cases {
+		t.Run(name, func(t *testing.T) {
+			b := NewRequestUserAgent()
+			stack := middleware.NewStack("testStack", smithyhttp.NewStackRequest)
+			err := stack.Build.Add(b, middleware.After)
+			if err != nil {
+				t.Fatalf("expect no error, got %v", err)
+			}
+
+			for _, f := range c.Features {
+				b.AddUserAgentFeature(f)
+			}
+
+			in := middleware.BuildInput{
+				Request: &smithyhttp.Request{
+					Request: &http.Request{
+						Header: map[string][]string{},
+					},
+				},
+			}
+			_, _, err = b.HandleBuild(context.Background(), in, middleware.BuildHandlerFunc(func(ctx context.Context, input middleware.BuildInput) (o middleware.BuildOutput, m middleware.Metadata, err error) {
+				return o, m, err
+			}))
+			if err != nil {
+				t.Fatalf("expect no error, got %v", err)
+			}
+			ua, ok := in.Request.(*smithyhttp.Request).Header["User-Agent"]
+			if !ok {
+				t.Fatalf("expect User-Agent to be present")
+			}
+			if ua[0] != c.Expect {
+				t.Errorf("User-Agent did not match expected, %v != %v", c.Expect, ua[0])
+			}
+		})
+	}
+}
+
 func TestAddSDKAgentKey(t *testing.T) {
 	restoreEnv := clearEnv()
 	defer restoreEnv()

aws/middleware/user_agent_test.go

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.smithy.aws.go.codegen.customization;
+
+import java.util.Map;
+import java.util.Set;
+
+import software.amazon.smithy.aws.go.codegen.AwsGoDependency;
+import software.amazon.smithy.codegen.core.Symbol;
+import software.amazon.smithy.codegen.core.SymbolProvider;
+import software.amazon.smithy.go.codegen.GoDelegator;
+import software.amazon.smithy.go.codegen.GoSettings;
+import software.amazon.smithy.go.codegen.SmithyGoDependency;
+import software.amazon.smithy.go.codegen.integration.Paginators;
+import software.amazon.smithy.model.Model;
+
+import static software.amazon.smithy.go.codegen.GoWriter.goTemplate;
+import static software.amazon.smithy.go.codegen.SymbolUtils.buildPackageSymbol;
+
+/**
+ * Extends the base smithy Paginators integration to track in the User-Agent string.
+ */
+public class AwsPaginators extends Paginators {
+    @Override
+    public Set<Symbol> getAdditionalClientOptions() {
+        return Set.of(buildPackageSymbol("addIsPaginatorUserAgent"));
+    }
+
+    @Override
+    public void writeAdditionalFiles(GoSettings settings, Model model, SymbolProvider symbolProvider, GoDelegator goDelegator) {
+        super.writeAdditionalFiles(settings, model, symbolProvider, goDelegator);
+
+        goDelegator.useFileWriter("api_client.go", settings.getModuleName(), goTemplate("""
+                func addIsPaginatorUserAgent(o *Options) {
+                    o.APIOptions = append(o.APIOptions, func(stack $stack:P) error {
+                        ua, err := getOrAddRequestUserAgent(stack)
+                        if err != nil {
+                            return err
+                        }
+
+                        ua.AddUserAgentFeature($featurePaginator:T)
+                        return nil
+                    })
+                }""",
+                Map.of(
+                        "stack", SmithyGoDependency.SMITHY_MIDDLEWARE.struct("Stack"),
+                        "featurePaginator", AwsGoDependency.AWS_MIDDLEWARE.valueSymbol("UserAgentFeaturePaginator")
+                )));
+    }
+}

codegen/smithy-aws-go-codegen/src/main/java/software/amazon/smithy/aws/go/codegen/customization/AwsPaginators.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.smithy.aws.go.codegen.customization;
+
+import java.util.Map;
+import java.util.Set;
+
+import software.amazon.smithy.aws.go.codegen.AwsGoDependency;
+import software.amazon.smithy.codegen.core.Symbol;
+import software.amazon.smithy.codegen.core.SymbolProvider;
+import software.amazon.smithy.go.codegen.GoDelegator;
+import software.amazon.smithy.go.codegen.GoSettings;
+import software.amazon.smithy.go.codegen.SmithyGoDependency;
+import software.amazon.smithy.go.codegen.integration.Waiters;
+import software.amazon.smithy.model.Model;
+
+import static software.amazon.smithy.go.codegen.GoWriter.goTemplate;
+import static software.amazon.smithy.go.codegen.SymbolUtils.buildPackageSymbol;
+
+/**
+ * Extends the base smithy Waiters integration to track in the User-Agent string.
+ */
+public class AwsWaiters extends Waiters {
+    @Override
+    public Set<Symbol> getAdditionalClientOptions() {
+        return Set.of(buildPackageSymbol("addIsWaiterUserAgent"));
+    }
+
+    @Override
+    public void writeAdditionalFiles(GoSettings settings, Model model, SymbolProvider symbolProvider, GoDelegator goDelegator) {
+        super.writeAdditionalFiles(settings, model, symbolProvider, goDelegator);
+
+        goDelegator.useFileWriter("api_client.go", settings.getModuleName(), goTemplate("""
+                func addIsWaiterUserAgent(o *Options) {
+                    o.APIOptions = append(o.APIOptions, func(stack $stack:P) error {
+                        ua, err := getOrAddRequestUserAgent(stack)
+                        if err != nil {
+                            return err
+                        }
+
+                        ua.AddUserAgentFeature($featureWaiter:T)
+                        return nil
+                    })
+                }""",
+                Map.of(
+                        "stack", SmithyGoDependency.SMITHY_MIDDLEWARE.struct("Stack"),
+                        "featureWaiter", AwsGoDependency.AWS_MIDDLEWARE.valueSymbol("UserAgentFeatureWaiter")
+                )));
+    }
+}