Disabling sdk feature (#219)

* DummySegment and DummySubSegment creation

* Revert "DummySegment and DummySubSegment creation"

This reverts commit 4edb4b5.

* Revert "Revert "DummySegment and DummySubSegment creation""

This reverts commit 0442d84.

* ran gofmt command to remove golint warnings

* remove redundent declaration

* Added Dummysegment creation logic

* Added reviewdog suggestions

* modified code based on feedback

* New Implementation of Dummy Segment and Subsegments

* Remove unnecessary code

* Review changes

* removed isDummy method

* Added return statement end of the block

* added SDK disabling feature

* Added README instructions

* Added minor fix

* Added reviewdog fix

* Added disablekey in DownStreamHeader API

* Review changes

* Fixed tests

* added function for checking if sdk is disabled

* Modified BeginSubSegment disable logic

* Clean up code and minor chnages

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()