tracing: fix lazytag export for Jaegar/otel

pkg/util/tracing/crdbspan.go

@@ -923,34 +923,19 @@ func (s *crdbSpan) getRecordingNoChildrenLocked(
 			// We encode the tag values as strings.
 			addTag(string(kv.Key), kv.Value.Emit())
 		}
-		for _, kv := range s.mu.lazyTags {
-			switch v := kv.Value.(type) {
-			case LazyTag:
-				var tagGroup *tracingpb.TagGroup
-				for _, tag := range v.Render() {
-					if tagGroup == nil {
-						// Only create the tag group if we have at least one child tag.
-						tagGroup = addTagGroup(kv.Key)
-					}
-					childKey := string(tag.Key)
-					childValue := tag.Value.Emit()
-
-					if rs.Tags == nil {
-						rs.Tags = make(map[string]string)
-					}
-					rs.Tags[childKey] = childValue
-
-					tagGroup.Tags = append(tagGroup.Tags,
-						tracingpb.Tag{
-							Key:   childKey,
-							Value: childValue,
-						},
-					)
+		for _, tg := range s.getLazyTagGroupsLocked() {
+			if tg.Name == tracingpb.AnonymousTagGroupName {
+				for _, tag := range tg.Tags {
+					addTag(tag.Key, tag.Value)
+				}
+			} else {
+				rs.TagGroups = append(rs.TagGroups, *tg)
+				if rs.Tags == nil {
+					rs.Tags = make(map[string]string)
+				}
+				for _, tag := range tg.Tags {
+					rs.Tags[tag.Key] = tag.Value
 				}
-			case fmt.Stringer:
-				addTag(kv.Key, v.String())
-			default:
-				addTag(kv.Key, fmt.Sprintf("<can't render %T>", kv.Value))
 			}
 		}
 
@@ -973,6 +958,59 @@ func (s *crdbSpan) getRecordingNoChildrenLocked(
 	return rs
 }
 
+// getLazyTagGroupsLocked returns a list of lazy tags as a slice of
+// *tracingpb.TagGroup
+func (s *crdbSpan) getLazyTagGroupsLocked() []*tracingpb.TagGroup {
+	var lazyTags []*tracingpb.TagGroup
+	addTagGroup := func(name string) *tracingpb.TagGroup {
+		lazyTags = append(lazyTags,
+			&tracingpb.TagGroup{
+				Name: name,
+			})
+		return lazyTags[len(lazyTags)-1]
+	}
+
+	addTag := func(k, v string) {
+		var tagGroup *tracingpb.TagGroup
+		for _, tg := range lazyTags {
+			if tg.Name == tracingpb.AnonymousTagGroupName {
+				tagGroup = tg
+				break
+			}
+		}
+		if tagGroup == nil {
+			tagGroup = addTagGroup(tracingpb.AnonymousTagGroupName)
+		}
+		tagGroup.Tags = append(tagGroup.Tags, tracingpb.Tag{
+			Key:   k,
+			Value: v,
+		})
+	}
+	for _, kv := range s.mu.lazyTags {
+		switch v := kv.Value.(type) {
+		case LazyTag:
+			var tagGroup *tracingpb.TagGroup
+			for _, tag := range v.Render() {
+				if tagGroup == nil {
+					// Only create the tag group if we have at least one child tag.
+					tagGroup = addTagGroup(kv.Key)
+				}
+				tagGroup.Tags = append(tagGroup.Tags,
+					tracingpb.Tag{
+						Key:   string(tag.Key),
+						Value: tag.Value.Emit(),
+					},
+				)
+			}
+		case fmt.Stringer:
+			addTag(kv.Key, v.String())
+		default:
+			addTag(kv.Key, fmt.Sprintf("<can't render %T>", kv.Value))
+		}
+	}
+	return lazyTags
+}
+
 // addChildLocked adds a child to the receiver.
 //
 // The receiver's lock must be held.

pkg/util/tracing/span_inner.go

@@ -104,8 +104,25 @@ func (s *spanInner) Finish() {
 	}
 
 	if s.otelSpan != nil {
+		// Serialize the lazy tags.
+		s.crdb.mu.Lock()
+		defer s.crdb.mu.Unlock()
+		for _, lazyTagGroup := range s.crdb.getLazyTagGroupsLocked() {
+			for _, tag := range lazyTagGroup.Tags {
+				key := attribute.Key(tag.Key)
+				if lazyTagGroup.Name != tracingpb.AnonymousTagGroupName {
+					key = attribute.Key(fmt.Sprintf("%s-%s", lazyTagGroup.Name, tag.Key))
+				}
+				s.otelSpan.SetAttributes(attribute.KeyValue{
+					Key:   key,
+					Value: attribute.StringValue(tag.Value),
+				})
+			}
+		}
+
 		s.otelSpan.End()
 	}
+
 	if s.netTr != nil {
 		s.netTr.Finish()
 	}

pkg/util/tracing/tracer_test.go

@@ -397,14 +397,38 @@ func TestOtelTracer(t *testing.T) {
 	}
 
 	sp2 := tr.StartSpan("child", WithRemoteParentFromSpanMeta(wireSpanMeta))
-	defer sp2.Finish()
+	sp2.SetLazyTag("lazy expanding tag", testExpandingTag{})
+	sp2.SetLazyTag("lazy tag", testStringerLazyTag{})
 
 	rs := sr.Started()
 	require.Len(t, rs, 2)
 	require.Len(t, rs[0].Events(), 1)
 	require.Equal(t, "hello", rs[0].Events()[0].Name)
 	require.Equal(t, rs[0].SpanContext().TraceID(), rs[1].Parent().TraceID())
 	require.Equal(t, rs[0].SpanContext().SpanID(), rs[1].Parent().SpanID())
+
+	sp2.Finish()
+	rs2 := sr.Ended()
+
+	expectedAttributes := []attribute.KeyValue{
+		{
+			Key:   attribute.Key("lazy expanding tag-exp1"),
+			Value: attribute.StringValue("1"),
+		},
+		{
+			Key:   attribute.Key("lazy expanding tag-exp2"),
+			Value: attribute.StringValue("2"),
+		},
+		{
+			Key:   attribute.Key("lazy tag"),
+			Value: attribute.StringValue("lazy stringer"),
+		},
+	}
+	actualAttributes := rs2[0].Attributes()
+
+	require.Len(t, rs2, 1)
+	require.Len(t, actualAttributes, 3)
+	require.Equal(t, expectedAttributes, actualAttributes)
 }
 
 func TestTracer_RegistryMaxSize(t *testing.T) {