initial import

core/trace/carrier.go

@@ -0,0 +1,41 @@
+package trace
+
+import (
+	"errors"
+	"net/http"
+	"strings"
+)
+
+var ErrInvalidCarrier = errors.New("invalid carrier")
+
+type (
+	Carrier interface {
+		Get(key string) string
+		Set(key, value string)
+	}
+
+	httpCarrier http.Header
+	// grpc metadata takes keys as case insensitive
+	grpcCarrier map[string][]string
+)
+
+func (h httpCarrier) Get(key string) string {
+	return http.Header(h).Get(key)
+}
+
+func (h httpCarrier) Set(key, val string) {
+	http.Header(h).Set(key, val)
+}
+
+func (g grpcCarrier) Get(key string) string {
+	if vals, ok := g[strings.ToLower(key)]; ok && len(vals) > 0 {
+		return vals[0]
+	} else {
+		return ""
+	}
+}
+
+func (g grpcCarrier) Set(key, val string) {
+	key = strings.ToLower(key)
+	g[key] = append(g[key], val)
+}

core/trace/carrier_test.go

@@ -0,0 +1,59 @@
+package trace
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"zero/core/stringx"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHttpCarrier(t *testing.T) {
+	tests := []map[string]string{
+		{},
+		{
+			"first":  "a",
+			"second": "b",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(stringx.RandId(), func(t *testing.T) {
+			req := httptest.NewRequest(http.MethodGet, "http://localhost", nil)
+			carrier := httpCarrier(req.Header)
+			for k, v := range test {
+				carrier.Set(k, v)
+			}
+			for k, v := range test {
+				assert.Equal(t, v, carrier.Get(k))
+			}
+			assert.Equal(t, "", carrier.Get("none"))
+		})
+	}
+}
+
+func TestGrpcCarrier(t *testing.T) {
+	tests := []map[string]string{
+		{},
+		{
+			"first":  "a",
+			"second": "b",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(stringx.RandId(), func(t *testing.T) {
+			m := make(map[string][]string)
+			carrier := grpcCarrier(m)
+			for k, v := range test {
+				carrier.Set(k, v)
+			}
+			for k, v := range test {
+				assert.Equal(t, v, carrier.Get(k))
+			}
+			assert.Equal(t, "", carrier.Get("none"))
+		})
+	}
+}

core/trace/constants.go

@@ -0,0 +1,6 @@
+package trace
+
+const (
+	traceIdKey = "X-Trace-ID"
+	spanIdKey  = "X-Span-ID"
+)

core/trace/noop.go

@@ -0,0 +1,33 @@
+package trace
+
+import (
+	"context"
+
+	"zero/core/trace/tracespec"
+)
+
+var emptyNoopSpan = noopSpan{}
+
+type noopSpan struct{}
+
+func (s noopSpan) Finish() {
+}
+
+func (s noopSpan) Follow(ctx context.Context, serviceName, operationName string) (context.Context, tracespec.Trace) {
+	return ctx, emptyNoopSpan
+}
+
+func (s noopSpan) Fork(ctx context.Context, serviceName, operationName string) (context.Context, tracespec.Trace) {
+	return ctx, emptyNoopSpan
+}
+
+func (s noopSpan) SpanId() string {
+	return ""
+}
+
+func (s noopSpan) TraceId() string {
+	return ""
+}
+
+func (s noopSpan) Visit(fn func(key, val string) bool) {
+}

core/trace/noop_test.go

@@ -0,0 +1,32 @@
+package trace
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNoopSpan_Fork(t *testing.T) {
+	ctx, span := emptyNoopSpan.Fork(context.Background(), "", "")
+	assert.Equal(t, emptyNoopSpan, span)
+	assert.Equal(t, context.Background(), ctx)
+}
+
+func TestNoopSpan_Follow(t *testing.T) {
+	ctx, span := emptyNoopSpan.Follow(context.Background(), "", "")
+	assert.Equal(t, emptyNoopSpan, span)
+	assert.Equal(t, context.Background(), ctx)
+}
+
+func TestNoopSpan(t *testing.T) {
+	emptyNoopSpan.Visit(func(key, val string) bool {
+		assert.Fail(t, "should not go here")
+		return true
+	})
+
+	ctx, span := emptyNoopSpan.Follow(context.Background(), "", "")
+	assert.Equal(t, context.Background(), ctx)
+	assert.Equal(t, "", span.TraceId())
+	assert.Equal(t, "", span.SpanId())
+}

core/trace/propagator.go

@@ -0,0 +1,85 @@
+package trace
+
+import (
+	"net/http"
+
+	"google.golang.org/grpc/metadata"
+)
+
+const (
+	HttpFormat = iota
+	GrpcFormat
+)
+
+var (
+	emptyHttpPropagator httpPropagator
+	emptyGrpcPropagator grpcPropagator
+)
+
+type (
+	Propagator interface {
+		Extract(carrier interface{}) (Carrier, error)
+		Inject(carrier interface{}) (Carrier, error)
+	}
+
+	httpPropagator struct{}
+	grpcPropagator struct{}
+)
+
+func (h httpPropagator) Extract(carrier interface{}) (Carrier, error) {
+	if c, ok := carrier.(http.Header); !ok {
+		return nil, ErrInvalidCarrier
+	} else {
+		return httpCarrier(c), nil
+	}
+}
+
+func (h httpPropagator) Inject(carrier interface{}) (Carrier, error) {
+	if c, ok := carrier.(http.Header); ok {
+		return httpCarrier(c), nil
+	} else {
+		return nil, ErrInvalidCarrier
+	}
+}
+
+func (g grpcPropagator) Extract(carrier interface{}) (Carrier, error) {
+	if c, ok := carrier.(metadata.MD); ok {
+		return grpcCarrier(c), nil
+	} else {
+		return nil, ErrInvalidCarrier
+	}
+}
+
+func (g grpcPropagator) Inject(carrier interface{}) (Carrier, error) {
+	if c, ok := carrier.(metadata.MD); ok {
+		return grpcCarrier(c), nil
+	} else {
+		return nil, ErrInvalidCarrier
+	}
+}
+
+func Extract(format, carrier interface{}) (Carrier, error) {
+	switch v := format.(type) {
+	case int:
+		if v == HttpFormat {
+			return emptyHttpPropagator.Extract(carrier)
+		} else if v == GrpcFormat {
+			return emptyGrpcPropagator.Extract(carrier)
+		}
+	}
+
+	return nil, ErrInvalidCarrier
+}
+
+func Inject(format, carrier interface{}) (Carrier, error) {
+	switch v := format.(type) {
+	case int:
+		if v == HttpFormat {
+			return emptyHttpPropagator.Inject(carrier)
+		} else if v == GrpcFormat {
+			return emptyGrpcPropagator.Inject(carrier)
+		}
+	}
+
+	return nil, ErrInvalidCarrier
+}

core/trace/propagator_test.go

@@ -0,0 +1,68 @@
+package trace
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestHttpPropagator_Extract(t *testing.T) {
+	req := httptest.NewRequest(http.MethodGet, "http://localhost", nil)
+	req.Header.Set(traceIdKey, "trace")
+	req.Header.Set(spanIdKey, "span")
+	carrier, err := Extract(HttpFormat, req.Header)
+	assert.Nil(t, err)
+	assert.Equal(t, "trace", carrier.Get(traceIdKey))
+	assert.Equal(t, "span", carrier.Get(spanIdKey))
+
+	carrier, err = Extract(HttpFormat, req)
+	assert.Equal(t, ErrInvalidCarrier, err)
+}
+
+func TestHttpPropagator_Inject(t *testing.T) {
+	req := httptest.NewRequest(http.MethodGet, "http://localhost", nil)
+	req.Header.Set(traceIdKey, "trace")
+	req.Header.Set(spanIdKey, "span")
+	carrier, err := Inject(HttpFormat, req.Header)
+	assert.Nil(t, err)
+	assert.Equal(t, "trace", carrier.Get(traceIdKey))
+	assert.Equal(t, "span", carrier.Get(spanIdKey))
+
+	carrier, err = Inject(HttpFormat, req)
+	assert.Equal(t, ErrInvalidCarrier, err)
+}
+
+func TestGrpcPropagator_Extract(t *testing.T) {
+	md := metadata.New(map[string]string{
+		traceIdKey: "trace",
+		spanIdKey:  "span",
+	})
+	carrier, err := Extract(GrpcFormat, md)
+	assert.Nil(t, err)
+	assert.Equal(t, "trace", carrier.Get(traceIdKey))
+	assert.Equal(t, "span", carrier.Get(spanIdKey))
+
+	carrier, err = Extract(GrpcFormat, 1)
+	assert.Equal(t, ErrInvalidCarrier, err)
+	carrier, err = Extract(nil, 1)
+	assert.Equal(t, ErrInvalidCarrier, err)
+}
+
+func TestGrpcPropagator_Inject(t *testing.T) {
+	md := metadata.New(map[string]string{
+		traceIdKey: "trace",
+		spanIdKey:  "span",
+	})
+	carrier, err := Inject(GrpcFormat, md)
+	assert.Nil(t, err)
+	assert.Equal(t, "trace", carrier.Get(traceIdKey))
+	assert.Equal(t, "span", carrier.Get(spanIdKey))
+
+	carrier, err = Inject(GrpcFormat, 1)
+	assert.Equal(t, ErrInvalidCarrier, err)
+	carrier, err = Inject(nil, 1)
+	assert.Equal(t, ErrInvalidCarrier, err)
+}

core/trace/span.go

@@ -0,0 +1,144 @@
+package trace
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+	"strings"
+	"time"
+
+	"zero/core/stringx"
+	"zero/core/timex"
+	"zero/core/trace/tracespec"
+)
+
+const (
+	initSpanId  = "0"
+	clientFlag  = "client"
+	serverFlag  = "server"
+	spanSepRune = '.'
+	timeFormat  = "2006-01-02 15:04:05.000"
+)
+
+var spanSep = string([]byte{spanSepRune})
+
+type Span struct {
+	ctx           spanContext
+	serviceName   string
+	operationName string
+	startTime     time.Time
+	flag          string
+	children      int
+}
+
+func newServerSpan(carrier Carrier, serviceName, operationName string) tracespec.Trace {
+	traceId := stringx.TakeWithPriority(func() string {
+		if carrier != nil {
+			return carrier.Get(traceIdKey)
+		}
+		return ""
+	}, func() string {
+		return stringx.RandId()
+	})
+	spanId := stringx.TakeWithPriority(func() string {
+		if carrier != nil {
+			return carrier.Get(spanIdKey)
+		}
+		return ""
+	}, func() string {
+		return initSpanId
+	})
+
+	return &Span{
+		ctx: spanContext{
+			traceId: traceId,
+			spanId:  spanId,
+		},
+		serviceName:   serviceName,
+		operationName: operationName,
+		startTime:     timex.Time(),
+		flag:          serverFlag,
+	}
+}
+
+func (s *Span) Finish() {
+}
+
+func (s *Span) Follow(ctx context.Context, serviceName, operationName string) (context.Context, tracespec.Trace) {
+	span := &Span{
+		ctx: spanContext{
+			traceId: s.ctx.traceId,
+			spanId:  s.followSpanId(),
+		},
+		serviceName:   serviceName,
+		operationName: operationName,
+		startTime:     timex.Time(),
+		flag:          s.flag,
+	}
+	return context.WithValue(ctx, tracespec.TracingKey, span), span
+}
+
+func (s *Span) Fork(ctx context.Context, serviceName, operationName string) (context.Context, tracespec.Trace) {
+	span := &Span{
+		ctx: spanContext{
+			traceId: s.ctx.traceId,
+			spanId:  s.forkSpanId(),
+		},
+		serviceName:   serviceName,
+		operationName: operationName,
+		startTime:     timex.Time(),
+		flag:          clientFlag,
+	}
+	return context.WithValue(ctx, tracespec.TracingKey, span), span
+}
+
+func (s *Span) SpanId() string {
+	return s.ctx.SpanId()
+}
+
+func (s *Span) TraceId() string {
+	return s.ctx.TraceId()
+}
+
+func (s *Span) Visit(fn func(key, val string) bool) {
+	s.ctx.Visit(fn)
+}
+
+func (s *Span) forkSpanId() string {
+	s.children++
+	return fmt.Sprintf("%s.%d", s.ctx.spanId, s.children)
+}
+
+func (s *Span) followSpanId() string {
+	fields := strings.FieldsFunc(s.ctx.spanId, func(r rune) bool {
+		return r == spanSepRune
+	})
+	if len(fields) == 0 {
+		return s.ctx.spanId
+	}
+
+	last := fields[len(fields)-1]
+	val, err := strconv.Atoi(last)
+	if err != nil {
+		return s.ctx.spanId
+	}
+
+	last = strconv.Itoa(val + 1)
+	fields[len(fields)-1] = last
+
+	return strings.Join(fields, spanSep)
+}
+
+func StartClientSpan(ctx context.Context, serviceName, operationName string) (context.Context, tracespec.Trace) {
+	if span, ok := ctx.Value(tracespec.TracingKey).(*Span); ok {
+		return span.Fork(ctx, serviceName, operationName)
+	}
+
+	return ctx, emptyNoopSpan
+}
+
+func StartServerSpan(ctx context.Context, carrier Carrier, serviceName, operationName string) (
+	context.Context, tracespec.Trace) {
+	span := newServerSpan(carrier, serviceName, operationName)
+	return context.WithValue(ctx, tracespec.TracingKey, span), span
+}

core/trace/span_test.go

@@ -0,0 +1,140 @@
+package trace
+
+import (
+	"context"
+	"testing"
+
+	"zero/core/stringx"
+	"zero/core/trace/tracespec"
+
+	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestClientSpan(t *testing.T) {
+	span := newServerSpan(nil, "service", "operation")
+	ctx := context.WithValue(context.Background(), tracespec.TracingKey, span)
+	ctx, span = StartClientSpan(ctx, "entrance", "operation")
+	defer span.Finish()
+	assert.Equal(t, span, ctx.Value(tracespec.TracingKey))
+
+	const serviceName = "authorization"
+	const operationName = "verification"
+	ctx, childSpan := span.Fork(ctx, serviceName, operationName)
+	defer childSpan.Finish()
+
+	assert.Equal(t, childSpan, ctx.Value(tracespec.TracingKey))
+	assert.Equal(t, getSpan(span).TraceId(), getSpan(childSpan).TraceId())
+	assert.Equal(t, "0.1.1", getSpan(childSpan).SpanId())
+	assert.Equal(t, serviceName, childSpan.(*Span).serviceName)
+	assert.Equal(t, operationName, childSpan.(*Span).operationName)
+	assert.Equal(t, clientFlag, childSpan.(*Span).flag)
+}
+
+func TestClientSpan_WithoutTrace(t *testing.T) {
+	ctx, span := StartClientSpan(context.Background(), "entrance", "operation")
+	defer span.Finish()
+	assert.Equal(t, emptyNoopSpan, span)
+	assert.Equal(t, context.Background(), ctx)
+}
+
+func TestServerSpan(t *testing.T) {
+	ctx, span := StartServerSpan(context.Background(), nil, "service", "operation")
+	defer span.Finish()
+	assert.Equal(t, span, ctx.Value(tracespec.TracingKey))
+
+	const serviceName = "authorization"
+	const operationName = "verification"
+	ctx, childSpan := span.Fork(ctx, serviceName, operationName)
+	defer childSpan.Finish()
+
+	assert.Equal(t, childSpan, ctx.Value(tracespec.TracingKey))
+	assert.Equal(t, getSpan(span).TraceId(), getSpan(childSpan).TraceId())
+	assert.Equal(t, "0.1", getSpan(childSpan).SpanId())
+	assert.Equal(t, serviceName, childSpan.(*Span).serviceName)
+	assert.Equal(t, operationName, childSpan.(*Span).operationName)
+	assert.Equal(t, clientFlag, childSpan.(*Span).flag)
+}
+
+func TestServerSpan_WithCarrier(t *testing.T) {
+	md := metadata.New(map[string]string{
+		traceIdKey: "a",
+		spanIdKey:  "0.1",
+	})
+	ctx, span := StartServerSpan(context.Background(), grpcCarrier(md), "service", "operation")
+	defer span.Finish()
+	assert.Equal(t, span, ctx.Value(tracespec.TracingKey))
+
+	const serviceName = "authorization"
+	const operationName = "verification"
+	ctx, childSpan := span.Fork(ctx, serviceName, operationName)
+	defer childSpan.Finish()
+
+	assert.Equal(t, childSpan, ctx.Value(tracespec.TracingKey))
+	assert.Equal(t, getSpan(span).TraceId(), getSpan(childSpan).TraceId())
+	assert.Equal(t, "0.1.1", getSpan(childSpan).SpanId())
+	assert.Equal(t, serviceName, childSpan.(*Span).serviceName)
+	assert.Equal(t, operationName, childSpan.(*Span).operationName)
+	assert.Equal(t, clientFlag, childSpan.(*Span).flag)
+}
+
+func TestSpan_Follow(t *testing.T) {
+	tests := []struct {
+		span       string
+		expectSpan string
+	}{
+		{
+			"0.1",
+			"0.2",
+		},
+		{
+			"0",
+			"1",
+		},
+		{
+			"a",
+			"a",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(stringx.RandId(), func(t *testing.T) {
+			md := metadata.New(map[string]string{
+				traceIdKey: "a",
+				spanIdKey:  test.span,
+			})
+			ctx, span := StartServerSpan(context.Background(), grpcCarrier(md),
+				"service", "operation")
+			defer span.Finish()
+			assert.Equal(t, span, ctx.Value(tracespec.TracingKey))
+
+			const serviceName = "authorization"
+			const operationName = "verification"
+			ctx, childSpan := span.Follow(ctx, serviceName, operationName)
+			defer childSpan.Finish()
+
+			assert.Equal(t, childSpan, ctx.Value(tracespec.TracingKey))
+			assert.Equal(t, getSpan(span).TraceId(), getSpan(childSpan).TraceId())
+			assert.Equal(t, test.expectSpan, getSpan(childSpan).SpanId())
+			assert.Equal(t, serviceName, childSpan.(*Span).serviceName)
+			assert.Equal(t, operationName, childSpan.(*Span).operationName)
+			assert.Equal(t, span.(*Span).flag, childSpan.(*Span).flag)
+		})
+	}
+}
+
+func TestSpan_Visit(t *testing.T) {
+	var run bool
+	span := newServerSpan(nil, "service", "operation")
+	span.Visit(func(key, val string) bool {
+		assert.True(t, len(key) > 0)
+		assert.True(t, len(val) > 0)
+		run = true
+		return true
+	})
+	assert.True(t, run)
+}
+
+func getSpan(span tracespec.Trace) tracespec.Trace {
+	return span.(*Span)
+}

core/trace/spancontext.go

@@ -0,0 +1,19 @@
+package trace
+
+type spanContext struct {
+	traceId string
+	spanId  string
+}
+
+func (sc spanContext) TraceId() string {
+	return sc.traceId
+}
+
+func (sc spanContext) SpanId() string {
+	return sc.spanId
+}
+
+func (sc spanContext) Visit(fn func(key, val string) bool) {
+	fn(traceIdKey, sc.traceId)
+	fn(spanIdKey, sc.spanId)
+}

core/trace/tracespec/spancontext.go

@@ -0,0 +1,7 @@
+package tracespec
+
+type SpanContext interface {
+	TraceId() string
+	SpanId() string
+	Visit(fn func(key, val string) bool)
+}

core/trace/tracespec/trace.go

@@ -0,0 +1,10 @@
+package tracespec
+
+import "context"
+
+type Trace interface {
+	SpanContext
+	Finish()
+	Fork(ctx context.Context, serviceName, operationName string) (context.Context, Trace)
+	Follow(ctx context.Context, serviceName, operationName string) (context.Context, Trace)
+}

core/trace/tracespec/vars.go

@@ -0,0 +1,3 @@
+package tracespec
+
+const TracingKey = "X-Trace"