trace+span impl

trace_ctx.go → context.go

@@ -11,32 +11,25 @@ package trace_ctx
 
 import (
 	"context"
-	"strings"
-
-	"github.com/google/uuid"
 )
 
 var (
-	TraceKeyInCtx any = new(byte)
-	rp                = strings.NewReplacer("-", "")
+	TraceKeyInCtx any = new(TraceID)
+	TraceInCtx    any = new(Trace)
 )
 
 const (
 	TraceIDKeyName = "trace_id"
 )
 
-func genTraceID() string {
-	return rp.Replace(uuid.NewString())
-}
-
-func GetTraceID(ctx context.Context) string {
+func GetTraceID(ctx context.Context) TraceID {
 	if ctx == nil {
 		return genTraceID()
 	}
 
 	val := ctx.Value(TraceKeyInCtx)
 	if val != nil {
-		return val.(string)
+		return val.(TraceID)
 	}
 
 	return genTraceID()
@@ -50,10 +43,31 @@ func WithTraceID(ctx context.Context) context.Context {
 	return context.WithValue(ctx, TraceKeyInCtx, genTraceID())
 }
 
-func SetTraceID(ctx context.Context, traceID string) context.Context {
+func SetTraceID(ctx context.Context, traceID TraceID) context.Context {
 	if ctx == nil {
 		panic("ctx is nil")
 	}
 
 	return context.WithValue(ctx, TraceKeyInCtx, traceID)
 }
+
+func WithTrace(ctx context.Context, trace *Trace) context.Context {
+	if ctx == nil {
+		panic("ctx is nil")
+	}
+
+	return context.WithValue(ctx, TraceInCtx, trace)
+}
+
+func GetTrace(ctx context.Context) *Trace {
+	if ctx == nil {
+		panic("ctx is nil")
+	}
+
+	val := ctx.Value(TraceInCtx)
+	if val != nil {
+		return val.(*Trace)
+	}
+
+	return nil
+}

ctx_test.go

@@ -0,0 +1,14 @@
+package trace_ctx_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/rzaripov1990/trace_ctx"
+)
+
+func TestCtx(t *testing.T) {
+	ctx := trace_ctx.WithTraceID(context.Background())
+	fmt.Println(trace_ctx.GetTraceID(ctx))
+}

span.go

@@ -0,0 +1,72 @@
+package trace_ctx
+
+import (
+	"log/slog"
+	"sync"
+	"time"
+
+	"github.com/google/uuid"
+)
+
+type (
+	SpanID string
+
+	Span struct {
+		traceID      TraceID
+		id           SpanID
+		parentSpanID *SpanID
+		operation    string
+		startTime    int64
+		duration     int64
+		tags         *map[string]string
+		attrs        []slog.Attr
+		mtx          *sync.Mutex
+	}
+)
+
+func NewSpan(traceID TraceID, operation string, parentID *SpanID) *Span {
+	return &Span{
+		id:           genSpanID(),
+		traceID:      traceID,
+		parentSpanID: parentID,
+		operation:    operation,
+		startTime:    time.Now().UnixNano() / int64(time.Microsecond),
+		tags:         new(map[string]string),
+		attrs:        []slog.Attr{},
+		mtx:          &sync.Mutex{},
+	}
+}
+
+func (s *Span) End() {
+	s.duration = (time.Now().UnixNano() / int64(time.Microsecond)) - s.startTime
+}
+
+func (s *Span) AttrAdd(a slog.Attr) *Span {
+	s.mtx.Lock()
+	defer s.mtx.Unlock()
+
+	s.attrs = append(s.attrs, a)
+	return s
+}
+
+func (s *Span) GetAttrs() (a []slog.Attr) {
+	a = append(a, s.attrs...)
+	a = append(a,
+		slog.String(TraceIDKeyName, string(s.traceID)),
+		slog.String("span_id", string(s.id)),
+		slog.String("operation_name", s.operation),
+		slog.Int64("start_time", s.startTime),
+		slog.Int64("duration", s.duration),
+	)
+	if s.parentSpanID != nil {
+		a = append(a, slog.String("parent_span_id", string(*s.parentSpanID)))
+	}
+	if s.tags != nil {
+		a = append(a, slog.Any("tags", s.tags))
+	}
+	return a
+}
+
+func genSpanID() SpanID {
+	return SpanID(replace(uuid.New().String()))
+}

tools.go

@@ -0,0 +1,11 @@
+package trace_ctx
+
+import "strings"
+
+var (
+	rp = strings.NewReplacer("-", "")
+)
+
+func replace(s string) string {
+	return rp.Replace(s)
+}

trace.go

@@ -0,0 +1,60 @@
+package trace_ctx
+
+import (
+	"sync"
+
+	"github.com/google/uuid"
+)
+
+type (
+	TraceID string
+
+	Trace struct {
+		ID    TraceID
+		spans []*Span
+		mtx   *sync.Mutex
+	}
+)
+
+func NewTrace() *Trace {
+	return &Trace{
+		ID:    genTraceID(),
+		spans: []*Span{},
+		mtx:   &sync.Mutex{},
+	}
+}
+
+func (t *Trace) StartSpan(operation string, parent *Span) *Span {
+	t.mtx.Lock()
+	defer t.mtx.Unlock()
+
+	parentID := new(SpanID)
+	if parent != nil {
+		parentID = &parent.id
+	}
+
+	span := NewSpan(t.ID, operation, parentID)
+	t.spans = append(t.spans, span)
+	return span
+}
+
+func (t *Trace) WithSpan(operation string, parent *Span, f func()) *Span {
+	t.mtx.Lock()
+	defer t.mtx.Unlock()
+
+	parentID := new(SpanID)
+	if parent != nil {
+		parentID = &parent.id
+	}
+
+	span := NewSpan(t.ID, operation, parentID)
+	f()
+	span.End()
+	t.spans = append(t.spans, span)
+
+	return span
+}
+
+func genTraceID() TraceID {
+	return TraceID(replace(uuid.NewString()))
+}

tracing_test.go

@@ -0,0 +1,56 @@
+package trace_ctx_test
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/rzaripov1990/trace_ctx"
+)
+
+func TestTracing(t *testing.T) {
+	log := slog.New(
+		slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{
+			AddSource: false,
+			Level:     slog.LevelDebug,
+		}),
+	)
+
+	trace := trace_ctx.NewTrace()
+	ctx := trace_ctx.WithTrace(context.Background(), trace)
+
+	span1 := trace.StartSpan("operation-1", nil)
+	time.Sleep(100 * time.Millisecond)
+	span1.End()
+	span1.AttrAdd(slog.Bool("slow", true))
+	log.LogAttrs(ctx, slog.LevelDebug, "operation-1", span1.GetAttrs()...)
+
+	span2 := trace.StartSpan("operation-2", span1)
+	time.Sleep(200 * time.Millisecond)
+	span2.End()
+	log.LogAttrs(ctx, slog.LevelDebug, "operation-2", span2.GetAttrs()...)
+
+	var (
+		err   error
+		count int
+	)
+	log.LogAttrs(
+		ctx,
+		slog.LevelDebug,
+		"operation-3",
+		trace.WithSpan("operation-3", span2,
+			func() {
+				if count == 0 && err == nil {
+					count = 5
+				}
+				time.Sleep(300 * time.Millisecond)
+			},
+		).AttrAdd(slog.Int("count", count)).GetAttrs()...)
+
+	fmt.Println("Spans logged to standard output")
+
+	//fmt.Printf("%#v", trace_ctx.GetTrace(ctx))
+}