diff --git a/instrumentation/awsv2/awsv2.go b/instrumentation/awsv2/awsv2.go index ffe73e3a..9584361d 100644 --- a/instrumentation/awsv2/awsv2.go +++ b/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) diff --git a/instrumentation/awsv2/awsv2_test.go b/instrumentation/awsv2/awsv2_test.go index 038aa183..2b45a264 100644 --- a/instrumentation/awsv2/awsv2_test.go +++ b/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(` + + + Tried to create resource record set duplicate.example.com. type A, but it already exists + + b25f48e8-84fd-11e6-80d9-574e0c4664cb + `), + }, + + "error response": { + responseStatus: 404, + responseBody: []byte(` + + + Sender + MalformedXML + 1 validation error detected: Value null at 'route53#ChangeSet' failed to satisfy constraint: Member must not be null + + 1234567890A + + `), + }, + + "success response": { + responseStatus: 200, + responseBody: []byte(` + + + mockComment + mockID + + `), + }, + } + + 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) + } +}