go - 使用 Otel 和 Shopify/Sarama 跟踪 Kafka

package kafka import ( "crypto/tls" "encoding/json" "net" "os" "strings" "sync" "time" "github.com/Shopify/sarama" "github.com/prometheus/client_golang/prometheus" "go.opentelemetry.io/contrib/instrumentation/github.com/Shopify/sarama/otelsarama" "go.opentelemetry.io/otel/propagation" "go.opentelemetry.io/otel/trace" "golang.org/x/net/context" "abc.com/go/gofr/pkg" "abc.com/go/gofr/pkg/datastore/pubsub" "abc.com/go/gofr/pkg/datastore/pubsub/avro" "abc.com/go/gofr/pkg/errors" "abc.com/go/gofr/pkg/gofr/types" "abc.com/go/gofr/pkg/log" ) const ( ErrConsumeMsg = errors.Error("error while consuming the message") SASLTypeSCRAMSHA512 = "SCRAM-SHA-512" PLAIN = "PLAIN" errInvalidMechanism = errors.Error("Invalid SASL Mechanism") ) type Kafka struct { config *Config Producer sarama.SyncProducer Consumer *Consumer } type AvroWithKafkaConfig struct { KafkaConfig Config AvroConfig avro.Config } const ( // OffsetNewest stands for the log head offset, i.e. the offset that will be // assigned to the next message that will be produced to the partition. OffsetNewest int64 = sarama.OffsetNewest // OffsetOldest stands for the oldest offset available on the broker for a // partition. You can send this to a client's GetOffset method to get this // offset, or when calling ConsumePartition to start consuming from the // oldest offset that is still available on the broker. OffsetOldest int64 = sarama.OffsetOldest ) // Consumer is a wrapper on sarama ConsumerGroup. type Consumer struct { ConsumerGroup sarama.ConsumerGroup ConsumerGroupHandler *ConsumerHandler // initSessionRebalance is used to maintain the state of the consumer // in a consumer group to check if the partition rebalance has been // initiated on this consumer session. The Consume method must be called // only once in a go routine, as an infinite loop, which is ensured by // using this field. initSessionRebalance sync.Once errCh <-chan error } // Config provide values for kafka producer and consumer type Config struct { // Brokers comma separated kafka brokers Brokers string // SASL provide configs for authentication SASL SASLConfig // MaxRetry number of times to retry sending a failing message MaxRetry int // RetryFrequency backoff time in milliseconds before retrying RetryFrequency int // Topics multiple topics to subscribe messages from // first topic will be used for publishing the message Topics []string // GroupID consumer group id GroupID string Config *sarama.Config // ConnRetryDuration for specifying connection retry duration ConnRetryDuration int // Offsets is slice of TopicPartition in which "Topic","Partition" and "Offset" // are the field needed to be set to start consuming from specific offset Offsets []pubsub.TopicPartition InitialOffsets int64 // This config will allow application to disable kafka consumer auto commit DisableAutoCommit bool } type SASLConfig struct { // User username to connect to protected kafka instance User string // Password password to connect to protected kafka instance Password string // Mechanism SASL mechanism used for authentication Mechanism string // SecurityProtocol SSL or PLAINTEXT SecurityProtocol string // SSLVerify set it to true if certificate verification is required SSLVerify bool } /* New establishes connection to Kafka using the config provided in KafkaConfig */ func New(config *Config, logger log.Logger) (*Kafka, error) { _ = prometheus.Register(subscribeRecieveCount) _ = prometheus.Register(subscribeFailureCount) _ = prometheus.Register(subscribeSuccessCount) _ = prometheus.Register(publishFailureCount) _ = prometheus.Register(publishSuccessCount) _ = prometheus.Register(publishTotalCount) if config.SASL.Mechanism != SASLTypeSCRAMSHA512 && config.SASL.Mechanism != PLAIN && config.SASL.User != "" { return nil, errInvalidMechanism } populateOffsetTopic(config) convertKafkaConfig(config) sarama.Logger = kafkaLogger{logger: logger} brokers := strings.Split(config.Brokers, ",") p, err := sarama.NewSyncProducer(brokers, config.Config) if err != nil { return &Kafka{config: config}, err } c, err := NewKafkaConsumer(config) if err != nil { return &Kafka{config: config}, err } kafkaObj := &Kafka{config, p, c} return kafkaObj, nil } // populateOffsetTopic populate the Offsets Topic field with the topic provided in kafka config. func populateOffsetTopic(c *Config) { var topic string if c.Topics != nil { topic = c.Topics[0] } for i, v := range c.Offsets { if v.Topic == "" { c.Offsets[i].Topic = topic } } } func (k *Kafka) PublishEvent(key string, value interface{}, headers map[string]string) (err error) { return k.PublishEventWithOptions(key, value, headers, &pubsub.PublishOptions{ Topic: k.config.Topics[0], Timestamp: time.Now(), }) } func (k *Kafka) PublishEventWithOptions(key string, value interface{}, headers map[string]string, options *pubsub.PublishOptions) (err error) { if options == nil { options = &pubsub.PublishOptions{} } if options.Topic == "" { options.Topic = k.config.Topics[0] } if options.Timestamp.IsZero() { options.Timestamp = time.Now() } publishTotalCount.WithLabelValues(options.Topic, k.config.GroupID).Inc() valBytes, ok := value.([]byte) if !ok { valBytes, err = json.Marshal(value) if err != nil { publishFailureCount.WithLabelValues(options.Topic, k.config.GroupID).Inc() return err } } kafkaHeaders := make([]sarama.RecordHeader, 0, len(headers)) for key, value := range headers { kafkaHeaders = append(kafkaHeaders, sarama.RecordHeader{Key: []byte(key), Value: []byte(value)}) } message := &sarama.ProducerMessage{ Topic: options.Topic, Partition: int32(options.Partition), Value: sarama.ByteEncoder(valBytes), Key: sarama.StringEncoder(key), Timestamp: options.Timestamp, Headers: kafkaHeaders, } span := trace.SpanFromContext(context.Background()) producer := otelsarama.WrapSyncProducer(k.config.Config, k.Producer, otelsarama.WithTracerProvider(span.TracerProvider())) propagators := propagation.TraceContext{} propagators.Inject(context.Background(), otelsarama.NewProducerMessageCarrier(message)) _, _, err = producer.SendMessage(message) if err != nil { publishFailureCount.WithLabelValues(message.Topic, k.config.GroupID).Inc() return err } publishSuccessCount.WithLabelValues(message.Topic, k.config.GroupID).Inc() return nil } // rebalanceSession this method is responsible for rebalancing partitions // allocation whenever a new consumer is spawned, or an old one is closed. // It runs as an infinite for loop inside a goroutine. func (c *Consumer) rebalanceSession(ctx context.Context, config *Config) <-chan error { err := make(chan error) go func(consumerGroup sarama.ConsumerGroup, handler *ConsumerHandler) { for { propagators := propagation.TraceContext{} h := otelsarama.WrapConsumerGroupHandler(handler, otelsarama.WithPropagators(propagators)) if e := consumerGroup.Consume(ctx, config.Topics, h); e != nil { err <- e return } if e := ctx.Err(); e != nil { err <- e return } handler.ready = make(chan bool) } }(c.ConsumerGroup, c.ConsumerGroupHandler) return err } func (k *Kafka) subscribeMessage() (*pubsub.Message, error) { k.Consumer.initSessionRebalance.Do( func() { k.Consumer.errCh = k.Consumer.rebalanceSession(context.TODO(), k.config) }) select { case e := <-k.Consumer.errCh: return nil, e case <-k.Consumer.ConsumerGroupHandler.ready: } msg := <-k.Consumer.ConsumerGroupHandler.msg if msg == nil { return nil, ErrConsumeMsg } headers := make(map[string]string, len(msg.Headers)) for _, v := range msg.Headers { if len(v.Key) != 0 && len(v.Value) != 0 { headers[string(v.Key)] = string(v.Value) } } // Mark the message as read for autocommit to work k.Consumer.ConsumerGroupHandler.mu.Lock() k.Consumer.ConsumerGroupHandler.consumerGroupSession.MarkMessage(msg, "") k.Consumer.ConsumerGroupHandler.mu.Unlock() return &pubsub.Message{ Topic: msg.Topic, Partition: int(msg.Partition), Offset: msg.Offset, Key: string(msg.Key), Value: string(msg.Value), Headers: headers, }, nil } // Subscribe method is responsible for consuming a single message from a // sarama Consumer Group. When Subscribe is called first time we initiate // consumer group session rebalance, which handles the partition assignment // to multiple consumers in the group. func (k *Kafka) Subscribe() (*pubsub.Message, error) { message, err := k.subscribeMessage() if err != nil { // for unsuccessful subscribe topics := strings.Join(k.config.Topics, ",") subscribeFailureCount.WithLabelValues(topics, k.config.GroupID).Inc() subscribeRecieveCount.WithLabelValues(topics, k.config.GroupID).Inc() return nil, err } k.CommitOffset(pubsub.TopicPartition{ Topic: message.Topic, Partition: message.Partition, Offset: message.Offset, }) // for successful subscribe subscribeSuccessCount.WithLabelValues(message.Topic, k.config.GroupID).Inc() subscribeRecieveCount.WithLabelValues(message.Topic, k.config.GroupID).Inc() return message, nil } func (k *Kafka) Pause() error { k.Consumer.ConsumerGroup.PauseAll() return nil } func (k *Kafka) Resume() error { k.Consumer.ConsumerGroup.ResumeAll() return nil } type brokersErr struct{} func (b brokersErr) Error() string { return "invalid brokers connection failed" } // ConsumerHandler represents a Sarama consumer group message handler. // It is responsible for handling the messages in a specific topic/partition. // The handler implements methods called during the lifecycle of a consumer // group session. type ConsumerHandler struct { msg chan *sarama.ConsumerMessage ready chan bool consumerGroupSession sarama.ConsumerGroupSession offsets []pubsub.TopicPartition initialiseOffset sync.Once mu sync.Mutex } // setOffset set the starting offset of all the partition in the topic. func setOffset(s sarama.ConsumerGroupSession, consumer *ConsumerHandler) { for _, v := range consumer.offsets { p := int32(v.Partition) // To start subscribing uncommitted message from specific offset s.MarkOffset(v.Topic, p, v.Offset, "") // To start subscribing committed message from specific offset s.ResetOffset(v.Topic, p, v.Offset, "") } } // Setup is run at the beginning of a new session, before ConsumeClaim func (consumer *ConsumerHandler) Setup(s sarama.ConsumerGroupSession) error { // Set the initial custom offset value for subscribe if consumer.offsets != nil { consumer.initialiseOffset.Do( func() { setOffset(s, consumer) }) } // Mark the consumer as ready close(consumer.ready) return nil } // Cleanup is run at the end of a session, once all ConsumeClaim goroutines have exited func (consumer *ConsumerHandler) Cleanup(sarama.ConsumerGroupSession) error { return nil } // kafkaLogger is used only for debugging purposes // step-by-step logs are provided by sarama Kafka for connection, // publishing and subscribing messages type kafkaLogger struct { logger log.Logger }