fixup tracing

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>

event.go

@@ -1,12 +1,14 @@
 package kgo
 
 import (
+	"context"
 	"sync"
 
 	"go.unistack.org/micro/v3/broker"
 )
 
 type event struct {
+	ctx   context.Context
 	topic string
 	err   error
 	sync.RWMutex
@@ -14,6 +16,10 @@ type event struct {
 	ack bool
 }
 
+func (p *event) Context() context.Context {
+	return p.ctx
+}
+
 func (p *event) Topic() string {
 	return p.topic
 }

kgo.go

@@ -73,7 +73,7 @@ func (k *Broker) Name() string {
 	return k.opts.Name
 }
 
-func (k *Broker) connect(ctx context.Context, opts ...kgo.Opt) (*kgo.Client, error) {
+func (k *Broker) connect(ctx context.Context, opts ...kgo.Opt) (*kgo.Client, *hookTracer, error) {
 	var c *kgo.Client
 	var err error
 
@@ -90,28 +90,33 @@ func (k *Broker) connect(ctx context.Context, opts ...kgo.Opt) (*kgo.Client, err
 		}
 	}
 
+	htracer := &hookTracer{group: group, clientID: clientID, tracer: k.opts.Tracer}
 	opts = append(opts,
 		kgo.WithHooks(&hookMeter{meter: k.opts.Meter}),
-		kgo.WithHooks(&hookTracer{group: group, clientID: clientID, tracer: k.opts.Tracer}),
+		kgo.WithHooks(htracer),
 	)
 
 	select {
 	case <-ctx.Done():
 		if ctx.Err() != nil {
-			sp.SetStatus(tracer.SpanStatusError, ctx.Err().Error())
+			if sp != nil {
+				sp.SetStatus(tracer.SpanStatusError, ctx.Err().Error())
+			}
 		}
-		return nil, ctx.Err()
+		return nil, nil, ctx.Err()
 	default:
 		c, err = kgo.NewClient(opts...)
 		if err == nil {
 			err = c.Ping(ctx) // check connectivity to cluster
 		}
 		if err != nil {
-			sp.SetStatus(tracer.SpanStatusError, err.Error())
-			return nil, err
+			if sp != nil {
+				sp.SetStatus(tracer.SpanStatusError, err.Error())
+			}
+			return nil, nil, err
 		}
 	}
-	return c, nil
+	return c, htracer, nil
 }
 
 func (k *Broker) Connect(ctx context.Context) error {
@@ -127,7 +132,7 @@ func (k *Broker) Connect(ctx context.Context) error {
 		nctx = ctx
 	}
 
-	c, err := k.connect(nctx, k.kopts...)
+	c, _, err := k.connect(nctx, k.kopts...)
 	if err != nil {
 		return err
 	}
@@ -230,13 +235,9 @@ func (k *Broker) Publish(ctx context.Context, topic string, msg *broker.Message,
 }
 
 func (k *Broker) publish(ctx context.Context, msgs []*broker.Message, opts ...broker.PublishOption) error {
-	var span tracer.Span
-	ctx, span = k.opts.Tracer.Start(ctx, "Publish")
-	defer span.Finish()
-
 	k.Lock()
 	if !k.connected {
-		c, err := k.connect(ctx, k.kopts...)
+		c, _, err := k.connect(ctx, k.kopts...)
 		if err != nil {
 			k.Unlock()
 			return err
@@ -370,7 +371,7 @@ func (k *Broker) Subscribe(ctx context.Context, topic string, handler broker.Han
 		}
 	}
 
-	c, err := k.connect(ctx, kopts...)
+	c, htracer, err := k.connect(ctx, kopts...)
 	if err != nil {
 		return nil, err
 	}
@@ -388,6 +389,7 @@ func (k *Broker) Subscribe(ctx context.Context, topic string, handler broker.Han
 	}
 
 	sub.c = c
+	sub.htracer = htracer
 
 	go sub.poll(ctx)
 

subscriber.go

@@ -12,6 +12,7 @@ import (
 	"go.unistack.org/micro/v3/logger"
 	"go.unistack.org/micro/v3/metadata"
 	"go.unistack.org/micro/v3/semconv"
+	"go.unistack.org/micro/v3/tracer"
 )
 
 type tp struct {
@@ -23,6 +24,7 @@ type consumer struct {
 	c         *kgo.Client
 	topic     string
 	partition int32
+	htracer   *hookTracer
 	opts      broker.SubscribeOptions
 	kopts     broker.Options
 	handler   broker.Handler
@@ -34,6 +36,7 @@ type consumer struct {
 type subscriber struct {
 	c         *kgo.Client
 	topic     string
+	htracer   *hookTracer
 	opts      broker.SubscribeOptions
 	kopts     broker.Options
 	handler   broker.Handler
@@ -106,10 +109,8 @@ func (s *subscriber) poll(ctx context.Context) {
 				}
 
 				s.Lock()
-				for tp := range s.consumers {
-					if v, ok := lmap[tp.p]; ok {
-						s.kopts.Meter.Counter(semconv.BrokerGroupLag, "topic", s.topic, "group", s.opts.Group, "partition", strconv.Itoa(int(tp.p)), "lag", strconv.Itoa(int(v.Lag)))
-					}
+				for p, l := range lmap {
+					s.kopts.Meter.Counter(semconv.BrokerGroupLag, "topic", s.topic, "group", s.opts.Group, "partition", strconv.Itoa(int(p)), "lag", strconv.Itoa(int(l.Lag)))
 				}
 				s.Unlock()
 
@@ -180,13 +181,13 @@ func (s *subscriber) assigned(_ context.Context, c *kgo.Client, assigned map[str
 				c:         c,
 				topic:     topic,
 				partition: partition,
-
-				quit:    make(chan struct{}),
-				done:    make(chan struct{}),
-				recs:    make(chan kgo.FetchTopicPartition, 100),
-				handler: s.handler,
-				kopts:   s.kopts,
-				opts:    s.opts,
+				htracer:   s.htracer,
+				quit:      make(chan struct{}),
+				done:      make(chan struct{}),
+				recs:      make(chan kgo.FetchTopicPartition, 100),
+				handler:   s.handler,
+				kopts:     s.kopts,
+				opts:      s.opts,
 			}
 			s.Lock()
 			s.consumers[tp{topic, partition}] = pc
@@ -212,6 +213,7 @@ func (pc *consumer) consume() {
 			return
 		case p := <-pc.recs:
 			for _, record := range p.Records {
+				ctx, sp := pc.htracer.WithProcessSpan(record)
 				ts := time.Now()
 				pc.kopts.Meter.Counter(semconv.SubscribeMessageInflight, "endpoint", record.Topic, "topic", record.Topic).Inc()
 				p := eventPool.Get().(*event)
@@ -221,6 +223,7 @@ func (pc *consumer) consume() {
 				p.err = nil
 				p.ack = false
 				p.msg.Header = metadata.New(len(record.Headers))
+				p.ctx = ctx
 				for _, hdr := range record.Headers {
 					p.msg.Header.Set(hdr.Key, string(hdr.Value))
 				}
@@ -229,7 +232,11 @@ func (pc *consumer) consume() {
 				} else if pc.opts.BodyOnly {
 					p.msg.Body = record.Value
 				} else {
-					if err := pc.kopts.Codec.Unmarshal(record.Value, p.msg); err != nil {
+					sp.AddEvent("codec unmarshal start")
+					err := pc.kopts.Codec.Unmarshal(record.Value, p.msg)
+					sp.AddEvent("codec unmarshal stop")
+					if err != nil {
+						sp.SetStatus(tracer.SpanStatusError, err.Error())
 						pc.kopts.Meter.Counter(semconv.SubscribeMessageTotal, "endpoint", record.Topic, "topic", record.Topic, "status", "failure").Inc()
 						p.err = err
 						p.msg.Body = record.Value
@@ -259,13 +266,17 @@ func (pc *consumer) consume() {
 						pc.kopts.Meter.Histogram(semconv.SubscribeMessageDurationSeconds, "endpoint", record.Topic, "topic", record.Topic).Update(te.Seconds())
 						eventPool.Put(p)
 						pc.kopts.Logger.Fatalf(pc.kopts.Context, "[kgo] Unmarshal err not handled wtf?")
+						sp.Finish()
 						return
 					}
 				}
+				sp.AddEvent("handler start")
 				err := pc.handler(p)
+				sp.AddEvent("handler stop")
 				if err == nil {
 					pc.kopts.Meter.Counter(semconv.SubscribeMessageTotal, "endpoint", record.Topic, "topic", record.Topic, "status", "success").Inc()
 				} else {
+					sp.SetStatus(tracer.SpanStatusError, err.Error())
 					pc.kopts.Meter.Counter(semconv.SubscribeMessageTotal, "endpoint", record.Topic, "topic", record.Topic, "status", "failure").Inc()
 				}
 				pc.kopts.Meter.Counter(semconv.SubscribeMessageInflight, "endpoint", record.Topic, "topic", record.Topic).Dec()
@@ -274,7 +285,9 @@ func (pc *consumer) consume() {
 				} else if err != nil {
 					p.err = err
 					if eh != nil {
+						sp.AddEvent("error handler start")
 						_ = eh(p)
+						sp.AddEvent("error handler stop")
 					} else {
 						if pc.kopts.Logger.V(logger.ErrorLevel) {
 							pc.kopts.Logger.Errorf(pc.kopts.Context, "[kgo]: subscriber error: %v", err)
@@ -290,8 +303,11 @@ func (pc *consumer) consume() {
 				} else {
 					eventPool.Put(p)
 					pc.kopts.Logger.Fatalf(pc.kopts.Context, "[kgo] ErrLostMessage wtf?")
+					sp.SetStatus(tracer.SpanStatusError, "ErrLostMessage")
+					sp.Finish()
 					return
 				}
+				sp.Finish()
 			}
 		}
 	}

tracer.go

@@ -2,12 +2,11 @@ package kgo
 
 import (
 	"context"
-	"net"
-	"time"
 	"unicode/utf8"
 
 	"github.com/twmb/franz-go/pkg/kgo"
 	semconv "go.opentelemetry.io/otel/semconv/v1.18.0"
+	"go.unistack.org/micro/v3/metadata"
 	"go.unistack.org/micro/v3/tracer"
 )
 
@@ -18,40 +17,12 @@ type hookTracer struct {
 }
 
 var (
-	_ kgo.HookBrokerConnect       = &hookTracer{}
-	_ kgo.HookBrokerDisconnect    = &hookTracer{}
-	_ kgo.HookBrokerRead          = &hookTracer{}
-	_ kgo.HookBrokerThrottle      = &hookTracer{}
-	_ kgo.HookBrokerWrite         = &hookTracer{}
-	_ kgo.HookFetchBatchRead      = &hookTracer{}
-	_ kgo.HookProduceBatchWritten = &hookTracer{}
-	_ kgo.HookGroupManageError    = &hookTracer{}
+	_ kgo.HookProduceRecordBuffered   = (*hookTracer)(nil)
+	_ kgo.HookProduceRecordUnbuffered = (*hookTracer)(nil)
+	_ kgo.HookFetchRecordBuffered     = (*hookTracer)(nil)
+	_ kgo.HookFetchRecordUnbuffered   = (*hookTracer)(nil)
 )
 
-func (m *hookTracer) OnGroupManageError(err error) {
-}
-
-func (m *hookTracer) OnBrokerConnect(meta kgo.BrokerMetadata, _ time.Duration, _ net.Conn, err error) {
-}
-
-func (m *hookTracer) OnBrokerDisconnect(meta kgo.BrokerMetadata, _ net.Conn) {
-}
-
-func (m *hookTracer) OnBrokerWrite(meta kgo.BrokerMetadata, _ int16, bytesWritten int, writeWait, timeToWrite time.Duration, err error) {
-}
-
-func (m *hookTracer) OnBrokerRead(meta kgo.BrokerMetadata, _ int16, bytesRead int, readWait, timeToRead time.Duration, err error) {
-}
-
-func (m *hookTracer) OnBrokerThrottle(meta kgo.BrokerMetadata, throttleInterval time.Duration, _ bool) {
-}
-
-func (m *hookTracer) OnProduceBatchWritten(meta kgo.BrokerMetadata, topic string, _ int32, kmetrics kgo.ProduceBatchMetrics) {
-}
-
-func (m *hookTracer) OnFetchBatchRead(meta kgo.BrokerMetadata, topic string, _ int32, kmetrics kgo.FetchBatchMetrics) {
-}
-
 // OnProduceRecordBuffered starts a new span for the "publish" operation on a
 // buffered record.
 //
@@ -74,8 +45,21 @@ func (m *hookTracer) OnProduceRecordBuffered(r *kgo.Record) {
 		tracer.WithSpanLabels(attrs...),
 		tracer.WithSpanKind(tracer.SpanKindProducer),
 	}
+
+	if r.Context == nil {
+		r.Context = context.Background()
+	}
+
+	md, ok := metadata.FromOutgoingContext(r.Context)
+	if !ok {
+		md = metadata.New(len(r.Headers))
+	}
+	for _, h := range r.Headers {
+		md.Set(h.Key, string(h.Value))
+	}
+
 	// Start the "publish" span.
-	ctx, _ := m.tracer.Start(r.Context, r.Topic+" publish", opts...)
+	ctx, _ := m.tracer.Start(metadata.NewOutgoingContext(r.Context, md), r.Topic+" publish", opts...)
 	// Inject the span context into the record.
 	// t.propagators.Inject(ctx, NewRecordCarrier(r))
 	// Update the record context.
@@ -88,17 +72,14 @@ func (m *hookTracer) OnProduceRecordBuffered(r *kgo.Record) {
 // It sets attributes with values unset when producing and records any error
 // that occurred during the publish operation.
 func (m *hookTracer) OnProduceRecordUnbuffered(r *kgo.Record, err error) {
-	span, ok := tracer.SpanFromContext(r.Context)
-	if !ok {
-		return
-	}
-	defer span.Finish()
+	span, _ := tracer.SpanFromContext(r.Context)
 	span.AddLabels(
 		semconv.MessagingKafkaDestinationPartition(int(r.Partition)),
 	)
 	if err != nil {
 		span.SetStatus(tracer.SpanStatusError, err.Error())
 	}
+	span.Finish()
 }
 
 // OnFetchRecordBuffered starts a new span for the "receive" operation on a
@@ -132,10 +113,18 @@ func (m *hookTracer) OnFetchRecordBuffered(r *kgo.Record) {
 	if r.Context == nil {
 		r.Context = context.Background()
 	}
+	md, ok := metadata.FromIncomingContext(r.Context)
+	if !ok {
+		md = metadata.New(len(r.Headers))
+	}
+	for _, h := range r.Headers {
+		md.Set(h.Key, string(h.Value))
+	}
+
 	// Extract the span context from the record.
 	// ctx := t.propagators.Extract(r.Context, NewRecordCarrier(r))
 	// Start the "receive" span.
-	newCtx, _ := m.tracer.Start(r.Context, r.Topic+" receive", opts...)
+	newCtx, _ := m.tracer.Start(metadata.NewIncomingContext(r.Context, md), r.Topic+" receive", opts...)
 	// Update the record context.
 	r.Context = newCtx
 }
@@ -143,9 +132,8 @@ func (m *hookTracer) OnFetchRecordBuffered(r *kgo.Record) {
 // OnFetchRecordUnbuffered continues and ends the "receive" span for an
 // unbuffered record.
 func (m *hookTracer) OnFetchRecordUnbuffered(r *kgo.Record, _ bool) {
-	if span, ok := tracer.SpanFromContext(r.Context); ok {
-		defer span.Finish()
-	}
+	span, _ := tracer.SpanFromContext(r.Context)
+	span.Finish()
 }
 
 // WithProcessSpan starts a new span for the "process" operation on a consumer
@@ -182,6 +170,14 @@ func (m *hookTracer) WithProcessSpan(r *kgo.Record) (context.Context, tracer.Spa
 	if r.Context == nil {
 		r.Context = context.Background()
 	}
+	md, ok := metadata.FromIncomingContext(r.Context)
+	if !ok {
+		md = metadata.New(len(r.Headers))
+	}
+	for _, h := range r.Headers {
+		md.Set(h.Key, string(h.Value))
+	}
+
 	// Start a new span using the provided context and options.
 	return m.tracer.Start(r.Context, r.Topic+" process", opts...)
 }