refactor: logger with tid & improve open tracing in metadata (#651)

Author: woorui

core/client_test.go

@@ -129,7 +129,7 @@ func TestFrameRoundTrip(t *testing.T) {
 	exited = checkClientExited(sfn, time.Second)
 	assert.False(t, exited, "sfn stream should not exited")
 
-	sfnMd := NewDefaultMetadata(source.clientID, "sfn-tid", "sfn-sid", false)
+	sfnMd := NewMetadata(source.clientID, "tid", "trace-id", "span-id", false)
 
 	sfnMetaBytes, _ := sfnMd.Encode()
 
@@ -152,7 +152,7 @@ func TestFrameRoundTrip(t *testing.T) {
 	assert.ElementsMatch(t, nameList, []string{"source", "sfn-1"})
 
 	md := metadata.New(
-		NewDefaultMetadata(source.clientID, "tid", "sid", false),
+		NewMetadata(source.clientID, "tid", "trace-id", "span-id", false),
 		metadata.M{
 			"foo": "bar",
 		},

core/context.go

@@ -132,6 +132,9 @@ func (c *Context) WithFrame(f frame.Frame) error {
 	c.Frame = df
 	c.FrameMetadata = fmd
 
+	// log with tid
+	c.Logger = c.BaseLogger.With("tid", GetTIDFromMetadata(fmd))
+
 	return nil
 }
 

core/metadata.go

@@ -2,26 +2,31 @@ package core
 
 import (
 	"github.com/yomorun/yomo/core/metadata"
+	"github.com/yomorun/yomo/pkg/id"
+	"github.com/yomorun/yomo/pkg/trace"
+	oteltrace "go.opentelemetry.io/otel/trace"
+	"golang.org/x/exp/slog"
 )
 
 const (
+	// the keys for yomo working.
 	MetadataSourceIDKey = "yomo-source-id"
 	MetadataTIDKey      = "yomo-tid"
-	MetadataSIDKey      = "yomo-sid"
-	MetaTraced          = "yomo-traced"
+
+	// the keys for tracing.
+	MetadataTraceIDKey = "yomo-trace-id"
+	MetadataSpanIDKey  = "yomo-span-id"
+	MetaTracedKey      = "yomo-traced"
 )
 
-// NewDefaultMetadata returns a default metadata.
-func NewDefaultMetadata(sourceID string, tid string, sid string, traced bool) metadata.M {
-	tracedString := "false"
-	if traced {
-		tracedString = "true"
-	}
+// NewMetadata returns metadata for yomo working.
+func NewMetadata(sourceID, tid string, traceID string, spanID string, traced bool) metadata.M {
 	return metadata.M{
 		MetadataSourceIDKey: sourceID,
 		MetadataTIDKey:      tid,
-		MetadataSIDKey:      sid,
-		MetaTraced:          tracedString,
+		MetadataTraceIDKey:  traceID,
+		MetadataSpanIDKey:   spanID,
+		MetaTracedKey:       tracedString(traced),
 	}
 }
 
@@ -37,33 +42,124 @@ func GetTIDFromMetadata(m metadata.M) string {
 	return tid
 }
 
-// GetSIDFromMetadata gets SID from metadata.
-func GetSIDFromMetadata(m metadata.M) string {
-	sid, _ := m.Get(MetadataSIDKey)
-	return sid
-}
-
 // GetTracedFromMetadata gets traced from metadata.
 func GetTracedFromMetadata(m metadata.M) bool {
-	traced, _ := m.Get(MetaTraced)
-	return traced == "true"
+	tracedString, _ := m.Get(MetaTracedKey)
+	return tracedString == "true"
+}
+
+// SourceMetadata generates source metadata with trace information.
+func SourceMetadata(
+	sourceID, tid string,
+	spanName string, // the span name usually is the source name.
+	tp oteltrace.TracerProvider, logger *slog.Logger,
+) (metadata.M, func()) {
+	var (
+		traceID string
+		spanID  string
+		traced  bool
+		endFn   = func() {}
+	)
+	if tp != nil {
+		span, err := trace.NewSpan(tp, "Source", spanName, "", "")
+		if err != nil {
+			logger.Debug("trace error", "tracer_name", "Source", "span_name", spanName, "err", err)
+		} else {
+			endFn = func() { span.End() }
+			traceID = span.SpanContext().TraceID().String()
+			spanID = span.SpanContext().SpanID().String()
+			traced = true
+		}
+	}
+	if traceID == "" {
+		logger.Debug("create new traceID", "tracer_name", "Source", "span_name", spanName, "trace_id", traceID)
+		traceID = id.NewTraceID()
+	}
+	if spanID == "" {
+		logger.Debug("create new spanID", "tracer_name", "Source", "span_name", spanName, "span_id", spanID)
+		spanID = id.NewSpanID()
+	}
+	logger.Debug(
+		"trace metadata",
+		"tracer_name", "Source", "span_name", spanName,
+		"trace_id", traceID, "span_id", spanID, "traced", traced,
+	)
+	md := NewMetadata(sourceID, id.New(), traceID, spanID, traced)
+
+	return md, endFn
+}
+
+// ExtendTraceMetadata extends source metadata with trace information.
+func ExtendTraceMetadata(
+	md metadata.M,
+	tracerName string, // the tracer name is `StreamFunction` or `Zipper`.
+	spanName string, // the span name usually is the sfn name.
+	tp oteltrace.TracerProvider, logger *slog.Logger,
+) (metadata.M, func()) {
+	var (
+		traceID, _   = md.Get(MetadataTraceIDKey)
+		spanID, _    = md.Get(MetadataSpanIDKey)
+		parentTraced = GetTracedFromMetadata(md)
+		endFn        = func() {}
+	)
+	traced := false
+	if tp != nil {
+		var span oteltrace.Span
+		var err error
+		// set parent span, if not traced, use empty string
+		if parentTraced {
+			span, err = trace.NewSpan(tp, string(tracerName), spanName, traceID, spanID)
+		} else {
+			span, err = trace.NewSpan(tp, string(tracerName), spanName, "", "")
+		}
+		if err != nil {
+			logger.Debug("trace error", "tracer_name", tracerName, "span_name", spanName, "err", err)
+		} else {
+			endFn = func() { span.End() }
+			traceID = span.SpanContext().TraceID().String()
+			spanID = span.SpanContext().SpanID().String()
+			traced = true
+		}
+	}
+	if traceID == "" {
+		logger.Debug("create new traceID", "tracer_name", tracerName, "span_name", spanName, "trace_id", traceID)
+		traceID = id.NewTraceID()
+	}
+	if spanID == "" {
+		logger.Debug("create new spanID", "tracer_name", tracerName, "span_name", spanName, "span_id", spanID)
+		spanID = id.NewSpanID()
+	}
+	logger.Debug(
+		"trace metadata",
+		"tracer_name", tracerName, "span_name", spanName,
+		"trace_id", traceID, "span_id", spanID, "traced", traced, "parent_traced", parentTraced,
+	)
+
+	if tracerName == "Zipper" {
+		traced = traced || parentTraced
+	}
+
+	// reallocate metadata with new TraceID and SpanID
+	md.Set(MetadataTraceIDKey, traceID)
+	md.Set(MetadataSpanIDKey, spanID)
+	md.Set(MetaTracedKey, tracedString(traced))
+
+	return md, endFn
 }
 
-// SetTIDToMetadata sets tid to metadata.
-func SetTIDToMetadata(m metadata.M, tid string) {
-	m.Set(MetadataTIDKey, tid)
+// SfnTraceMetadata extends metadata for StreamFunction.
+func SfnTraceMetadata(md metadata.M, sfnName string, tp oteltrace.TracerProvider, logger *slog.Logger) (metadata.M, func()) {
+	return ExtendTraceMetadata(md, "StreamFunction", sfnName, tp, logger)
 }
 
-// SetSIDToMetadata sets sid to metadata.
-func SetSIDToMetadata(m metadata.M, sid string) {
-	m.Set(MetadataSIDKey, sid)
+// ZipperTraceMetadata extends metadata for Zipper.
+func ZipperTraceMetadata(md metadata.M, tp oteltrace.TracerProvider, logger *slog.Logger) (metadata.M, func()) {
+	return ExtendTraceMetadata(md, "Zipper", "zipper endpoint", tp, logger)
 }
 
-// SetTracedToMetadata sets traced to metadata.
-func SetTracedToMetadata(m metadata.M, traced bool) {
-	tracedString := "false"
+func tracedString(traced bool) string {
 	if traced {
-		tracedString = "true"
+		return "true"
 	}
-	m.Set(MetaTraced, tracedString)
+	return "false"
 }

core/metadata_test.go

@@ -7,19 +7,9 @@ import (
 )
 
 func TestMetadata(t *testing.T) {
-	md := NewDefaultMetadata("source", "xxxxxxx", "sssssss", true)
+	md := NewMetadata("source", "tid", "traceID", "spanID", true)
 
 	assert.Equal(t, "source", GetSourceIDFromMetadata(md))
-	assert.Equal(t, "xxxxxxx", GetTIDFromMetadata(md))
-	assert.Equal(t, "sssssss", GetSIDFromMetadata(md))
+	assert.Equal(t, "tid", GetTIDFromMetadata(md))
 	assert.Equal(t, true, GetTracedFromMetadata(md))
-
-	SetTIDToMetadata(md, "ccccccc")
-	assert.Equal(t, "ccccccc", GetTIDFromMetadata(md))
-
-	SetSIDToMetadata(md, "aaaaaaa")
-	assert.Equal(t, "aaaaaaa", GetSIDFromMetadata(md))
-
-	SetTracedToMetadata(md, false)
-	assert.Equal(t, false, GetTracedFromMetadata(md))
 }

core/server.go

@@ -20,9 +20,7 @@ import (
 	// authentication implements, Currently, only token authentication is implemented
 	_ "github.com/yomorun/yomo/pkg/auth"
 	"github.com/yomorun/yomo/pkg/frame-codec/y3codec"
-	"github.com/yomorun/yomo/pkg/id"
 	pkgtls "github.com/yomorun/yomo/pkg/tls"
-	"github.com/yomorun/yomo/pkg/trace"
 	oteltrace "go.opentelemetry.io/otel/trace"
 )
 
@@ -345,51 +343,18 @@ func (s *Server) handleDataFrame(c *Context) error {
 	// counter +1
 	atomic.AddInt64(&s.counterOfDataFrame, 1)
 
-	from := c.Connection
-	tid := GetTIDFromMetadata(c.FrameMetadata)
-	sid := GetSIDFromMetadata(c.FrameMetadata)
-	parentTraced := GetTracedFromMetadata(c.FrameMetadata)
-	traced := false
-	// trace
-	tp := s.TracerProvider()
-	if tp != nil {
-		// create span
-		var span oteltrace.Span
-		var err error
-		// set parent span, if not traced, use empty string
-		if parentTraced {
-			span, err = trace.NewSpan(tp, "zipper", "handle DataFrame", tid, sid)
-		} else {
-			span, err = trace.NewSpan(tp, "zipper", "handle DataFrame", "", "")
-		}
-		if err != nil {
-			s.logger.Error("zipper trace error", "err", err)
-		} else {
-			defer span.End()
-			tid = span.SpanContext().TraceID().String()
-			sid = span.SpanContext().SpanID().String()
-			traced = true
-		}
-	}
-	if tid == "" {
-		s.logger.Debug("zipper create new tid")
-		tid = id.TID()
-	}
-	if sid == "" {
-		s.logger.Debug("zipper create new sid")
-		sid = id.SID()
-	}
-	// reallocate metadata with new TID and SID
-	SetTIDToMetadata(c.FrameMetadata, tid)
-	SetSIDToMetadata(c.FrameMetadata, sid)
-	SetTracedToMetadata(c.FrameMetadata, traced || parentTraced)
-	md, err := c.FrameMetadata.Encode()
+	md, endFn := ZipperTraceMetadata(c.FrameMetadata, s.TracerProvider(), c.Logger)
+	defer endFn()
+
+	c.FrameMetadata = md
+
+	mdBytes, err := c.FrameMetadata.Encode()
 	if err != nil {
-		s.logger.Error("encode metadata error", "err", err)
+		c.Logger.Error("encode metadata error", "err", err)
 		return err
 	}
-	dataFrame.Metadata = md
-	s.logger.Debug("zipper metadata", "tid", tid, "sid", sid, "parentTraced", parentTraced, "traced", traced, "frome_stream_name", from.Name())
+	dataFrame.Metadata = mdBytes
+
 	// route
 	route := s.router.Route(c.FrameMetadata)
 	if route == nil {
@@ -415,7 +380,7 @@ func (s *Server) handleDataFrame(c *Context) error {
 			continue
 		}
 
-		c.Logger.Info("data routing", "tid", tid, "sid", sid, "tag", dataFrame.Tag, "data_length", data_length, "to_id", toID, "to_name", stream.Name())
+		c.Logger.Info("data routing", "tag", dataFrame.Tag, "data_length", data_length, "to_id", toID, "to_name", stream.Name())
 
 		// write data frame to stream
 		if err := stream.WriteFrame(dataFrame); err != nil {
@@ -512,10 +477,6 @@ func (s *Server) dispatchToDownstreams(c *Context) {
 		return
 	}
 
-	var (
-		tid = GetTIDFromMetadata(c.FrameMetadata)
-		sid = GetSIDFromMetadata(c.FrameMetadata)
-	)
 	mdBytes, err := c.FrameMetadata.Encode()
 	if err != nil {
 		c.Logger.Error("failed to dispatch to downstream", "err", err)
@@ -526,9 +487,9 @@ func (s *Server) dispatchToDownstreams(c *Context) {
 	for _, ds := range s.downstreams {
 		c.Logger.Info(
 			"dispatching to downstream",
-			"tid", tid, "sid", sid, "tag", dataFrame.Tag, "data_length", len(dataFrame.Payload),
-			"downstream_id", ds.ID(), "downstream_name", ds.LocalName(),
-		)
+			"tag", dataFrame.Tag, "data_length", len(dataFrame.Payload),
+			"downstream_id", ds.ID(), "downstream_name", ds.LocalName())
+
 		_ = ds.WriteFrame(dataFrame)
 	}
 }

pkg/id/id.go

@@ -19,17 +19,17 @@ func New(l ...int) string {
 	return tid
 }
 
-// TID generate trace id.
-func TID() string {
+// NewTraceID returns a trace id.
+func NewTraceID() string {
 	bytes := make([]byte, 16)
 	if _, err := rand.Read(bytes); err != nil {
 		return ""
 	}
 	return hex.EncodeToString(bytes)
 }
 
-// SID generate span id.
-func SID() string {
+// NewSpanID returns a span id.
+func NewSpanID() string {
 	bytes := make([]byte, 8)
 	if _, err := rand.Read(bytes); err != nil {
 		return ""

pkg/id/id_test.go

@@ -7,14 +7,20 @@ import (
 )
 
 func TestNew(t *testing.T) {
-	str := New()
-	assert.IsType(t, "", str)
+	t.Run("new a random id", func(t *testing.T) {
+		str := New()
+		assert.IsType(t, "", str)
+	})
 
-	tid := TID()
-	assert.IsType(t, "", tid)
-	assert.Equal(t, 32, len(tid))
+	t.Run("new trace id", func(t *testing.T) {
+		traceID := NewTraceID()
+		assert.IsType(t, "", traceID)
+		assert.Equal(t, 32, len(traceID))
+	})
 
-	sid := SID()
-	assert.IsType(t, "", sid)
-	assert.Equal(t, 16, len(sid))
+	t.Run("new span id", func(t *testing.T) {
+		spanID := NewSpanID()
+		assert.IsType(t, "", spanID)
+		assert.Equal(t, 16, len(spanID))
+	})
 }