func expectationConsumer(c sarama.Consumer, expectations <-chan *sarama.ProducerMessage, wg *sync.WaitGroup) {
defer wg.Done()
defer func() {
if err := c.Close(); err != nil {
logger.Println("Failed to close consumer:", err)
}
}()
var (
partitionVerifiers = make(map[int32]*partitionVerifier)
consumerWg sync.WaitGroup
)
for expectation := range expectations {
partition := expectation.Partition
if partitionVerifiers[partition] == nil {
logger.Printf("Starting message verifier for partition %d...\n", partition)
pc, err := c.ConsumePartition(*topic, partition, expectation.Offset)
if err != nil {
logger.Fatalf("Failed to open partition consumer for %s/%d: %s", *topic, expectation.Partition, err)
}
partitionExpectations := make(chan *sarama.ProducerMessage)
partitionVerifiers[partition] = &partitionVerifier{pc: pc, expectations: partitionExpectations}
consumerWg.Add(1)
go partitionExpectationConsumer(pc, partitionExpectations, &consumerWg)
}
partitionVerifiers[partition].expectations <- expectation
}
for _, pv := range partitionVerifiers {
close(pv.expectations)
}
consumerWg.Wait()
}
// GetKafkaPartitions is a helper function to look up which partitions are available
// via the given brokers for the given topic. This should be called only on startup.
func GetKafkaPartitions(brokerHosts []string, topic string) (partitions []int32, err error) {
if len(brokerHosts) == 0 {
return partitions, errors.New("at least 1 broker host is required")
}
if len(topic) == 0 {
return partitions, errors.New("topic name is required")
}
var cnsmr sarama.Consumer
cnsmr, err = sarama.NewConsumer(brokerHosts, sarama.NewConfig())
if err != nil {
return partitions, err
}
defer func() {
if cerr := cnsmr.Close(); cerr != nil && err == nil {
err = cerr
}
}()
return cnsmr.Partitions(topic)
}
// Connects to a consumer group, using Zookeeper for auto-discovery
func JoinConsumerGroup(name string, topics []string, zookeeper []string, config *Config) (cg *ConsumerGroup, err error) {
if name == "" {
return nil, sarama.ConfigurationError("Empty consumergroup name")
}
if len(topics) == 0 {
return nil, sarama.ConfigurationError("No topics provided")
}
if len(zookeeper) == 0 {
return nil, errors.New("You need to provide at least one zookeeper node address!")
}
if config == nil {
config = NewConfig()
}
config.ClientID = name
// Validate configuration
if err = config.Validate(); err != nil {
return
}
var kz *kazoo.Kazoo
if kz, err = kazoo.NewKazoo(zookeeper, config.Zookeeper); err != nil {
return
}
brokers, err := kz.BrokerList()
if err != nil {
kz.Close()
return
}
group := kz.Consumergroup(name)
if config.Offsets.ResetOffsets {
err = group.ResetOffsets()
if err != nil {
cg.Logf("FAILED to reset offsets of consumergroup: %s!\n", err)
kz.Close()
return
}
}
instance := group.NewInstance()
var consumer sarama.Consumer
if consumer, err = sarama.NewConsumer(brokers, config.Config); err != nil {
kz.Close()
return
}
cg = &ConsumerGroup{
config: config,
consumer: consumer,
kazoo: kz,
group: group,
instance: instance,
messages: make(chan *sarama.ConsumerMessage, config.ChannelBufferSize),
errors: make(chan *sarama.ConsumerError, config.ChannelBufferSize),
stopper: make(chan struct{}),
}
// Register consumer group
if exists, err := cg.group.Exists(); err != nil {
cg.Logf("FAILED to check for existence of consumergroup: %s!\n", err)
_ = consumer.Close()
_ = kz.Close()
return nil, err
} else if !exists {
cg.Logf("Consumergroup `%s` does not yet exists, creating...\n", cg.group.Name)
if err := cg.group.Create(); err != nil {
cg.Logf("FAILED to create consumergroup in Zookeeper: %s!\n", err)
_ = consumer.Close()
_ = kz.Close()
return nil, err
}
}
// Register itself with zookeeper
if err := cg.instance.Register(topics); err != nil {
cg.Logf("FAILED to register consumer instance: %s!\n", err)
return nil, err
} else {
cg.Logf("Consumer instance registered (%s).", cg.instance.ID)
}
offsetConfig := OffsetManagerConfig{CommitInterval: config.Offsets.CommitInterval}
cg.offsetManager = NewZookeeperOffsetManager(cg, &offsetConfig)
go cg.topicListConsumer(topics)
return
}
func (cg *ConsumerGroup) Load(logger zap.Logger) error {
var kz *kazoo.Kazoo
var err error
if kz, err = kazoo.NewKazoo(cg.zookeeper, cg.config.Zookeeper); err != nil {
return err
}
logger.Info("KAFKA: Getting broker list for replica",
zap.Int("replicaId", cg.replicaId),
)
brokers, err := kz.BrokerList()
if err != nil {
kz.Close()
return err
}
group := kz.Consumergroup(cg.config.ClientID)
instance := group.NewInstance()
var consumer sarama.Consumer
if consumer, err = sarama.NewConsumer(brokers, cg.config.Config); err != nil {
kz.Close()
return err
}
cg.kazoo = kz
cg.group = group
cg.instance = instance
cg.messages = make(chan *sarama.ConsumerMessage, cg.config.ChannelBufferSize)
cg.consumer = consumer
cg.singleShutdown = sync.Once{}
cg.errors = make(chan *sarama.ConsumerError, cg.config.ChannelBufferSize)
cg.stopper = make(chan struct{})
if exists, err := cg.group.Exists(); err != nil {
logger.Fatal("KAFKA: Replica failed to check existence of consumergroup",
zap.Int("replicaId", cg.replicaId),
zap.Error(err),
)
consumer.Close()
kz.Close()
return err
} else if !exists {
logger.Info("KAFKA: Consumergroup does not exist, creating it",
zap.Int("replicaId", cg.replicaId),
zap.String("consumerGroupName", cg.group.Name),
)
if err := cg.group.Create(); err != nil {
logger.Fatal("KAFKA: Failed to create consumergroup in Zookeeper",
zap.Int("replicaId", cg.replicaId),
zap.Error(err),
)
consumer.Close()
kz.Close()
return err
}
}
if err := cg.instance.Register(cg.topics); err != nil {
logger.Fatal("KAFKA: Failed to create consumer instance",
zap.Int("replicaId", cg.replicaId),
zap.Error(err),
)
return err
} else {
logger.Info("KAFKA: Consumer instance registered",
zap.Int("replicaId", cg.replicaId),
)
}
offsetConfig := OffsetManagerConfig{
CommitInterval: cg.config.Offsets.CommitInterval,
EnableAutoCommit: cg.config.EnableOffsetAutoCommit,
}
cg.offsetManager = NewZookeeperOffsetManager(cg, &offsetConfig, logger)
go cg.topicListConsumer(cg.topics, logger)
return nil
}
