From 0967e1df3f04499901913f80371198ed2c330a32 Mon Sep 17 00:00:00 2001
From: JBD <jbd@golang.org>
Date: Tue, 20 Mar 2018 12:58:05 -0700
Subject: [PATCH] Introduce SpanKind support (#610)

Updates #525.
---
 trace/export.go     |  1 +
 trace/trace.go      | 14 ++++++--
 trace/trace_test.go | 86 ++++++++++++++++++++++++++++++++++++++++-----
 3 files changed, 91 insertions(+), 10 deletions(-)

diff --git a/trace/export.go b/trace/export.go
index db5d275a2..086612c9a 100644
--- a/trace/export.go
+++ b/trace/export.go
@@ -58,6 +58,7 @@ func UnregisterExporter(e Exporter) {
 type SpanData struct {
 	SpanContext
 	ParentSpanID SpanID
+	SpanKind     int
 	Name         string
 	StartTime    time.Time
 	// The wall clock time of EndTime will be adjusted to always be offset
diff --git a/trace/trace.go b/trace/trace.go
index 27e0539cc..0340b6a83 100644
--- a/trace/trace.go
+++ b/trace/trace.go
@@ -100,6 +100,13 @@ func WithSpan(parent context.Context, s *Span) context.Context {
 	return context.WithValue(parent, contextKey{}, s)
 }
 
+// All available span kinds. Span kind must be either one of these values.
+const (
+	SpanKindUnspecified = iota
+	SpanKindServer
+	SpanKindClient
+)
+
 // StartOptions contains options concerning how a span is started.
 type StartOptions struct {
 	// Sampler to consult for this Span. If provided, it is always consulted.
@@ -111,9 +118,11 @@ type StartOptions struct {
 	// when there is a non-remote parent, no new sampling decision will be made:
 	// we will preserve the sampling of the parent.
 	Sampler Sampler
-}
 
-// TODO(jbd): Remove start options.
+	// SpanKind represents the kind of a span. If none is set,
+	// SpanKindUnspecified is used.
+	SpanKind int
+}
 
 // StartSpan starts a new child span of the current span in the context. If
 // there is no span in the context, creates a new trace and span.
@@ -180,6 +189,7 @@ func startSpanInternal(name string, hasParent bool, parent SpanContext, remotePa
 	span.data = &SpanData{
 		SpanContext:     span.spanContext,
 		StartTime:       time.Now(),
+		SpanKind:        o.SpanKind,
 		Name:            name,
 		HasRemoteParent: remoteParent,
 	}
diff --git a/trace/trace_test.go b/trace/trace_test.go
index 3dbb0f298..862b2c939 100644
--- a/trace/trace_test.go
+++ b/trace/trace_test.go
@@ -236,14 +236,15 @@ func TestStartSpanWithRemoteParent(t *testing.T) {
 }
 
 // startSpan returns a context with a new Span that is recording events and will be exported.
-func startSpan() *Span {
+func startSpan(o StartOptions) *Span {
 	return NewSpanWithRemoteParent("span0",
 		SpanContext{
 			TraceID:      tid,
 			SpanID:       sid,
 			TraceOptions: 1,
 		},
-		StartOptions{})
+		o,
+	)
 }
 
 type testExporter struct {
@@ -295,8 +296,77 @@ func checkTime(x *time.Time) bool {
 	return true
 }
 
+func TestSpanKind(t *testing.T) {
+	tests := []struct {
+		name         string
+		startOptions StartOptions
+		want         *SpanData
+	}{
+		{
+			name:         "zero StartOptions",
+			startOptions: StartOptions{},
+			want: &SpanData{
+				SpanContext: SpanContext{
+					TraceID:      tid,
+					SpanID:       SpanID{},
+					TraceOptions: 0x1,
+				},
+				ParentSpanID:    sid,
+				Name:            "span0",
+				SpanKind:        SpanKindUnspecified,
+				HasRemoteParent: true,
+			},
+		},
+		{
+			name: "client span",
+			startOptions: StartOptions{
+				SpanKind: SpanKindClient,
+			},
+			want: &SpanData{
+				SpanContext: SpanContext{
+					TraceID:      tid,
+					SpanID:       SpanID{},
+					TraceOptions: 0x1,
+				},
+				ParentSpanID:    sid,
+				Name:            "span0",
+				SpanKind:        SpanKindClient,
+				HasRemoteParent: true,
+			},
+		},
+		{
+			name: "server span",
+			startOptions: StartOptions{
+				SpanKind: SpanKindServer,
+			},
+			want: &SpanData{
+				SpanContext: SpanContext{
+					TraceID:      tid,
+					SpanID:       SpanID{},
+					TraceOptions: 0x1,
+				},
+				ParentSpanID:    sid,
+				Name:            "span0",
+				SpanKind:        SpanKindServer,
+				HasRemoteParent: true,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		span := startSpan(tt.startOptions)
+		got, err := endSpan(span)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !reflect.DeepEqual(got, tt.want) {
+			t.Errorf("exporting span: got %#v want %#v", got, tt.want)
+		}
+	}
+}
+
 func TestSetSpanAttributes(t *testing.T) {
-	span := startSpan()
+	span := startSpan(StartOptions{})
 	span.AddAttributes(StringAttribute("key1", "value1"))
 	got, err := endSpan(span)
 	if err != nil {
@@ -320,7 +390,7 @@ func TestSetSpanAttributes(t *testing.T) {
 }
 
 func TestAnnotations(t *testing.T) {
-	span := startSpan()
+	span := startSpan(StartOptions{})
 	span.Annotatef([]Attribute{StringAttribute("key1", "value1")}, "%f", 1.5)
 	span.Annotate([]Attribute{StringAttribute("key2", "value2")}, "Annotate")
 	got, err := endSpan(span)
@@ -354,7 +424,7 @@ func TestAnnotations(t *testing.T) {
 }
 
 func TestMessageEvents(t *testing.T) {
-	span := startSpan()
+	span := startSpan(StartOptions{})
 	span.AddMessageReceiveEvent(3, 400, 300)
 	span.AddMessageSendEvent(1, 200, 100)
 	got, err := endSpan(span)
@@ -388,7 +458,7 @@ func TestMessageEvents(t *testing.T) {
 }
 
 func TestSetSpanStatus(t *testing.T) {
-	span := startSpan()
+	span := startSpan(StartOptions{})
 	span.SetStatus(Status{Code: int32(1), Message: "request failed"})
 	got, err := endSpan(span)
 	if err != nil {
@@ -412,7 +482,7 @@ func TestSetSpanStatus(t *testing.T) {
 }
 
 func TestAddLink(t *testing.T) {
-	span := startSpan()
+	span := startSpan(StartOptions{})
 	span.AddLink(Link{
 		TraceID:    tid,
 		SpanID:     sid,
@@ -450,7 +520,7 @@ func TestUnregisterExporter(t *testing.T) {
 	RegisterExporter(&te)
 	UnregisterExporter(&te)
 
-	ctx := startSpan()
+	ctx := startSpan(StartOptions{})
 	endSpan(ctx)
 	if len(te.spans) != 0 {
 		t.Error("unregistered Exporter was called")