Add feature of refreshing service root spans

tuiexporter/internal/telemetry/cache.go

@@ -18,12 +18,17 @@ type TraceServiceSpanDataMap map[string]map[string][]*SpanData
 // the spans have any error status
 type TraceServiceHasErrorMap map[string]map[string]bool
 
+// TraceServiceParentIDMap is a map of trace id and service name to a parent span id
+// This is used to update service root spans in the trace list.
+type TraceServiceParentIDMap map[string]map[string]*SpanData
+
 // TraceCache is a cache of trace spans
 type TraceCache struct {
 	spanid2span       SpanDataMap
 	traceid2spans     TraceSpanDataMap
 	tracesvc2spans    TraceServiceSpanDataMap
 	tracesvc2haserror TraceServiceHasErrorMap
+	tracesvc2parent   TraceServiceParentIDMap
 }
 
 // NewTraceCache returns a new trace cache
@@ -33,34 +38,46 @@ func NewTraceCache() *TraceCache {
 		traceid2spans:     TraceSpanDataMap{},
 		tracesvc2spans:    TraceServiceSpanDataMap{},
 		tracesvc2haserror: TraceServiceHasErrorMap{},
+		tracesvc2parent:   TraceServiceParentIDMap{},
 	}
 }
 
 // UpdateCache updates the cache with a new span
-func (c *TraceCache) UpdateCache(sname string, data *SpanData) (newtracesvc bool) {
+func (c *TraceCache) UpdateCache(sname string, data *SpanData) (newtracesvc bool, replaceSpanID string) {
 	c.spanid2span[data.Span.SpanID().String()] = data
 	traceID := data.Span.TraceID().String()
 	hasError := spanHasError(data.Span)
 	if ts, ok := c.traceid2spans[traceID]; ok {
 		c.traceid2spans[traceID] = append(ts, data)
 		if _, ok := c.tracesvc2spans[traceID][sname]; ok {
 			c.tracesvc2spans[traceID][sname] = append(c.tracesvc2spans[traceID][sname], data)
+			if c.tracesvc2parent[traceID][sname].Span.ParentSpanID().String() == data.Span.SpanID().String() {
+				// This span is higher parent span
+				// NOTE: In this process, for performance reasons, only adjacent parent-child relationships
+				//   between spans are evaluated. For example, if the parent-child order of spans is 1, 2, 3, and
+				//   the arrival order is 3, 1, 2, span 2 will be recognized as the service root span. To recalculate
+				//   the specific parent-child relationship, use `R` key to trigger deep refreshing
+				replaceSpanID = c.tracesvc2parent[traceID][sname].Span.SpanID().String()
+				c.tracesvc2parent[traceID][sname] = data
+			}
 			if hasError {
 				c.tracesvc2haserror[traceID][sname] = hasError
 			}
 		} else {
 			c.tracesvc2spans[traceID][sname] = []*SpanData{data}
 			c.tracesvc2haserror[traceID][sname] = hasError
+			c.tracesvc2parent[traceID][sname] = data
 			newtracesvc = true
 		}
 	} else {
 		c.traceid2spans[traceID] = []*SpanData{data}
 		c.tracesvc2spans[traceID] = map[string][]*SpanData{sname: {data}}
 		c.tracesvc2haserror[traceID] = map[string]bool{sname: hasError}
+		c.tracesvc2parent[traceID] = map[string]*SpanData{sname: data}
 		newtracesvc = true
 	}
 
-	return newtracesvc
+	return newtracesvc, replaceSpanID
 }
 
 // DeleteCache deletes a list of spans from the cache
@@ -77,9 +94,11 @@ func (c *TraceCache) DeleteCache(serviceSpans []*SpanData) {
 		}
 		delete(c.tracesvc2spans[traceID], sname.AsString())
 		delete(c.tracesvc2haserror[traceID], sname.AsString())
+		delete(c.tracesvc2parent[traceID], sname.AsString())
 		if len(c.tracesvc2spans[traceID]) == 0 {
 			delete(c.tracesvc2spans, traceID)
 			delete(c.tracesvc2haserror, traceID)
+			delete(c.tracesvc2parent, traceID)
 			// delete spans in traceid2spans only if there are no spans left in tracesvc2spans
 			// for better performance
 			delete(c.traceid2spans, traceID)

tuiexporter/internal/telemetry/store.go

@@ -34,6 +34,15 @@ func (sd *SpanData) IsRoot() bool {
 // This is a slice of one span of a single service
 type SvcSpans []*SpanData
 
+func (ss *SvcSpans) replaceBySpanID(replaceSpanID string, data *SpanData) {
+	for i, s := range *ss {
+		if s.Span.SpanID().String() == replaceSpanID {
+			(*ss)[i] = data
+			return
+		}
+	}
+}
+
 // MetricData is a struct to represent a metric
 type MetricData struct {
 	Metric         *pmetric.Metric
@@ -249,6 +258,46 @@ func (s *Store) GetFilteredServiceSpansByIdx(idx int) []*SpanData {
 	return spans
 }
 
+// RecalculateServiceRootSpanByIdx recalculates service root span of the specified index
+func (s *Store) RecalculateServiceRootSpanByIdx(idx int) {
+	s.mut.Lock()
+	defer func() {
+		s.updatedAt = time.Now()
+		s.mut.Unlock()
+	}()
+
+	if idx < 0 || idx >= len(s.svcspansFiltered) {
+		return
+	}
+	traceID := s.svcspansFiltered[idx].Span.TraceID().String()
+	currentSpanID := s.svcspansFiltered[idx].Span.SpanID().String()
+	sname, ok := s.svcspansFiltered[idx].ResourceSpan.Resource().Attributes().Get("service.name")
+	if !ok {
+		return
+	}
+
+	type spanTreeNode struct {
+		span     *SpanData
+		children []*spanTreeNode
+	}
+
+	spans := s.tracecache.tracesvc2spans[traceID][sname.AsString()]
+	spanMemo := make(map[string]bool)
+	for _, span := range spans {
+		spanMemo[span.Span.SpanID().String()] = true
+	}
+	for _, span := range spans {
+		parentSpanID := span.Span.ParentSpanID().String()
+		spanID := span.Span.SpanID().String()
+		if _, ok := spanMemo[parentSpanID]; !ok {
+			// TODO: Condider orphan span?
+			sd := s.tracecache.spanid2span[spanID]
+			s.svcspansFiltered[idx] = sd
+			s.svcspans.replaceBySpanID(currentSpanID, sd)
+		}
+	}
+}
+
 // GetFilteredMetricByIdx returns the metric at the given index
 func (s *Store) GetFilteredMetricByIdx(idx int) *MetricData {
 	if idx < 0 || idx >= len(s.metricsFiltered) {
@@ -291,9 +340,12 @@ func (s *Store) AddSpan(traces *ptrace.Traces) {
 					ScopeSpans:   &ss,
 					ReceivedAt:   time.Now(),
 				}
-				newtracesvc := s.tracecache.UpdateCache(sname.AsString(), sd)
+				newtracesvc, replaceSpanID := s.tracecache.UpdateCache(sname.AsString(), sd)
 				if newtracesvc {
 					s.svcspans = append(s.svcspans, sd)
+				} else if len(replaceSpanID) > 0 {
+					// FIXME: More efficient logic is needed
+					s.svcspans.replaceBySpanID(replaceSpanID, sd)
 				}
 			}
 		}

tuiexporter/internal/telemetry/store_test.go

@@ -362,6 +362,124 @@ func TestStoreAddSpanWithRotation(t *testing.T) {
 	}
 }
 
+func TestStoreAddSpanServiceSpanCalculation(t *testing.T) {
+	// traceid: 1
+	//  └- resource: test-service-1
+	//    └- scope: test-scope-1-1
+	//      └- span: span-1-1-2
+	//      └- span: span-1-1-3
+	// traceid: 1 (the same trace)
+	//  └- resource: test-service-1
+	//    └- scope: test-scope-1-1
+	//      └- span: span-1-1-1
+	store := NewStore()
+	store.maxServiceSpanCount = 1
+	payload1, _ := test.GenerateOTLPTracesPayload(t, 1, 1, []int{3}, [][]int{{3, 0, 0}})
+	payload1.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(1).SetParentSpanID(
+		payload1.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0).SpanID(),
+	)
+	payload1.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(2).SetParentSpanID(
+		payload1.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(1).SpanID(),
+	)
+	payload1.ResourceSpans().At(0).ScopeSpans().At(0).Spans().RemoveIf(func(s ptrace.Span) bool {
+		return s.Name() == "span-0-0-0"
+	})
+	payload2, _ := test.GenerateOTLPTracesPayload(t, 1, 1, []int{3}, [][]int{{3, 0, 0}})
+	payload2.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(1).SetParentSpanID(
+		payload2.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0).SpanID(),
+	)
+	payload2.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(2).SetParentSpanID(
+		payload2.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(1).SpanID(),
+	)
+	payload2.ResourceSpans().At(0).ScopeSpans().At(0).Spans().RemoveIf(func(s ptrace.Span) bool {
+		return s.Name() == "span-0-0-1" || s.Name() == "span-0-0-2"
+	})
+
+	assert.Equal(t, 2, payload1.ResourceSpans().At(0).ScopeSpans().At(0).Spans().Len())
+	assert.Equal(t, 1, payload2.ResourceSpans().At(0).ScopeSpans().At(0).Spans().Len())
+
+	store.AddSpan(&payload1)
+
+	// The service root span should be span-1-1-2
+	assert.Equal(t, 1, len(store.svcspans))
+	assert.Equal(t, "span-0-0-1", store.svcspans[0].Span.Name())
+
+	store.AddSpan(&payload2)
+
+	// Now, The service root span should be span-1-1-1
+	assert.Equal(t, 1, len(store.svcspans))
+	assert.Equal(t, "span-0-0-0", store.svcspans[0].Span.Name())
+}
+
+func TestStoreAddSpanServiceSpanCalculationLimitation(t *testing.T) {
+	// traceid: 1
+	//  └- resource: test-service-1
+	//    └- scope: test-scope-1-1
+	//      └- span: span-1-1-3
+	// traceid: 1 (the same trace)
+	//  └- resource: test-service-1
+	//    └- scope: test-scope-1-1
+	//      └- span: span-1-1-1
+	// traceid: 1 (the same trace)
+	//  └- resource: test-service-1
+	//    └- scope: test-scope-1-1
+	//      └- span: span-1-1-2
+	store := NewStore()
+	store.maxServiceSpanCount = 1
+	payload1, _ := test.GenerateOTLPTracesPayload(t, 1, 1, []int{3}, [][]int{{3, 0, 0}})
+	payload1.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(1).SetParentSpanID(
+		payload1.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0).SpanID(),
+	)
+	payload1.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(2).SetParentSpanID(
+		payload1.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(1).SpanID(),
+	)
+	payload1.ResourceSpans().At(0).ScopeSpans().At(0).Spans().RemoveIf(func(s ptrace.Span) bool {
+		return s.Name() == "span-0-0-0" || s.Name() == "span-0-0-1"
+	})
+	payload2, _ := test.GenerateOTLPTracesPayload(t, 1, 1, []int{3}, [][]int{{3, 0, 0}})
+	payload2.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(1).SetParentSpanID(
+		payload2.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0).SpanID(),
+	)
+	payload2.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(2).SetParentSpanID(
+		payload2.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(1).SpanID(),
+	)
+	payload2.ResourceSpans().At(0).ScopeSpans().At(0).Spans().RemoveIf(func(s ptrace.Span) bool {
+		return s.Name() == "span-0-0-1" || s.Name() == "span-0-0-2"
+	})
+	payload3, _ := test.GenerateOTLPTracesPayload(t, 1, 1, []int{3}, [][]int{{3, 0, 0}})
+	payload3.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(1).SetParentSpanID(
+		payload3.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0).SpanID(),
+	)
+	payload3.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(2).SetParentSpanID(
+		payload3.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(1).SpanID(),
+	)
+	payload3.ResourceSpans().At(0).ScopeSpans().At(0).Spans().RemoveIf(func(s ptrace.Span) bool {
+		return s.Name() == "span-0-0-0" || s.Name() == "span-0-0-2"
+	})
+
+	assert.Equal(t, 1, payload1.ResourceSpans().At(0).ScopeSpans().At(0).Spans().Len())
+	assert.Equal(t, 1, payload2.ResourceSpans().At(0).ScopeSpans().At(0).Spans().Len())
+	assert.Equal(t, 1, payload3.ResourceSpans().At(0).ScopeSpans().At(0).Spans().Len())
+
+	store.AddSpan(&payload1)
+
+	// The service root span should be span-1-1-3
+	assert.Equal(t, 1, len(store.svcspans))
+	assert.Equal(t, "span-0-0-2", store.svcspans[0].Span.Name())
+
+	store.AddSpan(&payload2)
+
+	// The service root span should still be span-1-1-3
+	assert.Equal(t, 1, len(store.svcspans))
+	assert.Equal(t, "span-0-0-2", store.svcspans[0].Span.Name())
+
+	store.AddSpan(&payload3)
+
+	// Finally, The service root span should be span-1-1-2
+	assert.Equal(t, 1, len(store.svcspans))
+	assert.Equal(t, "span-0-0-1", store.svcspans[0].Span.Name())
+}
+
 func TestStoreAddMetricWithoutRotation(t *testing.T) {
 	// metric: 1
 	//  └- resource: test-service-1

tuiexporter/internal/tui/component/page.go

@@ -195,6 +195,14 @@ func (p *TUIPages) createTracePage(store *telemetry.Store) *tview.Flex {
 			}
 			log.Printf("sortType: %s", sortType)
 			store.ApplyFilterTraces(inputConfirmed, sortType)
+			return nil
+		} else if event.Rune() == 'r' {
+			row, _ := table.GetSelection()
+			if row == 0 {
+				return nil
+			}
+			store.RecalculateServiceRootSpanByIdx(row - 1)
+
 			return nil
 		}
 		return event
@@ -208,6 +216,10 @@ func (p *TUIPages) createTracePage(store *telemetry.Store) *tview.Flex {
 			key:         tcell.NewEventKey(tcell.KeyRune, 'S', tcell.ModCtrl),
 			description: "Toggle sort (Latency)",
 		},
+		{
+			key:         tcell.NewEventKey(tcell.KeyRune, 'R', tcell.ModNone),
+			description: "Recalculate service root span",
+		},
 		{
 			key:         tcell.NewEventKey(tcell.KeyRune, 'L', tcell.ModCtrl),
 			description: "Clear all data",