Allow users to provide context. (#28)

Users can now provide a custom context that will be used instead of
context.Background().

This fixes #27.

Author: knyar

stackdriver.go

@@ -168,6 +168,15 @@ type Options struct {
 	// default "opencensus_task" label. You should only do this if you know that
 	// the Resource you set uniquely identifies this Go process.
 	DefaultMonitoringLabels *Labels
+
+	// Context allows users to provide a custom context for API calls.
+	//
+	// This context will be used several times: first, to create Stackdriver
+	// trace and metric clients, and then every time a new batch of traces or
+	// stats needs to be uploaded.
+	//
+	// If unset, context.Background() will be used.
+	Context context.Context
 }
 
 // Exporter is a stats.Exporter and trace.Exporter
@@ -180,8 +189,11 @@ 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(context.Background(), traceapi.DefaultAuthScopes()...)
+		creds, err := google.FindDefaultCredentials(o.Context, traceapi.DefaultAuthScopes()...)
 		if err != nil {
 			return nil, fmt.Errorf("stackdriver: %v", err)
 		}

stats.go

@@ -78,7 +78,7 @@ func newStatsExporter(o Options, enforceProjectUniqueness bool) (*statsExporter,
 	}
 
 	opts := append(o.MonitoringClientOptions, option.WithUserAgent(userAgent))
-	client, err := monitoring.NewMetricClient(context.Background(), opts...)
+	client, err := monitoring.NewMetricClient(o.Context, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -150,7 +150,7 @@ func (e *statsExporter) Flush() {
 
 func (e *statsExporter) uploadStats(vds []*view.Data) error {
 	ctx, span := trace.StartSpan(
-		context.Background(),
+		e.o.Context,
 		"contrib.go.opencensus.io/exporter/stackdriver.uploadStats",
 		trace.WithSampler(trace.NeverSample()),
 	)
@@ -163,7 +163,7 @@ func (e *statsExporter) uploadStats(vds []*view.Data) error {
 		}
 	}
 	for _, req := range e.makeReq(vds, maxTimeSeriesPerUpload) {
-		if err := e.c.CreateTimeSeries(ctx, req); err != nil {
+		if err := createTimeSeries(ctx, e.c, req); err != nil {
 			span.SetStatus(trace.Status{Code: 2, Message: err.Error()})
 			// TODO(jbd): Don't fail fast here, batch errors?
 			return err
@@ -465,3 +465,7 @@ var createMetricDescriptor = func(ctx context.Context, c *monitoring.MetricClien
 var getMetricDescriptor = func(ctx context.Context, c *monitoring.MetricClient, mdr *monitoringpb.GetMetricDescriptorRequest) (*metric.MetricDescriptor, error) {
 	return c.GetMetricDescriptor(ctx, mdr)
 }
+
+var createTimeSeries = func(ctx context.Context, c *monitoring.MetricClient, ts *monitoringpb.CreateTimeSeriesRequest) error {
+	return c.CreateTimeSeries(ctx, ts)
+}

stats_test.go

@@ -16,6 +16,7 @@ package stackdriver
 
 import (
 	"context"
+	"fmt"
 	"reflect"
 	"testing"
 	"time"
@@ -1014,6 +1015,62 @@ func TestExporter_makeReq_withCustomMonitoredResource(t *testing.T) {
 	}
 }
 
+func TestExporter_customContext(t *testing.T) {
+	oldCreateMetricDescriptor := createMetricDescriptor
+	oldCreateTimeSeries := createTimeSeries
+
+	defer func() {
+		createMetricDescriptor = oldCreateMetricDescriptor
+		createTimeSeries = oldCreateTimeSeries
+	}()
+
+	var timedOut = 0
+	createMetricDescriptor = func(ctx context.Context, c *monitoring.MetricClient, mdr *monitoringpb.CreateMetricDescriptorRequest) (*metric.MetricDescriptor, error) {
+		select {
+		case <-time.After(15 * time.Millisecond):
+			fmt.Println("createMetricDescriptor did not time out")
+		case <-ctx.Done():
+			timedOut++
+		}
+		return &metric.MetricDescriptor{}, nil
+	}
+	createTimeSeries = func(ctx context.Context, c *monitoring.MetricClient, ts *monitoringpb.CreateTimeSeriesRequest) error {
+		select {
+		case <-time.After(15 * time.Millisecond):
+			fmt.Println("createTimeSeries did not time out")
+		case <-ctx.Done():
+			timedOut++
+		}
+		return nil
+	}
+
+	v := &view.View{
+		Name:        "test_view_count",
+		Description: "view_description",
+		Measure:     stats.Float64("test-measure/TestExporter_createMeasure", "measure desc", stats.UnitMilliseconds),
+		Aggregation: view.Count(),
+	}
+
+	data := &view.CountData{Value: 0}
+	vd := newTestViewData(v, time.Now(), time.Now(), data, data)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
+	defer cancel()
+	e := &statsExporter{
+		createdViews: make(map[string]*metricpb.MetricDescriptor),
+		o:            Options{ProjectID: "test_project", Context: ctx},
+	}
+	if err := e.uploadStats([]*view.Data{vd}); err != nil {
+		t.Errorf("Exporter.uploadStats() error = %v", err)
+	}
+	if ctx.Err() != context.DeadlineExceeded {
+		t.Errorf("expected context to time out; got %v", ctx.Err())
+	}
+	if timedOut != 2 {
+		t.Errorf("expected two functions to time out; got %d", timedOut)
+	}
+}
+
 func newTestViewData(v *view.View, start, end time.Time, data1, data2 view.AggregationData) *view.Data {
 	key, _ := tag.NewKey("test-key")
 	tag1 := tag.Tag{Key: key, Value: "test-value-1"}

trace.go

@@ -15,7 +15,6 @@
 package stackdriver
 
 import (
-	"context"
 	"fmt"
 	"log"
 	"sync"
@@ -43,7 +42,7 @@ type traceExporter struct {
 var _ trace.Exporter = (*traceExporter)(nil)
 
 func newTraceExporter(o Options) (*traceExporter, error) {
-	client, err := tracingclient.NewClient(context.Background(), o.TraceClientOptions...)
+	client, err := tracingclient.NewClient(o.Context, o.TraceClientOptions...)
 	if err != nil {
 		return nil, fmt.Errorf("stackdriver: couldn't initialize trace client: %v", err)
 	}
@@ -118,7 +117,7 @@ func (e *traceExporter) uploadSpans(spans []*trace.SpanData) {
 	}
 	// Create a never-sampled span to prevent traces associated with exporter.
 	ctx, span := trace.StartSpan( // TODO: add timeouts
-		context.Background(),
+		e.o.Context,
 		"contrib.go.opencensus.io/exporter/stackdriver.uploadSpans",
 		trace.WithSampler(trace.NeverSample()),
 	)