├── .gitignore ├── .travis.yml ├── Changelog ├── Gopkg.lock ├── Gopkg.toml ├── LICENSE ├── Makefile ├── README.md ├── README.zh-CN.md ├── config.go ├── config_test.go ├── consumer_group.go ├── consumer_group_test.go ├── example └── example.go ├── group_storage.go ├── partition.go ├── tests ├── docker │ ├── docker-compose.yml │ ├── setup.sh │ └── teardown.sh ├── integration │ ├── init.go │ ├── offset_test.go │ ├── rebalance_test.go │ └── util.go └── test.sh ├── topic.go ├── util.go ├── util_test.go ├── zk_group_storage.go └── zk_group_storage_test.go /.gitignore: -------------------------------------------------------------------------------- 1 | # Compiled Object files, Static and Dynamic libs (Shared Objects) 2 | *.o 3 | *.a 4 | *.so 5 | 6 | # Folders 7 | _obj 8 | _test 9 | 10 | # Architecture specific extensions/prefixes 11 | *.[568vq] 12 | [568vq].out 13 | 14 | *.cgo1.go 15 | *.cgo2.c 16 | _cgo_defun.c 17 | _cgo_gotypes.go 18 | _cgo_export.* 19 | 20 | _testmain.go 21 | 22 | *.exe 23 | *.test 24 | 25 | *.swp 26 | *.swo 27 | go-consumergroup 28 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | sudo: required 2 | language: go 3 | dist: trusty 4 | 5 | go: 6 | - 1.11 7 | 8 | services: 9 | - docker 10 | 11 | env: 12 | - DOCKER_COMPOSE_VERSION=1.14.0 13 | 14 | install: 15 | - sudo rm /usr/local/bin/docker-compose 16 | - curl -L https://github.com/docker/compose/releases/download/${DOCKER_COMPOSE_VERSION}/docker-compose-`uname -s`-`uname -m` > docker-compose 17 | - chmod +x docker-compose 18 | - sudo mv docker-compose /usr/local/bin 19 | - go get -u golang.org/x/lint/golint 20 | - curl https://raw.githubusercontent.com/golang/dep/master/install.sh | sh 21 | - dep ensure 22 | - make build 23 | 24 | script: 25 | - make test 26 | - make lint 27 | -------------------------------------------------------------------------------- /Changelog: -------------------------------------------------------------------------------- 1 | VERSION 0.1.0 2 | * init 3 | 4 | VERSION 0.2.0 5 | * Refactor consumer group log message 6 | * Add more annotation 7 | * Refactor GetTopicNextMessageChannel => GetMessages, GetTopicErrorsChannel 8 | => GetErrors 9 | * Refactor consumer-group into topic-consumer and partition-consumer 10 | 11 | VERSION 0.2.6 12 | 13 | * Fix wrong way to use defer function cause the failure of partition consumer 14 | close 15 | 16 | VERSION 0.2.7 17 | 18 | * Set the consumer group initial state as stopped 19 | -------------------------------------------------------------------------------- /Gopkg.lock: -------------------------------------------------------------------------------- 1 | # This file is autogenerated, do not edit; changes may be undone by the next 'dep ensure'. 2 | 3 | 4 | [[projects]] 5 | name = "github.com/Shopify/sarama" 6 | packages = ["."] 7 | revision = "f7be6aa2bc7b2e38edf816b08b582782194a1c02" 8 | version = "v1.16.0" 9 | 10 | [[projects]] 11 | name = "github.com/davecgh/go-spew" 12 | packages = ["spew"] 13 | revision = "346938d642f2ec3594ed81d874461961cd0faa76" 14 | version = "v1.1.0" 15 | 16 | [[projects]] 17 | name = "github.com/eapache/go-resiliency" 18 | packages = ["breaker"] 19 | revision = "ea41b0fad31007accc7f806884dcdf3da98b79ce" 20 | version = "v1.1.0" 21 | 22 | [[projects]] 23 | branch = "master" 24 | name = "github.com/eapache/go-xerial-snappy" 25 | packages = ["."] 26 | revision = "bb955e01b9346ac19dc29eb16586c90ded99a98c" 27 | 28 | [[projects]] 29 | name = "github.com/eapache/queue" 30 | packages = ["."] 31 | revision = "44cc805cf13205b55f69e14bcb69867d1ae92f98" 32 | version = "v1.1.0" 33 | 34 | [[projects]] 35 | branch = "master" 36 | name = "github.com/golang/snappy" 37 | packages = ["."] 38 | revision = "2e65f85255dbc3072edf28d6b5b8efc472979f5a" 39 | 40 | [[projects]] 41 | branch = "master" 42 | name = "github.com/meitu/go-zookeeper" 43 | packages = ["zk"] 44 | revision = "d516a4489ac1841674df100f5ed449ea3f1e915a" 45 | 46 | [[projects]] 47 | branch = "master" 48 | name = "github.com/meitu/zk_wrapper" 49 | packages = ["."] 50 | revision = "60b89077c5737b37b81454d28ac2b7adc70fafeb" 51 | 52 | [[projects]] 53 | name = "github.com/pierrec/lz4" 54 | packages = ["."] 55 | revision = "2fcda4cb7018ce05a25959d2fe08c83e3329f169" 56 | version = "v1.1" 57 | 58 | [[projects]] 59 | name = "github.com/pierrec/xxHash" 60 | packages = ["xxHash32"] 61 | revision = "f051bb7f1d1aaf1b5a665d74fb6b0217712c69f7" 62 | version = "v0.1.1" 63 | 64 | [[projects]] 65 | branch = "master" 66 | name = "github.com/rcrowley/go-metrics" 67 | packages = ["."] 68 | revision = "e2704e165165ec55d062f5919b4b29494e9fa790" 69 | 70 | [[projects]] 71 | name = "github.com/sirupsen/logrus" 72 | packages = ["."] 73 | revision = "c155da19408a8799da419ed3eeb0cb5db0ad5dbc" 74 | version = "v1.0.5" 75 | 76 | [[projects]] 77 | branch = "master" 78 | name = "golang.org/x/crypto" 79 | packages = ["ssh/terminal"] 80 | revision = "a49355c7e3f8fe157a85be2f77e6e269a0f89602" 81 | 82 | [[projects]] 83 | branch = "master" 84 | name = "golang.org/x/sys" 85 | packages = [ 86 | "unix", 87 | "windows" 88 | ] 89 | revision = "7138fd3d9dc8335c567ca206f4333fb75eb05d56" 90 | 91 | [solve-meta] 92 | analyzer-name = "dep" 93 | analyzer-version = 1 94 | inputs-digest = "3894696adaf114974713eeaf39b02e4ccc383c2aae9e798819c3eb0acaa39cd7" 95 | solver-name = "gps-cdcl" 96 | solver-version = 1 97 | -------------------------------------------------------------------------------- /Gopkg.toml: -------------------------------------------------------------------------------- 1 | # Gopkg.toml example 2 | # 3 | # Refer to https://golang.github.io/dep/docs/Gopkg.toml.html 4 | # for detailed Gopkg.toml documentation. 5 | # 6 | # required = ["github.com/user/thing/cmd/thing"] 7 | # ignored = ["github.com/user/project/pkgX", "bitbucket.org/user/project/pkgA/pkgY"] 8 | # 9 | # [[constraint]] 10 | # name = "github.com/user/project" 11 | # version = "1.0.0" 12 | # 13 | # [[constraint]] 14 | # name = "github.com/user/project2" 15 | # branch = "dev" 16 | # source = "github.com/myfork/project2" 17 | # 18 | # [[override]] 19 | # name = "github.com/x/y" 20 | # version = "2.4.0" 21 | # 22 | # [prune] 23 | # non-go = false 24 | # go-tests = true 25 | # unused-packages = true 26 | 27 | 28 | [[constraint]] 29 | name = "github.com/Shopify/sarama" 30 | version = "1.16.0" 31 | 32 | [[constraint]] 33 | branch = "master" 34 | name = "github.com/meitu/zk_wrapper" 35 | 36 | [[constraint]] 37 | branch = "master" 38 | name = "github.com/meitu/go-zookeeper" 39 | 40 | [prune] 41 | go-tests = true 42 | unused-packages = true 43 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2017 meitu 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | export GO15VENDOREXPERIMENT=1 2 | 3 | # Many Go tools take file globs or directories as arguments instead of packages. 4 | PKG_FILES ?= *.go 5 | 6 | # The linting tools evolve with each Go version, so run them only on the latest 7 | # stable release. 8 | GO_VERSION := $(shell go version | cut -d " " -f 3) 9 | GO_MINOR_VERSION := $(word 2,$(subst ., ,$(GO_VERSION))) 10 | ifneq ($(filter $(LINTABLE_MINOR_VERSIONS),$(GO_MINOR_VERSION)),) 11 | SHOULD_LINT := true 12 | endif 13 | 14 | .PHONY: build 15 | build: 16 | go build 17 | 18 | .PHONY: lint 19 | lint: 20 | @rm -rf lint.log 21 | @echo "Checking formatting..." 22 | @gofmt -d -s $(PKG_FILES) 2>&1 | tee lint.log 23 | @echo "Installing test dependencies for vet..." 24 | @go test -i $(PKGS) 25 | @echo "Checking vet..." 26 | @go tool vet $(PKG_FILES) 2>&1 | tee -a lint.log; 27 | @echo "Checking lint..." 28 | @golint $(PKG_FILES) 2>&1 | tee -a lint.log; 29 | @echo "Checking for unresolved FIXMEs..." 30 | @git grep -i fixme | grep -v -e vendor -e Makefile | tee -a lint.log 31 | @[ ! -s lint.log ] 32 | 33 | .PHONY: test 34 | test: 35 | cd tests && sh test.sh 36 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # go-consumergroup [![Build Status](https://travis-ci.org/meitu/go-consumergroup.svg?branch=master)](https://travis-ci.org/meitu/go-consumergroup) [![Go Report Card](https://goreportcard.com/badge/github.com/meitu/go-consumergroup)](https://goreportcard.com/report/github.com/meitu/go-consumergroup) 2 | 3 | Go-consumergroup is a kafka consumer library written in golang with rebalance and chroot supports. 4 | 5 | [Chinese Doc](./README.zh-CN.md) 6 | 7 | ## Requirements 8 | * Apache Kafka 0.8.x, 0.9.x, 0.10.x, 0.11.x, 1.0.x 9 | 10 | ## Dependencies 11 | * [go-zookeeper](https://github.com/samuel/go-zookeeper) 12 | * [sarama](https://github.com/Shopify/sarama) 13 | * [zk_wrapper](https://github.com/meitu/zk_wrapper) 14 | 15 | ## Getting started 16 | 17 | * API documentation and examples are available via [godoc](https://godoc.org/github.com/meitu/go-consumergroup). 18 | * The example directory contains more elaborate [example](example/example.go) applications. 19 | 20 | ## User Defined Logger 21 | 22 | ``` 23 | logger := logrus.New() 24 | cg.SetLogger(logger) 25 | ``` 26 | 27 | ## Run Tests 28 | 29 | ```shell 30 | $ make test 31 | ``` 32 | 33 | ***NOTE:*** `docker-compse` is required to run tests 34 | -------------------------------------------------------------------------------- /README.zh-CN.md: -------------------------------------------------------------------------------- 1 | # go-consumergroup [![Build Status](https://travis-ci.org/meitu/go-consumergroup.svg?branch=master)](https://travis-ci.org/meitu/go-consumergroup) [![Go Report Card](https://goreportcard.com/badge/github.com/meitu/go-consumergroup)](https://goreportcard.com/report/github.com/meitu/go-consumergroup) 2 | 3 | ### 简介 4 | go-consumergroup是一款提供集群功能的kafka客户端，支持 rebalance，offset 自动或者手动管理以及 chroot 功能。 5 | 6 | ### 依赖 7 | 8 | 9 | 10 | 11 | ## 快速上手 12 | 13 | * API 文档请参照 [godoc](https://godoc.org/github.com/meitu/go-consumergroup). 14 | * 使用例子参照 example 目录的 example.go 实现 [example](example/example.go) 15 | 16 | ## 测试 17 | 18 | ```shell 19 | $ make test 20 | ``` 21 | 22 | ***NOTE: *** 跑测试用例需要预先安装 docker-compose 23 | -------------------------------------------------------------------------------- /config.go: -------------------------------------------------------------------------------- 1 | package consumergroup 2 | 3 | import ( 4 | "errors" 5 | "time" 6 | 7 | "github.com/Shopify/sarama" 8 | ) 9 | 10 | // Config is used to pass multiple configuration options to ConsumerGroup's constructors 11 | type Config struct { 12 | Chroot string 13 | // ZkList is required, zookeeper address's list 14 | ZkList []string 15 | // Zookeeper session timeout, default is 6s 16 | ZkSessionTimeout time.Duration 17 | // GroupID is required, identifer to determin which ConsumerGroup would be joined 18 | GroupID string 19 | // ConsumerID is optional, identifer to sign partition's owner 20 | ConsumerID string 21 | // TopicList is required, topics that ConsumerGroup would be consumed 22 | TopicList []string 23 | // Just export Sarama Config 24 | SaramaConfig *sarama.Config 25 | // Size of error channel, default is 1024 26 | ErrorChannelBufferSize int 27 | // Whether auto commit the offset or not, default is true 28 | OffsetAutoCommitEnable bool 29 | // Offset auto commit interval, default is 10s 30 | OffsetAutoCommitInterval time.Duration 31 | // Where to fetch messages when offset was not found, default is newest 32 | OffsetAutoReset int64 33 | // Claim the partition would give up after ClaimPartitionRetryTimes(>0) retires, 34 | // ClaimPartitionRetryTimes <= 0 would retry until success or receive stop signal 35 | ClaimPartitionRetryTimes int 36 | // Retry interval when fail to clain the partition 37 | ClaimPartitionRetryInterval time.Duration 38 | } 39 | 40 | // NewConfig return the new config with default value. 41 | func NewConfig() *Config { 42 | config := new(Config) 43 | config.SaramaConfig = sarama.NewConfig() 44 | config.ErrorChannelBufferSize = 1024 45 | config.OffsetAutoCommitEnable = true 46 | config.OffsetAutoCommitInterval = 10 * time.Second 47 | config.OffsetAutoReset = sarama.OffsetNewest 48 | config.ClaimPartitionRetryTimes = 10 49 | config.ClaimPartitionRetryInterval = 3 * time.Second 50 | config.SaramaConfig.Consumer.Return.Errors = true 51 | return config 52 | } 53 | 54 | func (c *Config) validate() error { 55 | if c.ZkList == nil || len(c.ZkList) <= 0 { 56 | return errors.New("ZkList can't be empty") 57 | } 58 | if c.GroupID == "" { 59 | return errors.New("GroupID can't be empty") 60 | } 61 | if c.TopicList == nil || len(c.TopicList) <= 0 { 62 | return errors.New("GroupId can't be empty") 63 | } 64 | c.TopicList = sliceRemoveDuplicates(c.TopicList) 65 | c.ZkList = sliceRemoveDuplicates(c.ZkList) 66 | return nil 67 | } 68 | -------------------------------------------------------------------------------- /config_test.go: -------------------------------------------------------------------------------- 1 | package consumergroup 2 | 3 | import "testing" 4 | 5 | func TestConfigValidate(t *testing.T) { 6 | conf := NewConfig() 7 | conf.ZkList = []string{} 8 | conf.TopicList = []string{} 9 | if err := conf.validate(); err == nil { 10 | t.Fatal("config invalidate is expected") 11 | } 12 | conf.ZkList = []string{"127.0.0.1:2181", "127.0.0.1:2181"} 13 | if err := conf.validate(); err == nil { 14 | t.Fatal("config invalidate is expected") 15 | } 16 | conf.TopicList = []string{"a", "a", "b", "c", "a"} 17 | if err := conf.validate(); err == nil { 18 | t.Fatal("config validate is expected") 19 | } 20 | conf.GroupID = "go-test-group" 21 | if err := conf.validate(); err != nil { 22 | t.Fatalf("validate is expected, but got error %s", err) 23 | } 24 | if len(conf.TopicList) != 3 { 25 | t.Fatal("config validate should remove duplicate topics") 26 | } 27 | if len(conf.ZkList) != 1 { 28 | t.Fatal("config validate should remove duplicate zk addresses") 29 | } 30 | } 31 | -------------------------------------------------------------------------------- /consumer_group.go: -------------------------------------------------------------------------------- 1 | package consumergroup 2 | 3 | import ( 4 | "errors" 5 | "fmt" 6 | "reflect" 7 | "runtime/debug" 8 | "sync" 9 | "time" 10 | 11 | "github.com/Shopify/sarama" 12 | "github.com/meitu/go-zookeeper/zk" 13 | "github.com/sirupsen/logrus" 14 | ) 15 | 16 | const ( 17 | cgStarted = iota + 1 18 | cgStopped 19 | ) 20 | 21 | const ( 22 | restartEvent = 1 23 | ) 24 | 25 | // ConsumerGroup consume message from Kafka with rebalancing supports 26 | type ConsumerGroup struct { 27 | name string 28 | storage groupStorage 29 | topicConsumers map[string]*topicConsumer 30 | saramaClients map[string]sarama.Client 31 | saramaConsumers map[string]sarama.Consumer 32 | 33 | id string 34 | state int 35 | wg sync.WaitGroup 36 | stopCh chan struct{} 37 | triggerCh chan int 38 | triggerOnce *sync.Once 39 | owners map[string]map[int32]string 40 | 41 | config *Config 42 | logger *logrus.Logger 43 | 44 | onLoad, onClose []func() 45 | } 46 | 47 | // NewConsumerGroup create the ConsumerGroup instance with config 48 | func NewConsumerGroup(config *Config) (*ConsumerGroup, error) { 49 | if config == nil { 50 | return nil, errors.New("config can't be empty") 51 | } 52 | err := config.validate() 53 | if err != nil { 54 | return nil, fmt.Errorf("vaildate config failed, as %s", err) 55 | } 56 | 57 | cg := new(ConsumerGroup) 58 | cg.state = cgStopped 59 | cg.config = config 60 | cg.id = config.ConsumerID 61 | if cg.id == "" { 62 | cg.id = genConsumerID() 63 | } 64 | cg.name = config.GroupID 65 | cg.triggerCh = make(chan int) 66 | cg.topicConsumers = make(map[string]*topicConsumer) 67 | cg.saramaClients = make(map[string]sarama.Client) 68 | cg.saramaConsumers = make(map[string]sarama.Consumer) 69 | cg.onLoad = make([]func(), 0) 70 | cg.onClose = make([]func(), 0) 71 | cg.storage = newZKGroupStorage(config.ZkList, config.ZkSessionTimeout) 72 | cg.logger = logrus.New() 73 | if _, ok := cg.storage.(*zkGroupStorage); ok { 74 | cg.storage.(*zkGroupStorage).Chroot(config.Chroot) 75 | } 76 | 77 | err = cg.initSaramaConsumer() 78 | if err != nil { 79 | return nil, fmt.Errorf("init sarama consumer, as %s", err) 80 | } 81 | cg.owners = make(map[string]map[int32]string) 82 | for _, topic := range config.TopicList { 83 | cg.topicConsumers[topic] = newTopicConsumer(cg, topic) 84 | cg.owners[topic] = make(map[int32]string) 85 | } 86 | return cg, nil 87 | } 88 | 89 | func (cg *ConsumerGroup) initSaramaConsumer() error { 90 | brokerList, err := cg.storage.getBrokerList() 91 | if err != nil { 92 | return err 93 | } 94 | if len(brokerList) == 0 { 95 | return errors.New("no broker alive") 96 | } 97 | for _, topic := range cg.config.TopicList { 98 | saramaClient, err := sarama.NewClient(brokerList, cg.config.SaramaConfig) 99 | if err != nil { 100 | return err 101 | } 102 | saramaConsumer, err := sarama.NewConsumerFromClient(saramaClient) 103 | if err != nil { 104 | return err 105 | } 106 | cg.saramaClients[topic] = saramaClient 107 | cg.saramaConsumers[topic] = saramaConsumer 108 | } 109 | return nil 110 | } 111 | 112 | // Start would register ConsumerGroup, and rebalance would be triggered. 113 | // ConsumerGroup computes the partitions which should be consumed by consumer's num, and start fetching message. 114 | func (cg *ConsumerGroup) Start() error { 115 | // exit when failed to register the consumer 116 | err := cg.storage.registerConsumer(cg.name, cg.id, nil) 117 | if err != nil && err != zk.ErrNodeExists { 118 | return err 119 | } 120 | cg.wg.Add(1) 121 | go cg.start() 122 | return nil 123 | } 124 | 125 | // Stop would unregister ConsumerGroup, and rebalance would be triggered. 126 | // The partitions which consumed by this ConsumerGroup would be assigned to others. 127 | func (cg *ConsumerGroup) Stop() { 128 | cg.stop() 129 | cg.wg.Wait() 130 | } 131 | 132 | // SetLogger use to set the user's logger the consumer group 133 | func (cg *ConsumerGroup) SetLogger(l *logrus.Logger) { 134 | if l != nil { 135 | cg.logger = l 136 | } 137 | } 138 | 139 | // IsStopped return whether the ConsumerGroup was stopped or not. 140 | func (cg *ConsumerGroup) IsStopped() bool { 141 | return cg.state == cgStopped 142 | } 143 | 144 | func (cg *ConsumerGroup) callRecover() { 145 | if err := recover(); err != nil { 146 | cg.logger.WithFields(logrus.Fields{ 147 | "group": cg.name, 148 | "err": err, 149 | "stack": string(debug.Stack()), 150 | }).Error("Recover panic") 151 | cg.stop() 152 | } 153 | } 154 | 155 | func (cg *ConsumerGroup) start() { 156 | var wg sync.WaitGroup 157 | 158 | defer cg.callRecover() 159 | defer func() { 160 | cg.state = cgStopped 161 | err := cg.storage.deleteConsumer(cg.name, cg.id) 162 | if err != nil { 163 | cg.logger.WithFields(logrus.Fields{ 164 | "group": cg.name, 165 | "err": err, 166 | }).Error("Failed to delete consumer from zk") 167 | } 168 | for _, tc := range cg.topicConsumers { 169 | close(tc.messages) 170 | close(tc.errors) 171 | } 172 | cg.wg.Done() 173 | }() 174 | 175 | CONSUME_TOPIC_LOOP: 176 | for { 177 | cg.logger.WithField("group", cg.name).Info("Consumer group started") 178 | cg.triggerOnce = new(sync.Once) 179 | cg.stopCh = make(chan struct{}) 180 | 181 | err := cg.watchRebalance() 182 | if err != nil { 183 | cg.logger.WithFields(logrus.Fields{ 184 | "group": cg.name, 185 | "err": err, 186 | }).Error("Failed to watch rebalance") 187 | cg.stop() 188 | return 189 | } 190 | wg.Add(1) 191 | go func() { 192 | defer cg.callRecover() 193 | defer wg.Done() 194 | cg.autoReconnect(cg.storage.(*zkGroupStorage).sessionTimeout / 3) 195 | }() 196 | for _, consumer := range cg.topicConsumers { 197 | wg.Add(1) 198 | consumer.start() 199 | go func(tc *topicConsumer) { 200 | defer cg.callRecover() 201 | defer wg.Done() 202 | tc.wg.Wait() 203 | cg.logger.WithFields(logrus.Fields{ 204 | "group": tc.group, 205 | "topic": tc.name, 206 | }).Info("Stop the topic consumer") 207 | }(consumer) 208 | } 209 | cg.state = cgStarted 210 | for _, onLoadFunc := range cg.onLoad { 211 | onLoadFunc() 212 | } 213 | msg := <-cg.triggerCh 214 | for _, onCloseFunc := range cg.onClose { 215 | onCloseFunc() 216 | } 217 | switch msg { 218 | case restartEvent: 219 | close(cg.stopCh) 220 | // The stop channel was used to notify partition's consumer to stop consuming when rebalance is triggered. 221 | // So we should reinit when rebalance was triggered, as it would be closed. 222 | wg.Wait() 223 | continue CONSUME_TOPIC_LOOP 224 | default: 225 | close(cg.stopCh) 226 | cg.logger.WithField("group", cg.name).Info("ConsumerGroup is stopping") 227 | wg.Wait() 228 | cg.logger.WithField("group", cg.name).Info("ConsumerGroup was stopped") 229 | return 230 | } 231 | } 232 | } 233 | 234 | func (cg *ConsumerGroup) stop() { 235 | cg.triggerOnce.Do(func() { close(cg.triggerCh) }) 236 | } 237 | 238 | func (cg *ConsumerGroup) triggerRebalance() { 239 | cg.triggerOnce.Do(func() { cg.triggerCh <- restartEvent }) 240 | } 241 | 242 | func (cg *ConsumerGroup) getPartitionConsumer(topic string, partition int32, nextOffset int64) (sarama.PartitionConsumer, error) { 243 | saramaConsumer := cg.saramaConsumers[topic] 244 | consumer, err := saramaConsumer.ConsumePartition(topic, partition, nextOffset) 245 | if err == sarama.ErrOffsetOutOfRange { 246 | cg.logger.WithFields(logrus.Fields{ 247 | "group": cg.name, 248 | "topic": topic, 249 | "partition": partition, 250 | "offset": nextOffset, 251 | }).Error("Partition's offset was out of range, use auto-reset") 252 | nextOffset = cg.config.OffsetAutoReset 253 | consumer, err = saramaConsumer.ConsumePartition(topic, partition, nextOffset) 254 | } 255 | return consumer, err 256 | } 257 | 258 | // GetMessages was used to get a unbuffered message's channel from specified topic 259 | func (cg *ConsumerGroup) GetMessages(topic string) (<-chan *sarama.ConsumerMessage, bool) { 260 | if topicConsumer, ok := cg.topicConsumers[topic]; ok { 261 | return topicConsumer.messages, true 262 | } 263 | return nil, false 264 | } 265 | 266 | // GetErrors was used to get a unbuffered error's channel from specified topic 267 | func (cg *ConsumerGroup) GetErrors(topic string) (<-chan *sarama.ConsumerError, bool) { 268 | if topicConsumer, ok := cg.topicConsumers[topic]; ok { 269 | return topicConsumer.errors, true 270 | } 271 | return nil, false 272 | } 273 | 274 | // OnLoad load callback function that runs after startup 275 | func (cg *ConsumerGroup) OnLoad(cb func()) { 276 | cg.onLoad = append(cg.onLoad, cb) 277 | } 278 | 279 | // OnClose load callback function that runs before the end 280 | func (cg *ConsumerGroup) OnClose(cb func()) { 281 | cg.onClose = append(cg.onClose, cb) 282 | } 283 | 284 | func (cg *ConsumerGroup) autoReconnect(interval time.Duration) { 285 | timer := time.NewTimer(interval) 286 | cg.logger.WithField("group", cg.name).Info("The auto-reconnect consumer thread was started") 287 | defer cg.logger.WithField("group", cg.name).Info("The auto-reconnect consumer thread was stopped") 288 | for { 289 | select { 290 | case <-cg.stopCh: 291 | return 292 | case <-timer.C: 293 | timer.Reset(interval) 294 | exist, err := cg.storage.existsConsumer(cg.name, cg.id) 295 | if err != nil { 296 | cg.logger.WithFields(logrus.Fields{ 297 | "group": cg.name, 298 | "err": err, 299 | }).Error("Failed to check consumer existence") 300 | break 301 | } 302 | if exist { 303 | break 304 | } 305 | err = cg.storage.registerConsumer(cg.name, cg.id, nil) 306 | if err != nil { 307 | cg.logger.WithFields(logrus.Fields{ 308 | "group": cg.name, 309 | "err": err, 310 | }).Error("Failed to re-register consumer") 311 | } 312 | } 313 | } 314 | } 315 | 316 | func (cg *ConsumerGroup) watchRebalance() error { 317 | consumersWatcher, err := cg.storage.watchConsumerList(cg.name) 318 | if err != nil { 319 | return err 320 | } 321 | topicsChange, topicWatchers, err := cg.watchTopics(cg.config.TopicList) 322 | if err != nil { 323 | return err 324 | } 325 | go func() { 326 | defer cg.callRecover() 327 | cg.logger.WithField("group", cg.name).Info("Rebalance watcher thread was started") 328 | select { 329 | case <-consumersWatcher.EvCh: 330 | cg.triggerRebalance() 331 | cg.logger.WithField("group", cg.name).Info("Trigger rebalance while consumers was changed") 332 | for _, tw := range topicWatchers { 333 | cg.storage.removeWatcher(tw) 334 | } 335 | case topic := <-topicsChange: 336 | cg.triggerRebalance() 337 | cg.logger.WithFields(logrus.Fields{ 338 | "group": cg.name, 339 | "topic": topic, 340 | }).Info("Trigger rebalance while partitions was changed") 341 | cg.storage.removeWatcher(consumersWatcher) 342 | case <-cg.stopCh: 343 | } 344 | cg.logger.WithField("group", cg.name).Info("Rebalance watcher thread was exited") 345 | }() 346 | return nil 347 | } 348 | 349 | // CommitOffset is used to commit offset when auto commit was disabled. 350 | func (cg *ConsumerGroup) CommitOffset(topic string, partition int32, offset int64) error { 351 | if cg.config.OffsetAutoCommitEnable { 352 | return errors.New("commit offset take effect when offset auto commit was disabled") 353 | } 354 | return cg.storage.commitOffset(cg.name, topic, partition, offset) 355 | } 356 | 357 | // GetOffsets return the offset in memory for debug 358 | func (cg *ConsumerGroup) GetOffsets() map[string]interface{} { 359 | topics := make(map[string]interface{}) 360 | for topic, tc := range cg.topicConsumers { 361 | topics[topic] = tc.getOffsets() 362 | } 363 | return topics 364 | } 365 | 366 | // Owners return owners of all partitions 367 | func (cg *ConsumerGroup) Owners() map[string]map[int32]string { 368 | return cg.owners 369 | } 370 | 371 | func (cg *ConsumerGroup) watchTopics(topics []string) (<-chan string, []*zk.Watcher, error) { 372 | ch := make(chan string) 373 | cases := make([]reflect.SelectCase, len(topics)) 374 | watchers := make([]*zk.Watcher, len(topics)) 375 | for i, topic := range topics { 376 | w, err := cg.storage.watchTopic(topic) 377 | if err != nil { 378 | return nil, nil, fmt.Errorf("encounter error when watch topic: %s, err: %s", topic, err) 379 | } 380 | watchers[i] = w 381 | cases[i] = reflect.SelectCase{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(w.EvCh)} 382 | } 383 | go func(cases []reflect.SelectCase, ch chan string, topics []string) { 384 | chosen, _, ok := reflect.Select(cases) 385 | if !ok { 386 | //the chosen channel has been closed. 387 | return 388 | } 389 | topic := topics[chosen] 390 | num, err := cg.storage.getPartitionsNum(topic) 391 | if err != nil { 392 | cg.logger.WithFields(logrus.Fields{ 393 | "topic": topic, 394 | "err": err, 395 | }).Error("Failed to get partitions in zookeeper after topic metadata change") 396 | return 397 | } 398 | for { 399 | saramaClient := cg.saramaClients[topic] 400 | saramaClient.RefreshMetadata(topic) 401 | partitions, err := saramaClient.Partitions(topic) 402 | if err != nil { 403 | cg.logger.WithFields(logrus.Fields{ 404 | "topic": topic, 405 | "err": err, 406 | }).Error("Failed to get partitions in broker after topic metadata change") 407 | return 408 | } 409 | if len(partitions) == num { 410 | break 411 | } 412 | time.Sleep(100 * time.Millisecond) 413 | } 414 | 415 | ch <- topics[chosen] 416 | }(cases, ch, topics) 417 | return ch, watchers, nil 418 | } 419 | -------------------------------------------------------------------------------- /consumer_group_test.go: -------------------------------------------------------------------------------- 1 | package consumergroup 2 | 3 | import ( 4 | "fmt" 5 | "testing" 6 | "time" 7 | ) 8 | 9 | func TestConsumerGroup(t *testing.T) { 10 | conf := NewConfig() 11 | conf.TopicList = []string{"go-test-topic"} 12 | conf.GroupID = "go-test-group" 13 | conf.ZkList = []string{"127.0.0.1:2181"} 14 | conf.ZkSessionTimeout = 6 * time.Second 15 | cg, err := NewConsumerGroup(conf) 16 | if err != nil { 17 | t.Error(err) 18 | } 19 | fmt.Println(cg.Start()) 20 | } 21 | -------------------------------------------------------------------------------- /example/example.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "os" 6 | "os/signal" 7 | "syscall" 8 | "time" 9 | 10 | consumergroup "github.com/meitu/go-consumergroup" 11 | ) 12 | 13 | func handleSignal(sig os.Signal, cg *consumergroup.ConsumerGroup) { 14 | switch sig { 15 | case syscall.SIGINT: 16 | cg.Stop() 17 | case syscall.SIGTERM: 18 | cg.Stop() 19 | default: 20 | } 21 | } 22 | 23 | func registerSignal(cg *consumergroup.ConsumerGroup) { 24 | go func() { 25 | c := make(chan os.Signal) 26 | sigs := []os.Signal{ 27 | syscall.SIGINT, 28 | syscall.SIGTERM, 29 | } 30 | signal.Notify(c, sigs...) 31 | sig := <-c 32 | handleSignal(sig, cg) 33 | }() 34 | } 35 | 36 | func main() { 37 | conf := consumergroup.NewConfig() 38 | conf.ZkList = []string{"127.0.0.1:2181"} 39 | conf.ZkSessionTimeout = 6 * time.Second 40 | topic := "test" 41 | conf.TopicList = []string{topic} 42 | conf.GroupID = "go-test-group-id" 43 | 44 | cg, err := consumergroup.NewConsumerGroup(conf) 45 | if err != nil { 46 | fmt.Println("Failed to create consumer group, err ", err.Error()) 47 | os.Exit(1) 48 | } 49 | 50 | registerSignal(cg) 51 | 52 | err = cg.Start() 53 | if err != nil { 54 | fmt.Println("Failed to join group, err ", err.Error()) 55 | os.Exit(1) 56 | } 57 | 58 | // Retrieve the error and log 59 | go func() { 60 | if topicErrChan, ok := cg.GetErrors(topic); ok { 61 | for err := range topicErrChan { 62 | if err != nil { 63 | fmt.Printf("Toipic %s got err, %s\n", topic, err) 64 | } 65 | } 66 | } 67 | }() 68 | 69 | if msgChan, ok := cg.GetMessages(topic); ok { 70 | for message := range msgChan { 71 | fmt.Println(string(message.Value), message.Offset) 72 | time.Sleep(500 * time.Millisecond) 73 | } 74 | } else { 75 | fmt.Println("Topic was not found in consumergroup") 76 | } 77 | } 78 | -------------------------------------------------------------------------------- /group_storage.go: -------------------------------------------------------------------------------- 1 | package consumergroup 2 | 3 | import "github.com/meitu/go-zookeeper/zk" 4 | 5 | type groupStorage interface { 6 | claimPartition(group, topic string, partition int32, consumerID string) error 7 | releasePartition(group, topic string, partition int32) error 8 | getPartitionOwner(group, topic string, partition int32) (string, error) 9 | registerConsumer(group, consumerID string, data []byte) error 10 | existsConsumer(group, consumerID string) (bool, error) 11 | deleteConsumer(group, consumerID string) error 12 | getBrokerList() ([]string, error) 13 | getConsumerList(group string) ([]string, error) 14 | watchConsumerList(group string) (*zk.Watcher, error) 15 | watchTopic(topic string) (*zk.Watcher, error) 16 | getPartitionsNum(topic string) (int, error) 17 | commitOffset(group, topic string, partition int32, offset int64) error 18 | getOffset(group, topic string, partition int32) (int64, error) 19 | removeWatcher(watcher *zk.Watcher) bool 20 | } 21 | -------------------------------------------------------------------------------- /partition.go: -------------------------------------------------------------------------------- 1 | package consumergroup 2 | 3 | import ( 4 | "errors" 5 | "fmt" 6 | "sync" 7 | "time" 8 | 9 | "github.com/Shopify/sarama" 10 | "github.com/sirupsen/logrus" 11 | ) 12 | 13 | type partitionConsumer struct { 14 | owner *topicConsumer 15 | group string 16 | topic string 17 | partition int32 18 | offset int64 19 | prevOffset int64 20 | 21 | consumer sarama.PartitionConsumer 22 | } 23 | 24 | func newPartitionConsumer(owner *topicConsumer, partition int32) *partitionConsumer { 25 | return &partitionConsumer{ 26 | owner: owner, 27 | topic: owner.name, 28 | group: owner.owner.name, 29 | partition: partition, 30 | offset: 0, 31 | prevOffset: 0, 32 | } 33 | } 34 | 35 | func (pc *partitionConsumer) start() { 36 | var wg sync.WaitGroup 37 | 38 | cg := pc.owner.owner 39 | err := pc.claim() 40 | if err != nil { 41 | cg.logger.WithFields(logrus.Fields{ 42 | "group": pc.group, 43 | "topic": pc.topic, 44 | "partition": pc.partition, 45 | "err": err, 46 | }).Error("Failed to claim the partition and gave up") 47 | goto ERROR 48 | } 49 | defer func() { 50 | err = pc.release() 51 | if err != nil { 52 | cg.logger.WithFields(logrus.Fields{ 53 | "group": pc.group, 54 | "topic": pc.topic, 55 | "partition": pc.partition, 56 | "err": err, 57 | }).Error("Failed to release the partition") 58 | } else { 59 | cg.logger.WithFields(logrus.Fields{ 60 | "group": pc.group, 61 | "topic": pc.topic, 62 | "partition": pc.partition, 63 | }).Info("Success to release the partition") 64 | } 65 | }() 66 | 67 | err = pc.loadOffsetFromZk() 68 | if err != nil { 69 | cg.logger.WithFields(logrus.Fields{ 70 | "group": pc.group, 71 | "topic": pc.topic, 72 | "partition": pc.partition, 73 | "err": err, 74 | }).Error("Failed to fetch the partition's offset") 75 | goto ERROR 76 | } 77 | cg.logger.WithFields(logrus.Fields{ 78 | "group": pc.group, 79 | "topic": pc.topic, 80 | "partition": pc.partition, 81 | "offset": pc.offset, 82 | }).Info("Fetched the partition's offset from zk") 83 | pc.consumer, err = cg.getPartitionConsumer(pc.topic, pc.partition, pc.offset) 84 | if err != nil { 85 | cg.logger.WithFields(logrus.Fields{ 86 | "group": pc.group, 87 | "topic": pc.topic, 88 | "partition": pc.partition, 89 | "offset": pc.offset, 90 | "err": err, 91 | }).Error("Failed to create the partition's consumer") 92 | goto ERROR 93 | } 94 | defer func() { 95 | pc.consumer.Close() 96 | }() 97 | 98 | if cg.config.OffsetAutoCommitEnable { // start auto commit-offset thread when enable 99 | wg.Add(1) 100 | go func() { 101 | defer cg.callRecover() 102 | defer wg.Done() 103 | cg.logger.WithFields(logrus.Fields{ 104 | "group": pc.group, 105 | "topic": pc.topic, 106 | "partition": pc.partition, 107 | }).Info("Start the partition's offset auto-commit thread") 108 | pc.autoCommitOffset() 109 | }() 110 | } 111 | 112 | pc.fetch() 113 | if cg.config.OffsetAutoCommitEnable { 114 | err = pc.commitOffset() 115 | if err != nil { 116 | cg.logger.WithFields(logrus.Fields{ 117 | "group": pc.group, 118 | "topic": pc.topic, 119 | "partition": pc.partition, 120 | "offset": pc.offset, 121 | "err": err, 122 | }).Error("Failed to commit the partition's offset") 123 | } 124 | 125 | wg.Wait() // Wait for auto-commit-offset thread 126 | cg.logger.WithFields(logrus.Fields{ 127 | "group": pc.group, 128 | "topic": pc.topic, 129 | "partition": pc.partition, 130 | "offset": pc.offset, 131 | }).Info("The partition's offset auto-commit thread was stopped") 132 | } 133 | return 134 | 135 | ERROR: 136 | cg.stop() 137 | } 138 | 139 | func (pc *partitionConsumer) loadOffsetFromZk() error { 140 | cg := pc.owner.owner 141 | offset, err := cg.storage.getOffset(pc.group, pc.topic, pc.partition) 142 | if err != nil { 143 | return err 144 | } 145 | if offset == -1 { 146 | offset = cg.config.OffsetAutoReset 147 | } 148 | pc.offset = offset 149 | pc.prevOffset = offset 150 | return nil 151 | } 152 | 153 | func (pc *partitionConsumer) claim() error { 154 | cg := pc.owner.owner 155 | timer := time.NewTimer(cg.config.ClaimPartitionRetryInterval) 156 | defer timer.Stop() 157 | retry := cg.config.ClaimPartitionRetryTimes 158 | // Claim partition would retry until success 159 | for i := 0; i < retry+1 || retry <= 0; i++ { 160 | err := cg.storage.claimPartition(pc.group, pc.topic, pc.partition, cg.id) 161 | if err == nil { 162 | return nil 163 | } 164 | if i != 0 && (i%3 == 0 || retry > 0) { 165 | cg.logger.WithFields(logrus.Fields{ 166 | "group": pc.group, 167 | "topic": pc.topic, 168 | "partition": pc.partition, 169 | "retries": i, 170 | "err": err, 171 | }).Warn("Failed to claim the partition with retries") 172 | } 173 | select { 174 | case <-timer.C: 175 | timer.Reset(cg.config.ClaimPartitionRetryInterval) 176 | case <-cg.stopCh: 177 | return errors.New("stop signal was received when claim partition") 178 | } 179 | } 180 | return fmt.Errorf("failed to claim partition after %d retries", retry) 181 | } 182 | 183 | func (pc *partitionConsumer) release() error { 184 | cg := pc.owner.owner 185 | owner, err := cg.storage.getPartitionOwner(pc.group, pc.topic, pc.partition) 186 | if err != nil { 187 | return err 188 | } 189 | if cg.id == owner { 190 | return cg.storage.releasePartition(pc.group, pc.topic, pc.partition) 191 | } 192 | return fmt.Errorf("the owner of topic[%s] partition[%d] expected %s, but got %s", 193 | pc.topic, pc.partition, owner, cg.id) 194 | } 195 | 196 | func (pc *partitionConsumer) fetch() { 197 | cg := pc.owner.owner 198 | messageChan := pc.owner.messages 199 | errorChan := pc.owner.errors 200 | 201 | cg.logger.WithFields(logrus.Fields{ 202 | "group": pc.group, 203 | "topic": pc.topic, 204 | "partition": pc.partition, 205 | "offset": pc.offset, 206 | }).Info("Start to fetch the partition's messages") 207 | PARTITION_CONSUMER_LOOP: 208 | for { 209 | select { 210 | case <-cg.stopCh: 211 | break PARTITION_CONSUMER_LOOP 212 | case err := <-pc.consumer.Errors(): 213 | if err.Err == sarama.ErrOffsetOutOfRange { 214 | pc.restart() 215 | break 216 | } 217 | errorChan <- err 218 | case message, ok := <-pc.consumer.Messages(): 219 | //check if the channel is closed. message channel close while the offset out of range 220 | if !ok { 221 | pc.restart() 222 | break 223 | } 224 | if message == nil { 225 | cg.logger.WithFields(logrus.Fields{ 226 | "group": pc.group, 227 | "topic": pc.topic, 228 | "partition": pc.partition, 229 | "offset": pc.offset, 230 | }).Error("Sarama partition consumer encounter error, the consumer would be exited") 231 | cg.stop() 232 | break PARTITION_CONSUMER_LOOP 233 | } 234 | select { 235 | case messageChan <- message: 236 | pc.offset = message.Offset + 1 237 | case <-cg.stopCh: 238 | break PARTITION_CONSUMER_LOOP 239 | } 240 | } 241 | } 242 | } 243 | 244 | func (pc *partitionConsumer) autoCommitOffset() { 245 | cg := pc.owner.owner 246 | defer cg.callRecover() 247 | timer := time.NewTimer(cg.config.OffsetAutoCommitInterval) 248 | for { 249 | select { 250 | case <-cg.stopCh: 251 | return 252 | case <-timer.C: 253 | err := pc.commitOffset() 254 | if err != nil { 255 | cg.logger.WithFields(logrus.Fields{ 256 | "topic": pc.topic, 257 | "partition": pc.partition, 258 | "offset": pc.offset, 259 | "err": err, 260 | }).Error("Failed to auto-commit the partition's offset") 261 | } 262 | timer.Reset(cg.config.OffsetAutoCommitInterval) 263 | } 264 | } 265 | } 266 | 267 | func (pc *partitionConsumer) commitOffset() error { 268 | cg := pc.owner.owner 269 | offset := pc.offset 270 | if pc.prevOffset == offset { 271 | return nil 272 | } 273 | err := cg.storage.commitOffset(pc.group, pc.topic, pc.partition, offset) 274 | if err != nil { 275 | return err 276 | } 277 | pc.prevOffset = offset 278 | return nil 279 | } 280 | 281 | func (pc *partitionConsumer) restart() { 282 | cg := pc.owner.owner 283 | cg.logger.WithFields(logrus.Fields{ 284 | "group": pc.group, 285 | "topic": pc.topic, 286 | "partition": pc.partition, 287 | }).Infof("Restart partition consumer while the offset out of range") 288 | err := pc.consumer.Close() 289 | if err != nil { 290 | cg.logger.WithFields(logrus.Fields{ 291 | "group": pc.group, 292 | "topic": pc.topic, 293 | "partition": pc.partition, 294 | }).Error("Stop consumer group because the old partition consumer cannot be closed") 295 | cg.stop() 296 | return 297 | } 298 | pc.consumer, err = cg.getPartitionConsumer(pc.topic, pc.partition, sarama.OffsetOldest) 299 | if err != nil { 300 | cg.logger.WithFields(logrus.Fields{ 301 | "group": pc.group, 302 | "topic": pc.topic, 303 | "partition": pc.partition, 304 | "err": err, 305 | }).Error("Stop consumer group because the new partition consumer cannot be start") 306 | cg.stop() 307 | } 308 | } 309 | 310 | func (pc *partitionConsumer) getOffset() map[string]interface{} { 311 | offset := make(map[string]interface{}) 312 | offset["offset"] = pc.offset 313 | offset["prev_offset"] = pc.prevOffset 314 | return offset 315 | } 316 | -------------------------------------------------------------------------------- /tests/docker/docker-compose.yml: -------------------------------------------------------------------------------- 1 | version: '2' 2 | services: 3 | zookeeper: 4 | image: wurstmeister/zookeeper:3.4.6 5 | ports: 6 | - "2181:2181" 7 | kafka: 8 | image: wurstmeister/kafka:0.10.2.1 9 | ports: 10 | - "9092:9092" 11 | links: 12 | - zookeeper 13 | environment: 14 | KAFKA_BROKER_ID: 1 15 | # create topic with 10 partition, 1 replica 16 | KAFKA_CREATE_TOPICS: "test:10:1" 17 | KAFKA_ADVERTISED_HOST_NAME: 127.0.0.1 18 | KAFKA_ZOOKEEPER_CONNECT: zookeeper:2181 19 | KAFKA_AUTO_CREATE_TOPICS_ENABLE: 'false' 20 | volumes: 21 | - /var/run/docker.sock:/var/run/docker.sock 22 | -------------------------------------------------------------------------------- /tests/docker/setup.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | docker-compose -p go_consumergroup_test up --build -d 3 | 4 | docker-compose -p go_consumergroup_test exec -T zookeeper sh -c 'until nc -z 127.0.0.1 2181; do echo "zk is not ready"; sleep 1; done' 5 | docker-compose -p go_consumergroup_test exec -T kafka sh -c 'until nc -z 127.0.0.1 9092; do echo "kafka is not ready"; sleep 1; done' 6 | docker-compose -p go_consumergroup_test exec -T kafka sh -c 'until kafka-topics.sh --list --zookeeper zookeeper|grep test|grep -v grep; do echo "kafka topic is not ready"; sleep 1; done' 7 | -------------------------------------------------------------------------------- /tests/docker/teardown.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | docker-compose -p go_consumergroup_test down -v 3 | -------------------------------------------------------------------------------- /tests/integration/init.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | var ( 4 | zookeepers = []string{"127.0.0.1:2181"} 5 | brokers = []string{"127.0.0.1:9092"} 6 | topic = "test" 7 | ) 8 | -------------------------------------------------------------------------------- /tests/integration/offset_test.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "strconv" 6 | "testing" 7 | "time" 8 | 9 | "github.com/Shopify/sarama" 10 | "github.com/meitu/zk_wrapper" 11 | ) 12 | 13 | func TestOffsetAutoCommit(t *testing.T) { 14 | count := 10 15 | group := genRandomGroupID(10) 16 | c, err := createConsumerInstance(zookeepers, group, topic) 17 | if err != nil { 18 | t.Fatalf("Failed to create consumer instance, err %s", err) 19 | } 20 | defer c.Stop() 21 | time.Sleep(3000 * time.Millisecond) // we have no way to know if the consumer is ready 22 | produceMessages(brokers, topic, 0, count) 23 | messages, _ := c.GetMessages(topic) 24 | go func(mgsCh <-chan *sarama.ConsumerMessage) { 25 | for message := range messages { 26 | fmt.Println(message) 27 | } 28 | }(messages) 29 | time.Sleep(200 * time.Millisecond) // offset auto commit interval is 100ms 30 | 31 | zkCli, _, err := zk_wrapper.Connect(zookeepers, 6*time.Second) 32 | if err != nil { 33 | t.Fatal("Failed to connect zookeeper") 34 | } 35 | offsetPath := fmt.Sprintf("/consumers/%s/offsets/%s/%d", group, topic, 0) 36 | data, _, err := zkCli.Get(offsetPath) 37 | if err != nil { 38 | t.Fatalf("Failed to get partition offset, err %s", err) 39 | } 40 | offset, _ := strconv.Atoi(string(data)) 41 | if offset != count { 42 | t.Errorf("Auto commit offset expect %d, but got %d", count, offset) 43 | } 44 | } 45 | -------------------------------------------------------------------------------- /tests/integration/rebalance_test.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "testing" 6 | "time" 7 | 8 | consumergroup "github.com/meitu/go-consumergroup" 9 | 10 | "github.com/Shopify/sarama" 11 | "github.com/meitu/zk_wrapper" 12 | ) 13 | 14 | func TestRebalance(t *testing.T) { 15 | group := genRandomGroupID(10) 16 | consumers := make([]*consumergroup.ConsumerGroup, 0) 17 | for i := 0; i < 3; i++ { 18 | go func() { 19 | c, err := createConsumerInstance(zookeepers, group, topic) 20 | if err != nil { 21 | t.Errorf("Failed to create consumer instance, err %s", err) 22 | } 23 | consumers = append(consumers, c) 24 | }() 25 | } 26 | time.Sleep(3 * time.Second) // we have no way to know if the consumer is ready 27 | 28 | kafkaCli, _ := sarama.NewClient(brokers, nil) 29 | partitions, err := kafkaCli.Partitions(topic) 30 | if err != nil { 31 | t.Errorf("Failed to get partitons, err %s", err) 32 | } 33 | owners := make([]string, 0) 34 | zkCli, _, err := zk_wrapper.Connect(zookeepers, 6*time.Second) 35 | if err != nil { 36 | t.Fatal("Failed to connect zookeeper") 37 | } 38 | for i := 0; i < len(partitions); i++ { 39 | ownerPath := fmt.Sprintf("/consumers/%s/owners/%s/%d", group, topic, i) 40 | data, _, err := zkCli.Get(ownerPath) 41 | if err != nil { 42 | t.Errorf("Failed to get partition owner, err %s", err) 43 | } 44 | owners = append(owners, string(data)) 45 | } 46 | if len(owners) != len(partitions) { 47 | t.Errorf("Missing owner in some partitions expected %d, but got %d", 48 | len(partitions), len(owners)) 49 | } 50 | for i := 1; i < len(owners); i++ { 51 | if owners[i] == owners[i-1] { 52 | t.Fatal("Partition owner should be difference while consumer > 1") 53 | } 54 | } 55 | 56 | for _, c := range consumers { 57 | c.Stop() 58 | // no way to exit group when the consumer is rebalancing 59 | time.Sleep(3 * time.Second) 60 | } 61 | } 62 | -------------------------------------------------------------------------------- /tests/integration/util.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "math/rand" 6 | "time" 7 | 8 | consumergroup "github.com/meitu/go-consumergroup" 9 | 10 | "github.com/Shopify/sarama" 11 | ) 12 | 13 | func genRandomGroupID(n int) string { 14 | const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" 15 | suffix := "" 16 | for i := 0; i < n; i++ { 17 | suffix = suffix + string(letterBytes[rand.Int()%len(letterBytes)]) 18 | } 19 | return "integration-test-group-id-" + suffix 20 | } 21 | 22 | func produceMessages(addrs []string, topic string, partition int32, count int) { 23 | conf := sarama.NewConfig() 24 | conf.Producer.Return.Successes = true 25 | conf.Producer.Partitioner = sarama.NewManualPartitioner 26 | client, _ := sarama.NewClient(addrs, conf) 27 | producer, err := sarama.NewSyncProducerFromClient(client) 28 | if err != nil { 29 | panic(fmt.Sprintf("Failed to create producer, err %s", err)) 30 | } 31 | for i := 0; i < count; i++ { 32 | _, _, err := producer.SendMessage(&sarama.ProducerMessage{ 33 | Topic: topic, 34 | Partition: partition, 35 | Value: sarama.StringEncoder("test-value"), 36 | }) 37 | if err != nil { 38 | panic(fmt.Sprintf("Failed to send message, err %s", err)) 39 | } 40 | } 41 | } 42 | 43 | func createConsumerInstance(addrs []string, groupID, topic string) (*consumergroup.ConsumerGroup, error) { 44 | conf := consumergroup.NewConfig() 45 | conf.ZkList = addrs 46 | conf.ZkSessionTimeout = 6 * time.Second 47 | conf.TopicList = []string{topic} 48 | conf.GroupID = groupID 49 | conf.OffsetAutoCommitInterval = 100 * time.Millisecond 50 | 51 | cg, err := consumergroup.NewConsumerGroup(conf) 52 | if err != nil { 53 | return nil, err 54 | } 55 | err = cg.Start() 56 | if err != nil { 57 | return nil, err 58 | } 59 | return cg, nil 60 | } 61 | -------------------------------------------------------------------------------- /tests/test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | docker_path="`pwd`/docker" 4 | integration_path="`pwd`/integration" 5 | unittest_path="`pwd`/.." 6 | 7 | cd $docker_path && sh setup.sh 8 | # run unittest 9 | cd $unittest_path && go test . 10 | # run intergration test 11 | cd $integration_path && go test . 12 | cd $docker_path && sh teardown.sh 13 | -------------------------------------------------------------------------------- /topic.go: -------------------------------------------------------------------------------- 1 | package consumergroup 2 | 3 | import ( 4 | "errors" 5 | "sync" 6 | 7 | "github.com/Shopify/sarama" 8 | "github.com/sirupsen/logrus" 9 | ) 10 | 11 | type topicConsumer struct { 12 | group string 13 | name string 14 | owner *ConsumerGroup 15 | errors chan *sarama.ConsumerError 16 | messages chan *sarama.ConsumerMessage 17 | partitionConsumers map[int32]*partitionConsumer 18 | wg sync.WaitGroup 19 | } 20 | 21 | func newTopicConsumer(owner *ConsumerGroup, topic string) *topicConsumer { 22 | tc := new(topicConsumer) 23 | tc.owner = owner 24 | tc.group = owner.name 25 | tc.name = topic 26 | tc.errors = make(chan *sarama.ConsumerError) 27 | tc.messages = make(chan *sarama.ConsumerMessage) 28 | return tc 29 | } 30 | 31 | func (tc *topicConsumer) start() { 32 | 33 | cg := tc.owner 34 | topic := tc.name 35 | 36 | cg.logger.WithFields(logrus.Fields{ 37 | "group": tc.group, 38 | "topic": topic, 39 | }).Info("Start the topic consumer") 40 | 41 | partitions, err := tc.assignPartitions() 42 | if err != nil { 43 | cg.logger.WithFields(logrus.Fields{ 44 | "group": tc.group, 45 | "topic": topic, 46 | "err": err, 47 | }).Error("Failed to assign partitions to topic consumer") 48 | return 49 | } 50 | 51 | cg.logger.WithFields(logrus.Fields{ 52 | "group": tc.group, 53 | "topic": topic, 54 | "partitions": partitions, 55 | }).Info("The partitions was assigned to current topic consumer") 56 | tc.partitionConsumers = make(map[int32]*partitionConsumer) 57 | for _, partition := range partitions { 58 | tc.partitionConsumers[partition] = newPartitionConsumer(tc, partition) 59 | } 60 | for partition, consumer := range tc.partitionConsumers { 61 | tc.wg.Add(1) 62 | go func(pc *partitionConsumer) { 63 | defer cg.callRecover() 64 | defer tc.wg.Done() 65 | pc.start() 66 | }(consumer) 67 | cg.logger.WithFields(logrus.Fields{ 68 | "group": tc.group, 69 | "topic": topic, 70 | "partition": partition, 71 | }).Info("Topic consumer start to consume the partition") 72 | } 73 | } 74 | 75 | func (tc *topicConsumer) assignPartitions() ([]int32, error) { 76 | var partitions []int32 77 | 78 | cg := tc.owner 79 | partNum, err := tc.getPartitionNum() 80 | if err != nil || partNum == 0 { 81 | return nil, err 82 | } 83 | consumerList, err := cg.storage.getConsumerList(cg.name) 84 | if err != nil { 85 | return nil, err 86 | } 87 | consumerNum := len(consumerList) 88 | if consumerNum == 0 { 89 | return nil, errors.New("no consumer was found") 90 | } 91 | for i := int32(0); i < partNum; i++ { 92 | id := consumerList[i%int32(consumerNum)] 93 | cg.owners[tc.name][i] = id 94 | if id == cg.id { 95 | partitions = append(partitions, i) 96 | } 97 | } 98 | return partitions, nil 99 | } 100 | 101 | func (tc *topicConsumer) getPartitionNum() (int32, error) { 102 | saramaConsumer, ok := tc.owner.saramaConsumers[tc.name] 103 | if !ok { 104 | return 0, errors.New("sarama conumser was not found") 105 | } 106 | partitions, err := saramaConsumer.Partitions(tc.name) 107 | if err != nil { 108 | return 0, err 109 | } 110 | return int32(len(partitions)), nil 111 | } 112 | 113 | func (tc *topicConsumer) getOffsets() map[int32]interface{} { 114 | partitions := make(map[int32]interface{}) 115 | for partition, pc := range tc.partitionConsumers { 116 | partitions[partition] = pc.getOffset() 117 | } 118 | return partitions 119 | } 120 | -------------------------------------------------------------------------------- /util.go: -------------------------------------------------------------------------------- 1 | package consumergroup 2 | 3 | import ( 4 | "fmt" 5 | "math/rand" 6 | "os" 7 | "path" 8 | "sort" 9 | "time" 10 | 11 | "github.com/meitu/go-zookeeper/zk" 12 | "github.com/meitu/zk_wrapper" 13 | ) 14 | 15 | func sliceRemoveDuplicates(slice []string) []string { 16 | sort.Strings(slice) 17 | i := 0 18 | var j int 19 | for { 20 | if i >= len(slice)-1 { 21 | break 22 | } 23 | for j = i + 1; j < len(slice) && slice[i] == slice[j]; j++ { 24 | } 25 | slice = append(slice[:i+1], slice[j:]...) 26 | i++ 27 | } 28 | return slice 29 | } 30 | 31 | func genConsumerID() string { 32 | name, err := os.Hostname() 33 | if err != nil { 34 | name = "unknown" 35 | } 36 | currentMilliSec := time.Now().UnixNano() / int64(time.Millisecond) 37 | randBytes := make([]byte, 8) 38 | for i := 0; i < 8; i++ { 39 | randBytes[i] = byte(rand.Intn(26) + 65) 40 | } 41 | return fmt.Sprintf("%s-%d-%s", name, currentMilliSec, string(randBytes)) 42 | } 43 | 44 | func mkdirRecursive(c *zk_wrapper.Conn, zkPath string) error { 45 | var err error 46 | parent := path.Dir(zkPath) 47 | if parent != "/" { 48 | if err = mkdirRecursive(c, parent); err != nil { 49 | return err 50 | } 51 | } 52 | _, err = c.Create(zkPath, nil, 0, zk.WorldACL(zk.PermAll)) 53 | if err == zk.ErrNodeExists { 54 | err = nil 55 | } 56 | return err 57 | } 58 | 59 | func zkCreateEphemeralPath(c *zk_wrapper.Conn, zkPath string, data []byte) error { 60 | return zkCreateRecursive(c, zkPath, zk.FlagEphemeral, data) 61 | } 62 | 63 | func zkCreatePersistentPath(c *zk_wrapper.Conn, zkPath string, data []byte) error { 64 | return zkCreateRecursive(c, zkPath, 0, data) 65 | } 66 | 67 | func zkCreateRecursive(c *zk_wrapper.Conn, zkPath string, flags int32, data []byte) error { 68 | _, err := c.Create(zkPath, data, flags, zk.WorldACL(zk.PermAll)) 69 | if err != nil && err != zk.ErrNoNode { 70 | return err 71 | } 72 | if err == zk.ErrNoNode { 73 | mkdirRecursive(c, path.Dir(zkPath)) 74 | _, err = c.Create(zkPath, data, flags, zk.WorldACL(zk.PermAll)) 75 | } 76 | return err 77 | } 78 | 79 | func zkSetPersistentPath(c *zk_wrapper.Conn, zkPath string, data []byte) error { 80 | _, err := c.Set(zkPath, data, -1) 81 | if err != nil && err != zk.ErrNoNode { 82 | return err 83 | } 84 | if err == zk.ErrNoNode { 85 | mkdirRecursive(c, path.Dir(zkPath)) 86 | _, err = c.Create(zkPath, data, 0, zk.WorldACL(zk.PermAll)) 87 | } 88 | return err 89 | } 90 | -------------------------------------------------------------------------------- /util_test.go: -------------------------------------------------------------------------------- 1 | package consumergroup 2 | 3 | import ( 4 | "path" 5 | "reflect" 6 | "sort" 7 | "testing" 8 | "time" 9 | 10 | "github.com/meitu/go-zookeeper/zk" 11 | "github.com/meitu/zk_wrapper" 12 | ) 13 | 14 | const ( 15 | testPath = "/test1/test2/test3/test4" 16 | ) 17 | 18 | func TestUtilSliceRemoveDuplicates(t *testing.T) { 19 | slice := []string{"hello", "world", "i", "have", "i", "i", "i", "world", "have"} 20 | slice = sliceRemoveDuplicates(slice) 21 | s := []string{"have", "i", "world", "hello"} 22 | sort.Strings(s) 23 | sort.Strings(slice) 24 | if !reflect.DeepEqual(s, slice) { 25 | t.Error("slice remove duplicates failed") 26 | } 27 | } 28 | 29 | func deleteRecursive(c *zk_wrapper.Conn, zkPath string) error { 30 | p := zkPath 31 | for p != "/" { 32 | err := c.Delete(p, -1) 33 | if err != nil { 34 | return err 35 | } 36 | p = path.Dir(p) 37 | } 38 | return nil 39 | } 40 | 41 | func TestUtilMkdirRecursive(t *testing.T) { 42 | client, _, err := zk_wrapper.Connect([]string{"127.0.0.1:2181"}, time.Duration(6)*time.Second) 43 | if err != nil { 44 | t.Fatal(err) 45 | } 46 | err = mkdirRecursive(client, testPath) 47 | if err != nil { 48 | t.Fatal(err) 49 | } 50 | 51 | isExist, _, err := client.Exists(testPath) 52 | if err != nil { 53 | t.Fatal(err) 54 | } 55 | if !isExist { 56 | t.Fatal("make directory recursive failed, this path is not exist") 57 | } 58 | 59 | err = mkdirRecursive(client, testPath) 60 | if err != nil { 61 | if err == zk.ErrNodeExists { 62 | t.Error("expected function mkdirRecursive can ignore make directory repeatedly, but it didn't") 63 | } else { 64 | t.Error(err) 65 | } 66 | } 67 | 68 | err = deleteRecursive(client, testPath) 69 | if err != nil { 70 | t.Error("detele directory recursive failed, please delete by zk client") 71 | } 72 | } 73 | -------------------------------------------------------------------------------- /zk_group_storage.go: -------------------------------------------------------------------------------- 1 | package consumergroup 2 | 3 | import ( 4 | "encoding/json" 5 | "errors" 6 | "fmt" 7 | "sort" 8 | "strconv" 9 | "time" 10 | 11 | "github.com/meitu/go-zookeeper/zk" 12 | "github.com/meitu/zk_wrapper" 13 | ) 14 | 15 | // constants defining the fixed path format. 16 | const ( 17 | ownerPath = "/consumers/%s/owners/%s/%d" 18 | consumersDir = "/consumers/%s/ids" 19 | consumersPath = "/consumers/%s/ids/%s" 20 | offsetsPath = "/consumers/%s/offsets/%s/%d" 21 | brokersDir = "/brokers/ids" 22 | brokersPath = "/brokers/ids/%s" 23 | topicsPath = "/brokers/topics/%s" 24 | ) 25 | 26 | type zkGroupStorage struct { 27 | chroot string 28 | serverList []string 29 | client *zk_wrapper.Conn 30 | sessionTimeout time.Duration 31 | } 32 | 33 | var ( 34 | errInvalidGroup = errors.New("Invalid group") 35 | errInvalidTopic = errors.New("Invalid topic") 36 | errInvalidConsumerID = errors.New("Invalid consumer ID") 37 | errInvalidPartition = "Invalid partition %s" 38 | ) 39 | 40 | func newZKGroupStorage(serverList []string, sessionTimeout time.Duration) *zkGroupStorage { 41 | s := new(zkGroupStorage) 42 | if sessionTimeout <= 0 { 43 | sessionTimeout = 6 * time.Second 44 | } 45 | s.serverList = serverList 46 | s.sessionTimeout = sessionTimeout 47 | return s 48 | } 49 | 50 | func (s *zkGroupStorage) Chroot(chroot string) { 51 | s.chroot = chroot 52 | } 53 | 54 | // getClient returns a zookeeper connetion. 55 | func (s *zkGroupStorage) getClient() (*zk_wrapper.Conn, error) { 56 | var err error 57 | if s.client == nil { 58 | s.client, _, err = zk_wrapper.Connect(s.serverList, s.sessionTimeout) 59 | if s.client != nil && s.chroot != "" { 60 | if err = s.client.Chroot(s.chroot); err != nil { 61 | return nil, err 62 | } 63 | } 64 | } 65 | return s.client, err 66 | } 67 | 68 | func (s *zkGroupStorage) claimPartition(group, topic string, partition int32, consumerID string) error { 69 | if group == "" { 70 | return errInvalidGroup 71 | } 72 | if topic == "" { 73 | return errInvalidTopic 74 | } 75 | if consumerID == "" { 76 | return errInvalidConsumerID 77 | } 78 | if partition < 0 { 79 | return fmt.Errorf(errInvalidPartition, partition) 80 | } 81 | 82 | c, err := s.getClient() 83 | if err != nil { 84 | return err 85 | } 86 | zkPath := fmt.Sprintf(ownerPath, group, topic, partition) 87 | err = zkCreateEphemeralPath(c, zkPath, []byte(consumerID)) 88 | return err 89 | } 90 | 91 | func (s *zkGroupStorage) releasePartition(group, topic string, partition int32) error { 92 | if group == "" { 93 | return errInvalidGroup 94 | } 95 | if topic == "" { 96 | return errInvalidTopic 97 | } 98 | if partition < 0 { 99 | return fmt.Errorf(errInvalidPartition, partition) 100 | } 101 | 102 | c, err := s.getClient() 103 | if err != nil { 104 | return err 105 | } 106 | zkPath := fmt.Sprintf(ownerPath, group, topic, partition) 107 | err = c.Delete(zkPath, -1) 108 | return err 109 | } 110 | 111 | func (s *zkGroupStorage) getPartitionOwner(group, topic string, partition int32) (string, error) { 112 | if group == "" { 113 | return "", errInvalidGroup 114 | } 115 | if topic == "" { 116 | return "", errInvalidTopic 117 | } 118 | if partition < 0 { 119 | return "", fmt.Errorf(errInvalidPartition, partition) 120 | } 121 | 122 | c, err := s.getClient() 123 | if err != nil { 124 | return "", err 125 | } 126 | zkPath := fmt.Sprintf(ownerPath, group, topic, partition) 127 | value, _, err := c.Get(zkPath) 128 | if err != nil { 129 | return "", err 130 | } 131 | return string(value), nil 132 | } 133 | 134 | func (s *zkGroupStorage) registerConsumer(group, consumerID string, data []byte) error { 135 | if group == "" { 136 | return errInvalidGroup 137 | } 138 | if consumerID == "" { 139 | return errInvalidConsumerID 140 | } 141 | 142 | c, err := s.getClient() 143 | if err != nil { 144 | return err 145 | } 146 | zkPath := fmt.Sprintf(consumersPath, group, consumerID) 147 | err = zkCreateEphemeralPath(c, zkPath, data) 148 | return err 149 | } 150 | 151 | func (s *zkGroupStorage) existsConsumer(group, consumerID string) (bool, error) { 152 | if group == "" { 153 | return false, errInvalidGroup 154 | } 155 | if consumerID == "" { 156 | return false, errInvalidConsumerID 157 | } 158 | 159 | c, err := s.getClient() 160 | if err != nil { 161 | return false, err 162 | } 163 | zkPath := fmt.Sprintf(consumersPath, group, consumerID) 164 | exist, _, err := c.Exists(zkPath) 165 | return exist, err 166 | } 167 | 168 | func (s *zkGroupStorage) deleteConsumer(group, consumerID string) error { 169 | if group == "" { 170 | return errInvalidGroup 171 | } 172 | if consumerID == "" { 173 | return errInvalidConsumerID 174 | } 175 | 176 | c, err := s.getClient() 177 | if err != nil { 178 | return err 179 | } 180 | zkPath := fmt.Sprintf(consumersPath, group, consumerID) 181 | err = c.Delete(zkPath, -1) 182 | return err 183 | } 184 | 185 | func (s *zkGroupStorage) watchConsumerList(group string) (*zk.Watcher, error) { 186 | if group == "" { 187 | return nil, errInvalidGroup 188 | } 189 | 190 | c, err := s.getClient() 191 | if err != nil { 192 | return nil, err 193 | } 194 | 195 | zkPath := fmt.Sprintf(consumersDir, group) 196 | _, _, w, err := c.ChildrenW(zkPath) 197 | if err != nil { 198 | return nil, err 199 | } 200 | return w, nil 201 | } 202 | 203 | func (s *zkGroupStorage) watchTopic(topic string) (*zk.Watcher, error) { 204 | if topic == "" { 205 | return nil, errInvalidTopic 206 | } 207 | 208 | c, err := s.getClient() 209 | if err != nil { 210 | return nil, err 211 | } 212 | zkPath := fmt.Sprintf(topicsPath, topic) 213 | _, _, w, err := c.GetW(zkPath) 214 | if err != nil { 215 | return nil, err 216 | } 217 | return w, nil 218 | } 219 | 220 | func (s *zkGroupStorage) getBrokerList() ([]string, error) { 221 | var brokerList []string 222 | type broker struct { 223 | Host string 224 | Port int 225 | } 226 | var b broker 227 | 228 | c, err := s.getClient() 229 | if err != nil { 230 | return nil, err 231 | } 232 | 233 | idList, _, err := c.Children(brokersDir) 234 | if err != nil { 235 | return nil, err 236 | } 237 | 238 | for _, id := range idList { 239 | zkPath := fmt.Sprintf(brokersPath, id) 240 | value, _, err := c.Get(zkPath) 241 | if err != nil { 242 | return nil, err 243 | } 244 | err = json.Unmarshal(value, &b) 245 | if err != nil { 246 | return nil, err 247 | } 248 | brokerList = append(brokerList, fmt.Sprintf("%s:%d", b.Host, b.Port)) 249 | } 250 | return brokerList, nil 251 | } 252 | 253 | func (s *zkGroupStorage) getConsumerList(group string) ([]string, error) { 254 | if group == "" { 255 | return nil, errInvalidGroup 256 | } 257 | 258 | c, err := s.getClient() 259 | if err != nil { 260 | return nil, err 261 | } 262 | 263 | zkPath := fmt.Sprintf(consumersDir, group) 264 | consumerList, _, err := c.Children(zkPath) 265 | if err != nil { 266 | return nil, err 267 | } 268 | sort.Strings(consumerList) 269 | return consumerList, nil 270 | } 271 | 272 | func (s *zkGroupStorage) commitOffset(group, topic string, partition int32, offset int64) error { 273 | if group == "" { 274 | return errInvalidGroup 275 | } 276 | if topic == "" { 277 | return errInvalidTopic 278 | } 279 | if partition < 0 { 280 | return fmt.Errorf(errInvalidPartition, partition) 281 | } 282 | 283 | c, err := s.getClient() 284 | if err != nil { 285 | return err 286 | } 287 | data := []byte(strconv.FormatInt(offset, 10)) 288 | zkPath := fmt.Sprintf(offsetsPath, group, topic, partition) 289 | err = zkSetPersistentPath(c, zkPath, data) 290 | return err 291 | } 292 | 293 | func (s *zkGroupStorage) getOffset(group, topic string, partition int32) (int64, error) { 294 | if group == "" { 295 | return -1, errInvalidGroup 296 | } 297 | if topic == "" { 298 | return -1, errInvalidTopic 299 | } 300 | if partition < 0 { 301 | return -1, fmt.Errorf(errInvalidPartition, partition) 302 | } 303 | 304 | c, err := s.getClient() 305 | if err != nil { 306 | return -1, err 307 | } 308 | zkPath := fmt.Sprintf(offsetsPath, group, topic, partition) 309 | value, _, err := c.Get(zkPath) 310 | if err != nil { 311 | if err != zk.ErrNoNode { 312 | return -1, err 313 | } 314 | return -1, nil 315 | } 316 | return strconv.ParseInt(string(value), 10, 64) 317 | } 318 | 319 | func (s *zkGroupStorage) removeWatcher(w *zk.Watcher) bool { 320 | c, err := s.getClient() 321 | if err != nil { 322 | return false 323 | } 324 | return c.RemoveWatcher(w) 325 | } 326 | 327 | func (s *zkGroupStorage) getPartitionsNum(topic string) (int, error) { 328 | type meta struct { 329 | Partitions map[int][]int 330 | } 331 | c, err := s.getClient() 332 | if err != nil { 333 | return 0, err 334 | } 335 | zkPath := fmt.Sprintf(topicsPath, topic) 336 | metadata, _, err := c.Get(zkPath) 337 | if err != nil { 338 | return 0, err 339 | } 340 | var m meta 341 | err = json.Unmarshal([]byte(metadata), &m) 342 | if err != nil { 343 | return 0, err 344 | } 345 | return len(m.Partitions), nil 346 | } 347 | -------------------------------------------------------------------------------- /zk_group_storage_test.go: -------------------------------------------------------------------------------- 1 | package consumergroup 2 | 3 | import ( 4 | "fmt" 5 | "math/rand" 6 | "testing" 7 | "time" 8 | ) 9 | 10 | const ( 11 | testValue = "go_test_value" 12 | testTopic = "go_test_topic" 13 | testGroup = "go_test_group" 14 | testConsumerID = "go_test_consumer_id" 15 | ) 16 | 17 | func TestZKGroupStorageClaimAndGetAndReleasePartition(t *testing.T) { 18 | zk := newZKGroupStorage([]string{"127.0.0.1:2181"}, 6*time.Second) 19 | 20 | err := zk.claimPartition(testGroup, testTopic, 0, testConsumerID) 21 | if err != nil { 22 | t.Error(err) 23 | } 24 | 25 | err = zk.releasePartition(testGroup, testTopic, 0) 26 | if err != nil { 27 | t.Error(err) 28 | } 29 | 30 | zk.claimPartition(testGroup, testTopic, 0, testConsumerID) 31 | err = zk.claimPartition(testGroup, testTopic, 0, testConsumerID) 32 | if err == nil { 33 | zk.releasePartition(testGroup, testTopic, 0) 34 | t.Error("Expected it can't claim a partition twice, but it did") 35 | } 36 | 37 | cid, err := zk.getPartitionOwner(testGroup, testTopic, 0) 38 | if err != nil { 39 | zk.releasePartition(testGroup, testTopic, 0) 40 | t.Error("get partition owner failed, because: ", err) 41 | } 42 | if cid != testConsumerID { 43 | zk.releasePartition(testGroup, testTopic, 0) 44 | t.Error("partition owner get from zookeeper isn't unexpected") 45 | } 46 | 47 | zk.releasePartition(testGroup, testTopic, 0) 48 | } 49 | 50 | func TestZKGroupStorageRegisterAndGetAndDeleteConsumer(t *testing.T) { 51 | zk := newZKGroupStorage([]string{"127.0.0.1:2181"}, 6*time.Second) 52 | 53 | err := zk.registerConsumer(testGroup, testConsumerID, nil) 54 | if err != nil { 55 | t.Fatal(err) 56 | } 57 | 58 | err = zk.deleteConsumer(testGroup, testConsumerID) 59 | if err != nil { 60 | t.Fatal(err) 61 | } 62 | 63 | zk.registerConsumer(testGroup, testConsumerID, nil) 64 | err = zk.registerConsumer(testGroup, testConsumerID, nil) 65 | if err == nil { 66 | zk.deleteConsumer(testGroup, testConsumerID) 67 | t.Fatal("Expected it can't register consumer twice, but it did") 68 | } 69 | 70 | consumerList, err := zk.getConsumerList(testGroup) 71 | if err != nil { 72 | t.Fatal(err) 73 | } 74 | 75 | if consumerList[0] != testConsumerID { 76 | zk.deleteConsumer(testGroup, testConsumerID) 77 | t.Fatal("consumer id get from zookeeper isn't expected") 78 | } 79 | zk.deleteConsumer(testGroup, testConsumerID) 80 | } 81 | 82 | func TestZKGroupWatchConsumerList(t *testing.T) { 83 | zk := newZKGroupStorage([]string{"127.0.0.1:2181"}, 6*time.Second) 84 | 85 | consumer1 := fmt.Sprintf("%s-%d", testConsumerID, rand.Int()) 86 | consumer2 := fmt.Sprintf("%s-%d", testConsumerID, rand.Int()) 87 | consumer3 := fmt.Sprintf("%s-%d", testConsumerID, rand.Int()) 88 | consumerList := []string{consumer1, consumer2, consumer3} 89 | for _, consumer := range consumerList { 90 | zk.registerConsumer(testGroup, consumer, nil) 91 | } 92 | 93 | watcher, err := zk.watchConsumerList(testGroup) 94 | if err != nil { 95 | t.Error(err) 96 | } 97 | 98 | select { 99 | case <-watcher.EvCh: 100 | t.Error("channel receive message before consumer list change") 101 | default: 102 | } 103 | 104 | zk.deleteConsumer(testGroup, consumer1) 105 | 106 | select { 107 | case <-watcher.EvCh: 108 | default: 109 | t.Error("channel can't receive message after consumer list change") 110 | } 111 | 112 | for _, consumer := range consumerList { 113 | zk.deleteConsumer(testGroup, consumer) 114 | } 115 | } 116 | 117 | func TestZKGroupStorageCommitAndGetOffset(t *testing.T) { 118 | zk := newZKGroupStorage([]string{"127.0.0.1:2181"}, 6*time.Second) 119 | testOffset := rand.Int63() 120 | 121 | err := zk.commitOffset(testGroup, testTopic, 0, testOffset) 122 | if err != nil { 123 | t.Error(err) 124 | } 125 | 126 | offset, err := zk.getOffset(testGroup, testTopic, 0) 127 | if err != nil { 128 | t.Error(err) 129 | } 130 | 131 | if offset != testOffset { 132 | t.Error("offset get from zookeeper isn't unexpected") 133 | } 134 | 135 | err = zk.commitOffset(testGroup, testTopic, 0, testOffset+1) 136 | if err != nil { 137 | t.Error(err) 138 | } 139 | } 140 | --------------------------------------------------------------------------------