feat 初始化

db/kafka/consumer.go

@@ -12,6 +12,8 @@ import (
 	"go.opentelemetry.io/otel/attribute"
 	otelcodes "go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/trace"
+	"reflect"
+	"sync"
 	"time"
 )
 
@@ -31,10 +33,13 @@ type Consumer struct {
 func (c *Consumer) Consume(t []Topic) {
 	for {
 		topicArr := make([]string, 0)
-		handlerMap := make(map[string]Topic)
+		handlerMap := make(map[string]*sync.Pool)
 		for _, v := range t {
 			topicArr = append(topicArr, v.Name())
-			handlerMap[v.Name()] = v
+			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,
@@ -48,13 +53,13 @@ func (c *Consumer) Consume(t []Topic) {
 
 func (c *Consumer) Stop() error {
 	c.cancel()
-	return nil
+	return client.consumer.Close()
 }
 
 type Handler func(context.Context, []byte) error
 
 type handler struct {
-	handler map[string]Topic
+	handler map[string]*sync.Pool
 }
 
 func (h *handler) Setup(_ sarama.ConsumerGroupSession) error {
@@ -67,39 +72,72 @@ func (h *handler) Cleanup(sess sarama.ConsumerGroupSession) error {
 }
 
 func (h *handler) ConsumeClaim(sess sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
-	for message := range claim.Messages() {
-		ctx := NewCarrier().ExtractConsumer(message.Headers)
-		_, span := otel.Tracer("common.db.kafka").Start(ctx, "kafka.consume")
+	for {
+		select {
+		case message, ok := <-claim.Messages():
+			if ok {
+				ctx := NewCarrier().ExtractConsumer(message.Headers)
+				_, span := otel.Tracer("common.db.kafka").Start(ctx, "kafka.consume")
 
-		cb := h.handler[message.Topic]
-		if cb == nil {
-			span.SetStatus(otelcodes.Error, "handler not found")
-			span.End()
-			return utils.ErrorsWrap(errors.New("handler not found"))
-		}
+				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)))
 
-		if err := json.Unmarshal(message.Value, cb); err != nil {
-			span.SetStatus(otelcodes.Error, "handler json.Unmarshal error")
-			span.End()
-			return utils.ErrorsWrap(err, "handler json.Unmarshal error")
-		}
+				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 := cb.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())),
-				))
+				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()
 			}
-			span.SetStatus(otelcodes.Error, err.Error())
-			span.End()
-			return utils.ErrorsWrap(err, "kafka handler error")
+		case <-sess.Context().Done():
+			log.Warnf("kafka session done, reason: %v", sess.Context().Err())
+			sess.Commit()
+			return sess.Context().Err()
 		}
-
-		sess.MarkMessage(message, "")
-		span.End()
 	}
-	sess.Commit()
-	return nil
+}
+
+// 回收Topic对象
+func releaseTopicObject(pool *sync.Pool, obj Topic) {
+	v := reflect.ValueOf(obj).Elem()
+	v.Set(reflect.Zero(v.Type()))
+	pool.Put(obj)
 }
 
 func consumerError() {