Initial commit

This commit is contained in:
Donny
2019-04-22 20:46:32 +08:00
commit 49ab8aadd1
25441 changed files with 4055000 additions and 0 deletions

View File

@@ -0,0 +1,31 @@
load("@io_bazel_rules_go//go:def.bzl", "go_library")
go_library(
name = "go_default_library",
srcs = [
"consumer_group.go",
"offset_manager.go",
"utils.go",
],
importmap = "go-common/vendor/github.com/wvanbergen/kafka/consumergroup",
importpath = "github.com/wvanbergen/kafka/consumergroup",
visibility = ["//visibility:public"],
deps = [
"//vendor/github.com/Shopify/sarama:go_default_library",
"//vendor/github.com/wvanbergen/kazoo-go:go_default_library",
],
)
filegroup(
name = "package-srcs",
srcs = glob(["**"]),
tags = ["automanaged"],
visibility = ["//visibility:private"],
)
filegroup(
name = "all-srcs",
srcs = [":package-srcs"],
tags = ["automanaged"],
visibility = ["//visibility:public"],
)

View File

@@ -0,0 +1,514 @@
package consumergroup
import (
"errors"
"fmt"
"sync"
"time"
"github.com/Shopify/sarama"
"github.com/wvanbergen/kazoo-go"
)
var (
AlreadyClosing = errors.New("The consumer group is already shutting down.")
)
type Config struct {
*sarama.Config
Zookeeper *kazoo.Config
Offsets struct {
Initial int64 // The initial offset method to use if the consumer has no previously stored offset. Must be either sarama.OffsetOldest (default) or sarama.OffsetNewest.
ProcessingTimeout time.Duration // Time to wait for all the offsets for a partition to be processed after stopping to consume from it. Defaults to 1 minute.
CommitInterval time.Duration // The interval between which the processed offsets are commited.
ResetOffsets bool // Resets the offsets for the consumergroup so that it won't resume from where it left off previously.
}
}
func NewConfig() *Config {
config := &Config{}
config.Config = sarama.NewConfig()
config.Zookeeper = kazoo.NewConfig()
config.Offsets.Initial = sarama.OffsetOldest
config.Offsets.ProcessingTimeout = 60 * time.Second
config.Offsets.CommitInterval = 10 * time.Second
return config
}
func (cgc *Config) Validate() error {
if cgc.Zookeeper.Timeout <= 0 {
return sarama.ConfigurationError("ZookeeperTimeout should have a duration > 0")
}
if cgc.Offsets.CommitInterval < 0 {
return sarama.ConfigurationError("CommitInterval should have a duration >= 0")
}
if cgc.Offsets.Initial != sarama.OffsetOldest && cgc.Offsets.Initial != sarama.OffsetNewest {
return errors.New("Offsets.Initial should be sarama.OffsetOldest or sarama.OffsetNewest.")
}
if cgc.Config != nil {
if err := cgc.Config.Validate(); err != nil {
return err
}
}
return nil
}
// The ConsumerGroup type holds all the information for a consumer that is part
// of a consumer group. Call JoinConsumerGroup to start a consumer.
type ConsumerGroup struct {
config *Config
consumer sarama.Consumer
kazoo *kazoo.Kazoo
group *kazoo.Consumergroup
instance *kazoo.ConsumergroupInstance
wg sync.WaitGroup
singleShutdown sync.Once
messages chan *sarama.ConsumerMessage
errors chan error
stopper chan struct{}
consumers kazoo.ConsumergroupInstanceList
offsetManager OffsetManager
}
// 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 {
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 error, 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
}
// Returns a channel that you can read to obtain events from Kafka to process.
func (cg *ConsumerGroup) Messages() <-chan *sarama.ConsumerMessage {
return cg.messages
}
// Returns a channel that you can read to obtain events from Kafka to process.
func (cg *ConsumerGroup) Errors() <-chan error {
return cg.errors
}
func (cg *ConsumerGroup) Closed() bool {
return cg.instance == nil
}
func (cg *ConsumerGroup) Close() error {
shutdownError := AlreadyClosing
cg.singleShutdown.Do(func() {
defer cg.kazoo.Close()
shutdownError = nil
close(cg.stopper)
cg.wg.Wait()
if err := cg.offsetManager.Close(); err != nil {
cg.Logf("FAILED closing the offset manager: %s!\n", err)
}
if shutdownError = cg.instance.Deregister(); shutdownError != nil {
cg.Logf("FAILED deregistering consumer instance: %s!\n", shutdownError)
} else {
cg.Logf("Deregistered consumer instance %s.\n", cg.instance.ID)
}
if shutdownError = cg.consumer.Close(); shutdownError != nil {
cg.Logf("FAILED closing the Sarama client: %s\n", shutdownError)
}
close(cg.messages)
close(cg.errors)
cg.instance = nil
})
return shutdownError
}
func (cg *ConsumerGroup) Logf(format string, args ...interface{}) {
var identifier string
if cg.instance == nil {
identifier = "(defunct)"
} else {
identifier = cg.instance.ID[len(cg.instance.ID)-12:]
}
sarama.Logger.Printf("[%s/%s] %s", cg.group.Name, identifier, fmt.Sprintf(format, args...))
}
func (cg *ConsumerGroup) InstanceRegistered() (bool, error) {
return cg.instance.Registered()
}
func (cg *ConsumerGroup) CommitUpto(message *sarama.ConsumerMessage) error {
cg.offsetManager.MarkAsProcessed(message.Topic, message.Partition, message.Offset)
return nil
}
func (cg *ConsumerGroup) FlushOffsets() error {
return cg.offsetManager.Flush()
}
func (cg *ConsumerGroup) topicListConsumer(topics []string) {
for {
select {
case <-cg.stopper:
return
default:
}
consumers, consumerChanges, err := cg.group.WatchInstances()
if err != nil {
cg.Logf("FAILED to get list of registered consumer instances: %s\n", err)
return
}
cg.consumers = consumers
cg.Logf("Currently registered consumers: %d\n", len(cg.consumers))
stopper := make(chan struct{})
for _, topic := range topics {
cg.wg.Add(1)
go cg.topicConsumer(topic, cg.messages, cg.errors, stopper)
}
select {
case <-cg.stopper:
close(stopper)
return
case <-consumerChanges:
registered, err := cg.instance.Registered()
if err != nil {
cg.Logf("FAILED to get register status: %s\n", err)
} else if !registered {
err = cg.instance.Register(topics)
if err != nil {
cg.Logf("FAILED to register consumer instance: %s!\n", err)
} else {
cg.Logf("Consumer instance registered (%s).", cg.instance.ID)
}
}
cg.Logf("Triggering rebalance due to consumer list change\n")
close(stopper)
cg.wg.Wait()
}
}
}
func (cg *ConsumerGroup) topicConsumer(topic string, messages chan<- *sarama.ConsumerMessage, errors chan<- error, stopper <-chan struct{}) {
defer cg.wg.Done()
select {
case <-stopper:
return
default:
}
cg.Logf("%s :: Started topic consumer\n", topic)
// Fetch a list of partition IDs
partitions, err := cg.kazoo.Topic(topic).Partitions()
if err != nil {
cg.Logf("%s :: FAILED to get list of partitions: %s\n", topic, err)
cg.errors <- &sarama.ConsumerError{
Topic: topic,
Partition: -1,
Err: err,
}
return
}
partitionLeaders, err := retrievePartitionLeaders(partitions)
if err != nil {
cg.Logf("%s :: FAILED to get leaders of partitions: %s\n", topic, err)
cg.errors <- &sarama.ConsumerError{
Topic: topic,
Partition: -1,
Err: err,
}
return
}
dividedPartitions := dividePartitionsBetweenConsumers(cg.consumers, partitionLeaders)
myPartitions := dividedPartitions[cg.instance.ID]
cg.Logf("%s :: Claiming %d of %d partitions", topic, len(myPartitions), len(partitionLeaders))
// Consume all the assigned partitions
var wg sync.WaitGroup
for _, pid := range myPartitions {
wg.Add(1)
go cg.partitionConsumer(topic, pid.ID, messages, errors, &wg, stopper)
}
wg.Wait()
cg.Logf("%s :: Stopped topic consumer\n", topic)
}
func (cg *ConsumerGroup) consumePartition(topic string, partition int32, nextOffset int64) (sarama.PartitionConsumer, error) {
consumer, err := cg.consumer.ConsumePartition(topic, partition, nextOffset)
if err == sarama.ErrOffsetOutOfRange {
cg.Logf("%s/%d :: Partition consumer offset out of Range.\n", topic, partition)
// if the offset is out of range, simplistically decide whether to use OffsetNewest or OffsetOldest
// if the configuration specified offsetOldest, then switch to the oldest available offset, else
// switch to the newest available offset.
if cg.config.Offsets.Initial == sarama.OffsetOldest {
nextOffset = sarama.OffsetOldest
cg.Logf("%s/%d :: Partition consumer offset reset to oldest available offset.\n", topic, partition)
} else {
nextOffset = sarama.OffsetNewest
cg.Logf("%s/%d :: Partition consumer offset reset to newest available offset.\n", topic, partition)
}
// retry the consumePartition with the adjusted offset
consumer, err = cg.consumer.ConsumePartition(topic, partition, nextOffset)
}
if err != nil {
cg.Logf("%s/%d :: FAILED to start partition consumer: %s\n", topic, partition, err)
return nil, err
}
return consumer, err
}
// Consumes a partition
func (cg *ConsumerGroup) partitionConsumer(topic string, partition int32, messages chan<- *sarama.ConsumerMessage, errors chan<- error, wg *sync.WaitGroup, stopper <-chan struct{}) {
defer wg.Done()
select {
case <-stopper:
return
default:
}
// Since ProcessingTimeout is the amount of time we'll wait for the final batch
// of messages to be processed before releasing a partition, we need to wait slightly
// longer than that before timing out here to ensure that another consumer has had
// enough time to release the partition. Hence, +2 seconds.
maxRetries := int(cg.config.Offsets.ProcessingTimeout/time.Second) + 2
for tries := 0; tries < maxRetries; tries++ {
if err := cg.instance.ClaimPartition(topic, partition); err == nil {
break
} else if tries+1 < maxRetries {
if err == kazoo.ErrPartitionClaimedByOther {
// Another consumer still owns this partition. We should wait longer for it to release it.
time.Sleep(1 * time.Second)
} else {
// An unexpected error occurred. Log it and continue trying until we hit the timeout.
cg.Logf("%s/%d :: FAILED to claim partition on attempt %v of %v; retrying in 1 second. Error: %v", topic, partition, tries+1, maxRetries, err)
time.Sleep(1 * time.Second)
}
} else {
cg.Logf("%s/%d :: FAILED to claim the partition: %s\n", topic, partition, err)
cg.errors <- &sarama.ConsumerError{
Topic: topic,
Partition: partition,
Err: err,
}
return
}
}
defer func() {
err := cg.instance.ReleasePartition(topic, partition)
if err != nil {
cg.Logf("%s/%d :: FAILED to release partition: %s\n", topic, partition, err)
cg.errors <- &sarama.ConsumerError{
Topic: topic,
Partition: partition,
Err: err,
}
}
}()
nextOffset, err := cg.offsetManager.InitializePartition(topic, partition)
if err != nil {
cg.Logf("%s/%d :: FAILED to determine initial offset: %s\n", topic, partition, err)
return
}
if nextOffset >= 0 {
cg.Logf("%s/%d :: Partition consumer starting at offset %d.\n", topic, partition, nextOffset)
} else {
nextOffset = cg.config.Offsets.Initial
if nextOffset == sarama.OffsetOldest {
cg.Logf("%s/%d :: Partition consumer starting at the oldest available offset.\n", topic, partition)
} else if nextOffset == sarama.OffsetNewest {
cg.Logf("%s/%d :: Partition consumer listening for new messages only.\n", topic, partition)
}
}
consumer, err := cg.consumePartition(topic, partition, nextOffset)
if err != nil {
cg.Logf("%s/%d :: FAILED to start partition consumer: %s\n", topic, partition, err)
return
}
defer consumer.Close()
err = nil
var lastOffset int64 = -1 // aka unknown
partitionConsumerLoop:
for {
select {
case <-stopper:
break partitionConsumerLoop
case err := <-consumer.Errors():
if err == nil {
cg.Logf("%s/%d :: Consumer encountered an invalid state: re-establishing consumption of partition.\n", topic, partition)
// Errors encountered (if any) are logged in the consumerPartition function
var cErr error
consumer, cErr = cg.consumePartition(topic, partition, lastOffset)
if cErr != nil {
break partitionConsumerLoop
}
continue partitionConsumerLoop
}
for {
select {
case errors <- err:
continue partitionConsumerLoop
case <-stopper:
break partitionConsumerLoop
}
}
case message := <-consumer.Messages():
if message == nil {
cg.Logf("%s/%d :: Consumer encountered an invalid state: re-establishing consumption of partition.\n", topic, partition)
// Errors encountered (if any) are logged in the consumerPartition function
var cErr error
consumer, cErr = cg.consumePartition(topic, partition, lastOffset)
if cErr != nil {
break partitionConsumerLoop
}
continue partitionConsumerLoop
}
for {
select {
case <-stopper:
break partitionConsumerLoop
case messages <- message:
lastOffset = message.Offset
continue partitionConsumerLoop
}
}
}
}
cg.Logf("%s/%d :: Stopping partition consumer at offset %d\n", topic, partition, lastOffset)
if err := cg.offsetManager.FinalizePartition(topic, partition, lastOffset, cg.config.Offsets.ProcessingTimeout); err != nil {
cg.Logf("%s/%d :: %s\n", topic, partition, err)
}
}

View File

@@ -0,0 +1,298 @@
package consumergroup
import (
"errors"
"fmt"
"sync"
"time"
)
// OffsetManager is the main interface consumergroup requires to manage offsets of the consumergroup.
type OffsetManager interface {
// InitializePartition is called when the consumergroup is starting to consume a
// partition. It should return the last processed offset for this partition. Note:
// the same partition can be initialized multiple times during a single run of a
// consumer group due to other consumer instances coming online and offline.
InitializePartition(topic string, partition int32) (int64, error)
// MarkAsProcessed tells the offset manager than a certain message has been successfully
// processed by the consumer, and should be committed. The implementation does not have
// to store this offset right away, but should return true if it intends to do this at
// some point.
//
// Offsets should generally be increasing if the consumer
// processes events serially, but this cannot be guaranteed if the consumer does any
// asynchronous processing. This can be handled in various ways, e.g. by only accepting
// offsets that are higehr than the offsets seen before for the same partition.
MarkAsProcessed(topic string, partition int32, offset int64) bool
// Flush tells the offset manager to immediately commit offsets synchronously and to
// return any errors that may have occured during the process.
Flush() error
// FinalizePartition is called when the consumergroup is done consuming a
// partition. In this method, the offset manager can flush any remaining offsets to its
// backend store. It should return an error if it was not able to commit the offset.
// Note: it's possible that the consumergroup instance will start to consume the same
// partition again after this function is called.
FinalizePartition(topic string, partition int32, lastOffset int64, timeout time.Duration) error
// Close is called when the consumergroup is shutting down. In normal circumstances, all
// offsets are committed because FinalizePartition is called for all the running partition
// consumers. You may want to check for this to be true, and try to commit any outstanding
// offsets. If this doesn't succeed, it should return an error.
Close() error
}
var (
UncleanClose = errors.New("Not all offsets were committed before shutdown was completed")
)
// OffsetManagerConfig holds configuration setting son how the offset manager should behave.
type OffsetManagerConfig struct {
CommitInterval time.Duration // Interval between offset flushes to the backend store.
VerboseLogging bool // Whether to enable verbose logging.
}
// NewOffsetManagerConfig returns a new OffsetManagerConfig with sane defaults.
func NewOffsetManagerConfig() *OffsetManagerConfig {
return &OffsetManagerConfig{
CommitInterval: 10 * time.Second,
}
}
type (
topicOffsets map[int32]*partitionOffsetTracker
offsetsMap map[string]topicOffsets
offsetCommitter func(int64) error
)
type partitionOffsetTracker struct {
l sync.Mutex
waitingForOffset int64
highestProcessedOffset int64
lastCommittedOffset int64
done chan struct{}
}
type zookeeperOffsetManager struct {
config *OffsetManagerConfig
l sync.RWMutex
offsets offsetsMap
cg *ConsumerGroup
closing, closed, flush chan struct{}
flushErr chan error
}
// NewZookeeperOffsetManager returns an offset manager that uses Zookeeper
// to store offsets.
func NewZookeeperOffsetManager(cg *ConsumerGroup, config *OffsetManagerConfig) OffsetManager {
if config == nil {
config = NewOffsetManagerConfig()
}
zom := &zookeeperOffsetManager{
config: config,
cg: cg,
offsets: make(offsetsMap),
closing: make(chan struct{}),
closed: make(chan struct{}),
flush: make(chan struct{}),
flushErr: make(chan error),
}
go zom.offsetCommitter()
return zom
}
func (zom *zookeeperOffsetManager) InitializePartition(topic string, partition int32) (int64, error) {
zom.l.Lock()
defer zom.l.Unlock()
if zom.offsets[topic] == nil {
zom.offsets[topic] = make(topicOffsets)
}
nextOffset, err := zom.cg.group.FetchOffset(topic, partition)
if err != nil {
return 0, err
}
zom.offsets[topic][partition] = &partitionOffsetTracker{
highestProcessedOffset: nextOffset - 1,
lastCommittedOffset: nextOffset - 1,
done: make(chan struct{}),
}
return nextOffset, nil
}
func (zom *zookeeperOffsetManager) FinalizePartition(topic string, partition int32, lastOffset int64, timeout time.Duration) error {
zom.l.RLock()
tracker := zom.offsets[topic][partition]
zom.l.RUnlock()
if lastOffset >= 0 {
if lastOffset-tracker.highestProcessedOffset > 0 {
zom.cg.Logf("%s/%d :: Last processed offset: %d. Waiting up to %ds for another %d messages to process...", topic, partition, tracker.highestProcessedOffset, timeout/time.Second, lastOffset-tracker.highestProcessedOffset)
if !tracker.waitForOffset(lastOffset, timeout) {
return fmt.Errorf("TIMEOUT waiting for offset %d. Last committed offset: %d", lastOffset, tracker.lastCommittedOffset)
}
}
if err := zom.commitOffset(topic, partition, tracker); err != nil {
return fmt.Errorf("FAILED to commit offset %d to Zookeeper. Last committed offset: %d", tracker.highestProcessedOffset, tracker.lastCommittedOffset)
}
}
zom.l.Lock()
delete(zom.offsets[topic], partition)
zom.l.Unlock()
return nil
}
func (zom *zookeeperOffsetManager) MarkAsProcessed(topic string, partition int32, offset int64) bool {
zom.l.RLock()
defer zom.l.RUnlock()
if p, ok := zom.offsets[topic][partition]; ok {
return p.markAsProcessed(offset)
} else {
return false
}
}
func (zom *zookeeperOffsetManager) Flush() error {
zom.flush <- struct{}{}
return <-zom.flushErr
}
func (zom *zookeeperOffsetManager) Close() error {
close(zom.closing)
<-zom.closed
zom.l.Lock()
defer zom.l.Unlock()
var closeError error
for _, partitionOffsets := range zom.offsets {
if len(partitionOffsets) > 0 {
closeError = UncleanClose
}
}
return closeError
}
func (zom *zookeeperOffsetManager) offsetCommitter() {
var tickerChan <-chan time.Time
if zom.config.CommitInterval != 0 {
commitTicker := time.NewTicker(zom.config.CommitInterval)
tickerChan = commitTicker.C
defer commitTicker.Stop()
}
for {
select {
case <-zom.closing:
close(zom.closed)
return
case <-tickerChan:
if err := zom.commitOffsets(); err != nil {
zom.cg.errors <- err
}
case <-zom.flush:
zom.flushErr <- zom.commitOffsets()
}
}
}
func (zom *zookeeperOffsetManager) commitOffsets() error {
zom.l.RLock()
defer zom.l.RUnlock()
var returnErr error
for topic, partitionOffsets := range zom.offsets {
for partition, offsetTracker := range partitionOffsets {
err := zom.commitOffset(topic, partition, offsetTracker)
switch err {
case nil:
// noop
default:
returnErr = err
}
}
}
return returnErr
}
func (zom *zookeeperOffsetManager) commitOffset(topic string, partition int32, tracker *partitionOffsetTracker) error {
err := tracker.commit(func(offset int64) error {
if offset >= 0 {
return zom.cg.group.CommitOffset(topic, partition, offset+1)
} else {
return nil
}
})
if err != nil {
zom.cg.Logf("FAILED to commit offset %d for %s/%d!", tracker.highestProcessedOffset, topic, partition)
} else if zom.config.VerboseLogging {
zom.cg.Logf("Committed offset %d for %s/%d!", tracker.lastCommittedOffset, topic, partition)
}
return err
}
// MarkAsProcessed marks the provided offset as highest processed offset if
// it's higher than any previous offset it has received.
func (pot *partitionOffsetTracker) markAsProcessed(offset int64) bool {
pot.l.Lock()
defer pot.l.Unlock()
if offset > pot.highestProcessedOffset {
pot.highestProcessedOffset = offset
if pot.waitingForOffset == pot.highestProcessedOffset {
close(pot.done)
}
return true
} else {
return false
}
}
// Commit calls a committer function if the highest processed offset is out
// of sync with the last committed offset.
func (pot *partitionOffsetTracker) commit(committer offsetCommitter) error {
pot.l.Lock()
defer pot.l.Unlock()
if pot.highestProcessedOffset > pot.lastCommittedOffset {
if err := committer(pot.highestProcessedOffset); err != nil {
return err
}
pot.lastCommittedOffset = pot.highestProcessedOffset
return nil
} else {
return nil
}
}
func (pot *partitionOffsetTracker) waitForOffset(offset int64, timeout time.Duration) bool {
pot.l.Lock()
if offset > pot.highestProcessedOffset {
pot.waitingForOffset = offset
pot.l.Unlock()
select {
case <-pot.done:
return true
case <-time.After(timeout):
return false
}
} else {
pot.l.Unlock()
return true
}
}

View File

@@ -0,0 +1,113 @@
package consumergroup
import (
"crypto/rand"
"fmt"
"io"
"os"
"sort"
"github.com/wvanbergen/kazoo-go"
)
func retrievePartitionLeaders(partitions kazoo.PartitionList) (partitionLeaders, error) {
pls := make(partitionLeaders, 0, len(partitions))
for _, partition := range partitions {
leader, err := partition.Leader()
if err != nil {
return nil, err
}
pl := partitionLeader{id: partition.ID, leader: leader, partition: partition}
pls = append(pls, pl)
}
return pls, nil
}
// Divides a set of partitions between a set of consumers.
func dividePartitionsBetweenConsumers(consumers kazoo.ConsumergroupInstanceList, partitions partitionLeaders) map[string][]*kazoo.Partition {
result := make(map[string][]*kazoo.Partition)
plen := len(partitions)
clen := len(consumers)
if clen == 0 {
return result
}
sort.Sort(partitions)
sort.Sort(consumers)
n := plen / clen
m := plen % clen
p := 0
for i, consumer := range consumers {
first := p
last := first + n
if m > 0 && i < m {
last++
}
if last > plen {
last = plen
}
for _, pl := range partitions[first:last] {
result[consumer.ID] = append(result[consumer.ID], pl.partition)
}
p = last
}
return result
}
type partitionLeader struct {
id int32
leader int32
partition *kazoo.Partition
}
// A sortable slice of PartitionLeader structs
type partitionLeaders []partitionLeader
func (pls partitionLeaders) Len() int {
return len(pls)
}
func (pls partitionLeaders) Less(i, j int) bool {
return pls[i].leader < pls[j].leader || (pls[i].leader == pls[j].leader && pls[i].id < pls[j].id)
}
func (s partitionLeaders) Swap(i, j int) {
s[i], s[j] = s[j], s[i]
}
func generateUUID() (string, error) {
uuid := make([]byte, 16)
n, err := io.ReadFull(rand.Reader, uuid)
if n != len(uuid) || err != nil {
return "", err
}
// variant bits; see section 4.1.1
uuid[8] = uuid[8]&^0xc0 | 0x80
// version 4 (pseudo-random); see section 4.1.3
uuid[6] = uuid[6]&^0xf0 | 0x40
return fmt.Sprintf("%x-%x-%x-%x-%x", uuid[0:4], uuid[4:6], uuid[6:8], uuid[8:10], uuid[10:]), nil
}
func generateConsumerID() (consumerID string, err error) {
var uuid, hostname string
uuid, err = generateUUID()
if err != nil {
return
}
hostname, err = os.Hostname()
if err != nil {
return
}
consumerID = fmt.Sprintf("%s:%s", hostname, uuid)
return
}