Improve oc-receiver and oc-exporter span throughput (#395)

* Improve oc-receiver and oc-exporter span throughput

OC collector reception and export of spans did not scale up. This change
removes the `bundler` from the oc span receiver in favor of go-routines,
and additionally adds a worker count to the oc-exporter.

* fix lint

* PR comments

* PR comments

* Finish removal of unused options

* Update cached test

Author: Steven Karis

exporter/exporterparser/opencensus.go

@@ -17,6 +17,7 @@ package exporterparser
 import (
 	"context"
 	"fmt"
+	"sync/atomic"
 
 	"contrib.go.opencensus.io/exporter/ocagent"
 	"github.com/spf13/viper"
@@ -34,12 +35,18 @@ type opencensusConfig struct {
 	Compression string            `mapstructure:"compression,omitempty"`
 	Headers     map[string]string `mapstructure:"headers,omitempty"`
 	// TODO: add insecure, service name options.
+	NumWorkers int `mapstructure:"num-workers,omitempty"`
 }
 
 type ocagentExporter struct {
-	exporter *ocagent.Exporter
+	counter   uint32
+	exporters []*ocagent.Exporter
 }
 
+const (
+	defaultNumWorkers int = 2
+)
+
 var _ exporter.TraceExporter = (*ocagentExporter)(nil)
 
 // OpenCensusTraceExportersFromViper unmarshals the viper and returns an exporter.TraceExporter targeting
@@ -57,43 +64,55 @@ func OpenCensusTraceExportersFromViper(v *viper.Viper) (tes []exporter.TraceExpo
 	}
 
 	if ocac.Endpoint == "" {
-		return nil, nil, nil, fmt.Errorf("OpenCensus config requires an Endpoint")
+		return nil, nil, nil, fmt.Errorf("openCensus config requires an Endpoint")
 	}
 
 	opts := []ocagent.ExporterOption{ocagent.WithAddress(ocac.Endpoint), ocagent.WithInsecure()}
 	if ocac.Compression != "" {
 		if compressionKey := grpc.GetGRPCCompressionKey(ocac.Compression); compressionKey != compression.Unsupported {
 			opts = append(opts, ocagent.UseCompressor(compressionKey))
 		} else {
-			return nil, nil, nil, fmt.Errorf("Unsupported compression type: %s", ocac.Compression)
+			return nil, nil, nil, fmt.Errorf("unsupported compression type: %s", ocac.Compression)
 		}
 	}
 	if len(ocac.Headers) > 0 {
 		opts = append(opts, ocagent.WithHeaders(ocac.Headers))
 	}
 
-	sde, serr := ocagent.NewExporter(opts...)
-	if serr != nil {
-		return nil, nil, nil, fmt.Errorf("Cannot configure OpenCensus Trace exporter: %v", serr)
+	numWorkers := defaultNumWorkers
+	if ocac.NumWorkers > 0 {
+		numWorkers = ocac.NumWorkers
+	}
+
+	exporters := make([]*ocagent.Exporter, 0, numWorkers)
+	for exporterIndex := 0; exporterIndex < numWorkers; exporterIndex++ {
+		exporter, serr := ocagent.NewExporter(opts...)
+		if serr != nil {
+			return nil, nil, nil, fmt.Errorf("cannot configure OpenCensus Trace exporter: %v", serr)
+		}
+		exporters = append(exporters, exporter)
+		doneFns = append(doneFns, func() error {
+			exporter.Flush()
+			return nil
+		})
 	}
 
-	oexp := &ocagentExporter{exporter: sde}
+	oexp := &ocagentExporter{exporters: exporters}
 	tes = append(tes, oexp)
 
 	// TODO: (@odeke-em, @songya23) implement ExportMetrics for OpenCensus.
 	// mes = append(mes, oexp)
-	doneFns = append(doneFns, func() error {
-		sde.Flush()
-		return nil
-	})
 	return tes, mes, doneFns, nil
 }
 
-func (sde *ocagentExporter) ExportSpans(ctx context.Context, td data.TraceData) error {
-	err := sde.exporter.ExportTraceServiceRequest(
+func (oce *ocagentExporter) ExportSpans(ctx context.Context, td data.TraceData) error {
+	// Get an exporter worker round-robin
+	exporter := oce.exporters[atomic.AddUint32(&oce.counter, 1)%uint32(len(oce.exporters))]
+	err := exporter.ExportTraceServiceRequest(
 		&agenttracepb.ExportTraceServiceRequest{
-			Spans: td.Spans,
-			Node:  td.Node,
+			Spans:    td.Spans,
+			Resource: td.Resource,
+			Node:     td.Node,
 		},
 	)
 	if err != nil {

receiver/opencensus/octrace/observability_test.go

@@ -33,7 +33,6 @@ import (
 	agenttracepb "github.com/census-instrumentation/opencensus-proto/gen-go/agent/trace/v1"
 	tracepb "github.com/census-instrumentation/opencensus-proto/gen-go/trace/v1"
 	"github.com/census-instrumentation/opencensus-service/internal"
-	"github.com/census-instrumentation/opencensus-service/receiver/opencensus/octrace"
 )
 
 // Ensure that if we add a metrics exporter that our target metrics
@@ -48,7 +47,7 @@ func TestEnsureRecordedMetrics(t *testing.T) {
 
 	sappender := newSpanAppender()
 
-	_, port, doneFn := ocReceiverOnGRPCServer(t, sappender, octrace.WithSpanBufferPeriod(2*time.Millisecond))
+	_, port, doneFn := ocReceiverOnGRPCServer(t, sappender)
 	defer doneFn()
 
 	// Now the opencensus-agent exporter.
@@ -119,7 +118,7 @@ func TestEnsureRecordedMetrics_zeroLengthSpansSender(t *testing.T) {
 	)
 	sappender := newSpanAppender()
 
-	_, port, doneFn := ocReceiverOnGRPCServer(t, sappender, octrace.WithSpanBufferPeriod(2*time.Millisecond))
+	_, port, doneFn := ocReceiverOnGRPCServer(t, sappender)
 	defer doneFn()
 
 	// Now the opencensus-agent exporter.
@@ -191,7 +190,7 @@ func TestExportSpanLinkingMaintainsParentLink(t *testing.T) {
 
 	spanSink := newSpanAppender()
 	spansBufferPeriod := 10 * time.Millisecond
-	_, port, doneFn := ocReceiverOnGRPCServer(t, spanSink, octrace.WithSpanBufferPeriod(spansBufferPeriod))
+	_, port, doneFn := ocReceiverOnGRPCServer(t, spanSink)
 	defer doneFn()
 
 	traceSvcClient, traceSvcDoneFn, err := makeTraceServiceClient(port)
@@ -201,7 +200,7 @@ func TestExportSpanLinkingMaintainsParentLink(t *testing.T) {
 	defer traceSvcDoneFn()
 
 	n := 5
-	for i := 0; i <= n; i++ {
+	for i := 0; i < n; i++ {
 		sl := []*tracepb.Span{{TraceId: []byte("abcdefghijklmnop"), SpanId: []byte{byte(i + 1), 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08}}}
 		_ = traceSvcClient.Send(&agenttracepb.ExportTraceServiceRequest{Spans: sl, Node: &commonpb.Node{}})
 	}
@@ -225,7 +224,7 @@ func TestExportSpanLinkingMaintainsParentLink(t *testing.T) {
 	}
 
 	gotSpanData := ocSpansSaver.spanData[:]
-	if g, w := len(gotSpanData), 2; g != w {
+	if g, w := len(gotSpanData), n+1; g != w {
 		blob, _ := json.MarshalIndent(gotSpanData, "  ", " ")
 		t.Fatalf("Spandata count: Got %d Want %d\n\nData: %s", g, w, blob)
 	}
@@ -235,7 +234,8 @@ func TestExportSpanLinkingMaintainsParentLink(t *testing.T) {
 		t.Fatalf("Links count: Got %d Want %d\nGotSpanData: %#v", g, w, receiverSpanData)
 	}
 
-	rpcSpanData := gotSpanData[1]
+	// The rpc span is always last in the list
+	rpcSpanData := gotSpanData[len(gotSpanData)-1]
 
 	// Ensure that the link matches up exactly!
 	wantLink := trace.Link{

receiver/opencensus/octrace/opencensus.go

@@ -18,9 +18,6 @@ import (
 	"context"
 	"errors"
 	"io"
-	"time"
-
-	"google.golang.org/api/support/bundler"
 
 	"go.opencensus.io/trace"
 
@@ -35,9 +32,7 @@ import (
 
 // Receiver is the type used to handle spans from OpenCensus exporters.
 type Receiver struct {
-	spanSink         receiver.TraceReceiverSink
-	spanBufferPeriod time.Duration
-	spanBufferCount  int
+	spanSink receiver.TraceReceiverSink
 }
 
 // New creates a new opencensus.Receiver reference.
@@ -69,25 +64,8 @@ const receiverName = "opencensus_trace"
 // Export is the gRPC method that receives streamed traces from
 // OpenCensus-traceproto compatible libraries/applications.
 func (oci *Receiver) Export(tes agenttracepb.TraceService_ExportServer) error {
-	// The bundler will receive batches of spans i.e. []*tracepb.Span
 	// We need to ensure that it propagates the receiver name as a tag
 	ctxWithReceiverName := internal.ContextWithReceiverName(tes.Context(), receiverName)
-	traceBundler := bundler.NewBundler((*data.TraceData)(nil), func(payload interface{}) {
-		oci.batchSpanExporting(ctxWithReceiverName, payload)
-	})
-
-	spanBufferPeriod := oci.spanBufferPeriod
-	if spanBufferPeriod <= 0 {
-		spanBufferPeriod = 2 * time.Second // Arbitrary value
-	}
-	spanBufferCount := oci.spanBufferCount
-	if spanBufferCount <= 0 {
-		// TODO: (@odeke-em) provide an option to disable any buffering
-		spanBufferCount = 50 // Arbitrary value
-	}
-
-	traceBundler.DelayThreshold = spanBufferPeriod
-	traceBundler.BundleCountThreshold = spanBufferCount
 
 	// The first message MUST have a non-nil Node.
 	recv, err := tes.Recv()
@@ -100,19 +78,6 @@ func (oci *Receiver) Export(tes agenttracepb.TraceService_ExportServer) error {
 		return errTraceExportProtocolViolation
 	}
 
-	spansMetricsFn := internal.NewReceivedSpansRecorderStreaming(tes.Context(), receiverName)
-
-	processReceivedSpans := func(ni *commonpb.Node, resource *resourcepb.Resource, spans []*tracepb.Span) {
-		// Firstly, we'll add them to the bundler.
-		if len(spans) > 0 {
-			bundlerPayload := &data.TraceData{Node: ni, Resource: resource, Spans: spans}
-			traceBundler.Add(bundlerPayload, len(bundlerPayload.Spans))
-		}
-
-		// We MUST unconditionally record metrics from this reception.
-		spansMetricsFn(ni, spans)
-	}
-
 	var lastNonNilNode *commonpb.Node
 	var resource *resourcepb.Resource
 	// Now that we've got the first message with a Node, we can start to receive streamed up spans.
@@ -128,7 +93,7 @@ func (oci *Receiver) Export(tes agenttracepb.TraceService_ExportServer) error {
 			resource = recv.Resource
 		}
 
-		processReceivedSpans(lastNonNilNode, resource, recv.Spans)
+		go oci.export(ctxWithReceiverName, tes, lastNonNilNode, resource, recv.Spans)
 
 		recv, err = tes.Recv()
 		if err != nil {
@@ -142,9 +107,19 @@ func (oci *Receiver) Export(tes agenttracepb.TraceService_ExportServer) error {
 	}
 }
 
-func (oci *Receiver) batchSpanExporting(longLivedRPCCtx context.Context, payload interface{}) {
-	tracedata := payload.([]*data.TraceData)
-	if len(tracedata) == 0 {
+func (oci *Receiver) export(
+	longLivedCtx context.Context,
+	tes agenttracepb.TraceService_ExportServer,
+	node *commonpb.Node,
+	resource *resourcepb.Resource,
+	spans []*tracepb.Span,
+) {
+	tracedata := data.TraceData{Node: node, Resource: resource, Spans: spans}
+	// We MUST unconditionally record metrics from this reception.
+	spansMetricsFn := internal.NewReceivedSpansRecorderStreaming(tes.Context(), receiverName)
+	spansMetricsFn(tracedata.Node, tracedata.Spans)
+
+	if len(tracedata.Spans) == 0 {
 		return
 	}
 
@@ -157,15 +132,11 @@ func (oci *Receiver) batchSpanExporting(longLivedRPCCtx context.Context, payload
 	// spansAndNode list unfurling then send spans grouped per node
 
 	// If the starting RPC has a parent span, then add it as a parent link.
-	internal.SetParentLink(longLivedRPCCtx, span)
+	internal.SetParentLink(longLivedCtx, span)
 
-	nSpans := int64(0)
-	for _, td := range tracedata {
-		oci.spanSink.ReceiveTraceData(ctx, *td)
-		nSpans += int64(len(td.Spans))
-	}
+	oci.spanSink.ReceiveTraceData(ctx, tracedata)
 
 	span.Annotate([]trace.Attribute{
-		trace.Int64Attribute("num_spans", nSpans),
+		trace.Int64Attribute("num_spans", int64(len(tracedata.Spans))),
 	}, "")
 }

receiver/opencensus/octrace/opencensus_test.go

@@ -49,7 +49,7 @@ func TestReceiver_endToEnd(t *testing.T) {
 
 	sappender := newSpanAppender()
 
-	_, port, doneFn := ocReceiverOnGRPCServer(t, sappender, octrace.WithSpanBufferPeriod(100*time.Millisecond))
+	_, port, doneFn := ocReceiverOnGRPCServer(t, sappender)
 	defer doneFn()
 
 	// Now the opencensus-agent exporter.
@@ -172,7 +172,7 @@ func TestReceiver_endToEnd(t *testing.T) {
 func TestExportMultiplexing(t *testing.T) {
 	spanSink := newSpanAppender()
 
-	_, port, doneFn := ocReceiverOnGRPCServer(t, spanSink, octrace.WithSpanBufferPeriod(90*time.Millisecond))
+	_, port, doneFn := ocReceiverOnGRPCServer(t, spanSink)
 	defer doneFn()
 
 	traceClient, traceClientDoneFn, err := makeTraceServiceClient(port)
@@ -279,10 +279,27 @@ func TestExportMultiplexing(t *testing.T) {
 		nodeToKey(node2):          append(sL2, append(sLn2a, sLn2b...)...),
 	}
 
-	gotBlob, _ := json.Marshal(resultsMapping)
-	wantBlob, _ := json.Marshal(wantContents)
-	if !bytes.Equal(gotBlob, wantBlob) {
-		t.Errorf("Unequal serialization results\nGot:\n\t%s\nWant:\n\t%s\n", gotBlob, wantBlob)
+	for nodeKey, wantSpans := range wantContents {
+		gotSpans, ok := resultsMapping[nodeKey]
+		if !ok {
+			t.Errorf("Wanted to find a node that was not found for key: %s", nodeKey)
+		}
+		if len(gotSpans) != len(wantSpans) {
+			t.Errorf("Unequal number of spans for nodeKey: %s", nodeKey)
+		}
+		for _, wantSpan := range wantSpans {
+			found := false
+			for _, gotSpan := range gotSpans {
+				wantStr, _ := json.Marshal(wantSpan)
+				gotStr, _ := json.Marshal(gotSpan)
+				if bytes.Equal(wantStr, gotStr) {
+					found = true
+				}
+			}
+			if !found {
+				t.Errorf("Unequal span serialization\nGot:\n\t%s\nWant:\n\t%s\n", gotSpans, wantSpans)
+			}
+		}
 	}
 }
 
@@ -291,7 +308,7 @@ func TestExportMultiplexing(t *testing.T) {
 func TestExportProtocolViolations_nodelessFirstMessage(t *testing.T) {
 	spanSink := newSpanAppender()
 
-	_, port, doneFn := ocReceiverOnGRPCServer(t, spanSink, octrace.WithSpanBufferPeriod(90*time.Millisecond))
+	_, port, doneFn := ocReceiverOnGRPCServer(t, spanSink)
 	defer doneFn()
 
 	traceClient, traceClientDoneFn, err := makeTraceServiceClient(port)
@@ -359,7 +376,7 @@ func TestExportProtocolViolations_nodelessFirstMessage(t *testing.T) {
 func TestExportProtocolConformation_spansInFirstMessage(t *testing.T) {
 	spanSink := newSpanAppender()
 
-	_, port, doneFn := ocReceiverOnGRPCServer(t, spanSink, octrace.WithSpanBufferPeriod(70*time.Millisecond))
+	_, port, doneFn := ocReceiverOnGRPCServer(t, spanSink)
 	defer doneFn()
 
 	traceClient, traceClientDoneFn, err := makeTraceServiceClient(port)

receiver/opencensus/octrace/options.go

@@ -14,43 +14,9 @@
 
 package octrace
 
-import "time"
-
 // Option interface defines for configuration settings to be applied to receivers.
 //
 // WithReceiver applies the configuration to the given receiver.
 type Option interface {
 	WithReceiver(*Receiver)
 }
-
-type spanBufferPeriod struct {
-	period time.Duration
-}
-
-var _ Option = (*spanBufferPeriod)(nil)
-
-func (sfd *spanBufferPeriod) WithReceiver(oci *Receiver) {
-	oci.spanBufferPeriod = sfd.period
-}
-
-// WithSpanBufferPeriod is an option that allows one to configure
-// the period that spans are buffered for before the Receiver
-// sends them to its TraceReceiverSink.
-func WithSpanBufferPeriod(period time.Duration) Option {
-	return &spanBufferPeriod{period: period}
-}
-
-type spanBufferCount int
-
-var _ Option = (*spanBufferCount)(nil)
-
-func (spc spanBufferCount) WithReceiver(oci *Receiver) {
-	oci.spanBufferCount = int(spc)
-}
-
-// WithSpanBufferCount is an option that allows one to configure
-// the number of spans that are buffered before the Receiver
-// send them to its TraceReceiverSink.
-func WithSpanBufferCount(count int) Option {
-	return spanBufferCount(count)
-}