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 }