Repository: lf-edge/ekuiper
Pull request #3187 (2024-09-11 02:08:01)

docs: fix trace doc


Patch:
diff --git a/docs/directory.json b/docs/directory.json
index d836b8db..94e4bf23 100644
--- a/docs/directory.json
+++ b/docs/directory.json
@@ -410,6 +410,10 @@
         {
           "title": "错误码",
           "path": "operation/usage/error_code"
+        },
+        {
+          "title": "数据链路追踪",
+          "path": "operation/usage/trace_data"
         }
       ]
     },
@@ -542,6 +546,10 @@
         {
           "title": "规则试运行",
           "path": "api/restapi/ruletest"
+        },
+        {
+          "title": "数据链路追踪",
+          "path": "api/restapi/trace"
         }
       ]
     },
@@ -1228,6 +1236,10 @@
               "path": "operation/manager-ui/plugins_in_manager"
             }
           ]
+        },
+        {
+          "title": "Tracing Data",
+          "path": "operation/usage/trace_data"
         }
       ]
     },
@@ -1299,6 +1311,10 @@
         {
           "title": "Rule Test",
           "path": "api/restapi/ruletest"
+        },
+        {
+          "title": "Trace Data",
+          "path": "api/restapi/trace"
         }
       ]
     },
diff --git a/internal/io/neuron/sink.go b/internal/io/neuron/sink.go
index a1b64e99..8b24efc3 100644
--- a/internal/io/neuron/sink.go
+++ b/internal/io/neuron/sink.go
@@ -22,10 +22,12 @@ import (
 	"github.com/lf-edge/ekuiper/contract/v2/api"
 
 	"github.com/lf-edge/ekuiper/v2/internal/conf"
+	"github.com/lf-edge/ekuiper/v2/internal/xsql"
 	"github.com/lf-edge/ekuiper/v2/pkg/cast"
 	"github.com/lf-edge/ekuiper/v2/pkg/connection"
 	"github.com/lf-edge/ekuiper/v2/pkg/errorx"
 	"github.com/lf-edge/ekuiper/v2/pkg/nng"
+	"github.com/lf-edge/ekuiper/v2/pkg/tracer"
 )
 
 type c struct {
@@ -100,6 +102,7 @@ func (s *sink) Collect(ctx api.StreamContext, data api.MessageTuple) error {
 		if err != nil {
 			return err
 		}
+		r = extractSpanContextIntoData(ctx, data, r)
 		return s.cli.Send(ctx, r)
 	} else {
 		return s.SendMapToNeuron(ctx, data)
@@ -160,7 +163,7 @@ func (s *sink) SendMapToNeuron(ctx api.StreamContext, tuple api.MessageTuple) er
 			for _, k := range keys {
 				t.TagName = k
 				t.Value = el[k]
-				err := doPublish(ctx, s.cli, t)
+				err := doPublish(ctx, s.cli, tuple, t)
 				if err != nil {
 					return err
 				}
@@ -169,7 +172,7 @@ func (s *sink) SendMapToNeuron(ctx api.StreamContext, tuple api.MessageTuple) er
 			for k, v := range el {
 				t.TagName = k
 				t.Value = v
-				err := doPublish(ctx, s.cli, t)
+				err := doPublish(ctx, s.cli, tuple, t)
 				if err != nil {
 					return err
 				}
@@ -188,7 +191,7 @@ func (s *sink) SendMapToNeuron(ctx api.StreamContext, tuple api.MessageTuple) er
 				ctx.GetLogger().Errorf("Error get the value of tag %s: %v", t.TagName, err)
 				continue
 			}
-			err := doPublish(ctx, s.cli, t)
+			err := doPublish(ctx, s.cli, tuple, t)
 			if err != nil {
 				return err
 			}
@@ -197,11 +200,12 @@ func (s *sink) SendMapToNeuron(ctx api.StreamContext, tuple api.MessageTuple) er
 	return nil
 }
 
-func doPublish(ctx api.StreamContext, cli *nng.Sock, t *neuronTemplate) error {
+func doPublish(ctx api.StreamContext, cli *nng.Sock, tuple api.MessageTuple, t *neuronTemplate) error {
 	r, err := json.Marshal(t)
 	if err != nil {
 		return fmt.Errorf("Error marshall the tag payload %v: %v", t, err)
 	}
+	r = extractSpanContextIntoData(ctx, tuple, r)
 	err = cli.Send(ctx, r)
 	if err != nil {
 		return errorx.NewIOErr(fmt.Sprintf(`Error publish the tag payload %s: %v`, t.TagName, err))
@@ -210,6 +214,21 @@ func doPublish(ctx api.StreamContext, cli *nng.Sock, t *neuronTemplate) error {
 	return nil
 }
 
+func extractSpanContextIntoData(ctx api.StreamContext, data interface{}, sendBytes []byte) []byte {
+	if tracerCtx, ok := data.(xsql.HasTracerCtx); ok && ctx.IsTraceEnabled() {
+		_, span := tracer.GetTracer().Start(tracerCtx.GetTracerCtx(), ctx.GetOpId())
+		traceID := span.SpanContext().TraceID()
+		spanID := span.SpanContext().SpanID()
+		defer span.End()
+		r := NeuronTraceHeader
+		r = append(r, traceID[:]...)
+		r = append(r, spanID[:]...)
+		r = append(r, sendBytes...)
+		return r
+	}
+	return sendBytes
+}
+
 func GetSink() api.Sink {
 	return &sink{}
 }
diff --git a/internal/io/neuron/source.go b/internal/io/neuron/source.go
index 5013e6d0..69dbafcc 100644
--- a/internal/io/neuron/source.go
+++ b/internal/io/neuron/source.go
@@ -15,13 +15,16 @@
 package neuron
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"time"
 
 	"github.com/lf-edge/ekuiper/contract/v2/api"
 	"go.nanomsg.org/mangos/v3"
+	"go.opentelemetry.io/otel/trace"
 
+	"github.com/lf-edge/ekuiper/v2/internal/topo/node/tracenode"
 	"github.com/lf-edge/ekuiper/v2/pkg/infra"
 	"github.com/lf-edge/ekuiper/v2/pkg/nng"
 	"github.com/lf-edge/ekuiper/v2/pkg/timex"
@@ -32,6 +35,15 @@ const (
 	PROTOCOL         = "pair"
 )
 
+var (
+	NeuronTraceHeader           = []byte{0x0A, 0xCE}
+	NeuronTraceIDStartIndex     = len(NeuronTraceHeader)
+	NeuronTraceIDEndIndex       = NeuronTraceIDStartIndex + 16
+	NeuronTraceSpanIDStartIndex = NeuronTraceIDEndIndex
+	NeuronTraceSpanIDEndIndex   = NeuronTraceSpanIDStartIndex + 8
+	NeuronTraceHeaderLen        = 2 + 16 + 8
+)
+
 type source struct {
 	c     *nng.SockConf
 	cli   *nng.Sock
@@ -84,7 +96,8 @@ func (s *source) Subscribe(ctx api.StreamContext, ingest api.BytesIngest, ingest
 				// no receiving deadline, will wait until the socket closed
 				if msg, err := s.cli.Recv(); err == nil {
 					ctx.GetLogger().Debugf("nng received message %s", string(msg))
-					ingest(ctx, msg, nil, timex.GetNow())
+					rawData, meta := extractTraceMeta(ctx, msg)
+					ingest(ctx, rawData, meta, timex.GetNow())
 				} else if err == mangos.ErrClosed {
 					ctx.GetLogger().Infof("neuron connection closed, retry after 1 second")
 					ingestErr(ctx, errors.New("neuron connection closed"))
@@ -112,3 +125,31 @@ func (s *source) Close(ctx api.StreamContext) error {
 func GetSource() api.Source {
 	return &source{}
 }
+
+func extractTraceMeta(ctx api.StreamContext, data []byte) ([]byte, map[string]interface{}) {
+	rawData := data
+	// extract rawData
+	if len(data) > NeuronTraceHeaderLen && bytes.Equal(data[:2], NeuronTraceHeader) {
+		rawData = data[NeuronTraceHeaderLen:]
+	}
+	if !ctx.IsTraceEnabled() {
+		return rawData, nil
+	}
+	meta := make(map[string]interface{})
+	var traced bool
+	var tracerCtx api.StreamContext
+	var span trace.Span
+	if len(data) > NeuronTraceHeaderLen && bytes.Equal(data[:2], NeuronTraceHeader) {
+		traceID := data[NeuronTraceIDStartIndex:NeuronTraceIDEndIndex]
+		spanID := data[NeuronTraceSpanIDStartIndex:NeuronTraceSpanIDEndIndex]
+		traced, tracerCtx, span = tracenode.StartTraceByID(ctx, [16]byte(traceID), [8]byte(spanID))
+	} else {
+		traced, tracerCtx, span = tracenode.StartTrace(ctx, ctx.GetOpId())
+	}
+	if traced {
+		meta["traceId"] = span.SpanContext().TraceID().String()
+		meta["traceCtx"] = tracerCtx
+		defer span.End()
+	}
+	return rawData, meta
+}
diff --git a/internal/topo/node/tracenode/trace_node.go b/internal/topo/node/tracenode/trace_node.go
index 8f5e2aaf..cf1f7580 100644
--- a/internal/topo/node/tracenode/trace_node.go
+++ b/internal/topo/node/tracenode/trace_node.go
@@ -82,6 +82,20 @@ func StartTrace(ctx api.StreamContext, opName string) (bool, api.StreamContext,
 	return true, ingestCtx, span
 }
 
+func StartTraceByID(ctx api.StreamContext, traceID [16]byte, spanID [8]byte) (bool, api.StreamContext, trace.Span) {
+	if !ctx.IsTraceEnabled() {
+		return false, nil, nil
+	}
+	spanCtx, span := tracer.GetTracer().Start(trace.ContextWithSpanContext(context.Background(), trace.NewSpanContext(trace.SpanContextConfig{
+		TraceID: traceID,
+		SpanID:  spanID,
+	})), ctx.GetOpId())
+	span.SetAttributes(attribute.String(RuleKey, ctx.GetRuleId()))
+	ingestCtx := topoContext.WithContext(spanCtx)
+	ingestCtx.IsTraceEnabled()
+	return true, ingestCtx, span
+}
+
 func ToStringRow(r xsql.Row) string {
 	d := r.Clone().ToMap()
 	b, _ := json.Marshal(d)


Test patch:
diff --git a/internal/io/neuron/sink_test.go b/internal/io/neuron/sink_test.go
index 4a6533e2..9cc39285 100644
--- a/internal/io/neuron/sink_test.go
+++ b/internal/io/neuron/sink_test.go
@@ -20,7 +20,11 @@ import (
 
 	"github.com/lf-edge/ekuiper/contract/v2/api"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/trace"
 
+	"github.com/lf-edge/ekuiper/v2/internal/topo/context"
+	topoContext "github.com/lf-edge/ekuiper/v2/internal/topo/context"
 	"github.com/lf-edge/ekuiper/v2/internal/xsql"
 	"github.com/lf-edge/ekuiper/v2/pkg/mock"
 	mockContext "github.com/lf-edge/ekuiper/v2/pkg/mock/context"
@@ -249,3 +253,40 @@ func TestSinkProvision(t *testing.T) {
 	err = s.Close(ctx)
 	assert.NoError(t, err)
 }
+
+func TestExtractSinkTraceData(t *testing.T) {
+	data := &mockData{
+		traceID: "1234567890abcdef",
+		spanID:  "12345678",
+	}
+	rawData := []byte(`{"a":1}`)
+	ctx := mockContext.NewMockContext("1", "2")
+	ctx.EnableTracer(true)
+	got := extractSpanContextIntoData(ctx, data, rawData)
+	require.Equal(t, NeuronTraceHeader, got[:len(NeuronTraceHeader)])
+	require.Equal(t, []byte(data.traceID), got[NeuronTraceIDStartIndex:NeuronTraceIDEndIndex])
+	require.Equal(t, rawData, got[NeuronTraceHeaderLen:])
+}
+
+type mockData struct {
+	traceID string
+	spanID  string
+}
+
+func (m *mockData) GetTracerCtx() api.StreamContext {
+	var traceID [16]byte
+	for i := 0; i < 16; i++ {
+		traceID[i] = m.traceID[i]
+	}
+	var spanID [8]byte
+	for i := 0; i < 8; i++ {
+		spanID[i] = m.spanID[i]
+	}
+	ctx := trace.ContextWithSpanContext(context.Background(), trace.NewSpanContext(trace.SpanContextConfig{
+		TraceID: traceID,
+		SpanID:  spanID,
+	}))
+	return topoContext.WithContext(ctx)
+}
+
+func (m *mockData) SetTracerCtx(ctx api.StreamContext) {}
diff --git a/internal/io/neuron/source_test.go b/internal/io/neuron/source_test.go
index f6049511..1727eb79 100644
--- a/internal/io/neuron/source_test.go
+++ b/internal/io/neuron/source_test.go
@@ -15,10 +15,12 @@
 package neuron
 
 import (
+	"encoding/hex"
 	"testing"
 
 	"github.com/lf-edge/ekuiper/contract/v2/api"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	_ "go.nanomsg.org/mangos/v3/transport/ipc"
 
 	"github.com/lf-edge/ekuiper/v2/pkg/mock"
@@ -82,3 +84,40 @@ func TestProvision(t *testing.T) {
 	err = s.Close(ctx)
 	assert.NoError(t, err)
 }
+
+func TestExtractSourceTraceData(t *testing.T) {
+	rawData := `{"a":1}`
+	traceID := "0123456789abcdef"
+	spanID := "12345678"
+	traceData := []byte{}
+	traceData = append(traceData, NeuronTraceHeader...)
+	traceData = append(traceData, []byte(traceID)...)
+	traceData = append(traceData, []byte(spanID)...)
+	traceData = append(traceData, []byte(rawData)...)
+	tests := []struct {
+		data    []byte
+		traceID string
+		spanID  string
+	}{
+		{
+			data:    traceData,
+			traceID: traceID,
+			spanID:  spanID,
+		},
+		{
+			data: []byte(rawData),
+		},
+	}
+	ctx := mockContext.NewMockContext("1", "2")
+	ctx.EnableTracer(true)
+	for _, tc := range tests {
+		gotData, meta := extractTraceMeta(ctx, tc.data)
+		require.Equal(t, []byte(rawData), gotData)
+		if len(tc.traceID) > 0 {
+			require.Equal(t, meta["traceId"], hex.EncodeToString([]byte(tc.traceID)))
+		} else {
+			_, ok := meta["traceId"]
+			require.True(t, ok)
+		}
+	}
+}