cloud-services/pkg/kafka/consumer.go

package kafka

import (
	"context"
	"sync"

	"fiskerinc.com/modules/common"
	"fiskerinc.com/modules/logger"
	"fiskerinc.com/modules/utils/envtool"

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

// ConsumerInterface interface for Consumer utility
type ConsumerInterface interface {
	Consume(topics []string, handler func([]byte, []byte) error) error
	ConsumeToChannel(topics []string, events chan *kafka.Message) error
	ConsumeToChannelJson(topics []string, events chan common.EventRawJSON) error
	ConsumePartitionsToChannel(partitions []kafka.TopicPartition, events chan *kafka.Message) error
	GetMetadata(topic string) (*kafka.Metadata, error)
	Check(ctx context.Context) error
	Stop()
}

// Consumer utility to produce messages
type Consumer struct {
	consumer      *kafka.Consumer
	running       bool
	timeout       int
	connected     bool
	connectedLock sync.RWMutex
}

// NewConsumer serves as factory method for consumer to kafka
func NewConsumer(serviceName string) (ConsumerInterface, error) {
	var consumer *kafka.Consumer
	config := loadKafkaConsumerConfig(serviceName)

	consumer, err := kafka.NewConsumer(&config)
	logger.Info().Msgf("NewConsumer  hosts: %s", envtool.GetEnv("KAFKA_HOSTS", "localhost:9093"))
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return &Consumer{
		consumer: consumer,
		timeout:  envtool.GetEnvInt("KAFKA_TIMEOUT", -1),
	}, nil
}

// SetConsumer sets the consumer
func (c *Consumer) SetConsumer(consumer *kafka.Consumer) {
	c.consumer = consumer
}

// Consume runs a poll loop which waits for messages to consume
func (c *Consumer) Consume(topics []string, handler func([]byte, []byte) error) error {
	if err := c.consumer.SubscribeTopics(topics, nil); err != nil {
		return errors.WithStack(err)
	}

	c.setConnected(true)
	defer c.setConnected(false)

	c.running = true
	connected := true
	for c.running {
		e := c.consumer.Poll(c.timeout)
		switch msg := e.(type) {
		case *kafka.Message:
			if !connected {
				connected = true
				c.setConnected(true)
			}

			if err := handler(msg.Key, msg.Value); err != nil {
				logger.Warn().Err(err).Send()
			}
		case kafka.Error:
			if connected {
				connected = false
				c.setConnected(false)
			}

			logger.Error().Msg(e.String())
		}
	}

	return nil
}

// ConsumeToChannel runs a poll loop which waits for messages to consume
func (c *Consumer) ConsumeToChannel(topics []string, events chan *kafka.Message) error {
	err := c.consumer.SubscribeTopics(topics, nil)
	if err != nil {
		return errors.WithStack(err)
	}

	c.setConnected(true)
	defer c.setConnected(false)

	c.running = true
	connected := true
	for c.running {
		e := c.consumer.Poll(c.timeout)
		switch msg := e.(type) {
		case *kafka.Message:
			if !connected {
				connected = true
				c.setConnected(true)
			}

			events <- msg
		case kafka.Error:
			if connected {
				connected = false
				c.setConnected(false)
			}

			logger.Error().Msg(e.String())
		}
	}
	return nil
}

func (c *Consumer) ConsumeToChannelJson(topics []string, events chan common.EventRawJSON) error {
	err := c.consumer.SubscribeTopics(topics, nil)
	if err != nil {
		return errors.WithStack(err)
	}

	c.setConnected(true)
	defer c.setConnected(false)

	connected := true
	c.running = true
	for c.running {
		e := c.consumer.Poll(c.timeout)
		switch msg := e.(type) {
		case *kafka.Message:
			if !connected {
				connected = true
				c.setConnected(true)
			}

			events <- common.EventRawJSON{
				Topic:   *msg.TopicPartition.Topic,
				Key:     string(msg.Key),
				Payload: msg.Value,
			}
		case kafka.Error:
			if connected {
				connected = false
				c.setConnected(false)
			}

			logger.Error().Err(errors.New(e.String())).Send()
		}
	}

	return nil
}

// ConsumePartitionsToChannel runs a poll loop which waits for messages to consume
func (c *Consumer) ConsumePartitionsToChannel(partitions []kafka.TopicPartition, events chan *kafka.Message) error {
	if len(partitions) == 0 {
		return errors.WithStack(errors.New("no partitions provided"))
	}
	topic := "attendant_service"
	m, err := c.consumer.GetMetadata(&topic, false, 200)
	_ = m
	err = c.consumer.Assign(partitions)
	if err != nil {
		return errors.WithStack(err)
	}

	c.setConnected(true)
	defer c.setConnected(false)

	c.running = true
	connected := true
	for c.running {
		e := <-c.consumer.Events()
		if e != nil {
			logger.Error().Msg(e.String())
		}
		switch msg := e.(type) {
		case *kafka.Message:
			if !connected {
				connected = true
				c.setConnected(true)
			}

			events <- msg
		case kafka.Error:
			if connected {
				connected = false
				c.setConnected(false)
			}

			logger.Error().Msg(e.String())
		}
	}
	return nil
}

func (c *Consumer) GetMetadata(topic string) (*kafka.Metadata, error) {
	return c.consumer.GetMetadata(&topic, false, 200)
}

// Stop stops the poll loop running
func (c *Consumer) Stop() {
	if c.consumer != nil {
		c.running = false
		if err := c.consumer.Close(); err != nil {
			logger.Warn().Err(err).Send()
		}
	}
}

func (c *Consumer) Check(ctx context.Context) error {
	if !c.isConnected() {
		return errors.New("consumer isn't connected to Kafka")
	}

	return nil
}

func (c *Consumer) setConnected(cntd bool) {
	c.connectedLock.Lock()
	defer c.connectedLock.Unlock()

	c.connected = cntd
}

func (c *Consumer) isConnected() bool {
	c.connectedLock.RLock()
	defer c.connectedLock.RUnlock()

	return c.connected
}