Skip to content

Commit 6ae9763

Browse files
robcarlan-datadoggenesorhannahkm
authored
feat(dsm): add kafka_cluster_id to confluent-kafka-go (#4470)
## Summary Most DSM tracer instrumentations for Kafka enrich a kafka_cluster_id. This is useful because numerous generated metrics such as `data_streams.latency` and `data_streams.kafka.lag_messages` will automatically add this metric if present. In situations where people have the same topic names across multiple kafka clusters, the above metric will become incorrect without the cluster id tag. This PR adds feature parity for Go and collects this metric tag when DSM is enabled. This is implemented by: * Launching a goroutine on consumer/producer creation to initiate the admin API and query for cluster id * Enriching metrics with this cluster id after we obtain it * Signalling a stop for an inflight cluster id request if the consumer/producer closes early * Enabling this only for DSM For context, the other tracers implement this as follows: 1. Java (doesn't block, intercepts the metadata response to enrich cluster id going forwards, see [here](https://github.com/DataDog/dd-trace-java/blob/master/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ProducerAdvice.java#L44) and [here](https://github.com/DataDog/dd-trace-java/blob/master/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/MetadataInstrumentation.java#L82)) 2. Node (blocks, see [here](https://github.com/DataDog/dd-trace-js/blob/master/packages/datadog-instrumentations/src/kafkajs.js#L234) and [here](https://github.com/DataDog/dd-trace-py/blob/main/ddtrace/contrib/internal/kafka/patch.py#L183)) 3. Python (blocks with a 1 second timeout, see [here](https://github.com/DataDog/dd-trace-py/blob/main/ddtrace/contrib/internal/kafka/patch.py#L360) and [here](https://github.com/DataDog/dd-trace-py/blob/main/ddtrace/contrib/internal/kafka/patch.py#L183)) ## Test plan - Add to existing tests - Test manually Screenshots: Spans have `kafka_cluster_id`: <img width="876" height="447" alt="Screenshot 2026-02-25 at 4 25 51 pm" src="https://github.com/user-attachments/assets/57cacf9d-ba8f-480b-97c6-60fffa9f66a2" /> <img width="822" height="483" alt="Screenshot 2026-02-25 at 4 26 04 pm" src="https://github.com/user-attachments/assets/3308be4b-ae79-4de4-8a30-6b7d3b0e6e26" /> Metric now adds `kafka_cluster_id`: <img width="2385" height="980" alt="Screenshot 2026-02-25 at 4 49 09 pm" src="https://github.com/user-attachments/assets/3ca95243-a3e2-4e2e-9a4b-8c97f25d3b3a" /> 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: genesor <[email protected]> Co-authored-by: hannahs.kim <[email protected]>
1 parent 6f4b819 commit 6ae9763

15 files changed

Lines changed: 360 additions & 30 deletions

File tree

‎contrib/confluentinc/confluent-kafka-go/kafka.v2/kafka.go‎

Lines changed: 62 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
package kafka // import "github.com/DataDog/dd-trace-go/contrib/confluentinc/confluent-kafka-go/kafka.v2/v2"
88

99
import (
10+
"context"
1011
"time"
1112

1213
"github.com/confluentinc/confluent-kafka-go/v2/kafka"
@@ -52,11 +53,39 @@ func NewProducer(conf *kafka.ConfigMap, opts ...Option) (*Producer, error) {
5253
return WrapProducer(p, opts...), nil
5354
}
5455

56+
// startClusterIDFetch launches a goroutine to fetch the cluster ID from the
57+
// given admin client. It returns a stop function that cancels the fetch and
58+
// waits for the goroutine to exit.
59+
func startClusterIDFetch(tr *kafkatrace.Tracer, admin *kafka.AdminClient) func() {
60+
ctx, cancel := context.WithCancel(context.Background())
61+
done := make(chan struct{})
62+
go func() {
63+
defer close(done)
64+
defer admin.Close()
65+
ctx, cancel := context.WithTimeout(ctx, 2*time.Second)
66+
defer cancel()
67+
clusterID, err := admin.ClusterID(ctx)
68+
if err != nil {
69+
if ctx.Err() == context.Canceled {
70+
return
71+
}
72+
instr.Logger().Warn("failed to fetch Kafka cluster ID: %s", err)
73+
return
74+
}
75+
tr.SetClusterID(clusterID)
76+
}()
77+
return func() {
78+
cancel()
79+
<-done
80+
}
81+
}
82+
5583
// A Consumer wraps a kafka.Consumer.
5684
type Consumer struct {
5785
*kafka.Consumer
58-
tracer *kafkatrace.Tracer
59-
events chan kafka.Event
86+
tracer *kafkatrace.Tracer
87+
events chan kafka.Event
88+
closeAsync []func() // async jobs to cancel and wait for on Close
6089
}
6190

6291
// WrapConsumer wraps a kafka.Consumer so that any consumed events are traced.
@@ -67,12 +96,27 @@ func WrapConsumer(c *kafka.Consumer, opts ...Option) *Consumer {
6796
}
6897
instr.Logger().Debug("%s: Wrapping Consumer: %#v", pkgPath, wrapped.tracer)
6998
wrapped.events = kafkatrace.WrapConsumeEventsChannel(wrapped.tracer, c.Events(), c, wrapEvent)
99+
if !wrapped.tracer.DSMEnabled() {
100+
return wrapped
101+
}
102+
// Create an admin client to fetch the cluster ID for data streams monitoring
103+
// The retrieval of the cluster ID is async and can be cancelled on Close to avoid blocking shutdown
104+
admin, err := kafka.NewAdminClientFromConsumer(c)
105+
if err != nil {
106+
instr.Logger().Warn("failed to create admin client for cluster ID, not adding cluster_id tags: %s", err)
107+
return wrapped
108+
}
109+
wrapped.closeAsync = append(wrapped.closeAsync, startClusterIDFetch(wrapped.tracer, admin))
70110
return wrapped
71111
}
72112

73113
// Close calls the underlying Consumer.Close and if polling is enabled, finishes
74114
// any remaining span.
75115
func (c *Consumer) Close() error {
116+
// Close any async jobs that might still be running
117+
for _, stopAsync := range c.closeAsync {
118+
stopAsync()
119+
}
76120
err := c.Consumer.Close()
77121
// we only close the previous span if consuming via the events channel is
78122
// not enabled, because otherwise there would be a data race from the
@@ -159,6 +203,7 @@ type Producer struct {
159203
tracer *kafkatrace.Tracer
160204
produceChannel chan *kafka.Message
161205
events chan kafka.Event
206+
closeAsync []func() // async jobs to cancel and wait for on Close
162207
}
163208

164209
// WrapProducer wraps a kafka.Producer so requests are traced.
@@ -170,9 +215,18 @@ func WrapProducer(p *kafka.Producer, opts ...Option) *Producer {
170215
}
171216
instr.Logger().Debug("%s: Wrapping Producer: %#v", pkgPath, wrapped.tracer)
172217
wrapped.produceChannel = kafkatrace.WrapProduceChannel(wrapped.tracer, p.ProduceChannel(), wrapMessage)
173-
if wrapped.tracer.DSMEnabled() {
174-
wrapped.events = kafkatrace.WrapProduceEventsChannel(wrapped.tracer, p.Events(), wrapEvent)
218+
if !wrapped.tracer.DSMEnabled() {
219+
return wrapped
175220
}
221+
wrapped.events = kafkatrace.WrapProduceEventsChannel(wrapped.tracer, p.Events(), wrapEvent)
222+
// Create an admin client to fetch the cluster ID for data streams monitoring
223+
// The retrieval of the cluster ID is async and can be cancelled on Close to avoid blocking shutdown
224+
admin, err := kafka.NewAdminClientFromProducer(p)
225+
if err != nil {
226+
instr.Logger().Warn("failed to create admin client for cluster ID, not adding cluster_id tags: %s", err)
227+
return wrapped
228+
}
229+
wrapped.closeAsync = append(wrapped.closeAsync, startClusterIDFetch(wrapped.tracer, admin))
176230
return wrapped
177231
}
178232

@@ -185,6 +239,10 @@ func (p *Producer) Events() chan kafka.Event {
185239
// Close calls the underlying Producer.Close and also closes the internal
186240
// wrapping producer channel.
187241
func (p *Producer) Close() {
242+
// Close any async jobs that might still be running
243+
for _, stop := range p.closeAsync {
244+
stop()
245+
}
188246
close(p.produceChannel)
189247
p.Producer.Close()
190248
}

‎contrib/confluentinc/confluent-kafka-go/kafka.v2/kafka_test.go‎

Lines changed: 46 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -170,9 +170,12 @@ func TestConsumerFunctional(t *testing.T) {
170170

171171
p, ok := datastreams.PathwayFromContext(datastreams.ExtractFromBase64Carrier(context.Background(), NewMessageCarrier(msg)))
172172
assert.True(t, ok)
173+
clusterID, ok := s0.Tag(ext.MessagingKafkaClusterID).(string)
174+
require.True(t, ok, "produce span should have a cluster ID tag")
175+
require.NotEmpty(t, clusterID)
173176
mt := mocktracer.Start()
174-
ctx, _ := tracer.SetDataStreamsCheckpoint(context.Background(), "direction:out", "topic:"+testTopic, "type:kafka")
175-
expectedCtx, _ := tracer.SetDataStreamsCheckpoint(ctx, "group:"+testGroupID, "direction:in", "topic:"+testTopic, "type:kafka")
177+
ctx, _ := tracer.SetDataStreamsCheckpoint(context.Background(), "direction:out", "topic:"+testTopic, "type:kafka", "kafka_cluster_id:"+clusterID)
178+
expectedCtx, _ := tracer.SetDataStreamsCheckpoint(ctx, "group:"+testGroupID, "direction:in", "topic:"+testTopic, "type:kafka", "kafka_cluster_id:"+clusterID)
176179
expected, _ := datastreams.PathwayFromContext(expectedCtx)
177180
mt.Stop()
178181
assert.NotEqual(t, expected.GetHash(), 0)
@@ -181,6 +184,39 @@ func TestConsumerFunctional(t *testing.T) {
181184
}
182185
}
183186

187+
func TestConsumerFunctionalWithClusterID(t *testing.T) {
188+
action := func(c *Consumer) (*kafka.Message, error) {
189+
return c.ReadMessage(3000 * time.Millisecond)
190+
}
191+
spans, msg := produceThenConsume(t, action,
192+
[]Option{WithDataStreams()},
193+
[]Option{WithDataStreams()},
194+
false,
195+
)
196+
require.Len(t, spans, 2)
197+
198+
// Verify cluster ID is set as a span tag on both produce and consume spans.
199+
// The cluster ID is auto-fetched from the broker, so we just verify it's
200+
// present and consistent across spans.
201+
s0 := spans[0] // produce
202+
s1 := spans[1] // consume
203+
clusterID, ok := s0.Tag(ext.MessagingKafkaClusterID).(string)
204+
require.True(t, ok, "produce span should have a cluster ID tag")
205+
assert.NotEmpty(t, clusterID)
206+
assert.Equal(t, clusterID, s1.Tag(ext.MessagingKafkaClusterID))
207+
208+
// Verify DSM pathway hash includes kafka_cluster_id in edge tags
209+
p, ok := datastreams.PathwayFromContext(datastreams.ExtractFromBase64Carrier(context.Background(), NewMessageCarrier(msg)))
210+
assert.True(t, ok)
211+
mt := mocktracer.Start()
212+
ctx, _ := tracer.SetDataStreamsCheckpoint(context.Background(), "direction:out", "topic:"+testTopic, "type:kafka", "kafka_cluster_id:"+clusterID)
213+
expectedCtx, _ := tracer.SetDataStreamsCheckpoint(ctx, "group:"+testGroupID, "direction:in", "topic:"+testTopic, "type:kafka", "kafka_cluster_id:"+clusterID)
214+
expected, _ := datastreams.PathwayFromContext(expectedCtx)
215+
mt.Stop()
216+
assert.NotEqual(t, expected.GetHash(), 0)
217+
assert.Equal(t, expected.GetHash(), p.GetHash())
218+
}
219+
184220
// This tests the deprecated behavior of using cfg.context as the context passed via kafka messages
185221
// instead of the one passed in the message.
186222
func TestDeprecatedContext(t *testing.T) {
@@ -361,6 +397,7 @@ func produceThenConsume(t *testing.T, consumerAction consumerActionFn, producerO
361397
"go.delivery.reports": true,
362398
}, producerOpts...)
363399
require.NoError(t, err)
400+
require.Eventually(t, func() bool { return p.tracer.ClusterID() != "" }, 5*time.Second, 10*time.Millisecond)
364401

365402
var delivery chan kafka.Event = nil
366403
if !useProducerEventsChannel {
@@ -399,6 +436,7 @@ func produceThenConsume(t *testing.T, consumerAction consumerActionFn, producerO
399436
"enable.auto.offset.store": false,
400437
}, consumerOpts...)
401438
require.NoError(t, err)
439+
require.Eventually(t, func() bool { return c.tracer.ClusterID() != "" }, 5*time.Second, 10*time.Millisecond)
402440

403441
err = c.Assign([]kafka.TopicPartition{
404442
{Topic: &testTopic, Partition: 0, Offset: msg1.TopicPartition.Offset},
@@ -428,9 +466,12 @@ func produceThenConsume(t *testing.T, consumerAction consumerActionFn, producerO
428466
return m
429467
}
430468
backlogsMap := toMap(backlogs)
431-
require.Contains(t, backlogsMap, "consumer_group:"+testGroupID+"partition:0"+"topic:"+testTopic+"type:kafka_commit")
432-
require.Contains(t, backlogsMap, "partition:0"+"topic:"+testTopic+"type:kafka_high_watermark")
433-
require.Contains(t, backlogsMap, "partition:0"+"topic:"+testTopic+"type:kafka_produce")
469+
clusterID := c.tracer.ClusterID()
470+
require.NotEmpty(t, clusterID)
471+
clusterTag := "kafka_cluster_id:" + clusterID
472+
require.Contains(t, backlogsMap, "consumer_group:"+testGroupID+"partition:0"+"topic:"+testTopic+"type:kafka_commit"+clusterTag)
473+
require.Contains(t, backlogsMap, "partition:0"+"topic:"+testTopic+"type:kafka_high_watermark"+clusterTag)
474+
require.Contains(t, backlogsMap, "partition:0"+"topic:"+testTopic+"type:kafka_produce"+clusterTag)
434475
}
435476
return spans, msg2
436477
}

‎contrib/confluentinc/confluent-kafka-go/kafka/kafka.go‎

Lines changed: 62 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
package kafka // import "github.com/DataDog/dd-trace-go/contrib/confluentinc/confluent-kafka-go/kafka/v2"
88

99
import (
10+
"context"
1011
"time"
1112

1213
"github.com/confluentinc/confluent-kafka-go/kafka"
@@ -52,11 +53,39 @@ func NewProducer(conf *kafka.ConfigMap, opts ...Option) (*Producer, error) {
5253
return WrapProducer(p, opts...), nil
5354
}
5455

56+
// startClusterIDFetch launches a goroutine to fetch the cluster ID from the
57+
// given admin client. It returns a stop function that cancels the fetch and
58+
// waits for the goroutine to exit.
59+
func startClusterIDFetch(tr *kafkatrace.Tracer, admin *kafka.AdminClient) func() {
60+
ctx, cancel := context.WithCancel(context.Background())
61+
done := make(chan struct{})
62+
go func() {
63+
defer close(done)
64+
defer admin.Close()
65+
ctx, cancel := context.WithTimeout(ctx, 2*time.Second)
66+
defer cancel()
67+
clusterID, err := admin.ClusterID(ctx)
68+
if err != nil {
69+
if ctx.Err() == context.Canceled {
70+
return
71+
}
72+
instr.Logger().Warn("failed to fetch Kafka cluster ID: %s", err)
73+
return
74+
}
75+
tr.SetClusterID(clusterID)
76+
}()
77+
return func() {
78+
cancel()
79+
<-done
80+
}
81+
}
82+
5583
// A Consumer wraps a kafka.Consumer.
5684
type Consumer struct {
5785
*kafka.Consumer
58-
tracer *kafkatrace.Tracer
59-
events chan kafka.Event
86+
tracer *kafkatrace.Tracer
87+
events chan kafka.Event
88+
closeAsync []func() // async jobs to cancel and wait for on Close
6089
}
6190

6291
// WrapConsumer wraps a kafka.Consumer so that any consumed events are traced.
@@ -67,12 +96,27 @@ func WrapConsumer(c *kafka.Consumer, opts ...Option) *Consumer {
6796
}
6897
instr.Logger().Debug("%s: Wrapping Consumer: %#v", pkgPath, wrapped.tracer)
6998
wrapped.events = kafkatrace.WrapConsumeEventsChannel(wrapped.tracer, c.Events(), c, wrapEvent)
99+
if !wrapped.tracer.DSMEnabled() {
100+
return wrapped
101+
}
102+
// Create an admin client to fetch the cluster ID for data streams monitoring
103+
// The retrieval of the cluster ID is async and can be cancelled on Close to avoid blocking shutdown
104+
admin, err := kafka.NewAdminClientFromConsumer(c)
105+
if err != nil {
106+
instr.Logger().Warn("failed to create admin client for cluster ID, not adding cluster_id tags: %s", err)
107+
return wrapped
108+
}
109+
wrapped.closeAsync = append(wrapped.closeAsync, startClusterIDFetch(wrapped.tracer, admin))
70110
return wrapped
71111
}
72112

73113
// Close calls the underlying Consumer.Close and if polling is enabled, finishes
74114
// any remaining span.
75115
func (c *Consumer) Close() error {
116+
// Close any async jobs that might still be running
117+
for _, stop := range c.closeAsync {
118+
stop()
119+
}
76120
err := c.Consumer.Close()
77121
// we only close the previous span if consuming via the events channel is
78122
// not enabled, because otherwise there would be a data race from the
@@ -159,6 +203,7 @@ type Producer struct {
159203
tracer *kafkatrace.Tracer
160204
produceChannel chan *kafka.Message
161205
events chan kafka.Event
206+
closeAsync []func() // async jobs to cancel and wait for on Close
162207
}
163208

164209
// WrapProducer wraps a kafka.Producer so requests are traced.
@@ -170,9 +215,18 @@ func WrapProducer(p *kafka.Producer, opts ...Option) *Producer {
170215
}
171216
instr.Logger().Debug("%s: Wrapping Producer: %#v", pkgPath, wrapped.tracer)
172217
wrapped.produceChannel = kafkatrace.WrapProduceChannel(wrapped.tracer, p.ProduceChannel(), wrapMessage)
173-
if wrapped.tracer.DSMEnabled() {
174-
wrapped.events = kafkatrace.WrapProduceEventsChannel(wrapped.tracer, p.Events(), wrapEvent)
218+
if !wrapped.tracer.DSMEnabled() {
219+
return wrapped
175220
}
221+
wrapped.events = kafkatrace.WrapProduceEventsChannel(wrapped.tracer, p.Events(), wrapEvent)
222+
// Create an admin client to fetch the cluster ID for data streams monitoring
223+
// The retrieval of the cluster ID is async and can be cancelled on Close to avoid blocking shutdown
224+
admin, err := kafka.NewAdminClientFromProducer(p)
225+
if err != nil {
226+
instr.Logger().Warn("failed to create admin client for cluster ID, not adding cluster_id tags: %s", err)
227+
return wrapped
228+
}
229+
wrapped.closeAsync = append(wrapped.closeAsync, startClusterIDFetch(wrapped.tracer, admin))
176230
return wrapped
177231
}
178232

@@ -185,6 +239,10 @@ func (p *Producer) Events() chan kafka.Event {
185239
// Close calls the underlying Producer.Close and also closes the internal
186240
// wrapping producer channel.
187241
func (p *Producer) Close() {
242+
// Close any async jobs that might still be running
243+
for _, stop := range p.closeAsync {
244+
stop()
245+
}
188246
close(p.produceChannel)
189247
p.Producer.Close()
190248
}

0 commit comments

Comments
 (0)