cloud-services/pkg/kafka/producer.go

package kafka

import (
	"context"
	"encoding/json"

	"fiskerinc.com/modules/logger"

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

// ProducerInterface interface for Producer utilty
type ProducerInterface interface {
	Produce(topic string, key string, payload interface{}, headers map[string][]byte) error
	ProduceBinary(topic string, key string, data []byte, headers map[string][]byte) error
	ProduceBinaryToChannel(topic string, key string, payload []byte) error
	ProduceToChannel(string, string, interface{}) error
	ProduceSignalBatch(string, string, interface{}) error
	ReadEvents()
	Len() int
	Flush(timeoutMs int) int
	Close()
}

// Producer utility to produce messages
type Producer struct {
	producer *kafka.Producer
}

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

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

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

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

// Produce sends a Kafka Message to Kafka
func (p *Producer) 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 *Producer) 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].Key = key
		result[i].Value = value
		i++
	}

	return result
}

// Produce sends a Kafka Message to Kafka
func (p *Producer) 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)
	}
	deliveryChan := make(chan kafka.Event)

	err := p.producer.Produce(&km, deliveryChan)
	if err != nil {
		return errors.WithStack(err)
	}
	for e := range deliveryChan {
		switch ev := e.(type) {
		case *kafka.Message:
			// The message delivery report, indicating success or
			// permanent failure after retries have been exhausted.
			// Application level retries won't help since the client
			// is already configured to do that.
			m := ev
			if m.TopicPartition.Error != nil {
				logger.Error().Msgf("Delivery failed: %v\n", m.TopicPartition.Error)
				return m.TopicPartition.Error
			} else {
				logger.Info().Msgf("Delivered message to topic %s [%d] at offset %v\n",
					*m.TopicPartition.Topic, m.TopicPartition.Partition, m.TopicPartition.Offset)
				return nil
			}
		default:
			logger.Debug().Msgf("Ignored event: %s\n", ev)
		}
	}
	return nil
}

// ProduceBatch is a stub for async producer
func (p *Producer) ProduceBatch(topic string, key string, payload [][]byte) error {
	// stub
	return nil
}

// ProduceToChannel is a stub for async producer
func (p *Producer) 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
}

// ProduceBinaryToChannel producing binary to a channel.
func (p *Producer) 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
}

// Produce sends a Kafka Message to Kafka
func (p *Producer) ProduceSignalBatch(topic string, key string, payload interface{}) error {
	// stub
	return nil
}

// ReadEvents is a stub for async producer
func (p *Producer) ReadEvents() {
	for e := range p.producer.Events() {
		switch ev := e.(type) {
		case kafka.Error:
			// Generic client instance-level errors, such as
			// broker connection failures, authentication issues, etc.
			//
			// These errors should generally be considered informational
			// as the underlying client will automatically try to
			// recover from any errors encountered, the application
			// does not need to take action on them.
			logger.Error().Msgf("Error: %v\n", ev)
		default:
			logger.Debug().Msgf("Ignored event: %s\n", ev)
		}
	}
}

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