cloud-services/pkg/kafka/async_producer.go

package kafka

import (
	"context"
	"encoding/json"

	"fiskerinc.com/modules/logger"

	"github.com/confluentinc/confluent-kafka-go/v2/kafka"
	"github.com/pkg/errors"
)

type AsyncProducer struct {
	producer *kafka.Producer
}

// NewProducer serves as factory method for producer to kafka
func NewAsyncProducer(ctx context.Context) (ProducerInterface, error) {
	var producer *kafka.Producer
	configuration := loadKafkaProducerConfig()
	producer, err := kafka.NewProducer(
		&configuration,
		// kafkaTrace.WithContext(ctx),
	)
	logger.Info().Msgf("NewAsyncProducer  hosts: %s", kafkaHosts)
	if err != nil {
		return nil, errors.WithStack(err)
	}
	return &AsyncProducer{producer: producer}, nil
}

// SetProducer sets the producer instance
func (p *AsyncProducer) SetProducer(producer *kafka.Producer) {
	p.producer = producer
}

// Len returns len of messages in queue.
func (p *AsyncProducer) Len() int {
	return p.producer.Len()
}

// Flush calls producer's Flush function.
func (p *AsyncProducer) Flush(timeoutMs int) int {
	return p.producer.Flush(timeoutMs)
}

// Produce sends a Kafka Message to Kafka
func (p *AsyncProducer) Produce(topic string, key string, payload interface{}, headers map[string][]byte) error {
	v, err := json.Marshal(payload)
	if err != nil {
		return errors.WithStack(err)
	}

	return p.ProduceBinary(topic, key, v, headers)
}

func (p *AsyncProducer) makeHeaders(headers map[string][]byte) []kafka.Header {
	if headers == nil {
		return nil
	}

	i := 0
	result := make([]kafka.Header, len(headers))

	for key, value := range headers {
		result[i] = kafka.Header{
			Key: key,
		}
		copy(result[i].Value, value)
		i++
	}

	return result
}

func (p *AsyncProducer) ProduceBinary(topic string, key string, data []byte, headers map[string][]byte) error {
	km := kafka.Message{
		TopicPartition: kafka.TopicPartition{Topic: &topic, Partition: kafka.PartitionAny},
		Key:            []byte(key),
		Value:          data,
	}
	if headers != nil {
		km.Headers = p.makeHeaders(headers)
	}

	err := p.producer.Produce(&km, nil)
	if err != nil {
		// error handle inability to connect to kafka
		return errors.WithStack(err)
	}

	return err
}

// ProduceToChannel sends a list of Kafka Messages to Kafka
func (p *AsyncProducer) ProduceToChannel(topic string, key string, payload interface{}) error {
	v, err := json.Marshal(payload)
	if err != nil {
		return errors.WithStack(err)
	}

	km := kafka.Message{
		TopicPartition: kafka.TopicPartition{Topic: &topic, Partition: kafka.PartitionAny},
		Key:            []byte(key),
		Value:          v,
	}

	p.producer.ProduceChannel() <- &km
	return nil
}

func (p *AsyncProducer) ProduceBinaryToChannel(topic string, key string, payload []byte) error {
	km := kafka.Message{
		TopicPartition: kafka.TopicPartition{Topic: &topic, Partition: kafka.PartitionAny},
		Key:            []byte(key),
		Value:          payload,
	}

	p.producer.ProduceChannel() <- &km
	return nil
}

// ProduceSignalBatch sends a specified batch of CAN signals
//
//	NOTE: does NOT produce proper JSON, Value is custom for Clickhouse
func (p *AsyncProducer) ProduceSignalBatch(topic string, key string, payload interface{}) error {
	v, err := json.Marshal(payload)
	if err != nil {
		return errors.WithStack(err)
	}

	km := kafka.Message{
		TopicPartition: kafka.TopicPartition{Topic: &topic, Partition: kafka.PartitionAny},
		Key:            []byte(key),
		Value:          v[1 : len(v)-1],
	}

	err = p.producer.Produce(&km, nil)
	if err != nil {
		// error handle inability to connect to kafka
		return errors.WithStack(err)
	}

	return nil
}

// ReadEvents listens to producer's events which are emitted as responses to producing events
func (p *AsyncProducer) ReadEvents() {
	for e := range p.producer.Events() {
		switch m := e.(type) {
		case *kafka.Message:
			if m.TopicPartition.Error != nil {
				logger.Error().
					Err(m.TopicPartition.Error).
					Send()
			}
		case kafka.Error:
			logger.Error().
				Str("err", m.String()).
				Send()
		}
	}
}

// Close the Kafka Producer
func (p *AsyncProducer) Close() {
	p.producer.Flush(0)
	p.producer.Close()
	p.producer = nil
}