contrib/aws/aws-sdk-go: open only a single span for retryable operations

contrib/aws/aws-sdk-go/aws/aws.go

@@ -49,6 +49,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)),

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()
+
+	root, ctx := tracer.StartSpanFromContext(context.Background(), "test")
+	s3api := s3.New(session)
+	req, _ := s3api.GetObjectRequest(&s3.GetObjectInput{
+		Bucket: aws.String("BUCKET"),
+		Key:    aws.String("KEY"),
+	})
+	req.SetContext(ctx)
+	err := req.Send()
+	root.Finish()
+
+	assert.Equal(t, 3, req.RetryCount)
+	assert.Same(t, expectedError, err)
+	assert.Len(t, mt.UnfinishedSpans(), 0)
+	assert.Len(t, mt.FinishedSpans(), 2)
+}

ddtrace/mocktracer/mockspan.go

@@ -211,6 +211,7 @@ func (s *mockspan) Finish(opts ...ddtrace.FinishOption) {
 	}
 	s.finished = true
 	s.finishTime = t
+	s.tracer.removeUnfinishedSpan(s)
 	s.tracer.addFinishedSpan(s)
 }
 

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 {
+	// UnfinishedSpans returns the set of started spans.
+	UnfinishedSpans() []Span
+
 	// FinishedSpans returns the set of finished spans.
 	FinishedSpans() []Span
 
@@ -52,8 +55,9 @@ func Start() Tracer {
 }
 
 type mocktracer struct {
-	sync.RWMutex  // guards below spans
-	finishedSpans []Span
+	sync.RWMutex    // guards below spans
+	unfinishedSpans []Span
+	finishedSpans   []Span
 }
 
 // Stop deactivates the mock tracer and sets the active tracer to a no-op.
@@ -67,7 +71,15 @@ 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.addUnfinishedSpan(span)
+	return span
+}
+
+func (t *mocktracer) UnfinishedSpans() []Span {
+	t.RLock()
+	defer t.RUnlock()
+	return t.unfinishedSpans
 }
 
 func (t *mocktracer) FinishedSpans() []Span {
@@ -79,9 +91,31 @@ func (t *mocktracer) FinishedSpans() []Span {
 func (t *mocktracer) Reset() {
 	t.Lock()
 	defer t.Unlock()
+	t.unfinishedSpans = nil
 	t.finishedSpans = nil
 }
 
+func (t *mocktracer) addUnfinishedSpan(s Span) {
+	t.Lock()
+	defer t.Unlock()
+	t.unfinishedSpans = append(t.unfinishedSpans, s)
+}
+
+func (t *mocktracer) removeUnfinishedSpan(s Span) {
+	t.Lock()
+	defer t.Unlock()
+	idx := -1
+	for i, si := range t.unfinishedSpans {
+		if si == s {
+			idx = i
+			break
+		}
+	}
+	if idx >= 0 {
+		t.unfinishedSpans = append(t.unfinishedSpans[:idx], t.unfinishedSpans[idx+1:]...)
+	}
+}
+
 func (t *mocktracer) addFinishedSpan(s Span) {
 	t.Lock()
 	defer t.Unlock()