unistack-org/micro-broker-kgo
4
0
Fork 2
Code Issues 3 Pull Requests 2 Actions Packages Projects Releases 73 Wiki Activity

update to v4
All checks were successful
test / test (push) Successful in 3m29s
Details

Browse Source 
Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
This commit is contained in:
Василий Толстов
2025-03-12 16:57:25 +03:00
parent f4f8793686
commit ac454693a6
12 changed files with 240 additions and 300 deletions
Download Patch File Download Diff File Expand all files Collapse all files

245
kgo.go
View File

@@ -11,19 +11,27 @@ import (
"sync/atomic"
"time"
"github.com/google/uuid"
"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"
"go.unistack.org/micro/v3/semconv"
"go.unistack.org/micro/v3/tracer"
mrand "go.unistack.org/micro/v3/util/rand"
"go.unistack.org/micro/v4/broker"
"go.unistack.org/micro/v4/codec"
"go.unistack.org/micro/v4/logger"
"go.unistack.org/micro/v4/metadata"
"go.unistack.org/micro/v4/options"
"go.unistack.org/micro/v4/semconv"
"go.unistack.org/micro/v4/tracer"
"go.unistack.org/micro/v4/util/id"
mrand "go.unistack.org/micro/v4/util/rand"
)
var _ broker.Broker = (*Broker)(nil)
var messagePool = sync.Pool{
New: func() interface{} {
return &kgoMessage{}
},
}
var ErrLostMessage = errors.New("message not marked for offsets commit and will be lost in next iteration")
var DefaultRetryBackoffFn = func() func(int) time.Duration {
@@ -56,8 +64,10 @@ var DefaultRetryBackoffFn = func() func(int) time.Duration {
}()
type Broker struct {
c *kgo.Client
connected *atomic.Uint32
funcPublish broker.FuncPublish
funcSubscribe broker.FuncSubscribe
c *kgo.Client
connected *atomic.Uint32
kopts []kgo.Opt
subs []*Subscriber
@@ -92,6 +102,68 @@ func (k *Broker) Client() *kgo.Client {
return k.c
}
type kgoMessage struct {
c codec.Codec
topic string
ctx context.Context
body []byte
hdr metadata.Metadata
opts broker.PublishOptions
ack bool
}
func (m *kgoMessage) Ack() error {
m.ack = true
return nil
}
func (m *kgoMessage) Body() []byte {
return m.body
}
func (m *kgoMessage) Header() metadata.Metadata {
return m.hdr
}
func (m *kgoMessage) Context() context.Context {
return m.ctx
}
func (m *kgoMessage) Topic() string {
return ""
}
func (m *kgoMessage) Unmarshal(dst interface{}, opts ...codec.Option) error {
return m.c.Unmarshal(m.body, dst)
}
func (b *Broker) newCodec(ct string) (codec.Codec, error) {
if idx := strings.IndexRune(ct, ';'); idx >= 0 {
ct = ct[:idx]
}
b.RLock()
c, ok := b.opts.Codecs[ct]
b.RUnlock()
if ok {
return c, nil
}
return nil, codec.ErrUnknownContentType
}
func (b *Broker) NewMessage(ctx context.Context, hdr metadata.Metadata, body interface{}, opts ...broker.PublishOption) (broker.Message, error) {
options := broker.NewPublishOptions(opts...)
m := &kgoMessage{ctx: ctx, hdr: hdr, opts: options}
c, err := b.newCodec(m.opts.ContentType)
if err == nil {
m.body, err = c.Marshal(body)
}
if err != nil {
return nil, err
}
return m, nil
}
func (k *Broker) connect(ctx context.Context, opts ...kgo.Opt) (*kgo.Client, *hookTracer, error) {
var c *kgo.Client
var err error
@@ -238,6 +310,18 @@ func (k *Broker) Init(opts ...broker.Option) error {
}
}
k.funcPublish = k.fnPublish
k.funcSubscribe = k.fnSubscribe
k.opts.Hooks.EachPrev(func(hook options.Hook) {
switch h := hook.(type) {
case broker.HookPublish:
k.funcPublish = h(k.funcPublish)
case broker.HookSubscribe:
k.funcSubscribe = h(k.funcSubscribe)
}
})
k.init = true
return nil
@@ -247,95 +331,88 @@ func (k *Broker) Options() broker.Options {
return k.opts
}
func (k *Broker) BatchPublish(ctx context.Context, msgs []*broker.Message, opts ...broker.PublishOption) error {
return k.publish(ctx, msgs, opts...)
func (b *Broker) Publish(ctx context.Context, topic string, messages ...broker.Message) error {
return b.funcPublish(ctx, topic, messages...)
}
func (k *Broker) Publish(ctx context.Context, topic string, msg *broker.Message, opts ...broker.PublishOption) error {
msg.Header.Set(metadata.HeaderTopic, topic)
return k.publish(ctx, []*broker.Message{msg}, opts...)
func (b *Broker) fnPublish(ctx context.Context, topic string, messages ...broker.Message) error {
return b.publish(ctx, topic, messages...)
}
func (k *Broker) publish(ctx context.Context, msgs []*broker.Message, opts ...broker.PublishOption) error {
k.Lock()
if k.connected.Load() == 0 {
c, _, err := k.connect(ctx, k.kopts...)
func (b *Broker) publish(ctx context.Context, topic string, messages ...broker.Message) error {
if b.connected.Load() == 0 {
c, _, err := b.connect(ctx, b.kopts...)
if err != nil {
k.Unlock()
return err
}
k.c = c
k.connected.Store(1)
b.Lock()
b.c = c
b.Unlock()
}
k.Unlock()
options := broker.NewPublishOptions(opts...)
records := make([]*kgo.Record, 0, len(msgs))
records := make([]*kgo.Record, 0, len(messages))
var errs []string
var err error
var key []byte
var promise func(*kgo.Record, error)
if options.Context != nil {
if k, ok := options.Context.Value(publishKey{}).([]byte); ok && k != nil {
key = k
}
if p, ok := options.Context.Value(publishPromiseKey{}).(func(*kgo.Record, error)); ok && p != nil {
promise = p
}
}
for _, msg := range messages {
for _, msg := range msgs {
rec := &kgo.Record{Context: ctx, Key: key}
rec.Topic, _ = msg.Header.Get(metadata.HeaderTopic)
msg.Header.Del(metadata.HeaderTopic)
k.opts.Meter.Counter(semconv.PublishMessageInflight, "endpoint", rec.Topic, "topic", rec.Topic).Inc()
if options.BodyOnly || k.opts.Codec.String() == "noop" {
rec.Value = msg.Body
setHeaders(rec, msg.Header)
} else {
rec.Value, err = k.opts.Codec.Marshal(msg)
if err != nil {
return err
if mctx := msg.Context(); mctx != nil {
if k, ok := mctx.Value(publishKey{}).([]byte); ok && k != nil {
key = k
}
if p, ok := mctx.Value(publishPromiseKey{}).(func(*kgo.Record, error)); ok && p != nil {
promise = p
}
}
rec := &kgo.Record{
Context: ctx,
Key: key,
Topic: topic,
Value: msg.Body(),
}
b.opts.Meter.Counter(semconv.PublishMessageInflight, "endpoint", rec.Topic, "topic", rec.Topic).Inc()
setHeaders(rec, msg.Header())
records = append(records, rec)
}
ts := time.Now()
if promise != nil {
ts := time.Now()
for _, rec := range records {
k.c.Produce(ctx, rec, func(r *kgo.Record, err error) {
b.c.Produce(ctx, rec, func(r *kgo.Record, err error) {
te := time.Since(ts)
k.opts.Meter.Counter(semconv.PublishMessageInflight, "endpoint", rec.Topic, "topic", rec.Topic).Dec()
k.opts.Meter.Summary(semconv.PublishMessageLatencyMicroseconds, "endpoint", rec.Topic, "topic", rec.Topic).Update(te.Seconds())
k.opts.Meter.Histogram(semconv.PublishMessageDurationSeconds, "endpoint", rec.Topic, "topic", rec.Topic).Update(te.Seconds())
b.opts.Meter.Counter(semconv.PublishMessageInflight, "endpoint", rec.Topic, "topic", rec.Topic).Dec()
b.opts.Meter.Summary(semconv.PublishMessageLatencyMicroseconds, "endpoint", rec.Topic, "topic", rec.Topic).Update(te.Seconds())
b.opts.Meter.Histogram(semconv.PublishMessageDurationSeconds, "endpoint", rec.Topic, "topic", rec.Topic).Update(te.Seconds())
if err != nil {
k.opts.Meter.Counter(semconv.PublishMessageTotal, "endpoint", rec.Topic, "topic", rec.Topic, "status", "failure").Inc()
b.opts.Meter.Counter(semconv.PublishMessageTotal, "endpoint", rec.Topic, "topic", rec.Topic, "status", "failure").Inc()
} else {
k.opts.Meter.Counter(semconv.PublishMessageTotal, "endpoint", rec.Topic, "topic", rec.Topic, "status", "success").Inc()
b.opts.Meter.Counter(semconv.PublishMessageTotal, "endpoint", rec.Topic, "topic", rec.Topic, "status", "success").Inc()
}
promise(r, err)
})
}
return nil
}
ts := time.Now()
results := k.c.ProduceSync(ctx, records...)
results := b.c.ProduceSync(ctx, records...)
te := time.Since(ts)
for _, result := range results {
k.opts.Meter.Summary(semconv.PublishMessageLatencyMicroseconds, "endpoint", result.Record.Topic, "topic", result.Record.Topic).Update(te.Seconds())
k.opts.Meter.Histogram(semconv.PublishMessageDurationSeconds, "endpoint", result.Record.Topic, "topic", result.Record.Topic).Update(te.Seconds())
k.opts.Meter.Counter(semconv.PublishMessageInflight, "endpoint", result.Record.Topic, "topic", result.Record.Topic).Dec()
b.opts.Meter.Summary(semconv.PublishMessageLatencyMicroseconds, "endpoint", result.Record.Topic, "topic", result.Record.Topic).Update(te.Seconds())
b.opts.Meter.Histogram(semconv.PublishMessageDurationSeconds, "endpoint", result.Record.Topic, "topic", result.Record.Topic).Update(te.Seconds())
b.opts.Meter.Counter(semconv.PublishMessageInflight, "endpoint", result.Record.Topic, "topic", result.Record.Topic).Dec()
if result.Err != nil {
k.opts.Meter.Counter(semconv.PublishMessageTotal, "endpoint", result.Record.Topic, "topic", result.Record.Topic, "status", "failure").Inc()
b.opts.Meter.Counter(semconv.PublishMessageTotal, "endpoint", result.Record.Topic, "topic", result.Record.Topic, "status", "failure").Inc()
errs = append(errs, result.Err.Error())
} else {
k.opts.Meter.Counter(semconv.PublishMessageTotal, "endpoint", result.Record.Topic, "topic", result.Record.Topic, "status", "success").Inc()
b.opts.Meter.Counter(semconv.PublishMessageTotal, "endpoint", result.Record.Topic, "topic", result.Record.Topic, "status", "success").Inc()
}
}
@@ -362,31 +439,44 @@ func (k *Broker) TopicExists(ctx context.Context, topic string) error {
return nil
}
func (k *Broker) BatchSubscribe(_ context.Context, _ string, _ broker.BatchHandler, _ ...broker.SubscribeOption) (broker.Subscriber, error) {
return nil, nil
func (b *Broker) Subscribe(ctx context.Context, topic string, handler interface{}, opts ...broker.SubscribeOption) (broker.Subscriber, error) {
return b.funcSubscribe(ctx, topic, handler, opts...)
}
func (k *Broker) Subscribe(ctx context.Context, topic string, handler broker.Handler, opts ...broker.SubscribeOption) (broker.Subscriber, error) {
func (b *Broker) fnSubscribe(ctx context.Context, topic string, handler interface{}, opts ...broker.SubscribeOption) (broker.Subscriber, error) {
if err := broker.IsValidHandler(handler); err != nil {
return nil, err
}
options := broker.NewSubscribeOptions(opts...)
switch handler.(type) {
default:
return nil, broker.ErrInvalidHandler
case func(broker.Message) error:
break
case func([]broker.Message) error:
break
}
if options.Group == "" {
uid, err := uuid.NewRandom()
uid, err := id.New()
if err != nil {
return nil, err
}
options.Group = uid.String()
options.Group = uid
}
commitInterval := DefaultCommitInterval
if k.opts.Context != nil {
if v, ok := k.opts.Context.Value(commitIntervalKey{}).(time.Duration); ok && v > 0 {
if b.opts.Context != nil {
if v, ok := b.opts.Context.Value(commitIntervalKey{}).(time.Duration); ok && v > 0 {
commitInterval = v
}
}
var fatalOnError bool
if k.opts.Context != nil {
if v, ok := k.opts.Context.Value(fatalOnErrorKey{}).(bool); ok && v {
if b.opts.Context != nil {
if v, ok := b.opts.Context.Value(fatalOnErrorKey{}).(bool); ok && v {
fatalOnError = v
}
}
@@ -401,14 +491,14 @@ func (k *Broker) Subscribe(ctx context.Context, topic string, handler broker.Han
topic: topic,
opts: options,
handler: handler,
kopts: k.opts,
kopts: b.opts,
consumers: make(map[tp]*consumer),
done: make(chan struct{}),
fatalOnError: fatalOnError,
connected: k.connected,
connected: b.connected,
}
kopts := append(k.kopts,
kopts := append(b.kopts,
kgo.ConsumerGroup(options.Group),
kgo.ConsumeTopics(topic),
kgo.ConsumeResetOffset(kgo.NewOffset().AtStart()),
@@ -428,7 +518,7 @@ func (k *Broker) Subscribe(ctx context.Context, topic string, handler broker.Han
}
}
c, htracer, err := k.connect(ctx, kopts...)
c, htracer, err := b.connect(ctx, kopts...)
if err != nil {
return nil, err
}
@@ -450,9 +540,10 @@ func (k *Broker) Subscribe(ctx context.Context, topic string, handler broker.Han
go sub.poll(ctx)
k.Lock()
k.subs = append(k.subs, sub)
k.Unlock()
b.Lock()
b.subs = append(b.subs, sub)
b.Unlock()
return sub, nil
}