tracing fixes

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

@@ -62,7 +62,7 @@ type Broker struct {
 	connected bool
 	sync.RWMutex
 	opts broker.Options
-	subs []*subscriber
+	subs []*Subscriber
 }
 
 func (k *Broker) Address() string {
@@ -73,7 +73,11 @@ 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) Client() *kgo.Client {
+	return k.c
+}
+
+func (k *Broker) connect(ctx context.Context, opts ...kgo.Opt) (*kgo.Client, *hookTracer, error) {
 	var c *kgo.Client
 	var err error
 
@@ -90,9 +94,10 @@ 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 {
@@ -102,7 +107,7 @@ func (k *Broker) connect(ctx context.Context, opts ...kgo.Opt) (*kgo.Client, err
 				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 {
@@ -112,10 +117,10 @@ func (k *Broker) connect(ctx context.Context, opts ...kgo.Opt) (*kgo.Client, err
 			if sp != nil {
 				sp.SetStatus(tracer.SpanStatusError, err.Error())
 			}
-			return nil, err
+			return nil, nil, err
 		}
 	}
-	return c, nil
+	return c, htracer, nil
 }
 
 func (k *Broker) Connect(ctx context.Context) error {
@@ -131,7 +136,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
 	}
@@ -234,13 +239,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
@@ -325,6 +326,22 @@ func (k *Broker) publish(ctx context.Context, msgs []*broker.Message, opts ...br
 	return nil
 }
 
+func (k *Broker) TopicExists(ctx context.Context, topic string) error {
+	mdreq := kmsg.NewMetadataRequest()
+	mdreq.Topics = []kmsg.MetadataRequestTopic{
+		{Topic: &topic},
+	}
+
+	mdrsp, err := mdreq.RequestWith(ctx, k.c)
+	if err != nil {
+		return err
+	} else if mdrsp.Topics[0].ErrorCode != 0 {
+		return fmt.Errorf("topic %s not exists or permission error", topic)
+	}
+
+	return nil
+}
+
 func (k *Broker) BatchSubscribe(ctx context.Context, topic string, handler broker.BatchHandler, opts ...broker.SubscribeOption) (broker.Subscriber, error) {
 	return nil, nil
 }
@@ -347,7 +364,7 @@ func (k *Broker) Subscribe(ctx context.Context, topic string, handler broker.Han
 		}
 	}
 
-	sub := &subscriber{
+	sub := &Subscriber{
 		topic:     topic,
 		opts:      options,
 		handler:   handler,
@@ -374,7 +391,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
 	}
@@ -392,6 +409,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
@@ -31,9 +33,10 @@ type consumer struct {
 	recs      chan kgo.FetchTopicPartition
 }
 
-type subscriber struct {
+type Subscriber struct {
 	c         *kgo.Client
 	topic     string
+	htracer   *hookTracer
 	opts      broker.SubscribeOptions
 	kopts     broker.Options
 	handler   broker.Handler
@@ -43,15 +46,19 @@ type subscriber struct {
 	sync.RWMutex
 }
 
-func (s *subscriber) Options() broker.SubscribeOptions {
+func (s *Subscriber) Client() *kgo.Client {
+	return s.c
+}
+
+func (s *Subscriber) Options() broker.SubscribeOptions {
 	return s.opts
 }
 
-func (s *subscriber) Topic() string {
+func (s *Subscriber) Topic() string {
 	return s.topic
 }
 
-func (s *subscriber) Unsubscribe(ctx context.Context) error {
+func (s *Subscriber) Unsubscribe(ctx context.Context) error {
 	if s.closed {
 		return nil
 	}
@@ -73,7 +80,7 @@ func (s *subscriber) Unsubscribe(ctx context.Context) error {
 	return nil
 }
 
-func (s *subscriber) poll(ctx context.Context) {
+func (s *Subscriber) poll(ctx context.Context) {
 	maxInflight := DefaultSubscribeMaxInflight
 	if s.opts.Context != nil {
 		if n, ok := s.opts.Context.Value(subscribeMaxInflightKey{}).(int); n > 0 && ok {
@@ -141,7 +148,7 @@ func (s *subscriber) poll(ctx context.Context) {
 	}
 }
 
-func (s *subscriber) killConsumers(ctx context.Context, lost map[string][]int32) {
+func (s *Subscriber) killConsumers(ctx context.Context, lost map[string][]int32) {
 	var wg sync.WaitGroup
 	defer wg.Wait()
 
@@ -158,12 +165,12 @@ func (s *subscriber) killConsumers(ctx context.Context, lost map[string][]int32)
 	}
 }
 
-func (s *subscriber) lost(ctx context.Context, _ *kgo.Client, lost map[string][]int32) {
+func (s *Subscriber) lost(ctx context.Context, _ *kgo.Client, lost map[string][]int32) {
 	s.kopts.Logger.Debugf(ctx, "[kgo] lost %#+v", lost)
 	s.killConsumers(ctx, lost)
 }
 
-func (s *subscriber) revoked(ctx context.Context, c *kgo.Client, revoked map[string][]int32) {
+func (s *Subscriber) revoked(ctx context.Context, c *kgo.Client, revoked map[string][]int32) {
 	s.kopts.Logger.Debugf(ctx, "[kgo] revoked %#+v", revoked)
 	s.killConsumers(ctx, revoked)
 	if err := c.CommitMarkedOffsets(ctx); err != nil {
@@ -171,20 +178,20 @@ func (s *subscriber) revoked(ctx context.Context, c *kgo.Client, revoked map[str
 	}
 }
 
-func (s *subscriber) assigned(_ context.Context, c *kgo.Client, assigned map[string][]int32) {
+func (s *Subscriber) assigned(_ context.Context, c *kgo.Client, assigned map[string][]int32) {
 	for topic, partitions := range assigned {
 		for _, partition := range partitions {
 			pc := &consumer{
 				c:         c,
 				topic:     topic,
 				partition: partition,
-
+				htracer:   s.htracer,
-				quit:    make(chan struct{}),
+				quit:      make(chan struct{}),
-				done:    make(chan struct{}),
+				done:      make(chan struct{}),
-				recs:    make(chan kgo.FetchTopicPartition, 100),
+				recs:      make(chan kgo.FetchTopicPartition, 100),
-				handler: s.handler,
+				handler:   s.handler,
-				kopts:   s.kopts,
+				kopts:     s.kopts,
-				opts:    s.opts,
+				opts:      s.opts,
 			}
 			s.Lock()
 			s.consumers[tp{topic, partition}] = pc
@@ -210,6 +217,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)
@@ -219,6 +227,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))
 				}
@@ -227,7 +236,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
@@ -257,13 +270,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()
@@ -272,7 +289,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)
@@ -288,8 +307,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.HookProduceRecordBuffered   = (*hookTracer)(nil)
-	_ kgo.HookBrokerDisconnect    = &hookTracer{}
+	_ kgo.HookProduceRecordUnbuffered = (*hookTracer)(nil)
-	_ kgo.HookBrokerRead          = &hookTracer{}
+	_ kgo.HookFetchRecordBuffered     = (*hookTracer)(nil)
-	_ kgo.HookBrokerThrottle      = &hookTracer{}
+	_ kgo.HookFetchRecordUnbuffered   = (*hookTracer)(nil)
-	_ kgo.HookBrokerWrite         = &hookTracer{}
-	_ kgo.HookFetchBatchRead      = &hookTracer{}
-	_ kgo.HookProduceBatchWritten = &hookTracer{}
-	_ kgo.HookGroupManageError    = &hookTracer{}
 )
 
-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,12 +45,28 @@ func (m *hookTracer) OnProduceRecordBuffered(r *kgo.Record) {
 		tracer.WithSpanLabels(attrs...),
 		tracer.WithSpanKind(tracer.SpanKindProducer),
 	}
-	// Start the "publish" span.
+
-	ctx, _ := m.tracer.Start(r.Context, r.Topic+" publish", opts...)
+	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))
+	}
+
 	// Inject the span context into the record.
 	// t.propagators.Inject(ctx, NewRecordCarrier(r))
 	// Update the record context.
-	r.Context = ctx
+
+	if !ok {
+		r.Context, _ = m.tracer.Start(metadata.NewOutgoingContext(r.Context, md), r.Topic+" publish", opts...)
+	} else {
+		r.Context, _ = m.tracer.Start(r.Context, r.Topic+" publish", opts...)
+	}
 }
 
 // OnProduceRecordUnbuffered continues and ends the "publish" span for an
@@ -88,17 +75,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)
+	span, _ := tracer.SpanFromContext(r.Context)
-	if !ok {
-		return
-	}
-	defer span.Finish()
 	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,20 +116,31 @@ 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...)
+	if !ok {
+		r.Context, _ = m.tracer.Start(metadata.NewIncomingContext(r.Context, md), r.Topic+" receive", opts...)
+	} else {
+		r.Context, _ = m.tracer.Start(r.Context, r.Topic+" receive", opts...)
+	}
+
 	// Update the record context.
-	r.Context = newCtx
 }
 
 // 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 {
+	span, _ := tracer.SpanFromContext(r.Context)
-		defer span.Finish()
+	span.Finish()
-	}
 }
 
 // WithProcessSpan starts a new span for the "process" operation on a consumer
@@ -182,6 +177,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...)
 }