proto modify

pkg/common/kafka/consumer_group.go

@@ -7,6 +7,8 @@
 package kafka
 
 import (
+	"OpenIM/pkg/common/constant"
+	"OpenIM/pkg/common/tracelog"
 	"context"
 	"github.com/Shopify/sarama"
 )
@@ -39,6 +41,19 @@ func NewMConsumerGroup(consumerConfig *MConsumerGroupConfig, topics, addrs []str
 		topics,
 	}
 }
+
+func (mc *MConsumerGroup) GetContextFromMsg(cMsg *sarama.ConsumerMessage) context.Context {
+	ctx := context.Background()
+	var operationID string
+	for _, v := range cMsg.Headers {
+		if string(v.Key) == constant.OperationID {
+			operationID = string(v.Value)
+		}
+	}
+	tracelog.SetOperationID(ctx, operationID)
+	return ctx
+}
+
 func (mc *MConsumerGroup) RegisterHandleAndConsumer(handler sarama.ConsumerGroupHandler) {
 	ctx := context.Background()
 	for {

pkg/common/kafka/producer.go

@@ -2,6 +2,7 @@ package kafka
 
 import (
 	"OpenIM/pkg/common/config"
+	"OpenIM/pkg/common/constant"
 	log "OpenIM/pkg/common/log"
 	"OpenIM/pkg/common/tracelog"
 	"OpenIM/pkg/utils"
@@ -67,6 +68,7 @@ func (p *Producer) SendMessage(ctx context.Context, m proto.Message, key string)
 		return -1, -1, errors.New("key or value == 0")
 	}
 	kMsg.Metadata = ctx
+	kMsg.Headers = []sarama.RecordHeader{{Key: []byte(constant.OperationID), Value: []byte(operationID)}}
 	partition, offset, err := p.producer.SendMessage(kMsg)
 	log.Info(operationID, "ByteEncoder SendMessage end", "key ", kMsg.Key.Length(), kMsg.Value.Length(), p.producer)
 	if err == nil {