diff --git a/contrib/aws/aws-sdk-go/aws/aws.go b/contrib/aws/aws-sdk-go/aws/aws.go index 85d75a5cd3..2efbdb5192 100644 --- a/contrib/aws/aws-sdk-go/aws/aws.go +++ b/contrib/aws/aws-sdk-go/aws/aws.go @@ -19,9 +19,10 @@ import ( ) const ( - tagAWSAgent = "aws.agent" - tagAWSOperation = "aws.operation" - tagAWSRegion = "aws.region" + tagAWSAgent = "aws.agent" + tagAWSOperation = "aws.operation" + tagAWSRegion = "aws.region" + tagAWSRetryCount = "aws.retry_count" ) type handlers struct { @@ -49,6 +50,9 @@ func WrapSession(s *session.Session, opts ...Option) *session.Session { } func (h *handlers) Send(req *request.Request) { + if req.RetryCount != 0 { + return + } opts := []ddtrace.StartSpanOption{ tracer.SpanType(ext.SpanTypeHTTP), tracer.ServiceName(h.serviceName(req)), @@ -71,6 +75,7 @@ func (h *handlers) Complete(req *request.Request) { if !ok { return } + span.SetTag(tagAWSRetryCount, req.RetryCount) if req.HTTPResponse != nil { span.SetTag(ext.HTTPCode, strconv.Itoa(req.HTTPResponse.StatusCode)) } diff --git a/contrib/aws/aws-sdk-go/aws/aws_test.go b/contrib/aws/aws-sdk-go/aws/aws_test.go index bd50dd7b43..00a2141425 100644 --- a/contrib/aws/aws-sdk-go/aws/aws_test.go +++ b/contrib/aws/aws-sdk-go/aws/aws_test.go @@ -7,10 +7,12 @@ package aws import ( "context" + "errors" "testing" "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws/credentials" + "github.com/aws/aws-sdk-go/aws/request" "github.com/aws/aws-sdk-go/aws/session" "github.com/aws/aws-sdk-go/service/ec2" "github.com/aws/aws-sdk-go/service/s3" @@ -142,3 +144,35 @@ func TestAnalyticsSettings(t *testing.T) { assertRate(t, mt, 0.23, WithAnalyticsRate(0.23)) }) } + +func TestRetries(t *testing.T) { + cfg := aws.NewConfig(). + WithRegion("us-west-2"). + WithDisableSSL(true). + WithCredentials(credentials.AnonymousCredentials) + + session := WrapSession(session.Must(session.NewSession(cfg))) + expectedError := errors.New("an error") + session.Handlers.Send.PushBack(func(r *request.Request) { + r.Error = expectedError + r.Retryable = aws.Bool(true) + }) + + mt := mocktracer.Start() + defer mt.Stop() + + ctx := context.Background() + s3api := s3.New(session) + req, _ := s3api.GetObjectRequest(&s3.GetObjectInput{ + Bucket: aws.String("BUCKET"), + Key: aws.String("KEY"), + }) + req.SetContext(ctx) + err := req.Send() + + assert.Equal(t, 3, req.RetryCount) + assert.Same(t, expectedError, err) + assert.Len(t, mt.OpenSpans(), 0) + assert.Len(t, mt.FinishedSpans(), 1) + assert.Equal(t, mt.FinishedSpans()[0].Tag(tagAWSRetryCount), 3) +} diff --git a/ddtrace/mocktracer/mockspan_test.go b/ddtrace/mocktracer/mockspan_test.go index 826b02e5a9..7998036133 100644 --- a/ddtrace/mocktracer/mockspan_test.go +++ b/ddtrace/mocktracer/mockspan_test.go @@ -193,7 +193,7 @@ func TestSpanFinishTwice(t *testing.T) { func TestSpanWithID(t *testing.T) { spanID := uint64(123456789) - span := new(mocktracer).StartSpan("", tracer.WithSpanID(spanID)) + span := newMockTracer().StartSpan("", tracer.WithSpanID(spanID)) assert := assert.New(t) assert.Equal(spanID, span.Context().SpanID()) diff --git a/ddtrace/mocktracer/mocktracer.go b/ddtrace/mocktracer/mocktracer.go index 8161e2c5d5..9ab6ec7bbf 100644 --- a/ddtrace/mocktracer/mocktracer.go +++ b/ddtrace/mocktracer/mocktracer.go @@ -27,6 +27,9 @@ var _ Tracer = (*mocktracer)(nil) // Tracer exposes an interface for querying the currently running mock tracer. type Tracer interface { + // OpenSpans returns the set of started spans that have not been finished yet. + OpenSpans() []Span + // FinishedSpans returns the set of finished spans. FinishedSpans() []Span @@ -45,15 +48,22 @@ type Tracer interface { // to activate the mock tracer. When your test runs, use the returned // interface to query the tracer's state. func Start() Tracer { - var t mocktracer - internal.SetGlobalTracer(&t) + t := newMockTracer() + internal.SetGlobalTracer(t) internal.Testing = true - return &t + return t } type mocktracer struct { sync.RWMutex // guards below spans finishedSpans []Span + openSpans map[uint64]Span +} + +func newMockTracer() *mocktracer { + var t mocktracer + t.openSpans = make(map[uint64]Span) + return &t } // Stop deactivates the mock tracer and sets the active tracer to a no-op. @@ -67,7 +77,23 @@ func (t *mocktracer) StartSpan(operationName string, opts ...ddtrace.StartSpanOp for _, fn := range opts { fn(&cfg) } - return newSpan(t, operationName, &cfg) + span := newSpan(t, operationName, &cfg) + + t.Lock() + t.openSpans[span.SpanID()] = span + t.Unlock() + + return span +} + +func (t *mocktracer) OpenSpans() []Span { + t.RLock() + defer t.RUnlock() + spans := make([]Span, 0, len(t.openSpans)) + for _, s := range t.openSpans { + spans = append(spans, s) + } + return spans } func (t *mocktracer) FinishedSpans() []Span { @@ -79,12 +105,16 @@ func (t *mocktracer) FinishedSpans() []Span { func (t *mocktracer) Reset() { t.Lock() defer t.Unlock() + for k := range t.openSpans { + delete(t.openSpans, k) + } t.finishedSpans = nil } func (t *mocktracer) addFinishedSpan(s Span) { t.Lock() defer t.Unlock() + delete(t.openSpans, s.SpanID()) if t.finishedSpans == nil { t.finishedSpans = make([]Span, 0, 1) } diff --git a/ddtrace/mocktracer/mocktracer_test.go b/ddtrace/mocktracer/mocktracer_test.go index 0cac20ca9e..7cc3cacddd 100644 --- a/ddtrace/mocktracer/mocktracer_test.go +++ b/ddtrace/mocktracer/mocktracer_test.go @@ -36,8 +36,8 @@ func TestTracerStartSpan(t *testing.T) { startTime := time.Now() t.Run("with-service", func(t *testing.T) { - var mt mocktracer - parent := newSpan(&mt, "http.request", &ddtrace.StartSpanConfig{Tags: parentTags}) + mt := newMockTracer() + parent := newSpan(mt, "http.request", &ddtrace.StartSpanConfig{Tags: parentTags}) s, ok := mt.StartSpan( "db.query", tracer.ServiceName("my-service"), @@ -57,8 +57,8 @@ func TestTracerStartSpan(t *testing.T) { }) t.Run("inherit", func(t *testing.T) { - var mt mocktracer - parent := newSpan(&mt, "http.request", &ddtrace.StartSpanConfig{Tags: parentTags}) + mt := newMockTracer() + parent := newSpan(mt, "http.request", &ddtrace.StartSpanConfig{Tags: parentTags}) s, ok := mt.StartSpan("db.query", tracer.ChildOf(parent.Context())).(*mockspan) assert := assert.New(t) @@ -73,9 +73,11 @@ func TestTracerStartSpan(t *testing.T) { } func TestTracerFinishedSpans(t *testing.T) { - var mt mocktracer - parent := newSpan(&mt, "http.request", &ddtrace.StartSpanConfig{}) + mt := newMockTracer() + assert.Empty(t, mt.FinishedSpans()) + parent := mt.StartSpan("http.request") child := mt.StartSpan("db.query", tracer.ChildOf(parent.Context())) + assert.Empty(t, mt.FinishedSpans()) child.Finish() parent.Finish() found := 0 @@ -92,21 +94,45 @@ func TestTracerFinishedSpans(t *testing.T) { assert.Equal(t, 2, found) } -func TestTracerReset(t *testing.T) { - var mt mocktracer - mt.StartSpan("db.query").Finish() +func TestTracerOpenSpans(t *testing.T) { + mt := newMockTracer() + assert.Empty(t, mt.OpenSpans()) + parent := mt.StartSpan("http.request") + child := mt.StartSpan("db.query", tracer.ChildOf(parent.Context())) + + assert.Len(t, mt.OpenSpans(), 2) + assert.Contains(t, mt.OpenSpans(), parent) + assert.Contains(t, mt.OpenSpans(), child) + + child.Finish() + assert.Len(t, mt.OpenSpans(), 1) + assert.NotContains(t, mt.OpenSpans(), child) + parent.Finish() + assert.Empty(t, mt.OpenSpans()) +} + +func TestTracerReset(t *testing.T) { assert := assert.New(t) + mt := newMockTracer() + + span := mt.StartSpan("parent") + _ = mt.StartSpan("child", tracer.ChildOf(span.Context())) + assert.Len(mt.openSpans, 2) + + span.Finish() assert.Len(mt.finishedSpans, 1) + assert.Len(mt.openSpans, 1) mt.Reset() - assert.Nil(mt.finishedSpans) + assert.Empty(mt.finishedSpans) + assert.Empty(mt.openSpans) } func TestTracerInject(t *testing.T) { t.Run("errors", func(t *testing.T) { - var mt mocktracer + mt := newMockTracer() assert := assert.New(t) err := mt.Inject(&spanContext{}, 2)