148 lines
3.7 KiB
Go
148 lines
3.7 KiB
Go
package kafka
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"git.hlsq.asia/alex/service-common/log"
|
|
"git.hlsq.asia/alex/service-common/utils"
|
|
"github.com/IBM/sarama"
|
|
"go.opentelemetry.io/otel"
|
|
"go.opentelemetry.io/otel/attribute"
|
|
otelcodes "go.opentelemetry.io/otel/codes"
|
|
"go.opentelemetry.io/otel/trace"
|
|
"reflect"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
func NewConsumer() *Consumer {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
return &Consumer{
|
|
ctx: ctx,
|
|
cancel: cancel,
|
|
}
|
|
}
|
|
|
|
type Consumer struct {
|
|
ctx context.Context
|
|
cancel context.CancelFunc
|
|
}
|
|
|
|
func (c *Consumer) Consume(t []Topic) {
|
|
for {
|
|
topicArr := make([]string, 0)
|
|
handlerMap := make(map[string]*sync.Pool)
|
|
for _, v := range t {
|
|
topicArr = append(topicArr, v.Name())
|
|
typ := reflect.ValueOf(v).Elem().Type()
|
|
handlerMap[v.Name()] = &sync.Pool{New: func() interface{} {
|
|
return reflect.New(typ).Interface()
|
|
}}
|
|
}
|
|
err := client.consumer.Consume(c.ctx, topicArr, &handler{
|
|
handler: handlerMap,
|
|
})
|
|
if errors.Is(err, context.Canceled) || errors.Is(c.ctx.Err(), context.Canceled) {
|
|
return
|
|
}
|
|
time.Sleep(time.Second)
|
|
}
|
|
}
|
|
|
|
func (c *Consumer) Stop() error {
|
|
c.cancel()
|
|
return client.consumer.Close()
|
|
}
|
|
|
|
type Handler func(context.Context, []byte) error
|
|
|
|
type handler struct {
|
|
handler map[string]*sync.Pool
|
|
}
|
|
|
|
func (h *handler) Setup(_ sarama.ConsumerGroupSession) error {
|
|
return nil
|
|
}
|
|
|
|
func (h *handler) Cleanup(sess sarama.ConsumerGroupSession) error {
|
|
sess.Commit()
|
|
return nil
|
|
}
|
|
|
|
func (h *handler) ConsumeClaim(sess sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
|
|
for {
|
|
select {
|
|
case message, ok := <-claim.Messages():
|
|
if ok {
|
|
ctx := NewCarrier().ExtractConsumer(message.Headers)
|
|
_, span := otel.Tracer("common.db.kafka").Start(ctx, "kafka.consume")
|
|
|
|
pool := h.handler[message.Topic]
|
|
if pool == nil {
|
|
span.SetStatus(otelcodes.Error, "handler not found")
|
|
span.End()
|
|
return utils.ErrorsWrap(errors.New("handler not found"))
|
|
}
|
|
span.SetAttributes(attribute.String("kafka.topic", message.Topic))
|
|
span.SetAttributes(attribute.String("kafka.message", string(message.Value)))
|
|
|
|
newCb := pool.Get().(Topic)
|
|
if err := json.Unmarshal(message.Value, newCb); err != nil {
|
|
releaseTopicObject(pool, newCb)
|
|
span.SetStatus(otelcodes.Error, "handler json.Unmarshal error")
|
|
span.End()
|
|
return utils.ErrorsWrap(err, "handler json.Unmarshal error")
|
|
}
|
|
|
|
if err := func() (err error) {
|
|
defer func() {
|
|
if e := recover(); e != nil {
|
|
span.SetStatus(otelcodes.Error, fmt.Sprintf("%v", e))
|
|
span.End()
|
|
err = fmt.Errorf("kafka handler panic: %v", e)
|
|
}
|
|
}()
|
|
if err = newCb.OnMessage(ctx); err != nil {
|
|
if stack, ok := err.(interface{ StackTrace() string }); ok {
|
|
span.AddEvent("Stack Trace", trace.WithAttributes(
|
|
attribute.String("stack.trace", fmt.Sprintf("%v", stack.StackTrace())),
|
|
))
|
|
}
|
|
span.SetStatus(otelcodes.Error, err.Error())
|
|
span.End()
|
|
log.Errorf("kafka handler error: %v", err)
|
|
err = utils.ErrorsWrap(err, "kafka handler error")
|
|
}
|
|
return
|
|
}(); err != nil {
|
|
releaseTopicObject(pool, newCb)
|
|
return err
|
|
}
|
|
|
|
sess.MarkMessage(message, "")
|
|
releaseTopicObject(pool, newCb)
|
|
span.End()
|
|
}
|
|
case <-sess.Context().Done():
|
|
log.Warnf("kafka session done, reason: %v", sess.Context().Err())
|
|
sess.Commit()
|
|
return sess.Context().Err()
|
|
}
|
|
}
|
|
}
|
|
|
|
// 回收Topic对象
|
|
func releaseTopicObject(pool *sync.Pool, obj Topic) {
|
|
v := reflect.ValueOf(obj).Elem()
|
|
v.Set(reflect.Zero(v.Type()))
|
|
pool.Put(obj)
|
|
}
|
|
|
|
func consumerError() {
|
|
for err := range client.consumer.Errors() {
|
|
log.Errorf("kafka consumer error: %v", err)
|
|
}
|
|
}
|