Make max annotations and events per span configurable (#85)

* Implement changes to make max annotations and events per span configurable

* Go fmt the files modified

example/main.go

@@ -32,7 +32,11 @@ func main() {
 		ocagent.WithInsecure(),
 		ocagent.WithReconnectionPeriod(5*time.Second),
 		ocagent.WithAddress("localhost:55678"), // Only included here for demo purposes.
-		ocagent.WithServiceName("ocagent-go-example"))
+		ocagent.WithServiceName("ocagent-go-example"),
+		ocagent.WithSpanConfig(ocagent.SpanConfig{
+			AnnotationEventsPerSpan: 256,
+			MessageEventsPerSpan:    256,
+		}))
 	if err != nil {
 		log.Fatalf("Failed to create ocagent-exporter: %v", err)
 	}

ocagent.go

@@ -93,6 +93,8 @@ type Exporter struct {
 	clientTransportCredentials credentials.TransportCredentials
 
 	grpcDialOptions []grpc.DialOption
+
+	spanConfig SpanConfig
 }
 
 func NewExporter(opts ...ExporterOption) (*Exporter, error) {
@@ -475,14 +477,14 @@ func (ae *Exporter) ExportMetricsServiceRequest(batch *agentmetricspb.ExportMetr
 	}
 }
 
-func ocSpanDataToPbSpans(sdl []*trace.SpanData) []*tracepb.Span {
+func ocSpanDataToPbSpans(sdl []*trace.SpanData, spanConfig SpanConfig) []*tracepb.Span {
 	if len(sdl) == 0 {
 		return nil
 	}
 	protoSpans := make([]*tracepb.Span, 0, len(sdl))
 	for _, sd := range sdl {
 		if sd != nil {
-			protoSpans = append(protoSpans, ocSpanToProtoSpan(sd))
+			protoSpans = append(protoSpans, ocSpanToProtoSpan(sd, spanConfig))
 		}
 	}
 	return protoSpans
@@ -498,7 +500,7 @@ func (ae *Exporter) uploadTraces(sdl []*trace.SpanData) {
 			return
 		}
 
-		protoSpans := ocSpanDataToPbSpans(sdl)
+		protoSpans := ocSpanDataToPbSpans(sdl, ae.spanConfig)
 		if len(protoSpans) == 0 {
 			return
 		}

options.go

@@ -172,3 +172,14 @@ func (p metricNamePrefixSetter) withExporter(e *Exporter) {
 func WithMetricNamePrefix(prefix string) ExporterOption {
 	return metricNamePrefixSetter(prefix)
 }
+
+func (spanConfig SpanConfig) withExporter(e *Exporter) {
+	e.spanConfig = spanConfig
+}
+
+var _ ExporterOption = (*SpanConfig)(nil)
+
+// WithSpanConfig allows one to set the AnnotationEventsPerSpan and MessageEventsPerSpan
+func WithSpanConfig(spanConfig SpanConfig) ExporterOption {
+	return spanConfig
+}

span_config.go

@@ -0,0 +1,25 @@
+package ocagent
+
+const (
+	maxAnnotationEventsPerSpan = 32
+	maxMessageEventsPerSpan    = 128
+)
+
+type SpanConfig struct {
+	AnnotationEventsPerSpan int
+	MessageEventsPerSpan    int
+}
+
+func (spanConfig SpanConfig) GetAnnotationEventsPerSpan() int {
+	if spanConfig.AnnotationEventsPerSpan <= 0 {
+		return maxAnnotationEventsPerSpan
+	}
+	return spanConfig.AnnotationEventsPerSpan
+}
+
+func (spanConfig SpanConfig) GetMessageEventsPerSpan() int {
+	if spanConfig.MessageEventsPerSpan <= 0 {
+		return maxMessageEventsPerSpan
+	}
+	return spanConfig.MessageEventsPerSpan
+}

span_config_test.go

@@ -0,0 +1,59 @@
+package ocagent
+
+import "testing"
+
+func TestSpanConfig_GetAnnotationEventsPerSpan(t *testing.T) {
+	tests := []struct {
+		name                   string
+		inputSpanConfig        SpanConfig
+		expectedMaxAnnotations int
+	}{
+		{
+			name:                   "WhenConfigNotProvided",
+			inputSpanConfig:        SpanConfig{},
+			expectedMaxAnnotations: 32,
+		},
+		{
+			name:                   "WhenConfigProvided",
+			inputSpanConfig:        SpanConfig{AnnotationEventsPerSpan: 256},
+			expectedMaxAnnotations: 256,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			got := test.inputSpanConfig.GetAnnotationEventsPerSpan()
+			if got != test.expectedMaxAnnotations {
+				t.Errorf("Got = %d; want %d", got, test.expectedMaxAnnotations)
+			}
+		})
+	}
+}
+
+func TestSpanConfig_GetMessageEventsPerSpan(t *testing.T) {
+	tests := []struct {
+		name                     string
+		inputSpanConfig          SpanConfig
+		expectedMaxMessageEvents int
+	}{
+		{
+			name:                     "WhenConfigNotProvided",
+			inputSpanConfig:          SpanConfig{},
+			expectedMaxMessageEvents: 128,
+		},
+		{
+			name:                     "WhenConfigProvided",
+			inputSpanConfig:          SpanConfig{MessageEventsPerSpan: 256},
+			expectedMaxMessageEvents: 256,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			got := test.inputSpanConfig.GetMessageEventsPerSpan()
+			if got != test.expectedMaxMessageEvents {
+				t.Errorf("Got = %d; want %d", got, test.expectedMaxMessageEvents)
+			}
+		})
+	}
+}

transform_spans.go

@@ -25,12 +25,7 @@ import (
 	"github.com/golang/protobuf/ptypes/timestamp"
 )
 
-const (
-	maxAnnotationEventsPerSpan = 32
-	maxMessageEventsPerSpan    = 128
-)
-
-func ocSpanToProtoSpan(sd *trace.SpanData) *tracepb.Span {
+func ocSpanToProtoSpan(sd *trace.SpanData, spanConfig SpanConfig) *tracepb.Span {
 	if sd == nil {
 		return nil
 	}
@@ -49,7 +44,7 @@ func ocSpanToProtoSpan(sd *trace.SpanData) *tracepb.Span {
 		Kind:         ocSpanKindToProtoSpanKind(sd.SpanKind),
 		Name:         namePtr,
 		Attributes:   ocAttributesToProtoAttributes(sd.Attributes),
-		TimeEvents:   ocTimeEventsToProtoTimeEvents(sd.Annotations, sd.MessageEvents),
+		TimeEvents:   ocTimeEventsToProtoTimeEvents(sd.Annotations, sd.MessageEvents, spanConfig),
 		Tracestate:   ocTracestateToProtoTracestate(sd.Tracestate),
 	}
 }
@@ -131,7 +126,7 @@ func ocAttributesToProtoAttributes(attrs map[string]interface{}) *tracepb.Span_A
 
 // This code is mostly copied from
 // https://github.com/census-ecosystem/opencensus-go-exporter-stackdriver/blob/master/trace_proto.go#L46
-func ocTimeEventsToProtoTimeEvents(as []trace.Annotation, es []trace.MessageEvent) *tracepb.Span_TimeEvents {
+func ocTimeEventsToProtoTimeEvents(as []trace.Annotation, es []trace.MessageEvent, spanConfig SpanConfig) *tracepb.Span_TimeEvents {
 	if len(as) == 0 && len(es) == 0 {
 		return nil
 	}
@@ -142,7 +137,7 @@ func ocTimeEventsToProtoTimeEvents(as []trace.Annotation, es []trace.MessageEven
 
 	// Transform annotations
 	for i, a := range as {
-		if annotations >= maxAnnotationEventsPerSpan {
+		if annotations >= spanConfig.GetAnnotationEventsPerSpan() {
 			droppedAnnotationsCount = len(as) - i
 			break
 		}
@@ -157,7 +152,7 @@ func ocTimeEventsToProtoTimeEvents(as []trace.Annotation, es []trace.MessageEven
 
 	// Transform message events
 	for i, e := range es {
-		if messageEvents >= maxMessageEventsPerSpan {
+		if messageEvents >= spanConfig.GetMessageEventsPerSpan() {
 			droppedMessageEventsCount = len(es) - i
 			break
 		}