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)
+ }
+}