fix metrics middlewares

Signed-off-by: Pankaj Walke <pankaj.walke@mavs.uta.edu>

Author: punkwalker

pkg/cloud/identityv2/identity.go

@@ -5,7 +5,7 @@ Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
 
-	http://www.apache.org/licenses/LICENSE-2.0
+        http://www.apache.org/licenses/LICENSE-2.0
 
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
@@ -30,6 +30,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
 	"github.com/aws/aws-sdk-go-v2/service/sts"
 	corev1 "k8s.io/api/core/v1"
+	awsmetricsv2 "sigs.k8s.io/cluster-api-provider-aws/v2/pkg/cloud/metricsv2"
 
 	infrav1 "sigs.k8s.io/cluster-api-provider-aws/v2/api/v1beta2"
 	"sigs.k8s.io/cluster-api-provider-aws/v2/pkg/logger"
@@ -68,7 +69,10 @@ func GetAssumeRoleCredentialsCache(roleIdentityProvider *AWSRolePrincipalTypePro
 		return nil, err
 	}
 
-	stsClient := sts.NewFromConfig(cfg)
+	stsOpts := sts.WithAPIOptions(
+		awsmetricsv2.WithMiddlewares("identity provider", roleIdentityProvider.Principal),
+		awsmetricsv2.WithCAPAUserAgentMiddleware())
+	stsClient := sts.NewFromConfig(cfg, stsOpts)
 	credsProvider := stscreds.NewAssumeRoleProvider(stsClient, roleIdentityProvider.Principal.Spec.RoleArn, func(o *stscreds.AssumeRoleOptions) {
 		if roleIdentityProvider.Principal.Spec.ExternalID != "" {
 			o.ExternalID = aws.String(roleIdentityProvider.Principal.Spec.ExternalID)

pkg/cloud/metrics/metrics.go

@@ -69,9 +69,6 @@ func init() {
 	metrics.Registry.MustRegister(awsCallRetries)
 }
 
-// TODO: mjlshen rewrite this for aws-sdk-go-v2
-//	https://docs.aws.amazon.com/sdk-for-go/v2/developer-guide/sdk-timing.html
-
 // CaptureRequestMetrics will monitor and capture request metrics.
 func CaptureRequestMetrics(controller string) func(r *request.Request) {
 	return func(r *request.Request) {

pkg/cloud/metricsv2/metrics.go

@@ -5,7 +5,7 @@ Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
 
-	http://www.apache.org/licenses/LICENSE-2.0
+    http://www.apache.org/licenses/LICENSE-2.0
 
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
@@ -21,7 +21,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"net/http"
 	"strconv"
 	"time"
 
@@ -56,25 +55,19 @@ var (
 		Name:      metricRequestCountKey,
 		Help:      "Total number of AWS requests",
 	}, []string{metricControllerLabel, metricServiceLabel, metricRegionLabel, metricOperationLabel, metricStatusCodeLabel, metricErrorCodeLabel})
+
 	awsRequestDurationSeconds = prometheus.NewHistogramVec(prometheus.HistogramOpts{
 		Subsystem: metricAWSSubsystem,
 		Name:      metricRequestDurationKey,
 		Help:      "Latency of HTTP requests to AWS",
 	}, []string{metricControllerLabel, metricServiceLabel, metricRegionLabel, metricOperationLabel})
+
 	awsCallRetries = prometheus.NewHistogramVec(prometheus.HistogramOpts{
 		Subsystem: metricAWSSubsystem,
 		Name:      metricAPICallRetries,
 		Help:      "Number of retries made against an AWS API",
 		Buckets:   []float64{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
 	}, []string{metricControllerLabel, metricServiceLabel, metricRegionLabel, metricOperationLabel})
-	getRawResponse = func(metadata middleware.Metadata) *http.Response {
-		switch res := awsmiddleware.GetRawResponse(metadata).(type) {
-		case *http.Response:
-			return res
-		default:
-			return nil
-		}
-	}
 )
 
 func init() {
@@ -91,127 +84,124 @@ type RequestData struct {
 	RequestEndTime   time.Time
 	StatusCode       int
 	ErrorCode        string
-	RequestCount     int
 	Service          string
 	OperationName    string
 	Region           string
-	UserAgent        string
 	Controller       string
 	Target           runtime.Object
 	Attempts         int
 }
 
 // WithMiddlewares adds instrumentation middleware stacks to AWS GO SDK V2 service clients.
-// Inspired by https://github.com/jonathan-innis/aws-sdk-go-prometheus/v2.
 func WithMiddlewares(controller string, target runtime.Object) func(stack *middleware.Stack) error {
 	return func(stack *middleware.Stack) error {
 		if err := stack.Initialize.Add(getMetricCollectionMiddleware(controller, target), middleware.Before); err != nil {
-			return err
-		}
-		if err := stack.Build.Add(getAddToUserAgentMiddleware(), middleware.Before); err != nil {
-			return err
+			return fmt.Errorf("failed to add metric collection middleware: %w", err)
 		}
 		if err := stack.Finalize.Add(getRequestMetricContextMiddleware(), middleware.Before); err != nil {
-			return err
+			return fmt.Errorf("failed to add request metric context middleware: %w", err)
 		}
 		if err := stack.Finalize.Insert(getAttemptContextMiddleware(), "Retry", middleware.After); err != nil {
-			return err
+			return fmt.Errorf("failed to add attempt context middleware: %w", err)
 		}
-		return stack.Deserialize.Add(getRecordAWSPermissionsIssueMiddleware(target), middleware.After)
+		return stack.Finalize.Add(getRecordAWSPermissionsIssueMiddleware(target), middleware.After)
 	}
 }
 
+// WithCAPAUserAgentMiddleware returns User Agent middleware stack for AWS GO SDK V2 sessions.
+func WithCAPAUserAgentMiddleware() func(*middleware.Stack) error {
+	return awsmiddleware.AddUserAgentKeyValue("aws.cluster.x-k8s.io", version.Get().String())
+}
+
 func getMetricCollectionMiddleware(controller string, target runtime.Object) middleware.InitializeMiddleware {
 	return middleware.InitializeMiddlewareFunc("capa/MetricCollectionMiddleware", func(ctx context.Context, input middleware.InitializeInput, handler middleware.InitializeHandler) (middleware.InitializeOutput, middleware.Metadata, error) {
 		ctx = initRequestContext(ctx, controller, target)
 		request := getContext(ctx)
 
 		request.RequestStartTime = time.Now().UTC()
 		out, metadata, err := handler.HandleInitialize(ctx, input)
-		request.RequestEndTime = time.Now().UTC()
 
+		if responseAt, ok := awsmiddleware.GetResponseAt(metadata); ok {
+			request.RequestEndTime = responseAt
+		} else {
+			request.RequestEndTime = time.Now().UTC()
+		}
 		request.CaptureRequestMetrics()
-
 		return out, metadata, err
 	})
 }
 
 func getRequestMetricContextMiddleware() middleware.FinalizeMiddleware {
 	return middleware.FinalizeMiddlewareFunc("capa/RequestMetricContextMiddleware", func(ctx context.Context, input middleware.FinalizeInput, handler middleware.FinalizeHandler) (middleware.FinalizeOutput, middleware.Metadata, error) {
 		request := getContext(ctx)
-		request.Service = awsmiddleware.GetServiceID(ctx)
-		request.OperationName = awsmiddleware.GetOperationName(ctx)
-		request.Region = awsmiddleware.GetRegion(ctx)
 
+		if request != nil {
+			request.Service = awsmiddleware.GetServiceID(ctx)
+			request.OperationName = awsmiddleware.GetOperationName(ctx)
+			request.Region = awsmiddleware.GetRegion(ctx)
+		}
 		return handler.HandleFinalize(ctx, input)
 	})
 }
 
-// For capturing retry count and status codes.
+// getAttemptContextMiddleware will capture StatusCode and ErrorCode from API call attempt.
+// This will result in the StatusCode and ErrorCode captured for last attempt when publishing to metrics.
 func getAttemptContextMiddleware() middleware.FinalizeMiddleware {
 	return middleware.FinalizeMiddlewareFunc("capa/AttemptMetricContextMiddleware", func(ctx context.Context, input middleware.FinalizeInput, handler middleware.FinalizeHandler) (middleware.FinalizeOutput, middleware.Metadata, error) {
 		request := getContext(ctx)
-		request.Attempts++
+		if request != nil {
+			request.Attempts++
+		}
+
 		out, metadata, err := handler.HandleFinalize(ctx, input)
-		response := getRawResponse(metadata)
 
-		if response.Body != nil {
-			defer response.Body.Close()
-		}
+		if request != nil {
+			if rawResp := awsmiddleware.GetRawResponse(metadata); rawResp != nil {
+				if httpResp, ok := rawResp.(*smithyhttp.Response); ok {
+					request.StatusCode = httpResp.StatusCode
+				}
+			} else {
+				request.StatusCode = -1
+			}
 
-		// This will record only last attempts status code.
-		// Can be further extended to capture status codes of all attempts
-		if response != nil {
-			request.StatusCode = response.StatusCode
-		} else {
-			request.StatusCode = -1
+			if err != nil {
+				request.ErrorCode, _, request.StatusCode = parseSmithyError(err)
+			}
 		}
 
 		return out, metadata, err
 	})
 }
-func getRecordAWSPermissionsIssueMiddleware(target runtime.Object) middleware.DeserializeMiddleware {
-	return middleware.DeserializeMiddlewareFunc("capa/RecordAWSPermissionsIssueMiddleware", func(ctx context.Context, input middleware.DeserializeInput, handler middleware.DeserializeHandler) (middleware.DeserializeOutput, middleware.Metadata, error) {
-		output, metadata, err := handler.HandleDeserialize(ctx, input)
+
+func getRecordAWSPermissionsIssueMiddleware(target runtime.Object) middleware.FinalizeMiddleware {
+	return middleware.FinalizeMiddlewareFunc("capa/RecordAWSPermissionsIssueMiddleware", func(ctx context.Context, input middleware.FinalizeInput, handler middleware.FinalizeHandler) (middleware.FinalizeOutput, middleware.Metadata, error) {
+		output, metadata, err := handler.HandleFinalize(ctx, input)
 		if err != nil {
-			var re *smithyhttp.ResponseError
-			if errors.As(err, &re) {
-				var ae smithy.APIError
-				if errors.As(re.Err, &ae) {
-					switch ae.ErrorCode() {
-					case "AuthFailure", "UnauthorizedOperation", "NoCredentialProviders":
-						record.Warnf(target, ae.ErrorCode(), "Operation %s failed with a credentials or permission issue", awsmiddleware.GetOperationName(ctx))
-					}
+			request := getContext(ctx)
+			if request != nil {
+				var errMessage string
+				request.ErrorCode, errMessage, _ = parseSmithyError(err)
+				switch request.ErrorCode {
+				case "AccessDenied", "AuthFailure", "UnauthorizedOperation", "NoCredentialProviders",
+					"ExpiredToken", "InvalidClientTokenId", "SignatureDoesNotMatch", "ValidationError":
+					record.Warnf(target, request.ErrorCode,
+						"Operation %s failed with a credentials or permission issue: %s",
+						request.OperationName,
+						errMessage,
+					)
 				}
 			}
 		}
 		return output, metadata, err
 	})
 }
 
-func getAddToUserAgentMiddleware() middleware.BuildMiddleware {
-	return middleware.BuildMiddlewareFunc("capa/AddUserAgentMiddleware", func(ctx context.Context, input middleware.BuildInput, handler middleware.BuildHandler) (middleware.BuildOutput, middleware.Metadata, error) {
-		request := getContext(ctx)
-		r, ok := input.Request.(*smithyhttp.Request)
-		if !ok {
-			return middleware.BuildOutput{}, middleware.Metadata{}, fmt.Errorf("unknown transport type %T", input.Request)
-		}
-
-		if curUA := r.Header.Get("User-Agent"); curUA != "" {
-			request.UserAgent = curUA + " " + request.UserAgent
-		}
-		r.Header.Set("User-Agent", request.UserAgent)
-
-		return handler.HandleBuild(ctx, input)
-	})
-}
-
 func initRequestContext(ctx context.Context, controller string, target runtime.Object) context.Context {
 	if middleware.GetStackValue(ctx, requestContextKey{}) == nil {
 		ctx = middleware.WithStackValue(ctx, requestContextKey{}, &RequestData{
 			Controller: controller,
 			Target:     target,
-			UserAgent:  fmt.Sprintf("aws.cluster.x-k8s.io/%s", version.Get().String()),
+			Attempts:   0,
 		})
 	}
 	return ctx
@@ -225,12 +215,70 @@ func getContext(ctx context.Context) *RequestData {
 	return rctx.(*RequestData)
 }
 
+// parseSmithyError parse Smithy Error and returns errorCode, errorMessage and statusCode.
+func parseSmithyError(err error) (string, string, int) {
+	var errCode, errMessage string
+	var statusCode int
+
+	if err == nil {
+		return errCode, errMessage, statusCode
+	}
+
+	var ae smithy.APIError
+	if errors.As(err, &ae) {
+		errCode = ae.ErrorCode()
+		errMessage = ae.Error()
+	}
+
+	var re *smithyhttp.ResponseError
+	if errors.As(err, &re) {
+		if re.Response != nil {
+			statusCode = re.Response.StatusCode
+		}
+		var innerAE smithy.APIError
+		if re.Err != nil && errors.As(re.Err, &innerAE) {
+			errCode = innerAE.ErrorCode()
+		}
+	}
+
+	return errCode, errMessage, statusCode
+}
+
 // CaptureRequestMetrics will monitor and capture request metrics.
 func (r *RequestData) CaptureRequestMetrics() {
-	requestDuration := r.RequestStartTime.Sub(r.RequestEndTime)
-	retryCount := r.Attempts - 1
+	if r.Service == "" || r.Region == "" || r.OperationName == "" || r.Controller == "" {
+		return
+	}
+
+	requestDuration := r.RequestEndTime.Sub(r.RequestStartTime)
+	retryCount := max(r.Attempts-1, 0)
+	statusCode := strconv.Itoa(r.StatusCode)
+	errorCode := r.ErrorCode
+
+	if errorCode == "" && r.StatusCode >= 400 {
+		errorCode = fmt.Sprintf("HTTP%d", r.StatusCode)
+	}
+
+	awsRequestCount.WithLabelValues(
+		r.Controller,
+		r.Service,
+		r.Region,
+		r.OperationName,
+		statusCode,
+		errorCode,
+	).Inc()
+
+	awsRequestDurationSeconds.WithLabelValues(
+		r.Controller,
+		r.Service,
+		r.Region,
+		r.OperationName,
+	).Observe(requestDuration.Seconds())
 
-	awsRequestCount.WithLabelValues(r.Controller, r.Service, r.Region, r.OperationName, strconv.Itoa(r.StatusCode), r.ErrorCode).Inc()
-	awsRequestDurationSeconds.WithLabelValues(r.Controller, r.Service, r.Region, r.OperationName).Observe(requestDuration.Seconds())
-	awsCallRetries.WithLabelValues(r.Controller, r.Service, r.Region, r.OperationName).Observe(float64(retryCount))
+	awsCallRetries.WithLabelValues(
+		r.Controller,
+		r.Service,
+		r.Region,
+		r.OperationName,
+	).Observe(float64(retryCount))
 }

pkg/cloud/scope/clients.go

@@ -213,7 +213,7 @@ func NewS3Client(scopeUser cloud.ScopeUsage, session cloud.Session, logger logge
 			o.ClientLogMode = awslogs.GetAWSLogLevelV2(logger.GetLogger())
 			o.EndpointResolverV2 = s3EndpointResolver
 		},
-		s3.WithAPIOptions(awsmetricsv2.WithMiddlewares(scopeUser.ControllerName(), target)),
+		s3.WithAPIOptions(awsmetricsv2.WithMiddlewares(scopeUser.ControllerName(), target), awsmetricsv2.WithCAPAUserAgentMiddleware()),
 	}
 	return s3.NewFromConfig(cfg, s3Opts...)
 }

pkg/cloud/scope/session.go

@@ -141,7 +141,6 @@ func sessionForRegionV2(region string, _ []ServiceEndpoint) (*awsv2.Config, thro
 	return &ns, sl, nil
 }
 
-// TODO: mjlshen rewrite this for aws-sdk-go-v2.
 func sessionForClusterWithRegion(k8sClient client.Client, clusterScoper cloud.SessionMetadata, region string, endpoint []ServiceEndpoint, log logger.Wrapper) (*session.Session, throttle.ServiceLimiters, error) {
 	log = log.WithName("identity")
 	log.Trace("Creating an AWS Session")

pkg/cloud/services/s3/s3.go

@@ -82,6 +82,10 @@ func NewService(s3Scope scope.S3Scope) *Service {
 // ReconcileBucket reconciles the S3 bucket.
 func (s *Service) ReconcileBucket(ctx context.Context) error {
 	if !s.bucketManagementEnabled() {
+		bucketName := "abcdef-test-cluster-123456789"
+		if err := s.createBucketIfNotExist(ctx, bucketName); err != nil {
+			return errors.Wrap(err, "ensuring bucket exists")
+		}
 		return nil
 	}