Disabling sdk feature

README.md

@@ -151,6 +151,17 @@ func init() {
   seg.Close(nil)
 ```
 
+**Disabling XRay Tracing**
+
+XRay tracing can be disabled by setting up environment variable `AWS_XRAY_SDK_DISABLED` . Disabling XRay can be useful for specific use case like if customer wants to stop tracing in their test environment they can do so just by setting up the environment variable.
+
+
+
+```go
+  // Set environment variable TRUE to disable XRay
+  os.Setenv("AWS_XRAY_SDK_DISABLED", "TRUE")
+```
+
 **Capture**
 
 ```go

xray/segment.go

@@ -15,6 +15,7 @@ import (
 	"net/http"
 	"os"
 	"runtime"
+	"strings"
 	"sync/atomic"
 	"time"
 
@@ -61,6 +62,12 @@ func BeginSegment(ctx context.Context, name string) (context.Context, *Segment)
 }
 
 func BeginSegmentWithSampling(ctx context.Context, name string, r *http.Request, traceHeader *header.Header) (context.Context, *Segment) {
+	// If SDK is disabled then return with an empty segment
+	if SdkDisabled() {
+		seg := &Segment{}
+		return context.WithValue(ctx, ContextKey, seg), seg
+	}
+
 	seg := basicSegment(name, nil)
 
 	cfg := GetRecorder(ctx)
@@ -204,6 +211,12 @@ func (seg *Segment) assignConfiguration(cfg *Config) {
 
 // BeginSubsegment creates a subsegment for a given name and context.
 func BeginSubsegment(ctx context.Context, name string) (context.Context, *Segment) {
+	// If SDK is disabled then return with an empty segment
+	if SdkDisabled() {
+		seg := &Segment{}
+		return context.WithValue(ctx, ContextKey, seg), seg
+	}
+
 	if len(name) > 200 {
 		name = name[:200]
 	}
@@ -271,8 +284,19 @@ func NewSegmentFromHeader(ctx context.Context, name string, r *http.Request, h *
 	return con, seg
 }
 
+// Check if SDK is disabled
+func SdkDisabled() bool {
+	disableKey := os.Getenv("AWS_XRAY_SDK_DISABLED")
+	return strings.ToLower(disableKey) == "true"
+}
+
 // Close a segment.
 func (seg *Segment) Close(err error) {
+	// If SDK is disabled then return
+	if SdkDisabled() {
+		return
+	}
+
 	seg.Lock()
 	if seg.parent != nil {
 		logger.Debugf("Closing subsegment named %s", seg.Name)
@@ -298,6 +322,11 @@ func (seg *Segment) Close(err error) {
 
 // CloseAndStream closes a subsegment and sends it.
 func (seg *Segment) CloseAndStream(err error) {
+	// If SDK is disabled then return
+	if SdkDisabled() {
+		return
+	}
+
 	seg.Lock()
 	defer seg.Unlock()
 
@@ -472,6 +501,11 @@ func (seg *Segment) beforeEmitSubsegment(s *Segment) {
 
 // AddAnnotation allows adding an annotation to the segment.
 func (seg *Segment) AddAnnotation(key string, value interface{}) error {
+	// If SDK is disabled then return
+	if SdkDisabled() {
+		return nil
+	}
+
 	seg.Lock()
 	defer seg.Unlock()
 
@@ -495,6 +529,11 @@ func (seg *Segment) AddAnnotation(key string, value interface{}) error {
 
 // AddMetadata allows adding metadata to the segment.
 func (seg *Segment) AddMetadata(key string, value interface{}) error {
+	// If SDK is disabled then return
+	if SdkDisabled() {
+		return nil
+	}
+
 	seg.Lock()
 	defer seg.Unlock()
 
@@ -515,6 +554,11 @@ func (seg *Segment) AddMetadata(key string, value interface{}) error {
 
 // AddMetadataToNamespace allows adding a namespace into metadata for the segment.
 func (seg *Segment) AddMetadataToNamespace(namespace string, key string, value interface{}) error {
+	// If SDK is disabled then return
+	if SdkDisabled() {
+		return nil
+	}
+
 	seg.Lock()
 	defer seg.Unlock()
 

xray/segment_model.go

@@ -129,6 +129,12 @@ type SQLData struct {
 // DownstreamHeader returns a header for passing to downstream calls.
 func (s *Segment) DownstreamHeader() *header.Header {
 	r := &header.Header{}
+
+	// If SDK is disabled then return with an empty header
+	if SdkDisabled() {
+		return r
+	}
+
 	if parent := s.ParentSegment.IncomingHeader; parent != nil {
 		*r = *parent // copy parent incoming header
 	}

xray/segment_test.go

@@ -13,6 +13,7 @@ import (
 	"errors"
 	"net/http"
 	"net/url"
+	"os"
 	"sync"
 	"testing"
 	"time"
@@ -157,6 +158,62 @@ func TestSegment_isDummy(t *testing.T) {
 	assert.False(t, subSeg2.Dummy)
 }
 
+func TestSDKDisable_inOrder(t *testing.T) {
+	os.Setenv("AWS_XRAY_SDK_DISABLED", "TRue")
+	ctx, td := NewTestDaemon()
+	defer td.Close()
+	ctx, root := BeginSegment(ctx, "Segment")
+	ctxSubSeg1, subSeg1 := BeginSubsegment(ctx, "Subegment1")
+	_, subSeg2 := BeginSubsegment(ctxSubSeg1, "Subsegment2")
+	subSeg2.Close(nil)
+	subSeg1.Close(nil)
+	root.Close(nil)
+
+	assert.Equal(t, root, &Segment{})
+	assert.Equal(t, subSeg1, &Segment{})
+	assert.Equal(t, subSeg2, &Segment{})
+
+	os.Setenv("AWS_XRAY_SDK_DISABLED", "FALSE")
+}
+
+func TestSDKDisable_outOrder(t *testing.T) {
+	os.Setenv("AWS_XRAY_SDK_DISABLED", "TRUE")
+	ctx, td := NewTestDaemon()
+	defer td.Close()
+	_, subSeg := BeginSubsegment(ctx, "Subegment1")
+	_, seg := BeginSegment(context.Background(), "Segment")
+
+	subSeg.Close(nil)
+	seg.Close(nil)
+
+	assert.Equal(t, subSeg, &Segment{})
+	assert.Equal(t, seg, &Segment{})
+	os.Setenv("AWS_XRAY_SDK_DISABLED", "FALSE")
+}
+
+func TestSDKDisable_otherMethods(t *testing.T) {
+	os.Setenv("AWS_XRAY_SDK_DISABLED", "true")
+	ctx, td := NewTestDaemon()
+	defer td.Close()
+	ctx, seg := BeginSegment(ctx, "Segment")
+	_, subSeg := BeginSubsegment(ctx, "Subegment1")
+
+	if err := seg.AddAnnotation("key", "value"); err != nil {
+		return
+	}
+	if err := seg.AddMetadata("key", "value"); err != nil {
+		return
+	}
+	seg.DownstreamHeader()
+
+	subSeg.Close(nil)
+	seg.Close(nil)
+
+	assert.Equal(t, seg, &Segment{})
+	assert.Equal(t, subSeg, &Segment{})
+	os.Setenv("AWS_XRAY_SDK_DISABLED", "FALSE")
+}
+
 // Benchmarks
 func BenchmarkBeginSegment(b *testing.B) {
 	ctx, td := NewTestDaemon()