Fix panic when no parent segment is found with AWS SDK v2 (#316)

To avoid triggering a panic when a parent segment is not present, we
must call next middlewares instead of returning nil in the AWS SDK v2
middleware for X-Ray tracing.

Fixes #315.

Author: adamantike

instrumentation/awsv2/awsv2.go

@@ -28,7 +28,7 @@ func initializeMiddlewareAfter(stack *middleware.Stack) error {
 		// Start the subsegment
 		ctx, subseg := xray.BeginSubsegment(ctx, serviceName)
 		if subseg == nil {
-			return
+			return next.HandleInitialize(ctx, in)
 		}
 		subseg.Namespace = "aws"
 		subseg.GetAWS()["region"] = v2Middleware.GetRegion(ctx)
@@ -52,7 +52,11 @@ func deserializeMiddleware(stack *middleware.Stack) error {
 		ctx context.Context, in middleware.DeserializeInput, next middleware.DeserializeHandler) (
 		out middleware.DeserializeOutput, metadata middleware.Metadata, err error) {
 
-		subseg := ctx.Value(awsV2SubsegmentKey{}).(*xray.Segment)
+		subseg, ok := ctx.Value(awsV2SubsegmentKey{}).(*xray.Segment)
+		if !ok {
+			return next.HandleDeserialize(ctx, in)
+		}
+
 		in.Request.(*smithyhttp.Request).Header.Set(xray.TraceIDHeaderKey, subseg.DownstreamHeader().String())
 
 		out, metadata, err = next.HandleDeserialize(ctx, in)

instrumentation/awsv2/awsv2_test.go

@@ -21,6 +21,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/route53"
 	"github.com/aws/aws-sdk-go-v2/service/route53/types"
+	"github.com/aws/aws-xray-sdk-go/strategy/ctxmissing"
 	"github.com/aws/aws-xray-sdk-go/xray"
 )
 
@@ -151,3 +152,92 @@ func TestAWSV2(t *testing.T) {
 		time.Sleep(1 * time.Second)
 	}
 }
+
+func TestAWSV2WithoutSegment(t *testing.T) {
+	cases := map[string]struct {
+		responseStatus int
+		responseBody   []byte
+	}{
+		"fault response": {
+			responseStatus: 500,
+			responseBody: []byte(`<?xml version="1.0" encoding="UTF-8"?>
+		<InvalidChangeBatch xmlns="https://route53.amazonaws.com/doc/2013-04-01/">
+		<Messages>
+		  <Message>Tried to create resource record set duplicate.example.com. type A, but it already exists</Message>
+		</Messages>
+		<RequestId>b25f48e8-84fd-11e6-80d9-574e0c4664cb</RequestId>
+		</InvalidChangeBatch>`),
+		},
+
+		"error response": {
+			responseStatus: 404,
+			responseBody: []byte(`<?xml version="1.0"?>
+		<ErrorResponse xmlns="http://route53.amazonaws.com/doc/2016-09-07/">
+		<Error>
+		  <Type>Sender</Type>
+		  <Code>MalformedXML</Code>
+		  <Message>1 validation error detected: Value null at 'route53#ChangeSet' failed to satisfy constraint: Member must not be null</Message>
+		</Error>
+		<RequestId>1234567890A</RequestId>
+		</ErrorResponse>
+		`),
+		},
+
+		"success response": {
+			responseStatus: 200,
+			responseBody: []byte(`<?xml version="1.0" encoding="UTF-8"?>
+		<ChangeResourceRecordSetsResponse>
+			<ChangeInfo>
+			<Comment>mockComment</Comment>
+			<Id>mockID</Id>
+		</ChangeInfo>
+		</ChangeResourceRecordSetsResponse>`),
+		},
+	}
+
+	for name, c := range cases {
+		server := httptest.NewServer(http.HandlerFunc(
+			func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(c.responseStatus)
+				_, err := w.Write(c.responseBody)
+				if err != nil {
+					t.Fatal(err)
+				}
+			}))
+		defer server.Close()
+
+		t.Run(name, func(t *testing.T) {
+			// Ignore errors when segment cannot be found.
+			ctx, err := xray.ContextWithConfig(
+				context.Background(),
+				xray.Config{ContextMissingStrategy: ctxmissing.NewDefaultIgnoreErrorStrategy()},
+			)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			svc := route53.NewFromConfig(aws.Config{
+				EndpointResolver: aws.EndpointResolverFunc(func(service, region string) (aws.Endpoint, error) {
+					return aws.Endpoint{
+						URL:         server.URL,
+						SigningName: "route53",
+					}, nil
+				}),
+				Retryer: func() aws.Retryer {
+					return aws.NopRetryer{}
+				},
+			})
+
+			_, _ = svc.ChangeResourceRecordSets(ctx, &route53.ChangeResourceRecordSetsInput{
+				ChangeBatch: &types.ChangeBatch{
+					Changes: []types.Change{},
+					Comment: aws.String("mock"),
+				},
+				HostedZoneId: aws.String("zone"),
+			}, func(options *route53.Options) {
+				AWSV2Instrumentor(&options.APIOptions)
+			})
+		})
+		time.Sleep(1 * time.Second)
+	}
+}