From cfecb4afd024d46da37c31799c51e9cba6c8f7e1 Mon Sep 17 00:00:00 2001 From: Vasiliy Tolstov Date: Tue, 25 Feb 2025 17:30:48 +0300 Subject: [PATCH] export broker state Signed-off-by: Vasiliy Tolstov --- broker.go | 56 +++++++++++++++++++++++++++++++++++++++++++++++++++ kgo.go | 3 +++ subscriber.go | 50 ++++++++++++++++++++++++++++----------------- 3 files changed, 90 insertions(+), 19 deletions(-) create mode 100644 broker.go diff --git a/broker.go b/broker.go new file mode 100644 index 0000000..b85eb54 --- /dev/null +++ b/broker.go @@ -0,0 +1,56 @@ +package kgo + +import ( + "net" + "sync/atomic" + "time" + + "github.com/twmb/franz-go/pkg/kgo" +) + +type hookEvent struct { + connected *atomic.Uint32 +} + +var ( + _ kgo.HookBrokerConnect = &hookEvent{} + _ kgo.HookBrokerDisconnect = &hookEvent{} + _ kgo.HookBrokerRead = &hookEvent{} + _ kgo.HookBrokerWrite = &hookEvent{} + _ kgo.HookGroupManageError = &hookEvent{} + _ kgo.HookProduceRecordUnbuffered = &hookEvent{} +) + +func (m *hookEvent) OnGroupManageError(err error) { + if err != nil { + m.connected.Store(0) + } +} + +func (m *hookEvent) OnBrokerConnect(_ kgo.BrokerMetadata, _ time.Duration, _ net.Conn, err error) { + if err != nil { + m.connected.Store(0) + } +} + +func (m *hookEvent) OnBrokerDisconnect(_ kgo.BrokerMetadata, _ net.Conn) { + m.connected.Store(0) +} + +func (m *hookEvent) OnBrokerWrite(_ kgo.BrokerMetadata, _ int16, _ int, _ time.Duration, _ time.Duration, err error) { + if err != nil { + m.connected.Store(0) + } +} + +func (m *hookEvent) OnBrokerRead(_ kgo.BrokerMetadata, _ int16, _ int, _ time.Duration, _ time.Duration, err error) { + if err != nil { + m.connected.Store(0) + } +} + +func (m *hookEvent) OnProduceRecordUnbuffered(_ *kgo.Record, err error) { + if err != nil { + m.connected.Store(0) + } +} diff --git a/kgo.go b/kgo.go index af0549e..ae62c35 100644 --- a/kgo.go +++ b/kgo.go @@ -113,6 +113,7 @@ func (k *Broker) connect(ctx context.Context, opts ...kgo.Opt) (*kgo.Client, *ho opts = append(opts, kgo.WithHooks(&hookMeter{meter: k.opts.Meter}), kgo.WithHooks(htracer), + kgo.WithHooks(&hookEvent{connected: k.connected}), ) select { @@ -390,7 +391,9 @@ func (k *Broker) Subscribe(ctx context.Context, topic string, handler broker.Han kgo.AutoCommitInterval(commitInterval), kgo.OnPartitionsAssigned(sub.assigned), kgo.OnPartitionsRevoked(sub.revoked), + kgo.StopProducerOnDataLossDetected(), kgo.OnPartitionsLost(sub.lost), + kgo.AutoCommitCallback(sub.autocommit), kgo.AutoCommitMarks(), ) diff --git a/subscriber.go b/subscriber.go index 36b4929..64f5d11 100644 --- a/subscriber.go +++ b/subscriber.go @@ -5,10 +5,12 @@ import ( "fmt" "strconv" "sync" + "sync/atomic" "time" "github.com/twmb/franz-go/pkg/kadm" "github.com/twmb/franz-go/pkg/kgo" + "github.com/twmb/franz-go/pkg/kmsg" "go.unistack.org/micro/v3/broker" "go.unistack.org/micro/v3/logger" "go.unistack.org/micro/v3/metadata" @@ -22,20 +24,17 @@ type tp struct { } type consumer struct { - topic string - - c *kgo.Client + topic string + c *kgo.Client htracer *hookTracer - - handler broker.Handler - quit chan struct{} - done chan struct{} - recs chan kgo.FetchTopicPartition - - kopts broker.Options - opts broker.SubscribeOptions - + quit chan struct{} + done chan struct{} + recs chan kgo.FetchTopicPartition + kopts broker.Options partition int32 + opts broker.SubscribeOptions + handler broker.Handler + connected *atomic.Uint32 } type Subscriber struct { @@ -49,6 +48,7 @@ type Subscriber struct { kopts broker.Options opts broker.SubscribeOptions + connected *atomic.Uint32 sync.RWMutex closed bool } @@ -144,8 +144,8 @@ func (s *Subscriber) poll(ctx context.Context) { }) fetches.EachPartition(func(p kgo.FetchTopicPartition) { - nTp := tp{p.Topic, p.Partition} - s.consumers[nTp].recs <- p + tps := tp{p.Topic, p.Partition} + s.consumers[tps].recs <- p }) s.c.AllowRebalance() } @@ -158,9 +158,9 @@ func (s *Subscriber) killConsumers(ctx context.Context, lost map[string][]int32) for topic, partitions := range lost { for _, partition := range partitions { - nTp := tp{topic, partition} - pc := s.consumers[nTp] - delete(s.consumers, nTp) + tps := tp{topic, partition} + pc := s.consumers[tps] + delete(s.consumers, tps) close(pc.quit) if s.kopts.Logger.V(logger.DebugLevel) { s.kopts.Logger.Debug(ctx, fmt.Sprintf("[kgo] waiting for work to finish topic %s partition %d", topic, partition)) @@ -171,11 +171,18 @@ func (s *Subscriber) killConsumers(ctx context.Context, lost map[string][]int32) } } +func (s *Subscriber) autocommit(_ *kgo.Client, _ *kmsg.OffsetCommitRequest, _ *kmsg.OffsetCommitResponse, err error) { + if err != nil { + s.connected.Store(0) + } +} + func (s *Subscriber) lost(ctx context.Context, _ *kgo.Client, lost map[string][]int32) { - if s.kopts.Logger.V(logger.DebugLevel) { - s.kopts.Logger.Debug(ctx, fmt.Sprintf("[kgo] lost %#+v", lost)) + if s.kopts.Logger.V(logger.ErrorLevel) { + s.kopts.Logger.Error(ctx, fmt.Sprintf("[kgo] lost %#+v", lost)) } s.killConsumers(ctx, lost) + // s.connected.Store(0) } func (s *Subscriber) revoked(ctx context.Context, c *kgo.Client, revoked map[string][]int32) { @@ -185,6 +192,7 @@ func (s *Subscriber) revoked(ctx context.Context, c *kgo.Client, revoked map[str s.killConsumers(ctx, revoked) if err := c.CommitMarkedOffsets(ctx); err != nil { s.kopts.Logger.Error(ctx, "[kgo] revoked CommitMarkedOffsets error", err) + // s.connected.Store(0) } } @@ -202,6 +210,7 @@ func (s *Subscriber) assigned(_ context.Context, c *kgo.Client, assigned map[str handler: s.handler, kopts: s.kopts, opts: s.opts, + connected: s.connected, } s.Lock() s.consumers[tp{topic, partition}] = pc @@ -263,6 +272,7 @@ func (pc *consumer) consume() { pc.c.MarkCommitRecords(record) } else { eventPool.Put(p) + pc.connected.Store(0) pc.kopts.Logger.Fatal(pc.kopts.Context, "[kgo] ErrLostMessage wtf?") return } @@ -279,6 +289,7 @@ func (pc *consumer) consume() { pc.kopts.Meter.Summary(semconv.SubscribeMessageLatencyMicroseconds, "endpoint", record.Topic, "topic", record.Topic).Update(te.Seconds()) pc.kopts.Meter.Histogram(semconv.SubscribeMessageDurationSeconds, "endpoint", record.Topic, "topic", record.Topic).Update(te.Seconds()) eventPool.Put(p) + pc.connected.Store(0) pc.kopts.Logger.Fatal(pc.kopts.Context, "[kgo] Unmarshal err not handled wtf?") sp.Finish() return @@ -316,6 +327,7 @@ func (pc *consumer) consume() { pc.c.MarkCommitRecords(record) } else { eventPool.Put(p) + pc.connected.Store(0) pc.kopts.Logger.Fatal(pc.kopts.Context, "[kgo] ErrLostMessage wtf?") sp.SetStatus(tracer.SpanStatusError, "ErrLostMessage") sp.Finish()