Add timeouts; make span size buffer limit precise (#58)

Add a timeout (customizeable by option) to the contexts used to
make API calls (default to 2 seconds).

Add a size limit on spans we buffer in memory and compute the size
of the span by converting to protobuf before buffering, rather
than estimating from the SpanData struct.

Author: Ramon Nogueira

stackdriver.go

@@ -89,16 +89,22 @@ type Options struct {
 	TraceClientOptions []option.ClientOption
 
 	// BundleDelayThreshold determines the max amount of time
-	// the exporter can wait before uploading view data to
+	// the exporter can wait before uploading view data or trace spans to
 	// the backend.
 	// Optional.
 	BundleDelayThreshold time.Duration
 
-	// BundleCountThreshold determines how many view data events
+	// BundleCountThreshold determines how many view data events or trace spans
 	// can be buffered before batch uploading them to the backend.
 	// Optional.
 	BundleCountThreshold int
 
+	// TraceSpansBufferMaxBytes is the maximum size (in bytes) of spans that
+	// will be buffered in memory before being dropped.
+	//
+	// If unset, a default of 8MB will be used.
+	TraceSpansBufferMaxBytes int
+
 	// Resource sets the MonitoredResource against which all views will be
 	// recorded by this exporter.
 	//
@@ -190,9 +196,14 @@ type Options struct {
 	// trace and metric clients, and then every time a new batch of traces or
 	// stats needs to be uploaded.
 	//
+	// Do not set a timeout on this context. Instead, set the Timeout option.
+	//
 	// If unset, context.Background() will be used.
 	Context context.Context
 
+	// Timeout for all API calls. If not set, defaults to 5 seconds.
+	Timeout time.Duration
+
 	// GetMonitoredResource may be provided to supply the details of the
 	// monitored resource dynamically based on the tags associated with each
 	// data point. Most users will not need to set this, but should instead
@@ -209,6 +220,8 @@ type Options struct {
 	GetMonitoredResource func(*view.View, []tag.Tag) ([]tag.Tag, monitoredresource.Interface)
 }
 
+const defaultTimeout = 5 * time.Second
+
 // Exporter is a stats and trace exporter that uploads data to Stackdriver.
 //
 // You can create a single Exporter and register it as both a trace exporter
@@ -222,11 +235,12 @@ type Exporter struct {
 // NewExporter creates a new Exporter that implements both stats.Exporter and
 // trace.Exporter.
 func NewExporter(o Options) (*Exporter, error) {
-	if o.Context == nil {
-		o.Context = context.Background()
-	}
 	if o.ProjectID == "" {
-		creds, err := google.FindDefaultCredentials(o.Context, traceapi.DefaultAuthScopes()...)
+		ctx := o.Context
+		if ctx == nil {
+			ctx = context.Background()
+		}
+		creds, err := google.FindDefaultCredentials(ctx, traceapi.DefaultAuthScopes()...)
 		if err != nil {
 			return nil, fmt.Errorf("stackdriver: %v", err)
 		}
@@ -297,6 +311,18 @@ func (o Options) handleError(err error) {
 	log.Printf("Failed to export to Stackdriver: %v", err)
 }
 
+func (o Options) newContextWithTimeout() (context.Context, func()) {
+	ctx := o.Context
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	timeout := o.Timeout
+	if timeout <= 0 {
+		timeout = defaultTimeout
+	}
+	return context.WithTimeout(ctx, timeout)
+}
+
 // convertMonitoredResourceToPB converts MonitoredResource data in to
 // protocol buffer.
 func convertMonitoredResourceToPB(mr monitoredresource.Interface) *monitoredrespb.MonitoredResource {

stats.go

@@ -78,7 +78,9 @@ func newStatsExporter(o Options) (*statsExporter, error) {
 	}
 
 	opts := append(o.MonitoringClientOptions, option.WithUserAgent(userAgent))
-	client, err := monitoring.NewMetricClient(o.Context, opts...)
+	ctx, cancel := o.newContextWithTimeout()
+	defer cancel()
+	client, err := monitoring.NewMetricClient(ctx, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -167,8 +169,10 @@ func (e *statsExporter) Flush() {
 }
 
 func (e *statsExporter) uploadStats(vds []*view.Data) error {
+	ctx, cancel := e.o.newContextWithTimeout()
+	defer cancel()
 	ctx, span := trace.StartSpan(
-		e.o.Context,
+		ctx,
 		"contrib.go.opencensus.io/exporter/stackdriver.uploadStats",
 		trace.WithSampler(trace.NeverSample()),
 	)

stats_test.go

@@ -16,11 +16,12 @@ package stackdriver
 
 import (
 	"context"
-	"contrib.go.opencensus.io/exporter/stackdriver/monitoredresource"
 	"fmt"
 	"testing"
 	"time"
 
+	"contrib.go.opencensus.io/exporter/stackdriver/monitoredresource"
+
 	"cloud.google.com/go/monitoring/apiv3"
 	"github.com/golang/protobuf/ptypes/timestamp"
 	"github.com/google/go-cmp/cmp"

trace.go

@@ -15,12 +15,14 @@
 package stackdriver
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"sync"
 	"time"
 
 	tracingclient "cloud.google.com/go/trace/apiv2"
+	"github.com/golang/protobuf/proto"
 	"go.opencensus.io/trace"
 	"google.golang.org/api/support/bundler"
 	tracepb "google.golang.org/genproto/googleapis/devtools/cloudtrace/v2"
@@ -34,63 +36,69 @@ type traceExporter struct {
 	projectID string
 	bundler   *bundler.Bundler
 	// uploadFn defaults to uploadSpans; it can be replaced for tests.
-	uploadFn func(spans []*trace.SpanData)
+	uploadFn func(spans []*tracepb.Span)
 	overflowLogger
 	client *tracingclient.Client
 }
 
 var _ trace.Exporter = (*traceExporter)(nil)
 
 func newTraceExporter(o Options) (*traceExporter, error) {
-	client, err := tracingclient.NewClient(o.Context, o.TraceClientOptions...)
+	ctx := o.Context
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	client, err := tracingclient.NewClient(ctx, o.TraceClientOptions...)
 	if err != nil {
 		return nil, fmt.Errorf("stackdriver: couldn't initialize trace client: %v", err)
 	}
 	return newTraceExporterWithClient(o, client), nil
 }
 
+const defaultBufferedByteLimit = 8 * 1024 * 1024
+
 func newTraceExporterWithClient(o Options, c *tracingclient.Client) *traceExporter {
 	e := &traceExporter{
 		projectID: o.ProjectID,
 		client:    c,
 		o:         o,
 	}
-	bundler := bundler.NewBundler((*trace.SpanData)(nil), func(bundle interface{}) {
-		e.uploadFn(bundle.([]*trace.SpanData))
+	b := bundler.NewBundler((*tracepb.Span)(nil), func(bundle interface{}) {
+		e.uploadFn(bundle.([]*tracepb.Span))
 	})
 	if o.BundleDelayThreshold > 0 {
-		bundler.DelayThreshold = o.BundleDelayThreshold
+		b.DelayThreshold = o.BundleDelayThreshold
 	} else {
-		bundler.DelayThreshold = 2 * time.Second
+		b.DelayThreshold = 2 * time.Second
 	}
 	if o.BundleCountThreshold > 0 {
-		bundler.BundleCountThreshold = o.BundleCountThreshold
+		b.BundleCountThreshold = o.BundleCountThreshold
 	} else {
-		bundler.BundleCountThreshold = 50
+		b.BundleCountThreshold = 50
 	}
 	// The measured "bytes" are not really bytes, see exportReceiver.
-	bundler.BundleByteThreshold = bundler.BundleCountThreshold * 200
-	bundler.BundleByteLimit = bundler.BundleCountThreshold * 1000
-	bundler.BufferedByteLimit = bundler.BundleCountThreshold * 2000
+	b.BundleByteThreshold = b.BundleCountThreshold * 200
+	b.BundleByteLimit = b.BundleCountThreshold * 1000
+	if o.TraceSpansBufferMaxBytes > 0 {
+		b.BufferedByteLimit = o.TraceSpansBufferMaxBytes
+	} else {
+		b.BufferedByteLimit = defaultBufferedByteLimit
+	}
 
-	e.bundler = bundler
+	e.bundler = b
 	e.uploadFn = e.uploadSpans
 	return e
 }
 
 // ExportSpan exports a SpanData to Stackdriver Trace.
 func (e *traceExporter) ExportSpan(s *trace.SpanData) {
-	// n is a length heuristic.
-	n := 1
-	n += len(s.Attributes)
-	n += len(s.Annotations)
-	n += len(s.MessageEvents)
-	err := e.bundler.Add(s, n)
+	protoSpan := protoFromSpanData(s, e.projectID, e.o.Resource)
+	protoSize := proto.Size(protoSpan)
+	err := e.bundler.Add(protoSpan, protoSize)
 	switch err {
 	case nil:
 		return
 	case bundler.ErrOversizedItem:
-		go e.uploadFn([]*trace.SpanData{s})
 	case bundler.ErrOverflow:
 		e.overflowLogger.log()
 	default:
@@ -107,17 +115,16 @@ func (e *traceExporter) Flush() {
 }
 
 // uploadSpans uploads a set of spans to Stackdriver.
-func (e *traceExporter) uploadSpans(spans []*trace.SpanData) {
+func (e *traceExporter) uploadSpans(spans []*tracepb.Span) {
 	req := tracepb.BatchWriteSpansRequest{
 		Name:  "projects/" + e.projectID,
-		Spans: make([]*tracepb.Span, 0, len(spans)),
-	}
-	for _, span := range spans {
-		req.Spans = append(req.Spans, protoFromSpanData(span, e.projectID, e.o.Resource))
+		Spans: spans,
 	}
 	// Create a never-sampled span to prevent traces associated with exporter.
-	ctx, span := trace.StartSpan( // TODO: add timeouts
-		e.o.Context,
+	ctx, cancel := e.o.newContextWithTimeout()
+	defer cancel()
+	ctx, span := trace.StartSpan(
+		ctx,
 		"contrib.go.opencensus.io/exporter/stackdriver.uploadSpans",
 		trace.WithSampler(trace.NeverSample()),
 	)

trace_test.go

@@ -16,10 +16,12 @@ package stackdriver
 
 import (
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
 	"go.opencensus.io/trace"
+	tracepb "google.golang.org/genproto/googleapis/devtools/cloudtrace/v2"
 )
 
 func TestBundling(t *testing.T) {
@@ -29,8 +31,8 @@ func TestBundling(t *testing.T) {
 		BundleCountThreshold: 10,
 	}, nil)
 
-	ch := make(chan []*trace.SpanData)
-	exporter.uploadFn = func(spans []*trace.SpanData) {
+	ch := make(chan []*tracepb.Span)
+	exporter.uploadFn = func(spans []*tracepb.Span) {
 		ch <- spans
 	}
 	trace.RegisterExporter(exporter)
@@ -60,3 +62,51 @@ func TestBundling(t *testing.T) {
 	case <-time.After(time.Second / 5):
 	}
 }
+
+func TestNewContext_Timeout(t *testing.T) {
+	e := newTraceExporterWithClient(Options{
+		Timeout: 10 * time.Millisecond,
+	}, nil)
+	ctx, cancel := e.o.newContextWithTimeout()
+	defer cancel()
+	select {
+	case <-time.After(60 * time.Second):
+		t.Fatal("should have timed out")
+	case <-ctx.Done():
+	}
+}
+
+func TestTraceSpansBufferMaxBytes(t *testing.T) {
+	e := newTraceExporterWithClient(Options{
+		Context:                  context.Background(),
+		Timeout:                  10 * time.Millisecond,
+		TraceSpansBufferMaxBytes: 20000,
+	}, nil)
+	waitCh := make(chan struct{})
+	exported := 0
+	e.uploadFn = func(spans []*tracepb.Span) {
+		<-waitCh
+		exported++
+	}
+	for i := 0; i < 10; i++ {
+		e.ExportSpan(makeSampleSpanData())
+	}
+	close(waitCh)
+	e.Flush()
+	if exported != 2 {
+		t.Errorf("exported = %d; want 2", exported)
+	}
+}
+
+func makeSampleSpanData() *trace.SpanData {
+	sd := &trace.SpanData{
+		Annotations:   make([]trace.Annotation, 32),
+		Links:         make([]trace.Link, 32),
+		MessageEvents: make([]trace.MessageEvent, 128),
+		Attributes:    make(map[string]interface{}),
+	}
+	for i := 0; i < 32; i++ {
+		sd.Attributes[fmt.Sprintf("attribute-%d", i)] = ""
+	}
+	return sd
+}