Update collector code to use [Trace|Metrics]Data. (#431)

Author: Bogdan Drutu

cmd/occollector/app/sender/jaeger_thrift_http_sender.go

@@ -25,7 +25,7 @@ import (
 	"github.com/apache/thrift/lib/go/thrift"
 	"go.uber.org/zap"
 
-	agenttracepb "github.com/census-instrumentation/opencensus-proto/gen-go/agent/trace/v1"
+	"github.com/census-instrumentation/opencensus-service/data"
 	jaegertranslator "github.com/census-instrumentation/opencensus-service/translator/trace/jaeger"
 )
 
@@ -80,15 +80,11 @@ func NewJaegerThriftHTTPSender(
 }
 
 // ProcessSpans sends the received data to the configured Jaeger Thrift end-point.
-func (s *JaegerThriftHTTPSender) ProcessSpans(batch *agenttracepb.ExportTraceServiceRequest, spanFormat string) (uint64, error) {
+func (s *JaegerThriftHTTPSender) ProcessSpans(td data.TraceData, spanFormat string) (uint64, error) {
 	// TODO: (@pjanotti) In case of failure the translation to Jaeger Thrift is going to be remade, cache it somehow.
-	if batch == nil {
-		return 0, fmt.Errorf("Jaeger sender received nil batch")
-	}
-
-	tBatch, err := jaegertranslator.OCProtoToJaegerThrift(batch)
+	tBatch, err := jaegertranslator.OCProtoToJaegerThrift(td)
 	if err != nil {
-		return uint64(len(batch.Spans)), err
+		return uint64(len(td.Spans)), err
 	}
 
 	mSpans := tBatch.Spans

cmd/occollector/app/sender/jaeger_thrift_tchannel_sender.go

@@ -19,7 +19,7 @@ import (
 
 	reporter "github.com/jaegertracing/jaeger/cmd/agent/app/reporter"
 
-	agenttracepb "github.com/census-instrumentation/opencensus-proto/gen-go/agent/trace/v1"
+	"github.com/census-instrumentation/opencensus-service/data"
 	"github.com/census-instrumentation/opencensus-service/internal/collector/processor"
 	jaegertranslator "github.com/census-instrumentation/opencensus-service/translator/trace/jaeger"
 )
@@ -45,9 +45,9 @@ func NewJaegerThriftTChannelSender(
 }
 
 // ProcessSpans sends the received data to the configured Jaeger Thrift end-point.
-func (s *JaegerThriftTChannelSender) ProcessSpans(batch *agenttracepb.ExportTraceServiceRequest, spanFormat string) (uint64, error) {
+func (s *JaegerThriftTChannelSender) ProcessSpans(td data.TraceData, spanFormat string) (uint64, error) {
 	// TODO: (@pjanotti) In case of failure the translation to Jaeger Thrift is going to be remade, cache it somehow.
-	tBatch, err := jaegertranslator.OCProtoToJaegerThrift(batch)
+	tBatch, err := jaegertranslator.OCProtoToJaegerThrift(td)
 	if err != nil {
 		return uint64(len(tBatch.Spans)), err
 	}

internal/collector/processor/exporter_processor.go

@@ -17,7 +17,6 @@ package processor
 import (
 	"context"
 
-	agenttracepb "github.com/census-instrumentation/opencensus-proto/gen-go/agent/trace/v1"
 	"github.com/census-instrumentation/opencensus-service/data"
 	"github.com/census-instrumentation/opencensus-service/processor"
 )
@@ -33,8 +32,8 @@ func NewTraceExporterProcessor(traceExporters ...processor.TraceDataProcessor) S
 	return &exporterSpanProcessor{tdp: processor.NewMultiTraceDataProcessor(traceExporters)}
 }
 
-func (sp *exporterSpanProcessor) ProcessSpans(batch *agenttracepb.ExportTraceServiceRequest, spanFormat string) (uint64, error) {
-	err := sp.tdp.ProcessTraceData(context.Background(), data.TraceData{Node: batch.Node, Resource: batch.Resource, Spans: batch.Spans})
+func (sp *exporterSpanProcessor) ProcessSpans(td data.TraceData, spanFormat string) (uint64, error) {
+	err := sp.tdp.ProcessTraceData(context.Background(), td)
 	if err != nil {
 		// TODO: determine if the number of dropped spans is needed because it was wrong anyway.
 		return 0, err

internal/collector/processor/multi_processor.go

@@ -15,17 +15,17 @@
 package processor
 
 import (
+	"github.com/census-instrumentation/opencensus-service/data"
 	"github.com/census-instrumentation/opencensus-service/internal"
 
 	"github.com/spf13/cast"
 
-	agenttracepb "github.com/census-instrumentation/opencensus-proto/gen-go/agent/trace/v1"
 	tracepb "github.com/census-instrumentation/opencensus-proto/gen-go/trace/v1"
 )
 
 // MultiProcessorOption represents options that can be applied to a MultiSpanProcessor.
 type MultiProcessorOption func(*multiSpanProcessor)
-type preProcessFn func(*agenttracepb.ExportTraceServiceRequest, string)
+type preProcessFn func(data.TraceData, string)
 
 // MultiSpanProcessor enables processing on multiple processors.
 // For each incoming span batch, it calls ProcessSpans method on each span
@@ -59,11 +59,11 @@ func WithPreProcessFn(preProcFn preProcessFn) MultiProcessorOption {
 // in each ExportTraceServiceRequest.
 func WithAddAttributes(attributes map[string]interface{}, overwrite bool) MultiProcessorOption {
 	return WithPreProcessFn(
-		func(batch *agenttracepb.ExportTraceServiceRequest, spanFormat string) {
+		func(td data.TraceData, spanFormat string) {
 			if len(attributes) == 0 {
 				return
 			}
-			for _, span := range batch.Spans {
+			for _, span := range td.Spans {
 				if span == nil {
 					// We will not create nil spans with just attributes on them
 					continue
@@ -107,14 +107,14 @@ func WithAddAttributes(attributes map[string]interface{}, overwrite bool) MultiP
 }
 
 // ProcessSpans implements the SpanProcessor interface
-func (msp *multiSpanProcessor) ProcessSpans(batch *agenttracepb.ExportTraceServiceRequest, spanFormat string) (uint64, error) {
+func (msp *multiSpanProcessor) ProcessSpans(td data.TraceData, spanFormat string) (uint64, error) {
 	for _, preProcessFn := range msp.preProcessFns {
-		preProcessFn(batch, spanFormat)
+		preProcessFn(td, spanFormat)
 	}
 	var maxFailures uint64
 	var errors []error
 	for _, sp := range msp.processors {
-		failures, err := sp.ProcessSpans(batch, spanFormat)
+		failures, err := sp.ProcessSpans(td, spanFormat)
 		if err != nil {
 			errors = append(errors, err)
 		}

internal/collector/processor/multi_processor_test.go

@@ -19,8 +19,8 @@ import (
 	"sync/atomic"
 	"testing"
 
-	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/data"
 )
 
 func TestMultiSpanProcessorMultiplexing(t *testing.T) {
@@ -30,14 +30,14 @@ func TestMultiSpanProcessorMultiplexing(t *testing.T) {
 	}
 
 	tt := NewMultiSpanProcessor(processors)
-	batch := &agenttracepb.ExportTraceServiceRequest{
+	td := data.TraceData{
 		Spans: make([]*tracepb.Span, 7),
 	}
 
 	var wantSpansCount = 0
 	for i := 0; i < 2; i++ {
-		wantSpansCount += len(batch.Spans)
-		tt.ProcessSpans(batch, "test")
+		wantSpansCount += len(td.Spans)
+		tt.ProcessSpans(td, "test")
 	}
 
 	for _, p := range processors {
@@ -65,14 +65,14 @@ func TestMultiSpanProcessorSomeNotOk(t *testing.T) {
 	for i := range spans {
 		spans[i] = &tracepb.Span{}
 	}
-	batch := &agenttracepb.ExportTraceServiceRequest{
+	td := data.TraceData{
 		Spans: spans,
 	}
 
 	var wantSpansCount = 0
 	for i := 0; i < 2; i++ {
-		failures, _ := tt.ProcessSpans(batch, "test")
-		batchSize := len(batch.Spans)
+		failures, _ := tt.ProcessSpans(td, "test")
+		batchSize := len(td.Spans)
 		wantSpansCount += batchSize
 		if wantFailures != failures {
 			t.Errorf("Wanted %d failures but got %d", wantFailures, failures)
@@ -99,18 +99,18 @@ func TestMultiSpanProcessorWhenOneErrors(t *testing.T) {
 	m.MustFail = true
 
 	tt := NewMultiSpanProcessor(processors)
-	batch := &agenttracepb.ExportTraceServiceRequest{
+	td := data.TraceData{
 		Spans: make([]*tracepb.Span, 5),
 	}
 
 	var wantSpansCount = 0
 	for i := 0; i < 2; i++ {
-		failures, err := tt.ProcessSpans(batch, "test")
+		failures, err := tt.ProcessSpans(td, "test")
 		if err == nil {
 			t.Errorf("Wanted error got nil")
 			return
 		}
-		batchSize := len(batch.Spans)
+		batchSize := len(td.Spans)
 		wantSpansCount += batchSize
 		if failures != uint64(batchSize) {
 			t.Errorf("Wanted all spans to fail, got a different value.")
@@ -133,12 +133,12 @@ func TestMultiSpanProcessorWithPreProcessFn(t *testing.T) {
 	}
 
 	calledFnCount := int32(0)
-	testPreProcessFn := func(*agenttracepb.ExportTraceServiceRequest, string) {
+	testPreProcessFn := func(data.TraceData, string) {
 		atomic.AddInt32(&calledFnCount, 1)
 	}
 
 	tt := NewMultiSpanProcessor(processors, WithPreProcessFn(testPreProcessFn))
-	batch := &agenttracepb.ExportTraceServiceRequest{
+	batch := data.TraceData{
 		Spans: make([]*tracepb.Span, 7),
 	}
 
@@ -187,9 +187,9 @@ func multiSpanProcessorWithAddAttributesTestHelper(t *testing.T, overwrite bool)
 		}, overwrite),
 	)
 
-	batch := &agenttracepb.ExportTraceServiceRequest{}
+	td := data.TraceData{}
 	for i := 0; i < 7; i++ {
-		batch.Spans = append(batch.Spans, &tracepb.Span{
+		td.Spans = append(td.Spans, &tracepb.Span{
 			Attributes: &tracepb.Span_Attributes{
 				AttributeMap: map[string]*tracepb.AttributeValue{
 					"some_int": {
@@ -200,10 +200,10 @@ func multiSpanProcessorWithAddAttributesTestHelper(t *testing.T, overwrite bool)
 		})
 	}
 
-	spans := make([]*tracepb.Span, 0, len(batch.Spans)*2)
+	spans := make([]*tracepb.Span, 0, len(td.Spans)*2)
 	for i := 0; i < 2; i++ {
-		tt.ProcessSpans(batch, "test")
-		spans = append(spans, batch.Spans...)
+		tt.ProcessSpans(td, "test")
+		spans = append(spans, td.Spans...)
 	}
 
 	expectedSomeIntValue := int64(4567)
@@ -240,8 +240,8 @@ type mockSpanProcessor struct {
 
 var _ SpanProcessor = &mockSpanProcessor{}
 
-func (p *mockSpanProcessor) ProcessSpans(batch *agenttracepb.ExportTraceServiceRequest, spanFormat string) (uint64, error) {
-	batchSize := len(batch.Spans)
+func (p *mockSpanProcessor) ProcessSpans(td data.TraceData, spanFormat string) (uint64, error) {
+	batchSize := len(td.Spans)
 	p.TotalSpans += batchSize
 	if p.MustFail {
 		return uint64(batchSize), fmt.Errorf("this processor must fail")

internal/collector/processor/nodebatcher/node_batcher.go

@@ -26,13 +26,13 @@ import (
 	"unsafe"
 
 	commonpb "github.com/census-instrumentation/opencensus-proto/gen-go/agent/common/v1"
-	agenttracepb "github.com/census-instrumentation/opencensus-proto/gen-go/agent/trace/v1"
 	resourcepb "github.com/census-instrumentation/opencensus-proto/gen-go/resource/v1"
 	tracepb "github.com/census-instrumentation/opencensus-proto/gen-go/trace/v1"
 	"github.com/golang/protobuf/proto"
 	"go.opencensus.io/stats"
 	"go.uber.org/zap"
 
+	"github.com/census-instrumentation/opencensus-service/data"
 	"github.com/census-instrumentation/opencensus-service/internal/collector/processor"
 )
 
@@ -107,10 +107,10 @@ func NewBatcher(name string, logger *zap.Logger, sender processor.SpanProcessor,
 
 // ProcessSpans implements batcher as a SpanProcessor and takes the provided spans and adds them to
 // batches
-func (b *batcher) ProcessSpans(request *agenttracepb.ExportTraceServiceRequest, spanFormat string) (uint64, error) {
-	bucketID := b.genBucketID(request.Node, request.Resource, spanFormat)
-	bucket := b.getOrAddBucket(bucketID, request.Node, request.Resource, spanFormat)
-	bucket.add(request.Spans)
+func (b *batcher) ProcessSpans(td data.TraceData, spanFormat string) (uint64, error) {
+	bucketID := b.genBucketID(td.Node, td.Resource, spanFormat)
+	bucket := b.getOrAddBucket(bucketID, td.Node, td.Resource, spanFormat)
+	bucket.add(td.Spans)
 	return 0, nil
 }
 
@@ -270,12 +270,12 @@ func (nb *nodeBatcher) sendBatch(batch *batch) {
 	if len(spans) == 0 {
 		return
 	}
-	request := &agenttracepb.ExportTraceServiceRequest{
+	td := data.TraceData{
 		Node:     nb.node,
 		Resource: nb.resource,
 		Spans:    spans,
 	}
-	_, err := nb.parent.sender.ProcessSpans(request, nb.spanFormat)
+	_, err := nb.parent.sender.ProcessSpans(td, nb.spanFormat)
 	// Assumed that the next processor always handles a batch, and doesn't error
 	if err != nil {
 		nb.logger.Error(

internal/collector/processor/nodebatcher/node_batcher_test.go

@@ -20,9 +20,9 @@ import (
 	"time"
 
 	commonpb "github.com/census-instrumentation/opencensus-proto/gen-go/agent/common/v1"
-	agenttracepb "github.com/census-instrumentation/opencensus-proto/gen-go/agent/trace/v1"
 	resourcepb "github.com/census-instrumentation/opencensus-proto/gen-go/resource/v1"
 	tracepb "github.com/census-instrumentation/opencensus-proto/gen-go/trace/v1"
+	"github.com/census-instrumentation/opencensus-service/data"
 	"go.uber.org/zap"
 )
 
@@ -144,13 +144,13 @@ func TestConcurrentNodeAdds(t *testing.T) {
 		for spanIndex := 0; spanIndex < spansPerRequest; spanIndex++ {
 			spans = append(spans, &tracepb.Span{Name: getTestSpanName(requestNum, spanIndex)})
 		}
-		request := &agenttracepb.ExportTraceServiceRequest{
+		td := data.TraceData{
 			Node: &commonpb.Node{
 				ServiceInfo: &commonpb.ServiceInfo{Name: fmt.Sprintf("svc-%d", requestNum)},
 			},
 			Spans: spans,
 		}
-		go batcher.ProcessSpans(request, "oc")
+		go batcher.ProcessSpans(td, "oc")
 	}
 
 	err := sender.waitFor(requestCount*spansPerRequest, 3*time.Second)
@@ -190,7 +190,7 @@ func TestBucketRemove(t *testing.T) {
 	for spanIndex := 0; spanIndex < spansPerRequest; spanIndex++ {
 		spans = append(spans, &tracepb.Span{Name: getTestSpanName(0, spanIndex)})
 	}
-	request := &agenttracepb.ExportTraceServiceRequest{
+	request := data.TraceData{
 		Node: &commonpb.Node{
 			ServiceInfo: &commonpb.ServiceInfo{Name: "svc"},
 		},
@@ -225,7 +225,7 @@ func TestConcurrentBatchAdds(t *testing.T) {
 		for spanIndex := 0; spanIndex < spansPerRequest; spanIndex++ {
 			spans = append(spans, &tracepb.Span{Name: getTestSpanName(requestNum, spanIndex)})
 		}
-		request := &agenttracepb.ExportTraceServiceRequest{
+		request := data.TraceData{
 			Node: &commonpb.Node{
 				ServiceInfo: &commonpb.ServiceInfo{Name: "svc"},
 			},
@@ -254,7 +254,7 @@ func TestConcurrentBatchAdds(t *testing.T) {
 func BenchmarkConcurrentBatchAdds(b *testing.B) {
 	sender := newTestSender()
 	batcher := NewBatcher("test", zap.NewNop(), sender).(*batcher)
-	request := &agenttracepb.ExportTraceServiceRequest{
+	request := data.TraceData{
 		Node: &commonpb.Node{
 			ServiceInfo: &commonpb.ServiceInfo{Name: "svc"},
 		},
@@ -281,21 +281,21 @@ func getTestSpanName(requestNum, index int) *tracepb.TruncatableString {
 }
 
 type testSender struct {
-	reqChan             chan *agenttracepb.ExportTraceServiceRequest
+	reqChan             chan data.TraceData
 	batchesReceived     int
 	spansReceived       int
 	spansReceivedByName map[string]*tracepb.Span
 }
 
 func newTestSender() *testSender {
 	return &testSender{
-		reqChan:             make(chan *agenttracepb.ExportTraceServiceRequest, 100),
+		reqChan:             make(chan data.TraceData, 100),
 		spansReceivedByName: make(map[string]*tracepb.Span),
 	}
 }
 
-func (ts *testSender) ProcessSpans(request *agenttracepb.ExportTraceServiceRequest, spanFormat string) (uint64, error) {
-	ts.reqChan <- request
+func (ts *testSender) ProcessSpans(td data.TraceData, spanFormat string) (uint64, error) {
+	ts.reqChan <- td
 	return 0, nil
 }