├── .codecov.yml ├── .gitignore ├── .travis.yml ├── CHANGELOG.md ├── Gopkg.lock ├── Gopkg.toml ├── LICENSE ├── Makefile ├── README.md ├── client.go ├── client_test.go ├── cmd └── tools │ └── copyright │ └── licensegen.go ├── consumerBuilder.go ├── consumerBuilder_test.go ├── consumerOptions.go ├── consumerOptions_test.go ├── docs └── DEVELOPMENT-GUIDE.md ├── glide.lock ├── glide.yaml ├── internal ├── backoff │ ├── retry.go │ ├── retry_test.go │ ├── retrypolicy.go │ └── retrypolicy_test.go ├── consumer │ ├── ackMgr.go │ ├── clusterConsumer.go │ ├── clusterConsumer_test.go │ ├── dlq.go │ ├── dlqMetadata.pb.go │ ├── dlqMetadata.proto │ ├── dlq_test.go │ ├── limit.go │ ├── message.go │ ├── message_test.go │ ├── mocks_test.go │ ├── multiClusterConsumer_test.go │ ├── multiclusterConsumer.go │ ├── options.go │ ├── partitionConsumer.go │ ├── partitionConsumer_test.go │ ├── topicConsumer.go │ ├── types.go │ └── types_test.go ├── list │ ├── list.go │ └── list_test.go ├── metrics │ └── defs.go └── util │ ├── lifecycle.go │ ├── lifecycle_test.go │ └── testutil.go ├── kafka ├── config.go ├── config_test.go ├── interfaces.go ├── resolver.go └── resolver_test.go ├── scripts └── cover.sh └── version /.codecov.yml: -------------------------------------------------------------------------------- 1 | coverage: 2 | range: 80..100 3 | round: down 4 | precision: 2 5 | 6 | status: 7 | project: # measuring the overall project coverage 8 | default: # context, you can create multiple ones with custom titles 9 | enabled: yes # must be yes|true to enable this status 10 | target: 50% # specify the target coverage for each commit status 11 | # option: "auto" (must increase from parent commit or pull request base) 12 | # option: "X%" a static target percentage to hit 13 | if_not_found: success # if parent is not found report status as success, error, or failure 14 | if_ci_failed: error # if ci fails report status as success, error, or failure 15 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Folders 2 | vendor 3 | *.prof 4 | *.pprof 5 | *.out 6 | *.log 7 | .idea 8 | .okbuck 9 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: go 2 | sudo: false 3 | go: 4 | - 1.11 5 | go_import_path: github.com/uber-go/kafka-client 6 | cache: 7 | directories: 8 | - vendor 9 | install: 10 | - make dependencies 11 | script: 12 | - make lint 13 | - make test 14 | - make bench 15 | after_success: 16 | - make cover 17 | - bash <(curl -s https://codecov.io/bash) 18 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Changelog 2 | 3 | v0.2.3 (unreleased) 4 | ------------------- 5 | 6 | - Nothing changed yet. 7 | 8 | 9 | v0.2.2 (2019-05-08) 10 | ------------------- 11 | 12 | - Handle retryQ and DLQ only consumer topics. 13 | - Update KafkaPartitionOwned metric to be a boolean metric for ownership. 14 | - Rename number of partitions owned by specific worker to KafkaPartitionOwnedCount. 15 | - Add client id option. 16 | 17 | 18 | v0.2.1 (2018-08-07) 19 | ------------------- 20 | 21 | - Add method to nack to DLQ. 22 | 23 | 24 | v0.2.0 (2018-07-06) 25 | ------------------- 26 | 27 | - Tune producer max message bytes to 10mb. 28 | - Tune consumer default fetch bytes to 10mb. 29 | - Remove Offset.Initial.Reset config since it is unused. 30 | - Add Offset.Commit.Enabled config to enable auto offset commit. 31 | - Tune partition consumer logs to debug level 32 | - Remove Topic.BrokerList since it was unused in favor of NameResolver to resolve broker list. 33 | 34 | 35 | v0.1.8 (2018-06-04) 36 | ------------------- 37 | 38 | - Change sarama log to warn level. 39 | 40 | 41 | v0.1.7 (2018-05-01) 42 | ------------------- 43 | 44 | - Add metric for freshness 45 | 46 | 47 | v0.1.6 (2018-04-30) 48 | ------------------- 49 | 50 | - Allow RetryCount = -1 to signal infinite retry. 51 | - Fix off by one error for offset-lag metric. 52 | 53 | v0.1.5 (2018-04-11) 54 | ------------------- 55 | 56 | - Fix reset of rangePartitionConsumer with existing reset does not trigger new merge. 57 | - Update sarama config version to use 0.10.2. 58 | 59 | 60 | v0.1.4 (2018-03-31) 61 | ------------------- 62 | 63 | - Fix DLQMetadata decoding to use DLQMetadataDecoder func instead of inferred decoding from TopicType. 64 | - Fix consumer to use noopDLQ if RetryQ or DLQ in config is empty. 65 | - Fix ResetOffset fails on partition rebalance. 66 | - Add delay to Topic configuration 67 | 68 | 69 | v0.1.3 (2018-03-09) 70 | ------------------- 71 | 72 | - Add WithRetryTopics and WithDLQTopics to inject additional consumers for additional retry or DLQ topics. 73 | 74 | 75 | v0.1.2 (2018-03-07) 76 | ------------------- 77 | 78 | - Pin sarama-cluster to 2.1.13. 79 | 80 | 81 | v0.1.1 (2018-03-05) 82 | ------------------- 83 | 84 | - Fixed sarama-cluster dependency pin to cf455bc755fe41ac9bb2861e7a961833d9c2ecc3 because we need ResetOffsets method with NPE fix. 85 | 86 | 87 | v0.1.0 (2018-03-05) 88 | ------------------- 89 | 90 | - Added initial release 91 | -------------------------------------------------------------------------------- /Gopkg.lock: -------------------------------------------------------------------------------- 1 | # This file is autogenerated, do not edit; changes may be undone by the next 'dep ensure'. 2 | 3 | 4 | [[projects]] 5 | digest = "1:1e283d7680f5395cddb88a422783c8877ff60c28fc1d8fde27c808f1720c512e" 6 | name = "github.com/Shopify/sarama" 7 | packages = ["."] 8 | pruneopts = "UT" 9 | revision = "f7be6aa2bc7b2e38edf816b08b582782194a1c02" 10 | version = "v1.16.0" 11 | 12 | [[projects]] 13 | digest = "1:4fdffd1724c105db8c394019cfc2444fd23466be04812850506437361ee5de55" 14 | name = "github.com/bsm/sarama-cluster" 15 | packages = ["."] 16 | pruneopts = "UT" 17 | revision = "cf455bc755fe41ac9bb2861e7a961833d9c2ecc3" 18 | version = "v2.1.13" 19 | 20 | [[projects]] 21 | digest = "1:ffe9824d294da03b391f44e1ae8281281b4afc1bdaa9588c9097785e3af10cec" 22 | name = "github.com/davecgh/go-spew" 23 | packages = ["spew"] 24 | pruneopts = "UT" 25 | revision = "8991bc29aa16c548c550c7ff78260e27b9ab7c73" 26 | version = "v1.1.1" 27 | 28 | [[projects]] 29 | digest = "1:1f0c7ab489b407a7f8f9ad16c25a504d28ab461517a971d341388a56156c1bd7" 30 | name = "github.com/eapache/go-resiliency" 31 | packages = ["breaker"] 32 | pruneopts = "UT" 33 | revision = "ef9aaa7ea8bd2448429af1a77cf41b2b3b34bdd6" 34 | 35 | [[projects]] 36 | digest = "1:2b55f08306f24f60dab51e006c1890db08d34ccc1b92c53a08a1a6e7792e3eb8" 37 | name = "github.com/eapache/go-xerial-snappy" 38 | packages = ["."] 39 | pruneopts = "UT" 40 | revision = "bb955e01b9346ac19dc29eb16586c90ded99a98c" 41 | 42 | [[projects]] 43 | branch = "master" 44 | digest = "1:42b78fd006fab08021f1006fdb781b15494ada3c44c6a0f3fdc9f0f73e6d4b5e" 45 | name = "github.com/eapache/queue" 46 | packages = ["."] 47 | pruneopts = "UT" 48 | revision = "093482f3f8ce946c05bcba64badd2c82369e084d" 49 | 50 | [[projects]] 51 | branch = "master" 52 | digest = "1:5fe3f6ede1c208a2efd3b78fe4df0306aa9624edd39476143d14f0326e5a8d29" 53 | name = "github.com/facebookgo/clock" 54 | packages = ["."] 55 | pruneopts = "UT" 56 | revision = "600d898af40aa09a7a93ecb9265d87b0504b6f03" 57 | 58 | [[projects]] 59 | digest = "1:ffc060c551980d37ee9e428ef528ee2813137249ccebb0bfc412ef83071cac91" 60 | name = "github.com/golang/protobuf" 61 | packages = ["proto"] 62 | pruneopts = "UT" 63 | revision = "925541529c1fa6821df4e44ce2723319eb2be768" 64 | version = "v1.0.0" 65 | 66 | [[projects]] 67 | digest = "1:29a5ab9fa9e845fd8e8726f31b187d710afd271ef1eb32085fe3d604b7e06382" 68 | name = "github.com/golang/snappy" 69 | packages = ["."] 70 | pruneopts = "UT" 71 | revision = "553a641470496b2327abcac10b36396bd98e45c9" 72 | 73 | [[projects]] 74 | digest = "1:34833843589f17bb48234a8cdd36b1038aed85a09d60e30fc73891b7d89b6c1d" 75 | name = "github.com/pierrec/lz4" 76 | packages = ["."] 77 | pruneopts = "UT" 78 | revision = "ed8d4cc3b461464e69798080a0092bd028910298" 79 | 80 | [[projects]] 81 | branch = "master" 82 | digest = "1:5272d5bcf91d1db46fc9528f706019533f1b390b94b122cff427162bc5009e40" 83 | name = "github.com/pierrec/xxHash" 84 | packages = ["xxHash32"] 85 | pruneopts = "UT" 86 | revision = "a0006b13c722f7f12368c00a3d3c2ae8a999a0c6" 87 | 88 | [[projects]] 89 | digest = "1:0028cb19b2e4c3112225cd871870f2d9cf49b9b4276531f03438a88e94be86fe" 90 | name = "github.com/pmezard/go-difflib" 91 | packages = ["difflib"] 92 | pruneopts = "UT" 93 | revision = "792786c7400a136282c1664665ae0a8db921c6c2" 94 | version = "v1.0.0" 95 | 96 | [[projects]] 97 | digest = "1:565770be35414989a808a4c2fd74a9b0b1d48ca0b6c615f2f5908ab99bf77eee" 98 | name = "github.com/rcrowley/go-metrics" 99 | packages = ["."] 100 | pruneopts = "UT" 101 | revision = "8732c616f52954686704c8645fe1a9d59e9df7c1" 102 | 103 | [[projects]] 104 | digest = "1:c52d48fe752736ae45b72959e1e4fc1a5b60d5b56f8fad42c58414daab5f309d" 105 | name = "github.com/stretchr/testify" 106 | packages = [ 107 | "assert", 108 | "require", 109 | "suite", 110 | ] 111 | pruneopts = "UT" 112 | revision = "12b6f73e6084dad08a7c6e575284b177ecafbc71" 113 | version = "v1.2.1" 114 | 115 | [[projects]] 116 | digest = "1:84e66ba70567c8e4563aee2ffaef251458cba0de67ae61f9c01deb1755c2ffeb" 117 | name = "github.com/uber-go/tally" 118 | packages = ["."] 119 | pruneopts = "UT" 120 | revision = "522328b48efad0c6034dba92bf39228694e9d31f" 121 | version = "v3.3.2" 122 | 123 | [[projects]] 124 | digest = "1:9d5cf6f23377cc24f00ba9a7b0c83fe070171c094e42eece262c6252a392ec33" 125 | name = "go.uber.org/atomic" 126 | packages = ["."] 127 | pruneopts = "UT" 128 | revision = "8474b86a5a6f79c443ce4b2992817ff32cf208b8" 129 | version = "v1.3.1" 130 | 131 | [[projects]] 132 | digest = "1:60bf2a5e347af463c42ed31a493d817f8a72f102543060ed992754e689805d1a" 133 | name = "go.uber.org/multierr" 134 | packages = ["."] 135 | pruneopts = "UT" 136 | revision = "3c4937480c32f4c13a875a1829af76c98ca3d40a" 137 | version = "v1.1.0" 138 | 139 | [[projects]] 140 | digest = "1:0d59cfb5fb7d699ef8442a0d049935f6d0034aa001ca8328d048a814e1661035" 141 | name = "go.uber.org/zap" 142 | packages = [ 143 | ".", 144 | "buffer", 145 | "internal/bufferpool", 146 | "internal/color", 147 | "internal/exit", 148 | "zapcore", 149 | ] 150 | pruneopts = "UT" 151 | revision = "35aad584952c3e7020db7b839f6b102de6271f89" 152 | version = "v1.7.1" 153 | 154 | [solve-meta] 155 | analyzer-name = "dep" 156 | analyzer-version = 1 157 | input-imports = [ 158 | "github.com/Shopify/sarama", 159 | "github.com/bsm/sarama-cluster", 160 | "github.com/golang/protobuf/proto", 161 | "github.com/stretchr/testify/assert", 162 | "github.com/stretchr/testify/require", 163 | "github.com/stretchr/testify/suite", 164 | "github.com/uber-go/tally", 165 | "go.uber.org/zap", 166 | "go.uber.org/zap/zapcore", 167 | ] 168 | solver-name = "gps-cdcl" 169 | solver-version = 1 170 | -------------------------------------------------------------------------------- /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.0.0" 31 | 32 | [[constraint]] 33 | name = "github.com/bsm/sarama-cluster" 34 | version = "~2.1.13" 35 | 36 | [[constraint]] 37 | name = "github.com/golang/protobuf" 38 | version = "1.0.0" 39 | 40 | [[constraint]] 41 | name = "github.com/stretchr/testify" 42 | version = "1.2.1" 43 | 44 | [[constraint]] 45 | name = "github.com/uber-go/tally" 46 | version = "3.0.0" 47 | 48 | [[constraint]] 49 | name = "go.uber.org/zap" 50 | version = "1.0.0" 51 | 52 | [prune] 53 | go-tests = true 54 | unused-packages = true 55 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2017 Uber Technologies, Inc. 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy 4 | of this software and associated documentation files (the "Software"), to deal 5 | in the Software without restriction, including without limitation the rights 6 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | copies of the Software, and to permit persons to whom the Software is 8 | furnished to do so, subject to the following conditions: 9 | 10 | The above copyright notice and this permission notice shall be included in 11 | all copies or substantial portions of the Software. 12 | 13 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | THE SOFTWARE. -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | export GO15VENDOREXPERIMENT=1 2 | 3 | BENCH_FLAGS ?= -cpuprofile=cpu.pprof -memprofile=mem.pprof -benchmem 4 | PKGS ?= ./cmd/... ./internal/... ./kafka/... . 5 | # Many Go tools take file globs or directories as arguments instead of packages. 6 | PKG_FILES ?= *.go kafka internal/consumer internal/backoff internal/list internal/metrics internal/util 7 | 8 | # The linting tools evolve with each Go version, so run them only on the latest 9 | # stable release. 10 | GO_VERSION := $(shell go version | cut -d " " -f 3) 11 | GO_MINOR_VERSION := $(word 2,$(subst ., ,$(GO_VERSION))) 12 | LINTABLE_MINOR_VERSIONS := 9 13 | ifneq ($(filter $(LINTABLE_MINOR_VERSIONS),$(GO_MINOR_VERSION)),) 14 | SHOULD_LINT := true 15 | endif 16 | 17 | 18 | .PHONY: all 19 | all: lint test 20 | 21 | .PHONY: dependencies 22 | dependencies: 23 | @echo "Installing Dep and locked dependencies..." 24 | command -v dep || curl https://raw.githubusercontent.com/golang/dep/master/install.sh | sh 25 | dep ensure 26 | @echo "Installing test dependencies..." 27 | go get -u -f github.com/axw/gocov/gocov 28 | go get -u -f github.com/mattn/goveralls 29 | ifdef SHOULD_LINT 30 | @echo "Installing golint..." 31 | go get -u -f github.com/golang/lint/golint 32 | else 33 | @echo "Not installing golint, since we don't expect to lint on" $(GO_VERSION) 34 | endif 35 | 36 | # Disable printf-like invocation checking due to testify.assert.Error() 37 | VET_RULES := -printf=false 38 | 39 | .PHONY: lint 40 | lint: dependencies 41 | ifdef SHOULD_LINT 42 | @rm -rf lint.log 43 | @echo "Checking formatting..." 44 | @gofmt -d -s $(PKG_FILES) 2>&1 | tee lint.log 45 | @echo "Installing test dependencies for vet..." 46 | @go test -i $(PKGS) 47 | @echo "Checking vet..." 48 | @$(foreach dir,$(PKG_FILES),go tool vet $(VET_RULES) $(dir) 2>&1 | tee -a lint.log;) 49 | @echo "Checking lint..." 50 | @$(foreach dir,$(PKGS),golint $(dir) | grep -v ".pb.go" 2>&1 | tee -a lint.log;) 51 | @echo "Checking for unresolved FIXMEs..." 52 | @git grep -i fixme | grep -v -e vendor -e Makefile | tee -a lint.log 53 | @echo "Checking for license headers..." 54 | @go run ./cmd/tools/copyright/licensegen.go --verifyOnly | tee -a lint.log 55 | @[ ! -s lint.log ] 56 | else 57 | @echo "Skipping linters on" $(GO_VERSION) 58 | endif 59 | 60 | .PHONY: test 61 | test: dependencies 62 | go test $(PKGS) 63 | 64 | .PHONY: cover 65 | cover: dependencies 66 | ./scripts/cover.sh $(PKGS) 67 | 68 | .PHONY: bench 69 | BENCH ?= . 70 | bench: dependencies 71 | @$(foreach pkg,$(PKGS),go test -bench=$(BENCH) -run="^$$" $(BENCH_FLAGS) $(pkg);) 72 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Go Kafka Client Library [![Mit License][mit-img]][mit] [![Build Status][ci-img]][ci] [![Coverage Status][cov-img]][cov] 2 | 3 | A high level Go client library for Apache Kafka that provides the following primitives on top of [sarama-cluster](https://github.com/bsm/sarama-cluster): 4 | 5 | * Competing consumer semantics with dead letter queue (DLQ) 6 | * Ability to process messages across multiple goroutines 7 | * Ability to Ack or Nack messages out of order (with optional DLQ) 8 | * Ability to consume from topics spread across different kafka clusters 9 | 10 | ## Stability 11 | 12 | This library is in alpha. APIs are subject to change, use at your own risk 13 | 14 | ## Contributing 15 | If you are interested in contributing, please sign the [License Agreement](https://cla-assistant.io/uber-go/kafka-client) and see our [development guide](https://github.com/uber-go/kafka-client/blob/master/docs/DEVELOPMENT-GUIDE.md) 16 | 17 | ## Installation 18 | 19 | `go get -u github.com/uber-go/kafka-client` 20 | 21 | ## Quick Start 22 | 23 | ```go 24 | package main 25 | 26 | import ( 27 | "os" 28 | "os/signal" 29 | 30 | "github.com/uber-go/kafka-client" 31 | "github.com/uber-go/kafka-client/kafka" 32 | "github.com/uber-go/tally" 33 | "go.uber.org/zap" 34 | ) 35 | 36 | func main() { 37 | // mapping from cluster name to list of broker ip addresses 38 | brokers := map[string][]string{ 39 | "sample_cluster": []string{"127.0.0.1:9092"}, 40 | "sample_dlq_cluster": []string{"127.0.0.1:9092"}, 41 | } 42 | // mapping from topic name to cluster that has that topic 43 | topicClusterAssignment := map[string][]string{ 44 | "sample_topic": []string{"sample_cluster"}, 45 | } 46 | 47 | // First create the kafkaclient, its the entry point for creating consumers or producers 48 | // It takes as input a name resolver that knows how to map topic names to broker ip addrs 49 | client := kafkaclient.New(kafka.NewStaticNameResolver(topicClusterAssignment, brokers), zap.NewNop(), tally.NoopScope) 50 | 51 | // Next, setup the consumer config for consuming from a set of topics 52 | config := &kafka.ConsumerConfig{ 53 | TopicList: kafka.ConsumerTopicList{ 54 | kafka.ConsumerTopic{ // Consumer Topic is a combination of topic + dead-letter-queue 55 | Topic: kafka.Topic{ // Each topic is a tuple of (name, clusterName) 56 | Name: "sample_topic", 57 | Cluster: "sample_cluster", 58 | }, 59 | DLQ: kafka.Topic{ 60 | Name: "sample_consumer_dlq", 61 | Cluster: "sample_dlq_cluster", 62 | }, 63 | }, 64 | }, 65 | GroupName: "sample_consumer", 66 | Concurrency: 100, // number of go routines processing messages in parallel 67 | } 68 | 69 | // Create the consumer through the previously created client 70 | consumer, err := client.NewConsumer(config) 71 | if err != nil { 72 | panic(err) 73 | } 74 | 75 | // Finally, start consuming 76 | if err := consumer.Start(); err != nil { 77 | panic(err) 78 | } 79 | 80 | sigCh := make(chan os.Signal, 1) 81 | signal.Notify(sigCh, os.Interrupt) 82 | 83 | for { 84 | select { 85 | case msg, ok := <-consumer.Messages(): 86 | if !ok { 87 | return // channel closed 88 | } 89 | if err := process(msg); err != nil { 90 | msg.Nack() 91 | } else { 92 | msg.Ack() 93 | } 94 | case <-sigCh: 95 | consumer.Stop() 96 | <-consumer.Closed() 97 | } 98 | } 99 | } 100 | ``` 101 | 102 | [mit-img]: http://img.shields.io/badge/License-MIT-blue.svg 103 | [mit]: https://github.com/uber-go/kafka-client/blob/master/LICENSE 104 | 105 | [ci-img]: https://img.shields.io/travis/uber-go/kafka-client/master.svg 106 | [ci]: https://travis-ci.org/uber-go/kafka-client/branches 107 | 108 | [cov-img]: https://codecov.io/gh/uber-go/kafka-client/branch/master/graph/badge.svg 109 | [cov]: https://codecov.io/gh/uber-go/kafka-client/branch/master 110 | -------------------------------------------------------------------------------- /client.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package kafkaclient 22 | 23 | import ( 24 | "os" 25 | 26 | "github.com/uber-go/kafka-client/kafka" 27 | "github.com/uber-go/tally" 28 | "go.uber.org/zap" 29 | ) 30 | 31 | type ( 32 | // Client refers to the kafka client. Serves as 33 | // the entry point to producing or consuming 34 | // messages from kafka 35 | Client interface { 36 | // NewConsumer returns a new instance of kafka consumer. 37 | NewConsumer(config *kafka.ConsumerConfig, consumerOpts ...ConsumerOption) (kafka.Consumer, error) 38 | } 39 | 40 | client struct { 41 | tally tally.Scope 42 | logger *zap.Logger 43 | resolver kafka.NameResolver 44 | } 45 | ) 46 | 47 | // New returns a new kafka client. 48 | func New(resolver kafka.NameResolver, logger *zap.Logger, scope tally.Scope) Client { 49 | return &client{ 50 | resolver: resolver, 51 | logger: logger, 52 | tally: scope, 53 | } 54 | } 55 | 56 | // NewConsumer returns a new instance of kafka consumer. 57 | // 58 | // It is possible for NewConsumer to start a consumer which consumes from a subset of topics if EnablePartialConsumption, 59 | // ConsumerOption is used. 60 | // If partial consumption is enabled, error will not be returned. 61 | func (c *client) NewConsumer(config *kafka.ConsumerConfig, consumerOpts ...ConsumerOption) (kafka.Consumer, error) { 62 | return newConsumerBuilder(config, c.resolver, c.tally, c.logger, consumerOpts...).build() 63 | } 64 | 65 | func clientID() string { 66 | name, err := os.Hostname() 67 | if err != nil { 68 | name = "unknown-kafka-client" 69 | } 70 | return name 71 | } 72 | -------------------------------------------------------------------------------- /client_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package kafkaclient 22 | 23 | import ( 24 | "testing" 25 | "time" 26 | 27 | "github.com/Shopify/sarama" 28 | cluster "github.com/bsm/sarama-cluster" 29 | "github.com/stretchr/testify/assert" 30 | "github.com/uber-go/kafka-client/internal/consumer" 31 | "github.com/uber-go/kafka-client/kafka" 32 | "github.com/uber-go/tally" 33 | "go.uber.org/zap" 34 | ) 35 | 36 | type ( 37 | resolverMock struct { 38 | errs map[string]error 39 | clusterToIP map[string][]string 40 | } 41 | 42 | saramaConsumerConstructorMock struct { 43 | errRet map[string]error 44 | } 45 | 46 | saramaProducerConstructorMock struct { 47 | errRet map[string]error 48 | } 49 | 50 | clusterConsumerConstructorMock struct { 51 | errRet map[string]error 52 | } 53 | ) 54 | 55 | func (m *saramaConsumerConstructorMock) f(brokers []string, _ string, _ []string, _ *cluster.Config) (sc consumer.SaramaConsumer, err error) { 56 | key := "" 57 | for _, broker := range brokers { 58 | key += broker 59 | } 60 | 61 | err = m.errRet[key] 62 | return 63 | } 64 | 65 | func (m *saramaProducerConstructorMock) f(brokers []string) (p sarama.SyncProducer, err error) { 66 | key := "" 67 | for _, broker := range brokers { 68 | key += broker 69 | } 70 | 71 | err = m.errRet[key] 72 | return 73 | } 74 | 75 | func (m *resolverMock) ResolveIPForCluster(cluster string) (ip []string, err error) { 76 | err = m.errs[cluster] 77 | ip = m.clusterToIP[cluster] 78 | return 79 | } 80 | 81 | func (m *resolverMock) ResolveClusterForTopic(topic string) (cluster []string, err error) { 82 | err = m.errs[topic] 83 | return 84 | } 85 | 86 | func (m *clusterConsumerConstructorMock) f(_ string, cluster string, _ *consumer.Options, _ kafka.ConsumerTopicList, _ chan kafka.Message, _ consumer.SaramaConsumer, _ map[string]consumer.DLQ, _ tally.Scope, _ *zap.Logger) (kc kafka.Consumer, err error) { 87 | err = m.errRet[cluster] 88 | return 89 | } 90 | 91 | func TestBuildSaramaConfig(t *testing.T) { 92 | opts := &consumer.Options{ 93 | RcvBufferSize: 128, 94 | PartitionRcvBufferSize: 64, 95 | OffsetCommitInterval: time.Minute, 96 | RebalanceDwellTime: time.Hour, 97 | MaxProcessingTime: time.Second, 98 | OffsetPolicy: sarama.OffsetNewest, 99 | ConsumerMode: cluster.ConsumerModePartitions, 100 | } 101 | config := buildSaramaConfig(opts) 102 | assert.Equal(t, opts.PartitionRcvBufferSize, config.ChannelBufferSize) 103 | assert.Equal(t, opts.OffsetPolicy, config.Consumer.Offsets.Initial) 104 | assert.Equal(t, opts.OffsetCommitInterval, config.Consumer.Offsets.CommitInterval) 105 | assert.Equal(t, opts.MaxProcessingTime, config.Consumer.MaxProcessingTime) 106 | assert.True(t, config.Consumer.Return.Errors) 107 | assert.Equal(t, opts.RebalanceDwellTime, config.Group.Offsets.Synchronization.DwellTime) 108 | assert.True(t, config.Group.Return.Notifications) 109 | assert.Equal(t, cluster.ConsumerModePartitions, config.Group.Mode) 110 | } 111 | -------------------------------------------------------------------------------- /cmd/tools/copyright/licensegen.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package main 22 | 23 | import ( 24 | "bufio" 25 | "bytes" 26 | "flag" 27 | "fmt" 28 | "io" 29 | "io/ioutil" 30 | "os" 31 | "path/filepath" 32 | "strings" 33 | ) 34 | 35 | type ( 36 | // task that adds license header to source 37 | // files, if they don't already exist 38 | addLicenseHeaderTask struct { 39 | license string // license header string to add 40 | config *config // root directory of the project source 41 | } 42 | 43 | // command line config params 44 | config struct { 45 | rootDir string 46 | verifyOnly bool 47 | } 48 | ) 49 | 50 | // licenseFileName is the name of the license file 51 | const licenseFileName = "LICENSE" 52 | 53 | // unique prefix that identifies a license header 54 | const licenseHeaderPrefix = "// Copyright (c)" 55 | 56 | var ( 57 | // directories to be excluded 58 | dirBlacklist = []string{"vendor/"} 59 | // default perms for the newly created files 60 | defaultFilePerms = os.FileMode(0644) 61 | ) 62 | 63 | // command line utility that adds license header 64 | // to the source files. Usage as follows: 65 | // 66 | // ./cmd/tools/copyright/licensegen.go 67 | func main() { 68 | 69 | var cfg config 70 | flag.StringVar(&cfg.rootDir, "rootDir", ".", "project root directory") 71 | flag.BoolVar(&cfg.verifyOnly, "verifyOnly", false, 72 | "don't automatically add headers, just verify all files") 73 | flag.Parse() 74 | 75 | task := newAddLicenseHeaderTask(&cfg) 76 | if err := task.run(); err != nil { 77 | fmt.Println(err) 78 | os.Exit(-1) 79 | } 80 | } 81 | 82 | func newAddLicenseHeaderTask(cfg *config) *addLicenseHeaderTask { 83 | return &addLicenseHeaderTask{ 84 | config: cfg, 85 | } 86 | } 87 | 88 | func (task *addLicenseHeaderTask) run() error { 89 | file, err := os.Open(task.config.rootDir + "/" + licenseFileName) 90 | if err != nil { 91 | return err 92 | } 93 | 94 | var data bytes.Buffer 95 | scanner := bufio.NewScanner(file) 96 | for scanner.Scan() { 97 | data.WriteString("//") 98 | line := scanner.Text() 99 | if line != "" { 100 | data.WriteString(" " + scanner.Text()) 101 | } 102 | data.WriteString("\n") 103 | } 104 | data.WriteString("\n") 105 | 106 | task.license = data.String() 107 | 108 | err = filepath.Walk(task.config.rootDir, task.handleFile) 109 | if err != nil { 110 | return fmt.Errorf("copyright header check failed, err=%v", err.Error()) 111 | } 112 | return nil 113 | } 114 | 115 | func (task *addLicenseHeaderTask) handleFile(path string, fileInfo os.FileInfo, err error) error { 116 | 117 | if err != nil { 118 | return err 119 | } 120 | 121 | if fileInfo.IsDir() { 122 | return nil 123 | } 124 | 125 | if !mustProcessPath(path) { 126 | return nil 127 | } 128 | 129 | if !strings.HasSuffix(fileInfo.Name(), ".go") { 130 | return nil 131 | } 132 | 133 | f, err := os.Open(path) 134 | if err != nil { 135 | return err 136 | } 137 | 138 | buf := make([]byte, len(licenseHeaderPrefix)) 139 | _, err = io.ReadFull(f, buf) 140 | f.Close() 141 | 142 | if err != nil && !isEOF(err) { 143 | return err 144 | } 145 | 146 | if string(buf) == licenseHeaderPrefix { 147 | return nil // file already has the copyright header 148 | } 149 | 150 | // at this point, src file is missing the header 151 | if task.config.verifyOnly { 152 | if !isFileAutogenerated(path) { 153 | return fmt.Errorf("%v missing license header", path) 154 | } 155 | } 156 | 157 | data, err := ioutil.ReadFile(path) 158 | if err != nil { 159 | return err 160 | } 161 | 162 | return ioutil.WriteFile(path, []byte(task.license+string(data)), defaultFilePerms) 163 | } 164 | 165 | func isFileAutogenerated(path string) bool { 166 | return strings.HasPrefix(path, ".gen") 167 | } 168 | 169 | func mustProcessPath(path string) bool { 170 | for _, d := range dirBlacklist { 171 | if strings.HasPrefix(path, d) { 172 | return false 173 | } 174 | } 175 | return true 176 | } 177 | 178 | // returns true if the error type is an EOF 179 | func isEOF(err error) bool { 180 | return err == io.EOF || err == io.ErrUnexpectedEOF 181 | } 182 | -------------------------------------------------------------------------------- /consumerBuilder_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package kafkaclient 22 | 23 | import ( 24 | "sort" 25 | "testing" 26 | "time" 27 | 28 | "github.com/Shopify/sarama" 29 | "github.com/bsm/sarama-cluster" 30 | "github.com/stretchr/testify/suite" 31 | "github.com/uber-go/kafka-client/internal/consumer" 32 | "github.com/uber-go/kafka-client/kafka" 33 | "github.com/uber-go/tally" 34 | "go.uber.org/zap" 35 | ) 36 | 37 | type ( 38 | ConsumerBuilderTestSuite struct { 39 | suite.Suite 40 | builder *consumerBuilder 41 | config *kafka.ConsumerConfig 42 | resolver kafka.NameResolver 43 | mockSaramaClientConstructor *mockSaramaClientConstructor 44 | mockSaramaConsumerConstructor *mockSaramaConsumerConstructor 45 | mockSaramaProducerConstructor *mockSaramaProducerConstructor 46 | } 47 | 48 | mockSaramaConsumerConstructor struct { 49 | clusterTopicMap map[string][]string 50 | } 51 | mockSaramaClientConstructor struct { 52 | clusters map[string]bool 53 | } 54 | mockSaramaProducerConstructor struct{} 55 | ) 56 | 57 | func TestConsumerBuilderTestSuite(t *testing.T) { 58 | suite.Run(t, new(ConsumerBuilderTestSuite)) 59 | } 60 | 61 | func (s *ConsumerBuilderTestSuite) SetupTest() { 62 | s.config = kafka.NewConsumerConfig("consumergroup", []kafka.ConsumerTopic{ 63 | { 64 | Topic: kafka.Topic{ 65 | Name: "topic", 66 | Cluster: "cluster", 67 | }, 68 | RetryQ: kafka.Topic{ 69 | Name: "retry-topic", 70 | Cluster: "dlq-cluster", 71 | Delay: time.Microsecond, 72 | }, 73 | DLQ: kafka.Topic{ 74 | Name: "dlq-topic", 75 | Cluster: "dlq-cluster", 76 | }, 77 | MaxRetries: 1, 78 | }, 79 | { 80 | Topic: kafka.Topic{ 81 | Name: "topic1", 82 | Cluster: "cluster", 83 | }, 84 | RetryQ: kafka.Topic{}, 85 | DLQ: kafka.Topic{}, 86 | MaxRetries: 1, 87 | }, 88 | }) 89 | s.resolver = kafka.NewStaticNameResolver( 90 | map[string][]string{ 91 | "topic": {"cluster"}, 92 | "retry-topic": {"dlq-cluster"}, 93 | "dlq-topic": {"dlq-cluster"}, 94 | }, 95 | map[string][]string{ 96 | "cluster": {"broker1:9092"}, 97 | "dlq-cluster": {"broker2:9092"}, 98 | }, 99 | ) 100 | s.builder = newConsumerBuilder( 101 | s.config, 102 | s.resolver, 103 | tally.NoopScope, 104 | zap.NewNop(), 105 | WithClientID("test-id"), 106 | ) 107 | s.mockSaramaConsumerConstructor = &mockSaramaConsumerConstructor{ 108 | clusterTopicMap: make(map[string][]string), 109 | } 110 | s.mockSaramaClientConstructor = &mockSaramaClientConstructor{ 111 | clusters: make(map[string]bool), 112 | } 113 | s.mockSaramaProducerConstructor = &mockSaramaProducerConstructor{} 114 | s.builder.constructors = consumer.Constructors{ 115 | NewSaramaProducer: s.mockSaramaProducerConstructor.f, 116 | NewSaramaConsumer: s.mockSaramaConsumerConstructor.f, 117 | NewSaramaClient: s.mockSaramaClientConstructor.f, 118 | } 119 | } 120 | 121 | func (s *ConsumerBuilderTestSuite) TestBuild() { 122 | kafkaConsumer, err := s.builder.build() 123 | s.NoError(err) 124 | s.NotNil(kafkaConsumer) 125 | // 3 clusters used in consumer 126 | s.Equal([]string{"cluster", "dlq-cluster", "dlq-cluster"}, func() []string { 127 | output := make([]string, 0, 3) 128 | for cluster := range s.builder.clusterTopicsMap { 129 | output = append(output, cluster.name) 130 | } 131 | sort.Strings(output) 132 | return output 133 | }()) 134 | // 3 consumer group names: 2 default for original and retryQ consumer and one dlq-merger 135 | s.Equal([]string{"consumergroup", "consumergroup", "consumergroup-dlq-merger"}, func() []string { 136 | output := make([]string, 0, 3) 137 | for cluster := range s.builder.clusterTopicsMap { 138 | output = append(output, cluster.groupName) 139 | } 140 | sort.Strings(output) 141 | return output 142 | }()) 143 | // 2 sarama clients corresponding DLQ producers for 2 clusters. Group name is unused. 144 | s.Equal([]string{"dlq-cluster"}, func() []string { 145 | output := make([]string, 0, 2) 146 | for cluster := range s.builder.clusterSaramaClientMap { 147 | output = append(output, cluster.Cluster) 148 | } 149 | sort.Strings(output) 150 | return output 151 | }()) 152 | // consuming from 3 clusters 153 | s.Equal([]string{"cluster", "dlq-cluster", "dlq-cluster"}, func() []string { 154 | output := make([]string, 0, 3) 155 | for cluster := range s.builder.clusterSaramaConsumerMap { 156 | output = append(output, cluster.Cluster) 157 | } 158 | sort.Strings(output) 159 | return output 160 | }()) 161 | // consuming from 2 consumer group names 162 | s.Equal([]string{"consumergroup", "consumergroup", "consumergroup-dlq-merger"}, func() []string { 163 | output := make([]string, 0, 3) 164 | for cluster := range s.builder.clusterSaramaConsumerMap { 165 | output = append(output, cluster.Group) 166 | } 167 | sort.Strings(output) 168 | return output 169 | }()) 170 | // DLQ producers for 2 topics in 1 cluster 171 | s.Equal([]string{"dlq-topic", "retry-topic"}, func() []string { 172 | output := make([]string, 0, 2) 173 | for topic := range s.builder.clusterTopicSaramaProducerMap["dlq-cluster"] { 174 | output = append(output, topic) 175 | } 176 | sort.Strings(output) 177 | return output 178 | }()) 179 | // make sure retry topic with delay gets populated to the read topic list 180 | s.Equal([]string{"retry-topic"}, func() []string { 181 | output := make([]string, 0, 1) 182 | for _, topicList := range s.builder.clusterTopicsMap { 183 | for _, topic := range topicList { 184 | if topic.Delay > 0 { 185 | output = append(output, topic.Name) 186 | } 187 | } 188 | } 189 | sort.Strings(output) 190 | return output 191 | }()) 192 | // saramaClient constructor called for dlq-cluster = broker2:9092 193 | s.Equal([]string{"broker2:9092"}, func() []string { 194 | output := make([]string, 0, 1) 195 | for broker := range s.mockSaramaClientConstructor.clusters { 196 | output = append(output, broker) 197 | } 198 | sort.Strings(output) 199 | return output 200 | }()) 201 | // Sarama consumer constructor called for broker1:9092 and broker2:9092 clusters 202 | s.Equal([]string{"broker1:9092", "broker2:9092"}, func() []string { 203 | output := make([]string, 0, 2) 204 | for broker := range s.mockSaramaConsumerConstructor.clusterTopicMap { 205 | output = append(output, broker) 206 | } 207 | sort.Strings(output) 208 | return output 209 | }()) 210 | } 211 | 212 | func (s *ConsumerBuilderTestSuite) TestBuildConsumersWithCommitDisabled() { 213 | s.builder.kafkaConfig.Offsets.Commits.Enabled = false 214 | consumer, err := s.builder.build() 215 | s.NoError(err) 216 | s.NotNil(consumer) 217 | } 218 | 219 | func (m *mockSaramaConsumerConstructor) f(brokers []string, _ string, topics []string, _ *cluster.Config) (consumer.SaramaConsumer, error) { 220 | m.clusterTopicMap[brokers[0]] = topics 221 | return nil, nil 222 | } 223 | 224 | func (m *mockSaramaClientConstructor) f(brokers []string, _ *sarama.Config) (sarama.Client, error) { 225 | m.clusters[brokers[0]] = true 226 | return nil, nil 227 | } 228 | 229 | func (m *mockSaramaProducerConstructor) f(sarama.Client) (sarama.AsyncProducer, error) { 230 | return nil, nil 231 | } 232 | -------------------------------------------------------------------------------- /consumerOptions.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package kafkaclient 22 | 23 | import ( 24 | "github.com/uber-go/kafka-client/internal/consumer" 25 | "github.com/uber-go/kafka-client/kafka" 26 | ) 27 | 28 | type ( 29 | // ConsumerOption is the type for optional arguments to the NewConsumer constructor. 30 | ConsumerOption interface { 31 | apply(*consumer.Options) 32 | } 33 | 34 | dlqTopicsOptions struct { 35 | topicList kafka.ConsumerTopicList 36 | } 37 | 38 | retryTopicsOptions struct { 39 | topicList kafka.ConsumerTopicList 40 | } 41 | 42 | clientIDOptions struct { 43 | clientID string 44 | } 45 | ) 46 | 47 | // WithDLQTopics creates a range consumer for the specified consumer DLQ topics. 48 | func WithDLQTopics(topicList kafka.ConsumerTopicList) ConsumerOption { 49 | return &dlqTopicsOptions{ 50 | topicList: topicList, 51 | } 52 | } 53 | 54 | func (o *dlqTopicsOptions) apply(opts *consumer.Options) { 55 | for _, topic := range o.topicList { 56 | opts.OtherConsumerTopics = append(opts.OtherConsumerTopics, consumer.Topic{ 57 | ConsumerTopic: topic, 58 | DLQMetadataDecoder: consumer.ProtobufDLQMetadataDecoder, 59 | PartitionConsumerFactory: consumer.NewRangePartitionConsumer, 60 | ConsumerGroupSuffix: consumer.DLQConsumerGroupNameSuffix, 61 | }) 62 | } 63 | } 64 | 65 | // WithRetryTopics creates a consumer for the specified consumer Retry topics. 66 | func WithRetryTopics(topicList kafka.ConsumerTopicList) ConsumerOption { 67 | return &retryTopicsOptions{ 68 | topicList: topicList, 69 | } 70 | } 71 | 72 | func (o *retryTopicsOptions) apply(opts *consumer.Options) { 73 | for _, topic := range o.topicList { 74 | opts.OtherConsumerTopics = append(opts.OtherConsumerTopics, consumer.Topic{ 75 | ConsumerTopic: topic, 76 | DLQMetadataDecoder: consumer.ProtobufDLQMetadataDecoder, 77 | PartitionConsumerFactory: consumer.NewPartitionConsumer, 78 | }) 79 | } 80 | } 81 | 82 | // WithClientID sets client id. 83 | func WithClientID(clientID string) ConsumerOption { 84 | return &clientIDOptions{ 85 | clientID: clientID, 86 | } 87 | } 88 | 89 | func (o *clientIDOptions) apply(opts *consumer.Options) { 90 | opts.ClientID = o.clientID 91 | } 92 | -------------------------------------------------------------------------------- /consumerOptions_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package kafkaclient 22 | 23 | import ( 24 | "testing" 25 | 26 | "github.com/stretchr/testify/assert" 27 | "github.com/uber-go/kafka-client/internal/consumer" 28 | "github.com/uber-go/kafka-client/kafka" 29 | ) 30 | 31 | func TestDLQConsumerOptions(t *testing.T) { 32 | testTopics := make([]kafka.ConsumerTopic, 0, 10) 33 | testTopics = append(testTopics, *new(kafka.ConsumerTopic)) 34 | testTopics = append(testTopics, *new(kafka.ConsumerTopic)) 35 | consumerOption := WithDLQTopics(testTopics) 36 | options := consumer.DefaultOptions() 37 | consumerOption.apply(options) 38 | assert.Equal(t, 2, len(options.OtherConsumerTopics)) 39 | } 40 | 41 | func TestRetryConsumerOptions(t *testing.T) { 42 | testTopics := make([]kafka.ConsumerTopic, 0, 10) 43 | testTopics = append(testTopics, *new(kafka.ConsumerTopic)) 44 | testTopics = append(testTopics, *new(kafka.ConsumerTopic)) 45 | consumerOption := WithRetryTopics(testTopics) 46 | options := consumer.DefaultOptions() 47 | consumerOption.apply(options) 48 | assert.Equal(t, 2, len(options.OtherConsumerTopics)) 49 | } 50 | -------------------------------------------------------------------------------- /docs/DEVELOPMENT-GUIDE.md: -------------------------------------------------------------------------------- 1 | # Development Guide 2 | 3 | ## Contributions to the library 4 | 5 | We happily welcome your contributions. Following is the general process we follow for contributions 6 | 7 | * If you have an idea for a contribution, start with a Github issue to discuss your proposal. Once you receive 8 | feedback on your proposal, move onto the coding stage 9 | * Create a git fork of the repository and make your changes in your personal fork 10 | * Make sure the test coverage for your change is 80+% before you are ready 11 | * Run `make` and verify linter / tests pass without errors 12 | * Create a pull request against the master repo 13 | * All changes go through a code review and you must also sign the [Contributor License Agreement](https://cla-assistant.io/uber-go/kafka-client) 14 | 15 | 16 | ## Coding Style and Conventions 17 | 18 | We follow the standard Go coding conventions defined [here](https://github.com/golang/go/wiki/CodeReviewComments) 19 | 20 | -------------------------------------------------------------------------------- /glide.lock: -------------------------------------------------------------------------------- 1 | hash: d6d7552888db6c06b96621ff164834bd6b1f6e588f43c815c0c07e0a1a0cfff0 2 | updated: 2018-03-07T10:32:31.733217-08:00 3 | imports: 4 | - name: github.com/bsm/sarama-cluster 5 | version: cf455bc755fe41ac9bb2861e7a961833d9c2ecc3 6 | - name: github.com/davecgh/go-spew 7 | version: 8991bc29aa16c548c550c7ff78260e27b9ab7c73 8 | subpackages: 9 | - spew 10 | - name: github.com/eapache/go-resiliency 11 | version: ef9aaa7ea8bd2448429af1a77cf41b2b3b34bdd6 12 | subpackages: 13 | - breaker 14 | - name: github.com/eapache/go-xerial-snappy 15 | version: bb955e01b9346ac19dc29eb16586c90ded99a98c 16 | - name: github.com/eapache/queue 17 | version: 093482f3f8ce946c05bcba64badd2c82369e084d 18 | - name: github.com/facebookgo/clock 19 | version: 600d898af40aa09a7a93ecb9265d87b0504b6f03 20 | - name: github.com/golang/protobuf 21 | version: 925541529c1fa6821df4e44ce2723319eb2be768 22 | subpackages: 23 | - proto 24 | - name: github.com/golang/snappy 25 | version: 553a641470496b2327abcac10b36396bd98e45c9 26 | - name: github.com/pierrec/lz4 27 | version: ed8d4cc3b461464e69798080a0092bd028910298 28 | - name: github.com/pierrec/xxHash 29 | version: a0006b13c722f7f12368c00a3d3c2ae8a999a0c6 30 | subpackages: 31 | - xxHash32 32 | - name: github.com/rcrowley/go-metrics 33 | version: 8732c616f52954686704c8645fe1a9d59e9df7c1 34 | - name: github.com/Shopify/sarama 35 | version: f7be6aa2bc7b2e38edf816b08b582782194a1c02 36 | - name: github.com/uber-go/tally 37 | version: 522328b48efad0c6034dba92bf39228694e9d31f 38 | - name: go.uber.org/atomic 39 | version: 8474b86a5a6f79c443ce4b2992817ff32cf208b8 40 | - name: go.uber.org/multierr 41 | version: 3c4937480c32f4c13a875a1829af76c98ca3d40a 42 | - name: go.uber.org/zap 43 | version: 35aad584952c3e7020db7b839f6b102de6271f89 44 | subpackages: 45 | - buffer 46 | - internal/bufferpool 47 | - internal/color 48 | - internal/exit 49 | - zapcore 50 | testImports: 51 | - name: github.com/pmezard/go-difflib 52 | version: 792786c7400a136282c1664665ae0a8db921c6c2 53 | subpackages: 54 | - difflib 55 | - name: github.com/stretchr/testify 56 | version: 12b6f73e6084dad08a7c6e575284b177ecafbc71 57 | subpackages: 58 | - assert 59 | - require 60 | - suite 61 | -------------------------------------------------------------------------------- /glide.yaml: -------------------------------------------------------------------------------- 1 | package: github.com/uber-go/kafka-client 2 | license: MIT 3 | import: 4 | - package: go.uber.org/zap 5 | version: ^1 6 | - package: github.com/uber-go/tally 7 | version: ^3 8 | - package: github.com/Shopify/sarama 9 | version: ^1 10 | - package: github.com/bsm/sarama-cluster 11 | version: ~2.1.13 12 | - package: github.com/golang/protobuf 13 | version: ^1 14 | testImport: 15 | - package: github.com/stretchr/testify 16 | subpackages: 17 | - assert 18 | - require 19 | -------------------------------------------------------------------------------- /internal/backoff/retry.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package backoff 22 | 23 | import ( 24 | "time" 25 | ) 26 | 27 | type ( 28 | // Operation to retry 29 | Operation func() error 30 | 31 | // IsRetryable handler can be used to exclude certain errors during retry 32 | IsRetryable func(error) bool 33 | ) 34 | 35 | // Retry function can be used to wrap any call with retry logic using the passed in policy 36 | func Retry(operation Operation, policy RetryPolicy, isRetryable IsRetryable) error { 37 | var err error 38 | var next time.Duration 39 | 40 | r := NewRetrier(policy, SystemClock) 41 | for { 42 | // operation completed successfully. No need to retry. 43 | if err = operation(); err == nil { 44 | return nil 45 | } 46 | 47 | if next = r.NextBackOff(); next == done { 48 | return err 49 | } 50 | // Check if the error is retryable 51 | if isRetryable != nil && !isRetryable(err) { 52 | return err 53 | } 54 | 55 | time.Sleep(next) 56 | } 57 | } 58 | 59 | // IgnoreErrors can be used as IsRetryable handler for Retry function to exclude certain errors from the retry list 60 | func IgnoreErrors(errorsToExclude []error) func(error) bool { 61 | return func(err error) bool { 62 | for _, errorToExclude := range errorsToExclude { 63 | if err == errorToExclude { 64 | return false 65 | } 66 | } 67 | 68 | return true 69 | } 70 | } 71 | -------------------------------------------------------------------------------- /internal/backoff/retry_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package backoff 22 | 23 | import ( 24 | "testing" 25 | "time" 26 | 27 | "github.com/stretchr/testify/require" 28 | "github.com/stretchr/testify/suite" 29 | ) 30 | 31 | type ( 32 | RetrySuite struct { 33 | *require.Assertions // override suite.Suite.Assertions with require.Assertions; this means that s.NotNil(nil) will stop the test, not merely log an error 34 | suite.Suite 35 | } 36 | 37 | someError struct{} 38 | ) 39 | 40 | func TestRetrySuite(t *testing.T) { 41 | suite.Run(t, new(RetrySuite)) 42 | } 43 | 44 | func (s *RetrySuite) SetupTest() { 45 | s.Assertions = require.New(s.T()) // Have to define our overridden assertions in the test setup. If we did it earlier, s.T() will return nil 46 | } 47 | 48 | func (s *RetrySuite) TestRetrySuccess() { 49 | i := 0 50 | op := func() error { 51 | i++ 52 | 53 | if i == 5 { 54 | return nil 55 | } 56 | 57 | return &someError{} 58 | } 59 | 60 | policy := NewExponentialRetryPolicy(1 * time.Millisecond) 61 | policy.SetMaximumInterval(5 * time.Millisecond) 62 | policy.SetMaximumAttempts(10) 63 | 64 | err := Retry(op, policy, nil) 65 | s.NoError(err) 66 | s.Equal(5, i) 67 | } 68 | 69 | func (s *RetrySuite) TestRetryFailed() { 70 | i := 0 71 | op := func() error { 72 | i++ 73 | 74 | if i == 7 { 75 | return nil 76 | } 77 | 78 | return &someError{} 79 | } 80 | 81 | policy := NewExponentialRetryPolicy(1 * time.Millisecond) 82 | policy.SetMaximumInterval(5 * time.Millisecond) 83 | policy.SetMaximumAttempts(5) 84 | 85 | err := Retry(op, policy, nil) 86 | s.Error(err) 87 | } 88 | 89 | func (s *RetrySuite) TestIsRetryableSuccess() { 90 | i := 0 91 | op := func() error { 92 | i++ 93 | 94 | if i == 5 { 95 | return nil 96 | } 97 | 98 | return &someError{} 99 | } 100 | 101 | isRetryable := func(err error) bool { 102 | if _, ok := err.(*someError); ok { 103 | return true 104 | } 105 | 106 | return false 107 | } 108 | 109 | policy := NewExponentialRetryPolicy(1 * time.Millisecond) 110 | policy.SetMaximumInterval(5 * time.Millisecond) 111 | policy.SetMaximumAttempts(10) 112 | 113 | err := Retry(op, policy, isRetryable) 114 | s.NoError(err, "Retry count: %v", i) 115 | s.Equal(5, i) 116 | } 117 | 118 | func (s *RetrySuite) TestIsRetryableFailure() { 119 | i := 0 120 | op := func() error { 121 | i++ 122 | 123 | if i == 5 { 124 | return nil 125 | } 126 | 127 | return &someError{} 128 | } 129 | 130 | policy := NewExponentialRetryPolicy(1 * time.Millisecond) 131 | policy.SetMaximumInterval(5 * time.Millisecond) 132 | policy.SetMaximumAttempts(10) 133 | 134 | err := Retry(op, policy, IgnoreErrors([]error{&someError{}})) 135 | s.Error(err) 136 | s.Equal(1, i) 137 | } 138 | 139 | func (e *someError) Error() string { 140 | return "Some Error" 141 | } 142 | -------------------------------------------------------------------------------- /internal/backoff/retrypolicy.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package backoff 22 | 23 | import ( 24 | "math" 25 | "math/rand" 26 | "time" 27 | ) 28 | 29 | const ( 30 | // NoInterval represents Maximim interval 31 | NoInterval = 0 32 | done time.Duration = -1 33 | noMaximumAttempts = 0 34 | 35 | defaultBackoffCoefficient = 2.0 36 | defaultMaximumInterval = 10 * time.Second 37 | defaultExpirationInterval = time.Minute 38 | defaultMaximumAttempts = noMaximumAttempts 39 | ) 40 | 41 | type ( 42 | // RetryPolicy is the API which needs to be implemented by various retry policy implementations 43 | RetryPolicy interface { 44 | ComputeNextDelay(elapsedTime time.Duration, numAttempts int) time.Duration 45 | } 46 | 47 | // Retrier manages the state of retry operation 48 | Retrier interface { 49 | NextBackOff() time.Duration 50 | Reset() 51 | } 52 | 53 | // Clock used by ExponentialRetryPolicy implementation to get the current time. Mainly used for unit testing 54 | Clock interface { 55 | Now() time.Time 56 | } 57 | 58 | // ExponentialRetryPolicy provides the implementation for retry policy using a coefficient to compute the next delay. 59 | // Formula used to compute the next delay is: initialInterval * math.Pow(backoffCoefficient, currentAttempt) 60 | ExponentialRetryPolicy struct { 61 | initialInterval time.Duration 62 | backoffCoefficient float64 63 | maximumInterval time.Duration 64 | expirationInterval time.Duration 65 | maximumAttempts int 66 | } 67 | 68 | systemClock struct{} 69 | 70 | retrierImpl struct { 71 | policy RetryPolicy 72 | clock Clock 73 | currentAttempt int 74 | startTime time.Time 75 | } 76 | ) 77 | 78 | // SystemClock implements Clock interface that uses time.Now(). 79 | var SystemClock = systemClock{} 80 | 81 | // NewExponentialRetryPolicy returns an instance of ExponentialRetryPolicy using the provided initialInterval 82 | func NewExponentialRetryPolicy(initialInterval time.Duration) *ExponentialRetryPolicy { 83 | p := &ExponentialRetryPolicy{ 84 | initialInterval: initialInterval, 85 | backoffCoefficient: defaultBackoffCoefficient, 86 | maximumInterval: defaultMaximumInterval, 87 | expirationInterval: defaultExpirationInterval, 88 | maximumAttempts: defaultMaximumAttempts, 89 | } 90 | 91 | return p 92 | } 93 | 94 | // NewRetrier is used for creating a new instance of Retrier 95 | func NewRetrier(policy RetryPolicy, clock Clock) Retrier { 96 | return &retrierImpl{ 97 | policy: policy, 98 | clock: clock, 99 | startTime: clock.Now(), 100 | currentAttempt: 0, 101 | } 102 | } 103 | 104 | // SetInitialInterval sets the initial interval used by ExponentialRetryPolicy for the very first retry 105 | // All later retries are computed using the following formula: 106 | // initialInterval * math.Pow(backoffCoefficient, currentAttempt) 107 | func (p *ExponentialRetryPolicy) SetInitialInterval(initialInterval time.Duration) { 108 | p.initialInterval = initialInterval 109 | } 110 | 111 | // SetBackoffCoefficient sets the coefficient used by ExponentialRetryPolicy to compute next delay for each retry 112 | // All retries are computed using the following formula: 113 | // initialInterval * math.Pow(backoffCoefficient, currentAttempt) 114 | func (p *ExponentialRetryPolicy) SetBackoffCoefficient(backoffCoefficient float64) { 115 | p.backoffCoefficient = backoffCoefficient 116 | } 117 | 118 | // SetMaximumInterval sets the maximum interval for each retry 119 | func (p *ExponentialRetryPolicy) SetMaximumInterval(maximumInterval time.Duration) { 120 | p.maximumInterval = maximumInterval 121 | } 122 | 123 | // SetExpirationInterval sets the absolute expiration interval for all retries 124 | func (p *ExponentialRetryPolicy) SetExpirationInterval(expirationInterval time.Duration) { 125 | p.expirationInterval = expirationInterval 126 | } 127 | 128 | // SetMaximumAttempts sets the maximum number of retry attempts 129 | func (p *ExponentialRetryPolicy) SetMaximumAttempts(maximumAttempts int) { 130 | p.maximumAttempts = maximumAttempts 131 | } 132 | 133 | // ComputeNextDelay returns the next delay interval. This is used by Retrier to delay calling the operation again 134 | func (p *ExponentialRetryPolicy) ComputeNextDelay(elapsedTime time.Duration, numAttempts int) time.Duration { 135 | // Check to see if we ran out of maximum number of attempts 136 | if p.maximumAttempts != noMaximumAttempts && numAttempts >= p.maximumAttempts { 137 | return done 138 | } 139 | 140 | // Stop retrying after expiration interval is elapsed 141 | if p.expirationInterval != NoInterval && elapsedTime > p.expirationInterval { 142 | return done 143 | } 144 | 145 | nextInterval := float64(p.initialInterval) * math.Pow(p.backoffCoefficient, float64(numAttempts)) 146 | // Disallow retries if initialInterval is negative or nextInterval overflows 147 | if nextInterval <= 0 { 148 | return done 149 | } 150 | if p.maximumInterval != NoInterval { 151 | nextInterval = math.Min(nextInterval, float64(p.maximumInterval)) 152 | } 153 | 154 | if p.expirationInterval != NoInterval { 155 | remainingTime := float64(math.Max(0, float64(p.expirationInterval-elapsedTime))) 156 | nextInterval = math.Min(remainingTime, nextInterval) 157 | } 158 | 159 | // Bail out if the next interval is smaller than initial retry interval 160 | nextDuration := time.Duration(nextInterval) 161 | if nextDuration < p.initialInterval { 162 | return done 163 | } 164 | 165 | // add jitter to avoid global synchronization 166 | jitterPortion := int(0.2 * nextInterval) 167 | // Prevent overflow 168 | if jitterPortion < 1 { 169 | jitterPortion = 1 170 | } 171 | nextInterval = nextInterval*0.8 + float64(rand.Intn(jitterPortion)) 172 | 173 | return time.Duration(nextInterval) 174 | } 175 | 176 | // Now returns the current time using the system clock 177 | func (t systemClock) Now() time.Time { 178 | return time.Now() 179 | } 180 | 181 | // Reset will set the Retrier into initial state 182 | func (r *retrierImpl) Reset() { 183 | r.startTime = r.clock.Now() 184 | r.currentAttempt = 0 185 | } 186 | 187 | // NextBackOff returns the next delay interval. This is used by Retry to delay calling the operation again 188 | func (r *retrierImpl) NextBackOff() time.Duration { 189 | nextInterval := r.policy.ComputeNextDelay(r.getElapsedTime(), r.currentAttempt) 190 | 191 | // Now increment the current attempt 192 | r.currentAttempt++ 193 | return nextInterval 194 | } 195 | 196 | func (r *retrierImpl) getElapsedTime() time.Duration { 197 | return r.clock.Now().Sub(r.startTime) 198 | } 199 | -------------------------------------------------------------------------------- /internal/backoff/retrypolicy_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package backoff 22 | 23 | import ( 24 | "testing" 25 | "time" 26 | 27 | "github.com/stretchr/testify/require" 28 | "github.com/stretchr/testify/suite" 29 | ) 30 | 31 | type ( 32 | RetryPolicySuite struct { 33 | *require.Assertions // override suite.Suite.Assertions with require.Assertions; this means that s.NotNil(nil) will stop the test, not merely log an error 34 | suite.Suite 35 | } 36 | 37 | TestClock struct { 38 | currentTime time.Time 39 | } 40 | ) 41 | 42 | func TestRetryPolicySuite(t *testing.T) { 43 | suite.Run(t, new(RetryPolicySuite)) 44 | } 45 | 46 | func (s *RetryPolicySuite) SetupTest() { 47 | s.Assertions = require.New(s.T()) // Have to define our overridden assertions in the test setup. If we did it earlier, s.T() will return nil 48 | } 49 | 50 | func (s *RetryPolicySuite) TestExponentialBackoff() { 51 | policy := createPolicy(time.Second) 52 | policy.SetMaximumInterval(10 * time.Second) 53 | 54 | expectedResult := []time.Duration{1, 2, 4, 8, 10} 55 | for i, d := range expectedResult { 56 | expectedResult[i] = d * time.Second 57 | } 58 | 59 | r, _ := createRetrier(policy) 60 | for _, expected := range expectedResult { 61 | min, max := getNextBackoffRange(expected) 62 | next := r.NextBackOff() 63 | s.True(next >= min, "NextBackoff too low") 64 | s.True(next < max, "NextBackoff too high") 65 | } 66 | } 67 | 68 | func (s *RetryPolicySuite) TestNumberOfAttempts() { 69 | policy := createPolicy(time.Second) 70 | policy.SetMaximumAttempts(5) 71 | 72 | r, _ := createRetrier(policy) 73 | var next time.Duration 74 | for i := 0; i < 6; i++ { 75 | next = r.NextBackOff() 76 | } 77 | 78 | s.Equal(done, next) 79 | } 80 | 81 | // Test to make sure relative maximum interval for each retry is honoured 82 | func (s *RetryPolicySuite) TestMaximumInterval() { 83 | policy := createPolicy(time.Second) 84 | policy.SetMaximumInterval(10 * time.Second) 85 | 86 | expectedResult := []time.Duration{1, 2, 4, 8, 10, 10, 10, 10, 10, 10} 87 | for i, d := range expectedResult { 88 | expectedResult[i] = d * time.Second 89 | } 90 | 91 | r, _ := createRetrier(policy) 92 | for _, expected := range expectedResult { 93 | min, max := getNextBackoffRange(expected) 94 | next := r.NextBackOff() 95 | s.True(next >= min, "NextBackoff too low") 96 | s.True(next < max, "NextBackoff too high") 97 | } 98 | } 99 | 100 | func (s *RetryPolicySuite) TestBackoffCoefficient() { 101 | policy := createPolicy(2 * time.Second) 102 | policy.SetBackoffCoefficient(1.0) 103 | 104 | r, _ := createRetrier(policy) 105 | min, max := getNextBackoffRange(2 * time.Second) 106 | for i := 0; i < 10; i++ { 107 | next := r.NextBackOff() 108 | s.True(next >= min, "NextBackoff too low") 109 | s.True(next < max, "NextBackoff too high") 110 | } 111 | } 112 | 113 | func (s *RetryPolicySuite) TestExpirationInterval() { 114 | policy := createPolicy(2 * time.Second) 115 | policy.SetExpirationInterval(5 * time.Minute) 116 | 117 | r, clock := createRetrier(policy) 118 | clock.moveClock(6 * time.Minute) 119 | next := r.NextBackOff() 120 | 121 | s.Equal(done, next) 122 | } 123 | 124 | func (s *RetryPolicySuite) TestExpirationOverflow() { 125 | policy := createPolicy(2 * time.Second) 126 | policy.SetExpirationInterval(5 * time.Second) 127 | 128 | r, clock := createRetrier(policy) 129 | next := r.NextBackOff() 130 | min, max := getNextBackoffRange(2 * time.Second) 131 | s.True(next >= min, "NextBackoff too low") 132 | s.True(next < max, "NextBackoff too high") 133 | 134 | clock.moveClock(2 * time.Second) 135 | 136 | next = r.NextBackOff() 137 | min, max = getNextBackoffRange(3 * time.Second) 138 | s.True(next >= min, "NextBackoff too low") 139 | s.True(next < max, "NextBackoff too high") 140 | } 141 | 142 | func (s *RetryPolicySuite) TestDefaultPublishRetryPolicy() { 143 | policy := NewExponentialRetryPolicy(50 * time.Millisecond) 144 | policy.SetExpirationInterval(time.Minute) 145 | policy.SetMaximumInterval(10 * time.Second) 146 | 147 | r, clock := createRetrier(policy) 148 | expectedResult := []time.Duration{ 149 | 50 * time.Millisecond, 150 | 100 * time.Millisecond, 151 | 200 * time.Millisecond, 152 | 400 * time.Millisecond, 153 | 800 * time.Millisecond, 154 | 1600 * time.Millisecond, 155 | 3200 * time.Millisecond, 156 | 6400 * time.Millisecond, 157 | 10000 * time.Millisecond, 158 | 10000 * time.Millisecond, 159 | 10000 * time.Millisecond, 160 | 10000 * time.Millisecond, 161 | 6000 * time.Millisecond, 162 | 1300 * time.Millisecond, 163 | done, 164 | } 165 | 166 | for _, expected := range expectedResult { 167 | next := r.NextBackOff() 168 | if expected == done { 169 | s.Equal(done, next, "backoff not done yet!!!") 170 | } else { 171 | min, _ := getNextBackoffRange(expected) 172 | s.True(next >= min, "NextBackoff too low: actual: %v, expected: %v", next, expected) 173 | // s.True(next < max, "NextBackoff too high: actual: %v, expected: %v", next, expected) 174 | clock.moveClock(expected) 175 | } 176 | } 177 | } 178 | 179 | func (s *RetryPolicySuite) TestNoMaxAttempts() { 180 | policy := createPolicy(50 * time.Millisecond) 181 | policy.SetExpirationInterval(time.Minute) 182 | policy.SetMaximumInterval(10 * time.Second) 183 | 184 | r, clock := createRetrier(policy) 185 | for i := 0; i < 100; i++ { 186 | next := r.NextBackOff() 187 | //print("Iter: ", i, ", Next Backoff: ", next.String(), "\n") 188 | s.True(next > 0 || next == done, "Unexpected value for next retry duration: %v", next) 189 | clock.moveClock(next) 190 | } 191 | } 192 | 193 | func (s *RetryPolicySuite) TestUnbounded() { 194 | policy := createPolicy(50 * time.Millisecond) 195 | 196 | r, clock := createRetrier(policy) 197 | for i := 0; i < 100; i++ { 198 | next := r.NextBackOff() 199 | //print("Iter: ", i, ", Next Backoff: ", next.String(), "\n") 200 | s.True(next > 0 || next == done, "Unexpected value for next retry duration: %v", next) 201 | clock.moveClock(next) 202 | } 203 | } 204 | 205 | func (c *TestClock) Now() time.Time { 206 | return c.currentTime 207 | } 208 | 209 | func (c *TestClock) moveClock(duration time.Duration) { 210 | c.currentTime = c.currentTime.Add(duration) 211 | } 212 | 213 | func createPolicy(initialInterval time.Duration) *ExponentialRetryPolicy { 214 | policy := NewExponentialRetryPolicy(initialInterval) 215 | policy.SetBackoffCoefficient(2) 216 | policy.SetMaximumInterval(NoInterval) 217 | policy.SetExpirationInterval(NoInterval) 218 | policy.SetMaximumAttempts(noMaximumAttempts) 219 | 220 | return policy 221 | } 222 | 223 | func createRetrier(policy RetryPolicy) (Retrier, *TestClock) { 224 | clock := &TestClock{currentTime: time.Time{}} 225 | return NewRetrier(policy, clock), clock 226 | } 227 | 228 | func getNextBackoffRange(duration time.Duration) (time.Duration, time.Duration) { 229 | rangeMin := time.Duration(0.8 * float64(duration)) 230 | return rangeMin, duration 231 | } 232 | -------------------------------------------------------------------------------- /internal/consumer/ackMgr.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package consumer 22 | 23 | import ( 24 | "fmt" 25 | "sync" 26 | "time" 27 | 28 | "github.com/uber-go/kafka-client/internal/list" 29 | "github.com/uber-go/kafka-client/internal/metrics" 30 | "github.com/uber-go/tally" 31 | "go.uber.org/zap" 32 | ) 33 | 34 | const ( 35 | resetCheckInterval = time.Second 36 | ) 37 | 38 | type ( 39 | ackID struct { 40 | listAddr list.Address 41 | msgSeq int64 42 | } 43 | ackManager struct { 44 | unackedSeqList *threadSafeSortedList 45 | logger *zap.Logger 46 | tally tally.Scope 47 | } 48 | threadSafeSortedList struct { 49 | sync.Mutex 50 | list *list.IntegerList 51 | lastValue int64 52 | tally tally.Scope 53 | } 54 | ) 55 | 56 | // newAckManager returns a new instance of AckManager. The returned 57 | // ackManager can be used by a kafka partitioned consumer to keep 58 | // track of the current commit level for the partition. The ackManager 59 | // is only needed when consuming and processing kafka messages in 60 | // parallel (i.e. multiple goroutines), which can cause out of order 61 | // message completion. 62 | // 63 | // In the description below, the terms seqNum and offset are used interchangeably 64 | // Within the ackManager, offsets are nothing but message seqNums 65 | // 66 | // Usage: 67 | // - For every new message, call ackMgr.GetAckID(msgSeqNum) before processing the message 68 | // - After processing, call ackMgr.Ack(ackID) to mark the message as processed 69 | // - If the message processing fails, call ackMgr.Nack(ackID) to skip the message 70 | // and move it into an error queue (not supported yet) 71 | // - Periodically, call ackMgr.CommitLevel() to retrieve / checkpoint the safe commit level 72 | // 73 | // Assumptions: 74 | // - The first msgSeq added is considered as the beginSeqNum for this ackManager 75 | // - Any message with seqNum less than the most recent seqNum is considered a duplicate and ignored 76 | // - There CANNOT be more than maxOutstanding messages that are unacked at any given point of time 77 | // - Call ackMgr.Ack is an acknowledgement to move the seqNum past the acked message 78 | // 79 | // Implementation Notes: 80 | // The implementation works by keeping track of unacked seqNums in a doubly linked list. When a 81 | // message is acked, its removed from the linked list. The head of the linked list is the unacked 82 | // message with the lowest seqNum. So, any offset less that that is a safe commit checkpoint. 83 | // 84 | // Params: 85 | // maxOutstanding - max number of unacked messages at any given point in time 86 | // scope / logger - metrics / logging client 87 | func newAckManager(maxOutstanding int, scope tally.Scope, logger *zap.Logger) *ackManager { 88 | return &ackManager{ 89 | tally: scope, 90 | logger: logger, 91 | unackedSeqList: newThreadSafeSortedList(maxOutstanding, scope), 92 | } 93 | } 94 | 95 | // GetAckID adds the given seqNum to the list of unacked seqNums 96 | // and returns a opaque AckID that can be used to identify this 97 | // message when its subsequently acked or nacked 98 | // Returns an error if the msgSeqNum is unexpected 99 | func (mgr *ackManager) GetAckID(msgSeq int64) (ackID, error) { 100 | addr, err := mgr.unackedSeqList.Add(msgSeq) 101 | if err != nil { 102 | mgr.tally.Counter(metrics.KafkaPartitionGetAckIDErrors).Inc(1) 103 | if err != list.ErrCapacity { 104 | // list.ErrCapacity is handled gracefully so no need to log error. 105 | mgr.logger.Error("GetAckID() error", zap.Int64("rcvdSeq", msgSeq), zap.Error(err)) 106 | } 107 | return ackID{}, err 108 | } 109 | return newAckID(addr, msgSeq), nil 110 | } 111 | 112 | // Ack marks the given msgSeqNum as processed. 113 | // Ack always returns nil because the errors are non-actionable and do not affect correctness 114 | // so we do not want to propagate back to user. 115 | // Non-actionable errors are visible via "kafka.partition.ackmgr.ack-error" metrics. 116 | func (mgr *ackManager) Ack(id ackID) error { 117 | err := mgr.unackedSeqList.Remove(id.listAddr, id.msgSeq) 118 | if err != nil { 119 | mgr.tally.Counter(metrics.KafkaPartitionAckErrors).Inc(1) 120 | mgr.logger.Error("ack error: list remove failed", zap.Error(err)) 121 | } else { 122 | mgr.tally.Counter(metrics.KafkaPartitionAck).Inc(1) 123 | } 124 | return nil 125 | } 126 | 127 | // Nack marks the given msgSeqNum as processed, the expectation 128 | // is for the caller to move the message to an error queue 129 | // before calling this 130 | // Nack always returns nil because the errors are non-actionable and do not affect correctness 131 | // so we do not want to propagate back to user. 132 | // Non-actionable errors are visible via "kafka.partition.ackmgr.nack-error" metrics. 133 | func (mgr *ackManager) Nack(id ackID) error { 134 | err := mgr.unackedSeqList.Remove(id.listAddr, id.msgSeq) 135 | if err != nil { 136 | mgr.tally.Counter(metrics.KafkaPartitionNackErrors).Inc(1) 137 | mgr.logger.Error("nack error: list remove failed", zap.Error(err)) 138 | } else { 139 | mgr.tally.Counter(metrics.KafkaPartitionNack).Inc(1) 140 | } 141 | return nil 142 | } 143 | 144 | // CommitLevel returns the seqNum that can be 145 | // used as a safe commit checkpoint. Returns value 146 | // less than zero if there is no safe checkpoint yet 147 | func (mgr *ackManager) CommitLevel() int64 { 148 | unacked, err := mgr.unackedSeqList.PeekHead() 149 | if err != nil { 150 | if err != list.ErrEmpty { 151 | mgr.logger.Fatal("commitLevel error: list peekHead failed", zap.Error(err)) 152 | } 153 | return mgr.unackedSeqList.LastValue() 154 | } 155 | return unacked - 1 156 | } 157 | 158 | // Reset blocks until the list is empty and the offsets have been reset. 159 | func (mgr *ackManager) Reset() { 160 | mgr.unackedSeqList.Reset() 161 | } 162 | 163 | // newAckID returns a an ackID with the given params 164 | func newAckID(addr list.Address, value int64) ackID { 165 | return ackID{listAddr: addr, msgSeq: value} 166 | } 167 | 168 | // newThreadSafeSortedList returns a new instance of thread safe 169 | // integer list that expects its input to be received in a sorted 170 | // order 171 | func newThreadSafeSortedList(maxOutstanding int, scope tally.Scope) *threadSafeSortedList { 172 | list := list.NewIntegerList(maxOutstanding) 173 | return &threadSafeSortedList{list: list, lastValue: -1, tally: scope} 174 | } 175 | 176 | // PeekHead returns the value at the head of the list, if it exist 177 | func (l *threadSafeSortedList) PeekHead() (int64, error) { 178 | l.Lock() 179 | defer l.Unlock() 180 | return l.list.PeekHead() 181 | } 182 | 183 | func (l *threadSafeSortedList) LastValue() int64 { 184 | l.Lock() 185 | value := l.lastValue 186 | l.Unlock() 187 | return value 188 | } 189 | 190 | // Remove removes the entry at the given address, if and if only if 191 | // the entry has value equal to the given value 192 | func (l *threadSafeSortedList) Remove(addr list.Address, value int64) error { 193 | l.Lock() 194 | defer l.Unlock() 195 | got, err := l.list.Get(addr) 196 | if err != nil { 197 | return err 198 | } 199 | if value != got { 200 | return fmt.Errorf("address / value mismatch, expected value of %v but got %v", value, got) 201 | } 202 | return l.list.Remove(addr) 203 | } 204 | 205 | // Add adds the value to the end of the list. The value MUST be 206 | // greater than the last value added to this list to maintain 207 | // the sorted order. If not, an error is returned 208 | func (l *threadSafeSortedList) Add(value int64) (list.Address, error) { 209 | l.Lock() 210 | defer l.Unlock() 211 | if value <= l.lastValue { 212 | l.tally.Counter(metrics.KafkaPartitionAckMgrDups).Inc(1) 213 | return list.Null, fmt.Errorf("new value %v is not greater than last stored value %v", value, l.lastValue) 214 | } 215 | skipped := value - l.lastValue - 1 216 | if skipped > 0 { 217 | l.tally.Counter(metrics.KafkaPartitionAckMgrSkipped).Inc(skipped) 218 | } 219 | addr, err := l.list.Add(value) 220 | if err == nil { 221 | l.lastValue = value 222 | } 223 | return addr, err 224 | } 225 | 226 | // Reset blocks until the list is empty then sets lastValue to -1. 227 | func (l *threadSafeSortedList) Reset() { 228 | doneC := make(chan struct{}) 229 | checkInterval := time.NewTicker(resetCheckInterval) 230 | go func() { 231 | for { 232 | select { 233 | case <-checkInterval.C: 234 | l.Lock() 235 | if l.list.Empty() { 236 | l.lastValue = -1 237 | close(doneC) 238 | } 239 | l.Unlock() 240 | case <-doneC: 241 | return 242 | } 243 | } 244 | }() 245 | <-doneC // block until the list is reset 246 | checkInterval.Stop() 247 | } 248 | -------------------------------------------------------------------------------- /internal/consumer/clusterConsumer.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package consumer 22 | 23 | import ( 24 | "errors" 25 | "fmt" 26 | "strconv" 27 | "sync" 28 | "time" 29 | 30 | "github.com/bsm/sarama-cluster" 31 | "github.com/uber-go/kafka-client/internal/metrics" 32 | "github.com/uber-go/kafka-client/internal/util" 33 | "github.com/uber-go/kafka-client/kafka" 34 | "github.com/uber-go/tally" 35 | "go.uber.org/zap" 36 | "go.uber.org/zap/zapcore" 37 | ) 38 | 39 | const ( 40 | metricsInterval = time.Minute 41 | ) 42 | 43 | type ( 44 | // ClusterConsumer is a consumer for a single Kafka cluster. 45 | ClusterConsumer struct { 46 | cluster string 47 | consumer SaramaConsumer 48 | topicConsumerMap map[string]*TopicConsumer 49 | scope tally.Scope 50 | logger *zap.Logger 51 | lifecycle *util.RunLifecycle 52 | metricsTicker *time.Ticker 53 | stopC chan struct{} 54 | doneC chan struct{} 55 | } 56 | 57 | ClusterGroup struct { 58 | Cluster string 59 | Group string 60 | } 61 | ) 62 | 63 | // NewClusterConsumer returns a new single cluster consumer. 64 | func NewClusterConsumer( 65 | cluster string, 66 | saramaConsumer SaramaConsumer, 67 | consumerMap map[string]*TopicConsumer, 68 | scope tally.Scope, 69 | logger *zap.Logger, 70 | ) *ClusterConsumer { 71 | return &ClusterConsumer{ 72 | cluster: cluster, 73 | consumer: saramaConsumer, 74 | topicConsumerMap: consumerMap, 75 | scope: scope.Tagged(map[string]string{"cluster": cluster}), 76 | logger: logger.With(zap.String("cluster", cluster)), 77 | lifecycle: util.NewRunLifecycle(cluster + "-consumer"), 78 | metricsTicker: time.NewTicker(metricsInterval), 79 | stopC: make(chan struct{}), 80 | doneC: make(chan struct{}), 81 | } 82 | } 83 | 84 | // Start starts the consumer 85 | func (c *ClusterConsumer) Start() error { 86 | return c.lifecycle.Start(func() error { 87 | logger := c.logger.With( 88 | zap.Array("topicList", zapcore.ArrayMarshalerFunc(func(e zapcore.ArrayEncoder) error { 89 | for topic := range c.topicConsumerMap { 90 | e.AppendString(topic) 91 | } 92 | return nil 93 | })), 94 | ) 95 | for _, topicConsumer := range c.topicConsumerMap { 96 | if err := topicConsumer.Start(); err != nil { 97 | logger.Error("cluster consumer start error", zap.Error(err)) 98 | return err 99 | } 100 | } 101 | go c.eventLoop() 102 | logger.Info("cluster consumer started") 103 | return nil 104 | }) 105 | } 106 | 107 | // Stop stops the consumer 108 | func (c *ClusterConsumer) Stop() { 109 | c.lifecycle.Stop(func() { 110 | c.logger.With( 111 | zap.Array("topicList", zapcore.ArrayMarshalerFunc(func(e zapcore.ArrayEncoder) error { 112 | for topic := range c.topicConsumerMap { 113 | e.AppendString(topic) 114 | } 115 | return nil 116 | })), 117 | ).Info("cluster consumer stopping") 118 | close(c.stopC) 119 | }) 120 | } 121 | 122 | // Closed returns a channel which will closed after this consumer is shutdown 123 | func (c *ClusterConsumer) Closed() <-chan struct{} { 124 | return c.doneC 125 | } 126 | 127 | // eventLoop is the main event loop for this consumer 128 | func (c *ClusterConsumer) eventLoop() { 129 | var n *cluster.Notification 130 | var ok bool 131 | for { 132 | select { 133 | case pc, ok := <-c.consumer.Partitions(): 134 | if ok { 135 | c.addPartitionConsumer(pc) 136 | } 137 | case n, ok = <-c.consumer.Notifications(): 138 | if ok { 139 | c.handleNotification(n) 140 | } 141 | case err, ok := <-c.consumer.Errors(): 142 | if ok { 143 | c.logger.Warn("cluster consumer error", zap.Error(err)) 144 | } 145 | case _, ok := <-c.metricsTicker.C: 146 | if ok && n != nil { 147 | for topic, partitions := range n.Current { 148 | for _, partition := range partitions { 149 | c.scope.Tagged(map[string]string{"topic": topic, "partition": strconv.Itoa(int(partition))}).Gauge(metrics.KafkaPartitionOwned).Update(1.0) 150 | } 151 | } 152 | } 153 | case <-c.stopC: 154 | c.shutdown() 155 | c.logger.Info("cluster consumer stopped") 156 | return 157 | } 158 | } 159 | } 160 | 161 | // addPartition adds a new partition. If the partition already exist, 162 | // it is first stopped before overwriting it with the new partition 163 | func (c *ClusterConsumer) addPartitionConsumer(pc cluster.PartitionConsumer) { 164 | topic := pc.Topic() 165 | topicConsumer, ok := c.topicConsumerMap[topic] 166 | if !ok { 167 | c.logger.Error("cluster consumer cannot consume messages for missing topic consumer", zap.String("topic", topic)) 168 | return 169 | } 170 | topicConsumer.addPartitionConsumer(pc) 171 | } 172 | 173 | // handleNotification is the handler that handles notifications 174 | // from the underlying library about partition rebalances. There 175 | // is no action taken in this handler except for logging. 176 | func (c *ClusterConsumer) handleNotification(n *cluster.Notification) { 177 | for topic, partitions := range n.Claimed { 178 | for _, partition := range partitions { 179 | c.logger.Debug("cluster consumer partition rebalance claimed", zap.String("topic", topic), zap.Int32("partition", partition)) 180 | } 181 | } 182 | 183 | for topic, partitions := range n.Released { 184 | for _, partition := range partitions { 185 | c.logger.Debug("cluster consumer partition rebalance released", zap.String("topic", topic), zap.Int32("partition", partition)) 186 | } 187 | } 188 | 189 | var current []string 190 | for topic, partitions := range n.Current { 191 | for _, partition := range partitions { 192 | current = append(current, fmt.Sprintf("%s-%s", topic, strconv.Itoa(int(partition)))) 193 | } 194 | } 195 | 196 | c.logger.Info("cluster consumer owned topic-partitions after rebalance", zap.Strings("topic-partitions", current)) 197 | c.scope.Counter(metrics.KafkaPartitionRebalance).Inc(1) 198 | } 199 | 200 | // shutdown stops the consumer and frees resources 201 | func (c *ClusterConsumer) shutdown() { 202 | // Close each TopicConsumer 203 | var wg sync.WaitGroup 204 | for _, tc := range c.topicConsumerMap { 205 | wg.Add(1) 206 | go func(tc *TopicConsumer) { 207 | tc.Stop() 208 | wg.Done() 209 | }(tc) 210 | } 211 | wg.Wait() 212 | c.consumer.Close() // close sarama cluster consumer 213 | c.metricsTicker.Stop() 214 | close(c.doneC) 215 | } 216 | 217 | // ResetOffset will reset the consumer offset for the specified topic, partition. 218 | func (c *ClusterConsumer) ResetOffset(topic string, partition int32, offsetRange kafka.OffsetRange) error { 219 | tc, ok := c.topicConsumerMap[topic] 220 | if !ok { 221 | return errors.New("no topic consumer found") 222 | } 223 | 224 | return tc.ResetOffset(partition, offsetRange) 225 | 226 | } 227 | -------------------------------------------------------------------------------- /internal/consumer/clusterConsumer_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package consumer 22 | 23 | import ( 24 | "sync" 25 | "testing" 26 | "time" 27 | 28 | "github.com/Shopify/sarama" 29 | "github.com/bsm/sarama-cluster" 30 | "github.com/stretchr/testify/suite" 31 | "github.com/uber-go/kafka-client/internal/util" 32 | "github.com/uber-go/kafka-client/kafka" 33 | "github.com/uber-go/tally" 34 | "go.uber.org/zap" 35 | ) 36 | 37 | type ( 38 | ClusterConsumerTestSuite struct { 39 | suite.Suite 40 | consumer *ClusterConsumer 41 | topicConsumer *TopicConsumer 42 | saramaConsumer *mockSaramaConsumer 43 | dlqProducer *mockDLQProducer 44 | msgCh chan kafka.Message 45 | topic string 46 | dlqTopic string 47 | options *Options 48 | logger *zap.Logger 49 | } 50 | ) 51 | 52 | func testConsumerOptions() *Options { 53 | return &Options{ 54 | Concurrency: 4, 55 | RcvBufferSize: 4, 56 | PartitionRcvBufferSize: 2, 57 | OffsetCommitInterval: 25 * time.Millisecond, 58 | RebalanceDwellTime: time.Second, 59 | MaxProcessingTime: 5 * time.Millisecond, 60 | OffsetPolicy: sarama.OffsetOldest, 61 | } 62 | } 63 | 64 | func TestClusterConsumerTestSuite(t *testing.T) { 65 | suite.Run(t, new(ClusterConsumerTestSuite)) 66 | } 67 | 68 | func (s *ClusterConsumerTestSuite) SetupTest() { 69 | topic := kafka.ConsumerTopic{ 70 | Topic: kafka.Topic{ 71 | Name: "unit-test", 72 | Cluster: "production-cluster", 73 | }, 74 | DLQ: kafka.Topic{ 75 | Name: "unit-test-dlq", 76 | Cluster: "dlq-cluster", 77 | }, 78 | } 79 | s.topic = topic.Topic.Name 80 | s.dlqTopic = topic.DLQ.Name 81 | s.options = testConsumerOptions() 82 | s.logger = zap.NewNop() 83 | s.msgCh = make(chan kafka.Message, 10) 84 | s.saramaConsumer = newMockSaramaConsumer() 85 | s.dlqProducer = newMockDLQProducer() 86 | s.topicConsumer = NewTopicConsumer(Topic{ConsumerTopic: topic, DLQMetadataDecoder: NoopDLQMetadataDecoder, PartitionConsumerFactory: NewPartitionConsumer}, s.msgCh, s.saramaConsumer, s.dlqProducer, s.options, tally.NoopScope, s.logger) 87 | s.consumer = &ClusterConsumer{ 88 | cluster: topic.Cluster, 89 | consumer: s.saramaConsumer, 90 | topicConsumerMap: map[string]*TopicConsumer{s.topic: s.topicConsumer}, 91 | scope: tally.NoopScope, 92 | logger: s.logger, 93 | lifecycle: util.NewRunLifecycle(topic.Cluster + "-consumer"), 94 | metricsTicker: time.NewTicker(100 * time.Millisecond), 95 | stopC: make(chan struct{}), 96 | doneC: make(chan struct{}), 97 | } 98 | } 99 | 100 | func (s *ClusterConsumerTestSuite) TearDownTest() { 101 | s.consumer.Stop() 102 | <-s.consumer.Closed() 103 | s.True(util.AwaitCondition(func() bool { return s.saramaConsumer.isClosed() }, time.Second)) 104 | s.True(s.dlqProducer.isClosed()) 105 | } 106 | 107 | func (s *ClusterConsumerTestSuite) startWorker(count int, concurrency int, nack bool) *sync.WaitGroup { 108 | var wg sync.WaitGroup 109 | for i := 0; i < concurrency; i++ { 110 | wg.Add(1) 111 | go func() { 112 | for i := 0; i < count/concurrency; i++ { 113 | m := <-s.msgCh 114 | if nack { 115 | m.Nack() 116 | continue 117 | } 118 | m.Ack() 119 | } 120 | wg.Done() 121 | }() 122 | } 123 | return &wg 124 | } 125 | 126 | func (s *ClusterConsumerTestSuite) TestWithOnePartition() { 127 | s.consumer.Start() 128 | workerWG := s.startWorker(100, 4, false) 129 | 130 | // send new partition to consumer 131 | p1 := newMockPartitionedConsumer(s.topic, 1, 0, s.options.PartitionRcvBufferSize) 132 | s.saramaConsumer.partitionC <- p1 133 | p1.start() 134 | 135 | // send notification about a rebalance 136 | n := &cluster.Notification{ 137 | Claimed: map[string][]int32{s.topic: {1}}, 138 | Current: map[string][]int32{s.topic: {1}}, 139 | } 140 | s.saramaConsumer.notifyC <- n 141 | 142 | s.True(util.AwaitWaitGroup(workerWG, time.Second)) // wait for messages to be consumed 143 | 144 | // do assertions 145 | s.Equal(0, len(s.msgCh), "channel expected to be empty") 146 | s.True(util.AwaitCondition(func() bool { return s.saramaConsumer.offset(1) == int64(100) }, time.Second)) 147 | s.Equal(0, s.dlqProducer.backlog()) 148 | 149 | cp, ok := s.consumer.topicConsumerMap["unit-test"].partitionConsumerMap[1].(*partitionConsumer) 150 | s.True(ok) 151 | s.Equal(int64(100), s.saramaConsumer.offset(1), "wrong commit offset") 152 | s.True(cp.ackMgr.unackedSeqList.list.Empty(), "unacked offset list must be empty") 153 | 154 | // test shutdown 155 | p1.stop() 156 | s.True(util.AwaitCondition(func() bool { return p1.isClosed() }, time.Second)) 157 | } 158 | 159 | func (s *ClusterConsumerTestSuite) TestWithManyPartitions() { 160 | nPartitions := 8 161 | s.consumer.Start() 162 | workerWG := s.startWorker(nPartitions*100, 4, false) 163 | // start all N partitions 164 | for i := 0; i < nPartitions; i++ { 165 | pc := newMockPartitionedConsumer(s.topic, int32(i), 0, s.options.PartitionRcvBufferSize) 166 | s.saramaConsumer.partitionC <- pc 167 | pc.start() 168 | if i%2 == 0 { 169 | // send notification about a rebalance 170 | n := &cluster.Notification{Claimed: make(map[string][]int32)} 171 | n.Claimed[s.topic] = []int32{int32(i), int32(i - 1)} 172 | s.saramaConsumer.notifyC <- n 173 | } 174 | } 175 | s.True(util.AwaitWaitGroup(workerWG, 2*time.Second)) // wait for all messages to be consumed 176 | s.Equal(0, len(s.msgCh)) 177 | for i := 0; i < nPartitions; i++ { 178 | s.True(util.AwaitCondition(func() bool { return s.saramaConsumer.offset(i) == int64(100) }, time.Second)) 179 | cp, ok := s.consumer.topicConsumerMap["unit-test"].partitionConsumerMap[1].(*partitionConsumer) 180 | s.True(ok) 181 | s.Equal(int64(100), s.saramaConsumer.offset(i), "wrong commit offset") 182 | s.True(cp.ackMgr.unackedSeqList.list.Empty(), "unacked offset list must be empty") 183 | } 184 | s.Equal(0, s.dlqProducer.backlog()) 185 | } 186 | 187 | func (s *ClusterConsumerTestSuite) TestPartitionRebalance() { 188 | nPartitions := 4 189 | nRebalances := 3 190 | s.consumer.Start() 191 | for r := 0; r < nRebalances; r++ { 192 | workerWG := s.startWorker(nPartitions*100, 4, false) 193 | partitions := make([]*mockPartitionedConsumer, nPartitions) 194 | // start all N partitions 195 | for i := 0; i < nPartitions; i++ { 196 | pc := newMockPartitionedConsumer(s.topic, int32(i), int64(r*100), s.options.PartitionRcvBufferSize) 197 | partitions[i] = pc 198 | s.saramaConsumer.partitionC <- pc 199 | pc.start() 200 | } 201 | s.True(util.AwaitWaitGroup(workerWG, 2*time.Second)) // wait for all messages to be consumed 202 | s.Equal(0, len(s.msgCh)) 203 | for i := 0; i < nPartitions; i++ { 204 | off := int64(100 * (r + 1)) 205 | s.True(util.AwaitCondition(func() bool { return s.saramaConsumer.offset(i) == off }, time.Minute)) 206 | s.Equal(off, s.saramaConsumer.offset(i), "wrong commit offset for partition %v", i) 207 | } 208 | } 209 | s.Equal(0, s.dlqProducer.backlog()) 210 | } 211 | 212 | func (s *ClusterConsumerTestSuite) TestDuplicates() { 213 | s.consumer.Start() 214 | workerWG := s.startWorker(100, 1, false) 215 | pc := newMockPartitionedConsumer(s.topic, 1, 0, s.options.PartitionRcvBufferSize) 216 | s.saramaConsumer.partitionC <- pc 217 | // start two parallel message producers for the same offsets 218 | pc.start() 219 | pc.start() 220 | s.True(util.AwaitWaitGroup(workerWG, 2*time.Second)) // wait for all messages to be consumed 221 | s.Equal(0, len(s.msgCh)) 222 | s.True(util.AwaitCondition(func() bool { return s.saramaConsumer.offset(1) == int64(100) }, time.Second)) 223 | s.Equal(int64(100), s.saramaConsumer.offset(1)) 224 | s.Equal(0, s.dlqProducer.backlog()) 225 | } 226 | 227 | func (s *ClusterConsumerTestSuite) TestDLQ() { 228 | nPartitions := 4 229 | s.consumer.Start() 230 | workerWG := s.startWorker(nPartitions*100, 4, true) 231 | // start all N partitions 232 | for i := 0; i < nPartitions; i++ { 233 | pc := newMockPartitionedConsumer(s.topic, int32(i), 0, s.options.PartitionRcvBufferSize) 234 | s.saramaConsumer.partitionC <- pc 235 | pc.start() 236 | } 237 | s.True(util.AwaitWaitGroup(workerWG, 2*time.Second)) // wait for all messages to be consumed 238 | s.Equal(0, len(s.msgCh)) 239 | for i := 0; i < nPartitions; i++ { 240 | s.True(util.AwaitCondition(func() bool { return s.saramaConsumer.offset(i) == int64(100) }, time.Second)) 241 | cp, ok := s.consumer.topicConsumerMap["unit-test"].partitionConsumerMap[int32(i)].(*partitionConsumer) 242 | s.True(ok) 243 | s.Equal(int64(100), s.saramaConsumer.offset(i), "wrong commit offset") 244 | s.True(cp.ackMgr.unackedSeqList.list.Empty(), "unacked offset list must be empty") 245 | } 246 | s.Equal(nPartitions*100, s.dlqProducer.backlog()) 247 | } 248 | -------------------------------------------------------------------------------- /internal/consumer/dlq.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package consumer 22 | 23 | import ( 24 | "errors" 25 | "fmt" 26 | 27 | "github.com/Shopify/sarama" 28 | "github.com/golang/protobuf/proto" 29 | "github.com/uber-go/kafka-client/internal/metrics" 30 | "github.com/uber-go/kafka-client/internal/util" 31 | "github.com/uber-go/kafka-client/kafka" 32 | "github.com/uber-go/tally" 33 | "go.uber.org/zap" 34 | ) 35 | 36 | var ( 37 | errShutdown = errors.New("error consumer shutdown") 38 | errNoDLQ = errors.New("no persistent dlq configured") 39 | 40 | // RetryQErrorQType is the error queue for the retryQ. 41 | RetryQErrorQType ErrorQType = "retryQ" 42 | // DLQErrorQType is the error queue for DLQ. 43 | DLQErrorQType ErrorQType = "DLQ" 44 | 45 | // DLQConsumerGroupNameSuffix is the consumer group name used by the DLQ merge process. 46 | DLQConsumerGroupNameSuffix = "-dlq-merger" 47 | ) 48 | 49 | type ( 50 | // DLQ is the interface for implementations that 51 | // can take a message and put them into some sort 52 | // of error queue for later processing 53 | DLQ interface { 54 | // Start the DLQ producer 55 | Start() error 56 | // Stop the DLQ producer and close resources it holds. 57 | Stop() 58 | // Add adds the given message to DLQ. 59 | // This is a synchronous call and will block until sending is successful. 60 | Add(m kafka.Message, qTypes ...ErrorQType) error 61 | } 62 | 63 | // ErrorQType is the queue type to send messages to when using the DLQ interface. 64 | ErrorQType string 65 | 66 | // dlqMultiplexer is an implementation of DLQ which sends messages to a RetryQ for a configured number of times before 67 | // sending messages to thq DLQ. 68 | // Messages send to the RetryQ will be automatically reconsumed by the library. 69 | dlqMultiplexer struct { 70 | // retryCountThreshold is the threshold that is used to determine whether a message 71 | // goes to retryTopic or dlqTopic 72 | retryCountThreshold int64 73 | // retryTopic is topic that acts as a retry queue. 74 | // messages with retry count < the retry count threshold in the multiplexer will be sent to this topic. 75 | retryTopic DLQ 76 | // dlqTopic is topic that acts as a dead letter queue. 77 | // messages with retry count >= the retry count threshold in the multiplexer will be sent to this topic. 78 | dlqTopic DLQ 79 | } 80 | 81 | // bufferedErrorTopic is a client side abstraction for an error topic on the Kafka cluster. 82 | // This library uses the error topic as a base abstraction for retry and dlq topics. 83 | // 84 | // bufferedErrorTopic internally batches messages, so calls to Add may block until the internal 85 | // buffer is flushed and a batch of messages is send to the cluster. 86 | bufferedErrorTopic struct { 87 | topic kafka.Topic 88 | producer sarama.AsyncProducer 89 | 90 | stopC chan struct{} 91 | doneC chan struct{} 92 | 93 | scope tally.Scope 94 | logger *zap.Logger 95 | lifecycle *util.RunLifecycle 96 | } 97 | 98 | noopDLQ struct{} 99 | ) 100 | 101 | // NewBufferedDLQ returns a DLQ that is backed by a buffered async sarama producer. 102 | func NewBufferedDLQ(topic kafka.Topic, producer sarama.AsyncProducer, scope tally.Scope, logger *zap.Logger) DLQ { 103 | return newBufferedDLQ(topic, producer, scope, logger.With(zap.String("topic", topic.Name), zap.String("cluster", topic.Cluster))) 104 | } 105 | 106 | func newBufferedDLQ(topic kafka.Topic, producer sarama.AsyncProducer, scope tally.Scope, logger *zap.Logger) *bufferedErrorTopic { 107 | scope = scope.Tagged(map[string]string{"topic": topic.Name, "cluster": topic.Cluster}) 108 | logger = logger.With( 109 | zap.String("topic", topic.Name), 110 | zap.String("cluster", topic.Cluster), 111 | ) 112 | return &bufferedErrorTopic{ 113 | topic: topic, 114 | producer: producer, 115 | stopC: make(chan struct{}), 116 | doneC: make(chan struct{}), 117 | scope: scope, 118 | logger: logger, 119 | lifecycle: util.NewRunLifecycle(fmt.Sprintf("dlqProducer-%s-%s", topic.Name, topic.Cluster)), 120 | } 121 | } 122 | 123 | func (d *bufferedErrorTopic) Start() error { 124 | return d.lifecycle.Start(func() error { 125 | go d.asyncProducerResponseLoop() 126 | d.logger.Info("DLQ started") 127 | d.scope.Counter(metrics.KafkaDLQStarted).Inc(1) 128 | return nil 129 | }) 130 | } 131 | 132 | func (d *bufferedErrorTopic) Stop() { 133 | d.lifecycle.Stop(func() { 134 | close(d.stopC) 135 | d.producer.Close() 136 | <-d.doneC 137 | d.logger.Info("DLQ stopped") 138 | d.scope.Counter(metrics.KafkaDLQStopped).Inc(1) 139 | }) 140 | } 141 | 142 | // Add a message to the buffer for the error topic. 143 | // ErrorQType is ignored. 144 | func (d *bufferedErrorTopic) Add(m kafka.Message, qTypes ...ErrorQType) error { 145 | metadata := &DLQMetadata{ 146 | RetryCount: m.RetryCount() + 1, 147 | Data: m.Key(), 148 | Topic: m.Topic(), 149 | Partition: m.Partition(), 150 | Offset: m.Offset(), 151 | TimestampNs: m.Timestamp().UnixNano(), 152 | } 153 | 154 | key, err := proto.Marshal(metadata) 155 | if err != nil { 156 | d.logger.Error("failed to encode DLQ metadata", zap.Error(err)) 157 | d.scope.Counter(metrics.KafkaDLQMetadataError).Inc(1) 158 | return err 159 | } 160 | value := m.Value() 161 | // TODO (gteo): Use a channel pool 162 | responseC := make(chan error) 163 | 164 | sm := d.newSaramaMessage(key, value, responseC) 165 | select { 166 | case d.producer.Input() <- sm: 167 | case <-d.stopC: 168 | return errShutdown 169 | } 170 | select { 171 | case err := <-responseC: // block until response is received 172 | if err == nil { 173 | d.scope.Counter(metrics.KafkaDLQMessagesOut).Inc(1) 174 | } else { 175 | d.scope.Counter(metrics.KafkaDLQErrors).Inc(1) 176 | } 177 | return err 178 | case <-d.stopC: 179 | return errShutdown 180 | } 181 | } 182 | 183 | func (d *bufferedErrorTopic) asyncProducerResponseLoop() { 184 | for { 185 | select { 186 | case msg := <-d.producer.Successes(): 187 | if msg == nil { 188 | continue 189 | } 190 | responseC, ok := msg.Metadata.(chan error) 191 | if !ok { 192 | d.logger.Error("DLQ failed to decode metadata protobuf") 193 | continue 194 | } 195 | responseC <- nil 196 | case perr := <-d.producer.Errors(): 197 | if perr == nil { 198 | continue 199 | } 200 | responseC, ok := perr.Msg.Metadata.(chan error) 201 | if !ok { 202 | continue 203 | } 204 | err := perr.Err 205 | responseC <- err 206 | case <-d.stopC: 207 | d.shutdown() 208 | return 209 | } 210 | } 211 | } 212 | 213 | func (d *bufferedErrorTopic) shutdown() { 214 | close(d.doneC) 215 | } 216 | 217 | func (d *bufferedErrorTopic) newSaramaMessage(key, value []byte, responseC chan error) *sarama.ProducerMessage { 218 | return &sarama.ProducerMessage{ 219 | Topic: d.topic.Name, 220 | Key: sarama.ByteEncoder(key), 221 | Value: sarama.ByteEncoder(value), 222 | Metadata: responseC, 223 | } 224 | } 225 | 226 | // NewRetryDLQMultiplexer returns a DLQ that will produce messages to retryTopic or dlqTopic depending on 227 | // the threshold. 228 | // 229 | // Messages that are added to this DLQ will be sent to retryTopic if the retry count of the message is 230 | // < the threshold. 231 | // Else, it will go to the dlqTopic. 232 | func NewRetryDLQMultiplexer(retryTopic, dlqTopic DLQ, threshold int64) DLQ { 233 | return &dlqMultiplexer{ 234 | retryCountThreshold: threshold, 235 | retryTopic: retryTopic, 236 | dlqTopic: dlqTopic, 237 | } 238 | } 239 | 240 | // Add sends a kafka message to the retry topic or dlq topic depending on the retry count in the message. 241 | // 242 | // If the message RetryCount is greater than or equal to the retryCountThreshold in the multiplexer, 243 | // the message will be sent to the retry topic. 244 | // Else, it will be sent to the dlq topic. 245 | func (d *dlqMultiplexer) Add(m kafka.Message, qtypes ...ErrorQType) error { 246 | // If qtypes is specified, use the first specified qtype as queue target 247 | if len(qtypes) > 0 { 248 | switch qtypes[0] { 249 | case DLQErrorQType: 250 | return d.dlqTopic.Add(m) 251 | default: 252 | return d.retryTopic.Add(m) 253 | } 254 | } 255 | 256 | // Queue type is not specified so use the retry count to determine queue target. 257 | if d.retryCountThreshold >= 0 && m.RetryCount() >= d.retryCountThreshold { 258 | return d.dlqTopic.Add(m) 259 | } 260 | return d.retryTopic.Add(m) 261 | } 262 | 263 | // Start retryDLQMultiplexer will start the retry and dlq buffered dlq producers. 264 | func (d *dlqMultiplexer) Start() error { 265 | if err := d.retryTopic.Start(); err != nil { 266 | return err 267 | } 268 | if err := d.dlqTopic.Start(); err != nil { 269 | d.retryTopic.Stop() 270 | return err 271 | } 272 | return nil 273 | } 274 | 275 | // Stop closes the resources held by the retryTopic and dlqTopic. 276 | func (d *dlqMultiplexer) Stop() { 277 | d.dlqTopic.Stop() 278 | d.retryTopic.Stop() 279 | } 280 | 281 | // NewNoopDLQ returns returns a noop DLQ. 282 | func NewNoopDLQ() DLQ { 283 | return noopDLQ{} 284 | } 285 | 286 | // Start does nothing. 287 | func (d noopDLQ) Start() error { 288 | return nil 289 | } 290 | 291 | // Stop does nothing. 292 | func (d noopDLQ) Stop() {} 293 | 294 | // Add returns errNoDLQ because there is no kafka topic backing it. 295 | func (d noopDLQ) Add(m kafka.Message, qtypes ...ErrorQType) error { 296 | return errNoDLQ 297 | } 298 | -------------------------------------------------------------------------------- /internal/consumer/dlqMetadata.pb.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | // Code generated by protoc-gen-go. DO NOT EDIT. 22 | // source: dlqMetadata.proto 23 | 24 | /* 25 | Package consumer is a generated protocol buffer package. 26 | 27 | It is generated from these files: 28 | dlqMetadata.proto 29 | 30 | It has these top-level messages: 31 | DLQMetadata 32 | */ 33 | package consumer 34 | 35 | import proto "github.com/golang/protobuf/proto" 36 | import fmt "fmt" 37 | import math "math" 38 | 39 | // Reference imports to suppress errors if they are not otherwise used. 40 | var _ = proto.Marshal 41 | var _ = fmt.Errorf 42 | var _ = math.Inf 43 | 44 | // This is a compile-time assertion to ensure that this generated file 45 | // is compatible with the proto package it is being compiled against. 46 | // A compilation error at this line likely means your copy of the 47 | // proto package needs to be updated. 48 | const _ = proto.ProtoPackageIsVersion2 // please upgrade the proto package 49 | 50 | // DLQMetadata contains metadata from the original kafka message. 51 | // The metadata will be encoded and decoded when sending or receiving 52 | // messages from the DLQ cluster in order to present the library 53 | // user a seamless logical topic. 54 | type DLQMetadata struct { 55 | // retry_count is an incrementing value denoting the number 56 | // of times a message has been redelivered. 57 | // It will be 0 on first delivery. 58 | RetryCount int64 `protobuf:"varint,1,opt,name=retry_count,json=retryCount" json:"retry_count,omitempty"` 59 | // topic is the original kafka topic the mesasge was received on. 60 | // This is analogous to the logical topic name. 61 | Topic string `protobuf:"bytes,2,opt,name=topic" json:"topic,omitempty"` 62 | // partition is the original kafka partition the message was received on. 63 | Partition int32 `protobuf:"varint,3,opt,name=partition" json:"partition,omitempty"` 64 | // offset is the record offset of the original message in the original topic-partition. 65 | Offset int64 `protobuf:"varint,4,opt,name=offset" json:"offset,omitempty"` 66 | // timestamp_ns is the original record timestamp of the original mesage. 67 | TimestampNs int64 `protobuf:"varint,5,opt,name=timestamp_ns,json=timestampNs" json:"timestamp_ns,omitempty"` 68 | // data is a byte buffer for storing arbitrary information. 69 | // This is useful if the Kafka Broker version used is < 0.11 70 | // and hence Kafka native record headers (KAFKA-4208) are unavaiable 71 | // so the DLQ metadata must be stored in the record Key or Value. 72 | Data []byte `protobuf:"bytes,6,opt,name=data,proto3" json:"data,omitempty"` 73 | } 74 | 75 | func (m *DLQMetadata) Reset() { *m = DLQMetadata{} } 76 | func (m *DLQMetadata) String() string { return proto.CompactTextString(m) } 77 | func (*DLQMetadata) ProtoMessage() {} 78 | func (*DLQMetadata) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{0} } 79 | 80 | func (m *DLQMetadata) GetRetryCount() int64 { 81 | if m != nil { 82 | return m.RetryCount 83 | } 84 | return 0 85 | } 86 | 87 | func (m *DLQMetadata) GetTopic() string { 88 | if m != nil { 89 | return m.Topic 90 | } 91 | return "" 92 | } 93 | 94 | func (m *DLQMetadata) GetPartition() int32 { 95 | if m != nil { 96 | return m.Partition 97 | } 98 | return 0 99 | } 100 | 101 | func (m *DLQMetadata) GetOffset() int64 { 102 | if m != nil { 103 | return m.Offset 104 | } 105 | return 0 106 | } 107 | 108 | func (m *DLQMetadata) GetTimestampNs() int64 { 109 | if m != nil { 110 | return m.TimestampNs 111 | } 112 | return 0 113 | } 114 | 115 | func (m *DLQMetadata) GetData() []byte { 116 | if m != nil { 117 | return m.Data 118 | } 119 | return nil 120 | } 121 | 122 | func init() { 123 | proto.RegisterType((*DLQMetadata)(nil), "DLQMetadata") 124 | } 125 | 126 | func init() { proto.RegisterFile("dlqMetadata.proto", fileDescriptor0) } 127 | 128 | var fileDescriptor0 = []byte{ 129 | // 180 bytes of a gzipped FileDescriptorProto 130 | 0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0x44, 0xce, 0xbd, 0x0a, 0xc2, 0x30, 131 | 0x14, 0x05, 0x60, 0x62, 0x7f, 0xa0, 0xb7, 0x5d, 0x0c, 0x22, 0x19, 0x04, 0xa3, 0x53, 0x26, 0x17, 132 | 0x1f, 0x41, 0x47, 0x15, 0xcc, 0x0b, 0x94, 0xd8, 0xa6, 0x10, 0xb0, 0x4d, 0x4c, 0xae, 0x83, 0x8f, 133 | 0xe5, 0x1b, 0x4a, 0xe3, 0xdf, 0x76, 0xcf, 0x77, 0xe1, 0x70, 0x60, 0xda, 0x5e, 0x6f, 0x47, 0x8d, 134 | 0xaa, 0x55, 0xa8, 0x36, 0xce, 0x5b, 0xb4, 0xeb, 0x27, 0x81, 0x72, 0x7f, 0x38, 0x7f, 0x95, 0x2e, 135 | 0xa1, 0xf4, 0x1a, 0xfd, 0xa3, 0x6e, 0xec, 0x7d, 0x40, 0x46, 0x38, 0x11, 0x89, 0x84, 0x48, 0xbb, 136 | 0x51, 0xe8, 0x0c, 0x32, 0xb4, 0xce, 0x34, 0x6c, 0xc2, 0x89, 0x28, 0xe4, 0x3b, 0xd0, 0x05, 0x14, 137 | 0x4e, 0x79, 0x34, 0x68, 0xec, 0xc0, 0x12, 0x4e, 0x44, 0x26, 0xff, 0x40, 0xe7, 0x90, 0xdb, 0xae, 138 | 0x0b, 0x1a, 0x59, 0x1a, 0xfb, 0x3e, 0x89, 0xae, 0xa0, 0x42, 0xd3, 0xeb, 0x80, 0xaa, 0x77, 0xf5, 139 | 0x10, 0x58, 0x16, 0xbf, 0xe5, 0xcf, 0x4e, 0x81, 0x52, 0x48, 0xc7, 0x5d, 0x2c, 0xe7, 0x44, 0x54, 140 | 0x32, 0xde, 0x97, 0x3c, 0x4e, 0xdf, 0xbe, 0x02, 0x00, 0x00, 0xff, 0xff, 0x23, 0x67, 0x20, 0xf9, 141 | 0xcf, 0x00, 0x00, 0x00, 142 | } 143 | -------------------------------------------------------------------------------- /internal/consumer/dlqMetadata.proto: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | syntax = "proto3"; 22 | 23 | /* DLQMetadata contains metadata from the original kafka message. 24 | * The metadata will be encoded and decoded when sending or receiving 25 | * messages from the DLQ cluster in order to present the library 26 | * user a seamless logical topic. */ 27 | message DLQMetadata { 28 | // retry_count is an incrementing value denoting the number 29 | // of times a message has been redelivered. 30 | // It will be 0 on first delivery. 31 | int64 retry_count = 1; 32 | // topic is the original kafka topic the mesasge was received on. 33 | // This is analogous to the logical topic name. 34 | string topic = 2; 35 | // partition is the original kafka partition the message was received on. 36 | int32 partition = 3; 37 | // offset is the record offset of the original message in the original topic-partition. 38 | int64 offset = 4; 39 | // timestamp_ns is the original record timestamp of the original mesage. 40 | int64 timestamp_ns = 5; 41 | // data is a byte buffer for storing arbitrary information. 42 | // This is useful if the Kafka Broker version used is < 0.11 43 | // and hence Kafka native record headers (KAFKA-4208) are unavaiable 44 | // so the DLQ metadata must be stored in the record Key or Value. 45 | bytes data = 6; 46 | } -------------------------------------------------------------------------------- /internal/consumer/dlq_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package consumer 22 | 23 | import ( 24 | "errors" 25 | "github.com/Shopify/sarama" 26 | "testing" 27 | "time" 28 | 29 | "github.com/stretchr/testify/suite" 30 | "github.com/uber-go/kafka-client/kafka" 31 | "github.com/uber-go/tally" 32 | "go.uber.org/zap" 33 | ) 34 | 35 | type DLQTestSuite struct { 36 | suite.Suite 37 | saramaProducer *mockSaramaProducer 38 | dlq *bufferedErrorTopic 39 | } 40 | 41 | func TestDLQTestSuite(t *testing.T) { 42 | suite.Run(t, new(DLQTestSuite)) 43 | } 44 | 45 | func (s *DLQTestSuite) SetupTest() { 46 | s.saramaProducer = newMockSaramaProducer() 47 | topic := new(kafka.Topic) 48 | topic.Name = "topic" 49 | topic.Cluster = "cluster" 50 | s.dlq = newBufferedDLQ( 51 | *topic, 52 | s.saramaProducer, 53 | tally.NoopScope, 54 | zap.NewNop(), 55 | ) 56 | } 57 | 58 | func (s *DLQTestSuite) TestBatchProducerSuccessResponse() { 59 | m1 := &mockMessage{ 60 | topic: "topic", 61 | partition: 0, 62 | offset: 0, 63 | } 64 | 65 | s.dlq.Start() 66 | 67 | thread1 := make(chan error) 68 | go func() { 69 | err := s.dlq.Add(m1) 70 | thread1 <- err 71 | }() 72 | 73 | select { 74 | case <-thread1: 75 | s.Fail("add should block until response") 76 | case <-time.After(time.Millisecond): 77 | break 78 | } 79 | 80 | // flush success 81 | select { 82 | case pm := <-s.saramaProducer.inputC: 83 | s.saramaProducer.successC <- pm 84 | case <-time.After(time.Millisecond): 85 | s.Fail("") 86 | } 87 | 88 | select { 89 | case <-time.After(time.Millisecond): 90 | s.Fail("Expected thread1 to return nil") 91 | case err := <-thread1: 92 | s.NoError(err) 93 | } 94 | 95 | s.dlq.Stop() 96 | } 97 | 98 | func (s *DLQTestSuite) TestBatchProducerErrorResponse() { 99 | m1 := &mockMessage{ 100 | topic: "topic", 101 | partition: 0, 102 | offset: 0, 103 | } 104 | 105 | s.dlq.Start() 106 | 107 | thread1 := make(chan error) 108 | go func() { 109 | err := s.dlq.Add(m1) 110 | thread1 <- err 111 | }() 112 | 113 | select { 114 | case <-thread1: 115 | s.Fail("add should block until response") 116 | case <-time.After(time.Millisecond): 117 | break 118 | } 119 | 120 | // flush success 121 | select { 122 | case pm := <-s.saramaProducer.inputC: 123 | s.saramaProducer.errorC <- &sarama.ProducerError{Err: errors.New("error"), Msg: pm} 124 | case <-time.After(time.Millisecond): 125 | s.Fail("") 126 | } 127 | 128 | select { 129 | case <-time.After(time.Millisecond): 130 | s.Fail("Expected thread1 to return nil") 131 | case err := <-thread1: 132 | s.Error(err) 133 | } 134 | 135 | s.dlq.Stop() 136 | } 137 | 138 | func (s *DLQTestSuite) TestBatchProducerWaitingForResponseDoesNotDeadlock() { 139 | m1 := &mockMessage{ 140 | topic: "topic", 141 | partition: 0, 142 | offset: 0, 143 | } 144 | 145 | s.dlq.Start() 146 | 147 | thread1 := make(chan error) 148 | go func() { 149 | err := s.dlq.Add(m1) 150 | thread1 <- err 151 | }() 152 | 153 | select { 154 | case <-thread1: 155 | s.Fail("add should block until response or close") 156 | case <-time.After(time.Millisecond): 157 | break 158 | } 159 | 160 | s.dlq.Stop() 161 | 162 | select { 163 | case <-time.After(time.Millisecond): 164 | s.Fail("Expected thread1 to return nil") 165 | case err := <-thread1: 166 | s.Equal(errShutdown, err) 167 | } 168 | } 169 | 170 | func (s *DLQTestSuite) TestBatchProducerWaitingForProducerDoesNotDeadlock() { 171 | m1 := &mockMessage{ 172 | topic: "topic", 173 | partition: 0, 174 | offset: 0, 175 | } 176 | 177 | s.saramaProducer.inputC = make(chan *sarama.ProducerMessage) 178 | s.dlq.Start() 179 | 180 | thread1 := make(chan error) 181 | go func() { 182 | err := s.dlq.Add(m1) 183 | thread1 <- err 184 | }() 185 | 186 | select { 187 | case <-thread1: 188 | s.Fail("add should block until response or close") 189 | case <-time.After(time.Millisecond): 190 | break 191 | } 192 | 193 | s.dlq.Stop() 194 | 195 | select { 196 | case <-time.After(time.Millisecond): 197 | s.Fail("Expected thread1 to return nil") 198 | case err := <-thread1: 199 | s.Equal(errShutdown, err) 200 | } 201 | } 202 | 203 | type DLQMultiplexerTestSuite struct { 204 | suite.Suite 205 | msg *mockMessage 206 | retry *mockDLQProducer 207 | dlq *mockDLQProducer 208 | multiplexer *dlqMultiplexer 209 | } 210 | 211 | func TestDLQMultiplexerTestSuite(t *testing.T) { 212 | suite.Run(t, new(DLQMultiplexerTestSuite)) 213 | } 214 | 215 | func (s *DLQMultiplexerTestSuite) SetupTest() { 216 | s.msg = new(mockMessage) 217 | s.retry = newMockDLQProducer() 218 | s.dlq = newMockDLQProducer() 219 | s.multiplexer = &dlqMultiplexer{ 220 | retryCountThreshold: 3, 221 | retryTopic: s.retry, 222 | dlqTopic: s.dlq, 223 | } 224 | } 225 | 226 | func (s *DLQMultiplexerTestSuite) TestAdd() { 227 | s.multiplexer.Add(s.msg) 228 | s.Equal(1, s.retry.backlog()) 229 | s.Equal(0, s.dlq.backlog()) 230 | 231 | s.msg.retryCount = 3 232 | s.multiplexer.Add(s.msg) 233 | s.Equal(1, s.dlq.backlog()) 234 | } 235 | 236 | func (s *DLQMultiplexerTestSuite) TestAddWithInfiniteRetry() { 237 | // Set multiplexer for infinity retry 238 | s.multiplexer.retryCountThreshold = -1 239 | 240 | s.multiplexer.Add(s.msg) 241 | s.Equal(1, s.retry.backlog()) 242 | s.Equal(0, s.dlq.backlog()) 243 | 244 | s.msg.retryCount = 3 245 | s.multiplexer.Add(s.msg) 246 | s.Equal(2, s.retry.backlog()) 247 | s.Equal(0, s.dlq.backlog()) 248 | } 249 | 250 | func (s *DLQMultiplexerTestSuite) TestAddToErrorQueue() { 251 | s.multiplexer.Add(s.msg, DLQErrorQType) 252 | s.Equal(0, s.retry.backlog()) 253 | s.Equal(1, s.dlq.backlog()) 254 | 255 | s.multiplexer.Add(s.msg, RetryQErrorQType) 256 | s.Equal(1, s.retry.backlog()) 257 | s.Equal(1, s.dlq.backlog()) 258 | } 259 | -------------------------------------------------------------------------------- /internal/consumer/limit.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package consumer 22 | 23 | import "time" 24 | 25 | const ( 26 | noLimit = -2 27 | // defaultLimit set to -1 so that if you use default limit, no messages will be processed. 28 | defaultLimit = -1 29 | defaultLimitCheckInterval = time.Second 30 | ) 31 | -------------------------------------------------------------------------------- /internal/consumer/message.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package consumer 22 | 23 | import ( 24 | "time" 25 | 26 | "github.com/Shopify/sarama" 27 | "go.uber.org/zap/zapcore" 28 | ) 29 | 30 | type ( 31 | // Message is a wrapper around kafka consumer message 32 | Message struct { 33 | msg *sarama.ConsumerMessage 34 | metadata DLQMetadata 35 | ctx msgContext // consumer metadata, invisible to the application 36 | } 37 | // context that gets piggybacked in the message 38 | // will be used when the message is Acked/Nackd 39 | msgContext struct { 40 | ackID ackID 41 | ackMgr *ackManager 42 | dlq DLQ 43 | } 44 | ) 45 | 46 | func newDLQMetadata() *DLQMetadata { 47 | return &DLQMetadata{ 48 | RetryCount: 0, 49 | Topic: "", 50 | Partition: -1, 51 | Offset: -1, 52 | TimestampNs: -1, 53 | Data: nil, 54 | } 55 | } 56 | 57 | // newMessage builds a new Message object from the given kafka message 58 | func newMessage(scm *sarama.ConsumerMessage, ackID ackID, ackMgr *ackManager, dlq DLQ, metadata DLQMetadata) *Message { 59 | return &Message{ 60 | msg: scm, 61 | ctx: msgContext{ 62 | ackID: ackID, 63 | ackMgr: ackMgr, 64 | dlq: dlq, 65 | }, 66 | metadata: metadata, 67 | } 68 | } 69 | 70 | // Key is a mutable reference to the message's key 71 | func (m *Message) Key() (key []byte) { 72 | if m.metadata.Data != nil { 73 | key = make([]byte, len(m.metadata.Data)) 74 | copy(key, m.metadata.Data) 75 | } else { 76 | key = make([]byte, len(m.msg.Key)) 77 | copy(key, m.msg.Key) 78 | } 79 | return 80 | } 81 | 82 | // Value is a mutable reference to the message's value 83 | func (m *Message) Value() []byte { 84 | result := make([]byte, len(m.msg.Value)) 85 | copy(result, m.msg.Value) 86 | return result 87 | } 88 | 89 | // Topic is the topic from which the message was read 90 | func (m *Message) Topic() string { 91 | if m.metadata.Topic != "" { 92 | return m.metadata.Topic 93 | } 94 | return m.msg.Topic 95 | } 96 | 97 | // Partition is the ID of the partition from which the message was read 98 | func (m *Message) Partition() int32 { 99 | if m.metadata.Partition != -1 { 100 | return m.metadata.Partition 101 | } 102 | return m.msg.Partition 103 | } 104 | 105 | // Offset is the message's offset. 106 | func (m *Message) Offset() int64 { 107 | if m.metadata.Offset != -1 { 108 | return m.metadata.Offset 109 | } 110 | return m.msg.Offset 111 | } 112 | 113 | // Timestamp returns the timestamp for this message 114 | func (m *Message) Timestamp() time.Time { 115 | if m.metadata.TimestampNs != -1 { 116 | return time.Unix(0, m.metadata.TimestampNs) 117 | } 118 | return m.msg.Timestamp 119 | } 120 | 121 | // RetryCount returns the number of times this message has be retried. 122 | func (m *Message) RetryCount() int64 { 123 | return m.metadata.RetryCount 124 | } 125 | 126 | // Ack acknowledges the message 127 | func (m *Message) Ack() error { 128 | ctx := &m.ctx 129 | return ctx.ackMgr.Ack(ctx.ackID) 130 | } 131 | 132 | // Nack negatively acknowledges the message 133 | // also moves the message to a DLQ if the 134 | // consumer has a dlq configured. This method 135 | // will *block* until enqueue to the dlq succeeds 136 | func (m *Message) Nack() error { 137 | ctx := &m.ctx 138 | if err := ctx.dlq.Add(m); err != nil { 139 | return err 140 | } 141 | return ctx.ackMgr.Nack(ctx.ackID) 142 | } 143 | 144 | // NackToDLQ negatively acknowledges the message by sending it directly to the DLQ. 145 | // This method will *block* until enqueue to the dlq succeeds 146 | func (m *Message) NackToDLQ() error { 147 | ctx := &m.ctx 148 | if err := ctx.dlq.Add(m, DLQErrorQType); err != nil { 149 | return err 150 | } 151 | return ctx.ackMgr.Nack(ctx.ackID) 152 | } 153 | 154 | // MarshalLogObject implements zapcore.ObjectMarshaler for structured logging. 155 | func (m *Message) MarshalLogObject(e zapcore.ObjectEncoder) error { 156 | e.AddString("topic", m.Topic()) 157 | e.AddInt32("partition", m.Partition()) 158 | e.AddInt64("offset", m.Offset()) 159 | e.AddTime("timestamp", m.Timestamp()) 160 | e.AddInt64("retryCount", m.RetryCount()) 161 | 162 | if m.metadata.RetryCount != 0 { 163 | e.AddObject("underlyingMsg", zapcore.ObjectMarshalerFunc(func(ee zapcore.ObjectEncoder) error { 164 | ee.AddString("topic", m.msg.Topic) 165 | ee.AddInt32("partition", m.msg.Partition) 166 | ee.AddInt64("offset", m.msg.Offset) 167 | ee.AddTime("timestamp", m.msg.Timestamp) 168 | return nil 169 | })) 170 | } 171 | return nil 172 | } 173 | -------------------------------------------------------------------------------- /internal/consumer/message_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package consumer 22 | 23 | import ( 24 | "testing" 25 | "time" 26 | 27 | "github.com/Shopify/sarama" 28 | "github.com/stretchr/testify/suite" 29 | "github.com/uber-go/kafka-client/kafka" 30 | ) 31 | 32 | type MessageTestSuite struct { 33 | suite.Suite 34 | message *Message 35 | dlqMessage *Message 36 | } 37 | 38 | var _ kafka.Message = (*Message)(nil) 39 | 40 | func TestMessageTestSuite(t *testing.T) { 41 | suite.Run(t, new(MessageTestSuite)) 42 | } 43 | 44 | func (s *MessageTestSuite) SetupTest() { 45 | saramaMessage := &sarama.ConsumerMessage{ 46 | Key: []byte("key1"), 47 | Value: []byte("value1"), 48 | Topic: "topic1", 49 | Partition: 1, 50 | Offset: 1, 51 | Timestamp: time.Unix(1, 0), 52 | } 53 | metadata := &DLQMetadata{ 54 | RetryCount: 1, 55 | Topic: "topic2", 56 | Partition: 2, 57 | Offset: 2, 58 | TimestampNs: 100, 59 | Data: []byte("key2"), 60 | } 61 | ackID := new(ackID) 62 | s.message = newMessage(saramaMessage, *ackID, nil, nil, *newDLQMetadata()) 63 | s.dlqMessage = newMessage(saramaMessage, *ackID, nil, nil, *metadata) 64 | } 65 | 66 | func (s *MessageTestSuite) TestKey() { 67 | s.EqualValues(s.message.msg.Key, s.message.Key()) 68 | s.EqualValues(s.dlqMessage.metadata.Data, s.dlqMessage.Key()) 69 | } 70 | 71 | func (s *MessageTestSuite) TestValue() { 72 | s.EqualValues(s.message.msg.Value, s.message.Value()) 73 | s.EqualValues(s.dlqMessage.msg.Value, s.dlqMessage.Value()) 74 | } 75 | 76 | func (s *MessageTestSuite) TestTopic() { 77 | s.EqualValues(s.message.msg.Topic, s.message.Topic()) 78 | s.EqualValues(s.dlqMessage.metadata.Topic, s.dlqMessage.Topic()) 79 | } 80 | 81 | func (s *MessageTestSuite) TestPartition() { 82 | s.EqualValues(s.message.msg.Partition, s.message.Partition()) 83 | s.EqualValues(s.dlqMessage.metadata.Partition, s.dlqMessage.Partition()) 84 | } 85 | 86 | func (s *MessageTestSuite) TestOffset() { 87 | s.EqualValues(s.message.msg.Offset, s.message.Offset()) 88 | s.EqualValues(s.dlqMessage.metadata.Offset, s.dlqMessage.Offset()) 89 | } 90 | 91 | func (s *MessageTestSuite) TestTimestamp() { 92 | s.EqualValues(s.message.msg.Timestamp, s.message.Timestamp()) 93 | s.EqualValues(time.Unix(0, s.dlqMessage.metadata.TimestampNs), s.dlqMessage.Timestamp()) 94 | } 95 | 96 | func (s *MessageTestSuite) TestRetryCount() { 97 | s.EqualValues(s.message.metadata.RetryCount, s.message.RetryCount()) 98 | s.EqualValues(s.dlqMessage.metadata.RetryCount, s.dlqMessage.RetryCount()) 99 | } 100 | -------------------------------------------------------------------------------- /internal/consumer/mocks_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package consumer 22 | 23 | import ( 24 | "fmt" 25 | "sync" 26 | "sync/atomic" 27 | "time" 28 | 29 | "github.com/Shopify/sarama" 30 | cluster "github.com/bsm/sarama-cluster" 31 | "github.com/uber-go/kafka-client/internal/util" 32 | "github.com/uber-go/kafka-client/kafka" 33 | "go.uber.org/zap/zapcore" 34 | ) 35 | 36 | type ( 37 | mockConsumer struct { 38 | sync.Mutex 39 | name string 40 | topics []string 41 | startErr error 42 | msgC chan kafka.Message 43 | doneC chan struct{} 44 | lifecycle *util.RunLifecycle 45 | } 46 | mockMessage struct { 47 | key []byte 48 | value []byte 49 | topic string 50 | partition int32 51 | offset int64 52 | timestamp time.Time 53 | retryCount int64 54 | ackErr error 55 | nackErr error 56 | } 57 | 58 | mockSaramaClient struct { 59 | closed int32 60 | } 61 | mockSaramaProducer struct { 62 | closed int32 63 | inputC chan *sarama.ProducerMessage 64 | successC chan *sarama.ProducerMessage 65 | errorC chan *sarama.ProducerError 66 | } 67 | mockSaramaConsumer struct { 68 | sync.Mutex 69 | closed int64 70 | offsets map[int32]int64 71 | errorC chan error 72 | notifyC chan *cluster.Notification 73 | partitionC chan cluster.PartitionConsumer 74 | messages chan *sarama.ConsumerMessage 75 | } 76 | mockPartitionedConsumer struct { 77 | id int32 78 | topic string 79 | closed int64 80 | beginOffset int64 81 | msgC chan *sarama.ConsumerMessage 82 | } 83 | mockDLQProducer struct { 84 | sync.Mutex 85 | stopped int32 86 | messages []kafka.Message 87 | } 88 | ) 89 | 90 | func newMockSaramaProducer() *mockSaramaProducer { 91 | return &mockSaramaProducer{ 92 | inputC: make(chan *sarama.ProducerMessage, 10), 93 | successC: make(chan *sarama.ProducerMessage, 10), 94 | errorC: make(chan *sarama.ProducerError, 10), 95 | } 96 | } 97 | 98 | func (m *mockSaramaProducer) AsyncClose() { 99 | atomic.AddInt32(&m.closed, 1) 100 | } 101 | 102 | func (m *mockSaramaProducer) Close() error { 103 | m.AsyncClose() 104 | return nil 105 | } 106 | 107 | func (m *mockSaramaProducer) Input() chan<- *sarama.ProducerMessage { 108 | return m.inputC 109 | } 110 | 111 | func (m *mockSaramaProducer) Successes() <-chan *sarama.ProducerMessage { 112 | return m.successC 113 | } 114 | 115 | func (m *mockSaramaProducer) Errors() <-chan *sarama.ProducerError { 116 | return m.errorC 117 | } 118 | 119 | func newMockConsumer(name string, topics []string, msgC chan kafka.Message) *mockConsumer { 120 | return &mockConsumer{ 121 | name: name, 122 | topics: topics, 123 | msgC: msgC, 124 | doneC: make(chan struct{}), 125 | lifecycle: util.NewRunLifecycle("mockConsumer-" + name), 126 | } 127 | } 128 | 129 | // Name returns the id for this mockConsumer. 130 | func (c *mockConsumer) Name() string { 131 | c.Lock() 132 | defer c.Unlock() 133 | return c.name 134 | } 135 | 136 | // Topics returns the list of topics this mock consumer was assigned to consumer. 137 | func (c *mockConsumer) Topics() []string { 138 | c.Lock() 139 | defer c.Unlock() 140 | return c.topics 141 | } 142 | 143 | // Start will start the mockConsumer on the initialized lifecycle and return startErr. 144 | func (c *mockConsumer) Start() error { 145 | return c.lifecycle.Start(func() error { 146 | return c.startErr 147 | }) 148 | } 149 | 150 | // Stop will stop the lifecycle. 151 | func (c *mockConsumer) Stop() { 152 | c.lifecycle.Stop(func() { 153 | close(c.doneC) 154 | }) 155 | } 156 | 157 | // Closed will return a channel that can be used to check if the consumer has been stopped. 158 | func (c *mockConsumer) Closed() <-chan struct{} { 159 | return c.doneC 160 | } 161 | 162 | // Messages return the message channel. 163 | func (c *mockConsumer) Messages() <-chan kafka.Message { 164 | return c.msgC 165 | } 166 | 167 | func (m *mockMessage) MarshalLogObject(zapcore.ObjectEncoder) error { 168 | return nil 169 | } 170 | 171 | func (m *mockMessage) Key() []byte { 172 | return m.key 173 | } 174 | 175 | func (m *mockMessage) Value() []byte { 176 | return m.value 177 | } 178 | 179 | func (m *mockMessage) Topic() string { 180 | return m.topic 181 | } 182 | 183 | func (m *mockMessage) Partition() int32 { 184 | return m.partition 185 | } 186 | 187 | func (m *mockMessage) Offset() int64 { 188 | return m.offset 189 | } 190 | 191 | func (m *mockMessage) Timestamp() time.Time { 192 | return m.timestamp 193 | } 194 | 195 | func (m *mockMessage) RetryCount() int64 { 196 | return m.retryCount 197 | } 198 | 199 | func (m *mockMessage) Ack() error { 200 | return m.ackErr 201 | } 202 | 203 | func (m *mockMessage) Nack() error { 204 | return m.nackErr 205 | } 206 | 207 | func (m *mockMessage) NackToDLQ() error { 208 | return m.nackErr 209 | } 210 | 211 | func newMockPartitionedConsumer(topic string, id int32, beginOffset int64, rcvBufSize int) *mockPartitionedConsumer { 212 | return &mockPartitionedConsumer{ 213 | id: id, 214 | topic: topic, 215 | beginOffset: beginOffset, 216 | msgC: make(chan *sarama.ConsumerMessage, rcvBufSize), 217 | } 218 | } 219 | 220 | func (m *mockPartitionedConsumer) start() *sync.WaitGroup { 221 | var wg sync.WaitGroup 222 | wg.Add(1) 223 | go func() { 224 | offset := m.beginOffset + 1 225 | for i := 0; i < 100; i++ { 226 | m.sendMsg(offset) 227 | offset++ 228 | } 229 | wg.Done() 230 | }() 231 | return &wg 232 | } 233 | 234 | func (m *mockPartitionedConsumer) stop() { 235 | close(m.msgC) 236 | } 237 | 238 | func (m *mockPartitionedConsumer) sendMsg(offset int64) { 239 | msg := &sarama.ConsumerMessage{ 240 | Topic: m.topic, 241 | Partition: m.id, 242 | Value: []byte(fmt.Sprintf("msg-%v", offset)), 243 | Offset: offset, 244 | Timestamp: time.Now(), 245 | } 246 | m.msgC <- msg 247 | } 248 | 249 | func (m *mockPartitionedConsumer) Close() error { 250 | atomic.StoreInt64(&m.closed, 1) 251 | return nil 252 | } 253 | 254 | func (m *mockPartitionedConsumer) isClosed() bool { 255 | return atomic.LoadInt64(&m.closed) == 1 256 | } 257 | 258 | func (m *mockPartitionedConsumer) AsyncClose() {} 259 | 260 | func (m *mockPartitionedConsumer) Errors() <-chan *sarama.ConsumerError { 261 | return nil 262 | } 263 | 264 | // Messages returns the read channel for the messages that are returned by 265 | // the broker. 266 | func (m *mockPartitionedConsumer) Messages() <-chan *sarama.ConsumerMessage { 267 | return m.msgC 268 | } 269 | 270 | // HighWaterMarkOffset returns the high water mark offset of the partition, 271 | // i.e. the offset that will be used for the next message that will be produced. 272 | // You can use this to determine how far behind the processing is. 273 | func (m *mockPartitionedConsumer) HighWaterMarkOffset() int64 { 274 | return 0 275 | } 276 | 277 | // Topic returns the consumed topic name 278 | func (m *mockPartitionedConsumer) Topic() string { 279 | return m.topic 280 | } 281 | 282 | // Partition returns the consumed partition 283 | func (m *mockPartitionedConsumer) Partition() int32 { 284 | return m.id 285 | } 286 | 287 | func newMockSaramaConsumer() *mockSaramaConsumer { 288 | return &mockSaramaConsumer{ 289 | errorC: make(chan error, 1), 290 | notifyC: make(chan *cluster.Notification, 1), 291 | partitionC: make(chan cluster.PartitionConsumer, 1), 292 | offsets: make(map[int32]int64), 293 | messages: make(chan *sarama.ConsumerMessage, 1), 294 | } 295 | } 296 | 297 | func (m *mockSaramaConsumer) offset(id int) int64 { 298 | m.Lock() 299 | off, ok := m.offsets[int32(id)] 300 | m.Unlock() 301 | if !ok { 302 | return 0 303 | } 304 | return off 305 | } 306 | 307 | func (m *mockSaramaConsumer) ResetPartitionOffset(topic string, partition int32, offset int64, metadata string) { 308 | } 309 | 310 | func (m *mockSaramaConsumer) Errors() <-chan error { 311 | return m.errorC 312 | } 313 | 314 | func (m *mockSaramaConsumer) Notifications() <-chan *cluster.Notification { 315 | return m.notifyC 316 | } 317 | 318 | func (m *mockSaramaConsumer) Partitions() <-chan cluster.PartitionConsumer { 319 | return m.partitionC 320 | } 321 | 322 | func (m *mockSaramaConsumer) CommitOffsets() error { 323 | return nil 324 | } 325 | 326 | func (m *mockSaramaConsumer) Messages() <-chan *sarama.ConsumerMessage { 327 | return m.messages 328 | } 329 | 330 | func (m *mockSaramaConsumer) MarkOffset(msg *sarama.ConsumerMessage, metadata string) { 331 | m.Lock() 332 | m.offsets[msg.Partition] = msg.Offset 333 | m.Unlock() 334 | } 335 | 336 | func (m *mockSaramaConsumer) MarkPartitionOffset(topic string, partition int32, offset int64, metadata string) { 337 | m.Lock() 338 | m.offsets[partition] = offset 339 | m.Unlock() 340 | } 341 | 342 | func (m *mockSaramaConsumer) HighWaterMarks() map[string]map[int32]int64 { 343 | result := make(map[string]map[int32]int64) 344 | result["test"] = make(map[int32]int64) 345 | m.Lock() 346 | for k, v := range m.offsets { 347 | result["test"][k] = v 348 | } 349 | m.Unlock() 350 | return result 351 | } 352 | 353 | func (m *mockSaramaConsumer) Close() error { 354 | atomic.AddInt64(&m.closed, 1) 355 | return nil 356 | } 357 | 358 | func (m *mockSaramaConsumer) isClosed() bool { 359 | return atomic.LoadInt64(&m.closed) == 1 360 | } 361 | 362 | func newMockDLQProducer() *mockDLQProducer { 363 | return &mockDLQProducer{ 364 | messages: make([]kafka.Message, 0, 100), 365 | stopped: 0, 366 | } 367 | } 368 | 369 | func (d *mockDLQProducer) Start() error { 370 | return nil 371 | } 372 | 373 | func (d *mockDLQProducer) Stop() { 374 | atomic.AddInt32(&d.stopped, 1) 375 | } 376 | 377 | func (d *mockDLQProducer) Add(m kafka.Message, qType ...ErrorQType) error { 378 | d.Lock() 379 | defer d.Unlock() 380 | d.messages = append(d.messages, m) 381 | return nil 382 | } 383 | 384 | func (d *mockDLQProducer) isClosed() bool { 385 | return atomic.LoadInt32(&d.stopped) > 0 386 | } 387 | func (d *mockDLQProducer) backlog() int { 388 | d.Lock() 389 | defer d.Unlock() 390 | return len(d.messages) 391 | } 392 | 393 | func newMockSaramaClient() *mockSaramaClient { 394 | return &mockSaramaClient{ 395 | closed: 0, 396 | } 397 | } 398 | 399 | func (m *mockSaramaClient) Config() *sarama.Config { 400 | panic("implement me") 401 | } 402 | 403 | func (m *mockSaramaClient) Brokers() []*sarama.Broker { 404 | panic("implement me") 405 | } 406 | 407 | func (m *mockSaramaClient) Controller() (*sarama.Broker, error) { 408 | panic("implement me") 409 | } 410 | 411 | func (m *mockSaramaClient) Topics() ([]string, error) { 412 | panic("implement me") 413 | } 414 | 415 | func (m *mockSaramaClient) Partitions(topic string) ([]int32, error) { 416 | panic("implement me") 417 | } 418 | 419 | func (m *mockSaramaClient) WritablePartitions(topic string) ([]int32, error) { 420 | panic("implement me") 421 | } 422 | 423 | func (m *mockSaramaClient) Leader(topic string, partitionID int32) (*sarama.Broker, error) { 424 | panic("implement me") 425 | } 426 | 427 | func (m *mockSaramaClient) Replicas(topic string, partitionID int32) ([]int32, error) { 428 | panic("implement me") 429 | } 430 | 431 | func (m *mockSaramaClient) InSyncReplicas(topic string, partitionID int32) ([]int32, error) { 432 | panic("implement me") 433 | } 434 | 435 | func (m *mockSaramaClient) RefreshMetadata(topics ...string) error { 436 | panic("implement me") 437 | } 438 | 439 | func (m *mockSaramaClient) GetOffset(topic string, partitionID int32, time int64) (int64, error) { 440 | panic("implement me") 441 | } 442 | 443 | func (m *mockSaramaClient) Coordinator(consumerGroup string) (*sarama.Broker, error) { 444 | panic("implement me") 445 | } 446 | 447 | func (m *mockSaramaClient) RefreshCoordinator(consumerGroup string) error { 448 | panic("implement me") 449 | } 450 | 451 | func (m *mockSaramaClient) Close() error { 452 | atomic.AddInt32(&m.closed, 1) 453 | return nil 454 | } 455 | 456 | func (m *mockSaramaClient) Closed() bool { 457 | return atomic.LoadInt32(&m.closed) > 0 458 | } 459 | -------------------------------------------------------------------------------- /internal/consumer/multiClusterConsumer_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package consumer 22 | 23 | import ( 24 | "testing" 25 | "time" 26 | 27 | "github.com/Shopify/sarama" 28 | "github.com/stretchr/testify/suite" 29 | "github.com/uber-go/kafka-client/kafka" 30 | "github.com/uber-go/tally" 31 | "go.uber.org/zap" 32 | ) 33 | 34 | type MultiClusterConsumerTestSuite struct { 35 | suite.Suite 36 | consumer *MultiClusterConsumer 37 | config *kafka.ConsumerConfig 38 | topics kafka.ConsumerTopicList 39 | options *Options 40 | msgCh chan kafka.Message 41 | } 42 | 43 | // verify *MultiClusterConsumer is implements kafka.Consumer 44 | var _ kafka.Consumer = (*MultiClusterConsumer)(nil) 45 | 46 | func (s *MultiClusterConsumerTestSuite) SetupTest() { 47 | topic := kafka.ConsumerTopic{ 48 | Topic: kafka.Topic{ 49 | Name: "unit-test", 50 | Cluster: "production-cluster", 51 | }, 52 | DLQ: kafka.Topic{ 53 | Name: "unit-test-dlq", 54 | Cluster: "dlq-cluster", 55 | }, 56 | } 57 | s.topics = []kafka.ConsumerTopic{topic} 58 | s.config = &kafka.ConsumerConfig{ 59 | TopicList: s.topics, 60 | GroupName: "unit-test-cg", 61 | Concurrency: 4, 62 | } 63 | s.options = testConsumerOptions() 64 | s.msgCh = make(chan kafka.Message) 65 | s.consumer = NewMultiClusterConsumer( 66 | s.config.GroupName, 67 | s.topics, 68 | make(map[ClusterGroup]*ClusterConsumer), 69 | make(map[ClusterGroup]sarama.Client), 70 | s.msgCh, 71 | tally.NoopScope, 72 | zap.L(), 73 | ) 74 | s.Equal(s.config.GroupName, s.consumer.Name()) 75 | s.Equal(s.config.TopicList.TopicNames(), s.consumer.Topics().TopicNames()) 76 | } 77 | 78 | func (s *MultiClusterConsumerTestSuite) TeardownTest() { 79 | s.consumer.Stop() 80 | } 81 | 82 | func TestMultiClusterConsumerSuite(t *testing.T) { 83 | suite.Run(t, new(MultiClusterConsumerTestSuite)) 84 | } 85 | 86 | func (s *MultiClusterConsumerTestSuite) TestStartSucceeds() { 87 | cc1 := NewClusterConsumer("cc1", newMockSaramaConsumer(), make(map[string]*TopicConsumer), tally.NoopScope, zap.NewNop()) 88 | cc2 := NewClusterConsumer("cc2", newMockSaramaConsumer(), make(map[string]*TopicConsumer), tally.NoopScope, zap.NewNop()) 89 | s.consumer.clusterConsumerMap[ClusterGroup{"cc1", "g"}] = cc1 90 | s.consumer.clusterConsumerMap[ClusterGroup{"cc2", "g"}] = cc2 91 | 92 | s.NoError(s.consumer.Start()) 93 | 94 | started, stopped := cc1.lifecycle.Status() 95 | s.True(started) 96 | s.False(stopped) 97 | started, stopped = cc2.lifecycle.Status() 98 | s.True(started) 99 | s.False(stopped) 100 | 101 | s.consumer.Stop() 102 | select { 103 | case <-s.consumer.Closed(): 104 | case <-time.After(time.Millisecond): 105 | s.Fail("Consumer should be closed") 106 | } 107 | } 108 | 109 | func (s *MultiClusterConsumerTestSuite) TestStartError() { 110 | cc1 := NewClusterConsumer("cc1", newMockSaramaConsumer(), make(map[string]*TopicConsumer), tally.NoopScope, zap.L()) 111 | cc2 := NewClusterConsumer("cc2", newMockSaramaConsumer(), make(map[string]*TopicConsumer), tally.NoopScope, zap.L()) 112 | s.consumer.clusterConsumerMap[ClusterGroup{"cc1", "g"}] = cc1 113 | s.consumer.clusterConsumerMap[ClusterGroup{"cc2", "g"}] = cc2 114 | 115 | cc1.Start() 116 | cc1.Stop() 117 | 118 | s.Error(s.consumer.Start()) 119 | } 120 | -------------------------------------------------------------------------------- /internal/consumer/multiclusterConsumer.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package consumer 22 | 23 | import ( 24 | "errors" 25 | "github.com/Shopify/sarama" 26 | "github.com/uber-go/kafka-client/internal/metrics" 27 | "github.com/uber-go/kafka-client/internal/util" 28 | "github.com/uber-go/kafka-client/kafka" 29 | "github.com/uber-go/tally" 30 | "go.uber.org/zap" 31 | ) 32 | 33 | type ( 34 | // MultiClusterConsumer is a map that contains multiple kafka consumers 35 | MultiClusterConsumer struct { 36 | groupName string 37 | topics kafka.ConsumerTopicList 38 | clusterConsumerMap map[ClusterGroup]*ClusterConsumer 39 | clusterToSaramaClientMap map[ClusterGroup]sarama.Client 40 | msgC chan kafka.Message 41 | doneC chan struct{} 42 | scope tally.Scope 43 | logger *zap.Logger 44 | lifecycle *util.RunLifecycle 45 | } 46 | ) 47 | 48 | // NewMultiClusterConsumer returns a new consumer that consumes messages from 49 | // multiple Kafka clusters. 50 | func NewMultiClusterConsumer( 51 | groupName string, 52 | topics kafka.ConsumerTopicList, 53 | clusterConsumerMap map[ClusterGroup]*ClusterConsumer, 54 | saramaClients map[ClusterGroup]sarama.Client, 55 | msgC chan kafka.Message, 56 | scope tally.Scope, 57 | logger *zap.Logger, 58 | ) *MultiClusterConsumer { 59 | return &MultiClusterConsumer{ 60 | groupName: groupName, 61 | topics: topics, 62 | clusterConsumerMap: clusterConsumerMap, 63 | clusterToSaramaClientMap: saramaClients, 64 | msgC: msgC, 65 | doneC: make(chan struct{}), 66 | scope: scope, 67 | logger: logger, 68 | lifecycle: util.NewRunLifecycle(groupName + "-consumer"), 69 | } 70 | } 71 | 72 | // Name returns the consumer group name used by this consumer. 73 | func (c *MultiClusterConsumer) Name() string { 74 | return c.groupName 75 | } 76 | 77 | // Topics returns a list of topics this consumer is consuming from. 78 | func (c *MultiClusterConsumer) Topics() kafka.ConsumerTopicList { 79 | return c.topics 80 | } 81 | 82 | // Start will fail to start if there is any clusterConsumer that fails. 83 | func (c *MultiClusterConsumer) Start() error { 84 | err := c.lifecycle.Start(func() (err error) { 85 | for clusterName, consumer := range c.clusterConsumerMap { 86 | if err = consumer.Start(); err != nil { 87 | c.logger.With( 88 | zap.Error(err), 89 | zap.String("cluster", clusterName.Cluster), 90 | ).Error("multicluster consumer start error") 91 | return 92 | } 93 | } 94 | return 95 | }) 96 | if err != nil { 97 | c.Stop() 98 | return err 99 | } 100 | c.logger.Info("multicluster consumer started", zap.String("groupName", c.groupName), zap.Array("topicList", c.topics)) 101 | c.scope.Counter(metrics.KafkaConsumerStarted).Inc(1) 102 | return nil 103 | } 104 | 105 | // Stop will stop the consumer. 106 | func (c *MultiClusterConsumer) Stop() { 107 | c.lifecycle.Stop(func() { 108 | for _, consumer := range c.clusterConsumerMap { 109 | consumer.Stop() 110 | } 111 | for _, client := range c.clusterToSaramaClientMap { 112 | client.Close() 113 | } 114 | close(c.doneC) 115 | c.logger.Info("multicluster consumer stopped", zap.String("groupName", c.groupName), zap.Array("topicList", c.topics)) 116 | c.scope.Counter(metrics.KafkaConsumerStopped).Inc(1) 117 | }) 118 | } 119 | 120 | // Closed returns a channel that will be closed when the consumer is closed. 121 | func (c *MultiClusterConsumer) Closed() <-chan struct{} { 122 | return c.doneC 123 | } 124 | 125 | // Messages returns a channel to receive messages on. 126 | func (c *MultiClusterConsumer) Messages() <-chan kafka.Message { 127 | return c.msgC 128 | } 129 | 130 | // ResetOffset will reset the consumer offset for the specified cluster, topic, partition. 131 | func (c *MultiClusterConsumer) ResetOffset(cluster, group, topic string, partition int32, offsetRange kafka.OffsetRange) error { 132 | cc, ok := c.clusterConsumerMap[ClusterGroup{Cluster: cluster, Group: group}] 133 | if !ok { 134 | return errors.New("no cluster consumer found") 135 | } 136 | return cc.ResetOffset(topic, partition, offsetRange) 137 | } 138 | 139 | // MergeDLQ will merge the offset range for each partition of the DLQ topic for the specified ConsumerTopic. 140 | // Topic should be the DLQ topic (with __dlq). 141 | func (c *MultiClusterConsumer) MergeDLQ(cluster, group, topic string, partition int32, offsetRange kafka.OffsetRange) error { 142 | return c.ResetOffset(cluster, group+DLQConsumerGroupNameSuffix, topic, partition, offsetRange) 143 | } 144 | -------------------------------------------------------------------------------- /internal/consumer/options.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package consumer 22 | 23 | import ( 24 | "crypto/tls" 25 | "time" 26 | 27 | "github.com/Shopify/sarama" 28 | "github.com/bsm/sarama-cluster" 29 | ) 30 | 31 | type ( 32 | // Options are the tunable and injectable options for the consumer 33 | Options struct { 34 | ClientID string // client ID 35 | RcvBufferSize int // aggregate message buffer size 36 | PartitionRcvBufferSize int // message buffer size for each partition 37 | Concurrency int // number of goroutines that will concurrently process messages 38 | OffsetPolicy int64 39 | OffsetCommitInterval time.Duration 40 | RebalanceDwellTime time.Duration 41 | MaxProcessingTime time.Duration // amount of time a partitioned consumer will wait during a drain 42 | ConsumerMode cluster.ConsumerMode 43 | ProducerMaxMessageByes int 44 | FetchDefaultBytes int32 45 | OtherConsumerTopics []Topic 46 | TLSConfig *tls.Config // TLSConfig is the configuration to use for secure connections, not nil -> enable, nil -> disabled 47 | } 48 | ) 49 | 50 | // DefaultOptions returns the default options 51 | func DefaultOptions() *Options { 52 | return &Options{ 53 | Concurrency: 1024, 54 | RcvBufferSize: 2 * 1024, // twice the concurrency for compute/io overlap 55 | PartitionRcvBufferSize: 32, 56 | OffsetCommitInterval: time.Second, 57 | RebalanceDwellTime: time.Second, 58 | MaxProcessingTime: 250 * time.Millisecond, 59 | OffsetPolicy: sarama.OffsetOldest, 60 | ConsumerMode: cluster.ConsumerModePartitions, 61 | FetchDefaultBytes: 30 * 1024 * 1024, // 30MB. 62 | ProducerMaxMessageByes: 30 * 1024 * 1024, // 30MB 63 | OtherConsumerTopics: make([]Topic, 0, 10), 64 | } 65 | } 66 | -------------------------------------------------------------------------------- /internal/consumer/partitionConsumer_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package consumer 22 | 23 | import ( 24 | "testing" 25 | "time" 26 | 27 | "github.com/Shopify/sarama" 28 | "github.com/stretchr/testify/suite" 29 | "github.com/uber-go/kafka-client/kafka" 30 | "github.com/uber-go/tally" 31 | "go.uber.org/zap" 32 | ) 33 | 34 | const ( 35 | testWaitDuration = 2 * time.Second 36 | ) 37 | 38 | type RangePartitionConsumerTestSuite struct { 39 | suite.Suite 40 | msgC chan kafka.Message 41 | partitionConsumer *partitionConsumer 42 | rangePartitionConsumer *rangePartitionConsumer 43 | saramaConsumer *mockSaramaConsumer 44 | saramaPartitionConsumer *mockPartitionedConsumer 45 | } 46 | 47 | func TestRangePartitionConsumerTestSuite(t *testing.T) { 48 | suite.Run(t, new(RangePartitionConsumerTestSuite)) 49 | } 50 | 51 | func (s *RangePartitionConsumerTestSuite) SetupTest() { 52 | s.msgC = make(chan kafka.Message, 5) 53 | s.saramaConsumer = newMockSaramaConsumer() 54 | s.saramaPartitionConsumer = newMockPartitionedConsumer("topic", 0, 100, 5) 55 | topic := new(Topic) 56 | topic.Name = "topic" 57 | topic.DLQMetadataDecoder = NoopDLQMetadataDecoder 58 | opts := DefaultOptions() 59 | l := zap.NewNop() 60 | s.partitionConsumer = newPartitionConsumer( 61 | *topic, 62 | s.saramaConsumer, 63 | s.saramaPartitionConsumer, 64 | opts, 65 | s.msgC, 66 | nil, 67 | true, 68 | tally.NoopScope, 69 | l, 70 | ) 71 | s.rangePartitionConsumer = newRangePartitionConsumer(s.partitionConsumer) 72 | } 73 | 74 | func (s *RangePartitionConsumerTestSuite) TestDelayMsg() { 75 | m := &sarama.ConsumerMessage{ 76 | Offset: 100, 77 | Timestamp: time.Now(), 78 | } 79 | t1 := time.Now() 80 | delay := time.Millisecond 81 | s.rangePartitionConsumer.topicPartition.Delay = delay 82 | s.rangePartitionConsumer.delayMsg(m) 83 | t2 := time.Now() 84 | if t2.Sub(t1) < delay { 85 | s.Fail("expect to sleep around " + delay.String()) 86 | } 87 | 88 | // lag > delay, almost return immediately 89 | t1 = time.Now() 90 | m.Timestamp = t1.Add(time.Millisecond * -3) 91 | s.rangePartitionConsumer.delayMsg(m) 92 | t2 = time.Now() 93 | if t2.Sub(t1) > time.Millisecond { 94 | s.Fail("expect no delay on msg, actual time cost is " + t2.Sub(t1).String()) 95 | } 96 | 97 | // delay = 0, almost return immediately 98 | t1 = time.Now() 99 | s.rangePartitionConsumer.topicPartition.Delay = 0 100 | m.Timestamp = t1.Add(time.Millisecond) 101 | s.rangePartitionConsumer.delayMsg(m) 102 | t2 = time.Now() 103 | if t2.Sub(t1) > time.Millisecond { 104 | s.Fail("expect no delay on msg") 105 | } 106 | } 107 | 108 | func (s *RangePartitionConsumerTestSuite) TestOffsetNotificationTriggersMessageConsuming() { 109 | s.rangePartitionConsumer.Start() 110 | 111 | s.saramaPartitionConsumer.sendMsg(100) // this should be consumed 112 | s.saramaPartitionConsumer.sendMsg(101) // this message should be ignored 113 | s.saramaPartitionConsumer.sendMsg(91) // this should be the first message after ResetOffset 114 | s.saramaPartitionConsumer.sendMsg(92) // this should be second message after ResetOffset 115 | s.saramaPartitionConsumer.sendMsg(93) // this message should be ignored b/c it is larger than HighOffset of OffsetRange 116 | 117 | // set offset range to 100-100, so receive one message 118 | s.NoError(s.rangePartitionConsumer.ResetOffset(kafka.OffsetRange{LowOffset: 100, HighOffset: 100})) 119 | select { 120 | case msg := <-s.msgC: 121 | s.EqualValues(100, msg.Offset()) 122 | msg.Ack() 123 | case <-time.After(testWaitDuration): 124 | s.Fail("expected message 100") 125 | } 126 | 127 | // ackmgr should be at 100 since we commit 100 128 | s.EqualValues(100, s.rangePartitionConsumer.ackMgr.CommitLevel()) 129 | 130 | // Trigger reset offset 131 | s.NoError(s.rangePartitionConsumer.ResetOffset(kafka.OffsetRange{LowOffset: 91, HighOffset: 92})) 132 | for i := 0; i < 2; i++ { 133 | select { 134 | case msg := <-s.msgC: 135 | s.EqualValues(91+i, msg.Offset()) 136 | msg.Ack() 137 | case <-time.After(testWaitDuration): 138 | s.Fail("expected 2 messages on msgC") 139 | } 140 | } 141 | 142 | select { 143 | case <-s.msgC: 144 | s.Fail("expect only 2 messages on msgC") 145 | case <-time.After(testWaitDuration): 146 | break 147 | } 148 | 149 | // ackMgr should reset to 90 150 | s.EqualValues(92, s.rangePartitionConsumer.ackMgr.CommitLevel()) 151 | s.rangePartitionConsumer.Stop() 152 | } 153 | 154 | func (s *RangePartitionConsumerTestSuite) TestResetOffsetNoHighOffsetReturnsError() { 155 | s.Error(s.rangePartitionConsumer.ResetOffset(kafka.OffsetRange{LowOffset: 90, HighOffset: -1})) 156 | } 157 | 158 | func (s *RangePartitionConsumerTestSuite) TestResetOffsetStoppedReturnsError() { 159 | close(s.rangePartitionConsumer.stopC) 160 | s.Error(s.rangePartitionConsumer.ResetOffset(kafka.OffsetRange{LowOffset: 90, HighOffset: 100})) 161 | } 162 | 163 | func (s *RangePartitionConsumerTestSuite) TestResetOffsetWithCurrentOffsetReturnsError() { 164 | s.rangePartitionConsumer.Start() 165 | s.NoError(s.rangePartitionConsumer.ResetOffset(kafka.OffsetRange{LowOffset: 95, HighOffset: 100})) 166 | // existing 95-100 merge, so this returns error. 167 | s.Error(s.rangePartitionConsumer.ResetOffset(kafka.OffsetRange{LowOffset: 100, HighOffset: 120})) 168 | s.rangePartitionConsumer.Stop() 169 | } 170 | -------------------------------------------------------------------------------- /internal/consumer/topicConsumer.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package consumer 22 | 23 | import ( 24 | "sync" 25 | 26 | "github.com/bsm/sarama-cluster" 27 | "github.com/uber-go/kafka-client/internal/metrics" 28 | "github.com/uber-go/kafka-client/kafka" 29 | "github.com/uber-go/tally" 30 | "go.uber.org/zap" 31 | ) 32 | 33 | type ( 34 | // TopicConsumer is a consumer for a specific topic. 35 | // TopicConsumer is an abstraction that runs on the same 36 | // goroutine as the cluster consumer. 37 | TopicConsumer struct { 38 | topic Topic 39 | msgC chan kafka.Message 40 | partitionConsumerMap map[int32]PartitionConsumer 41 | dlq DLQ 42 | saramaConsumer SaramaConsumer // SaramaConsumer is a shared resource that is owned by ClusterConsumer. 43 | options *Options 44 | scope tally.Scope 45 | logger *zap.Logger 46 | } 47 | ) 48 | 49 | // NewTopicConsumer returns a new TopicConsumer for consuming from a single topic. 50 | func NewTopicConsumer( 51 | topic Topic, 52 | msgC chan kafka.Message, 53 | consumer SaramaConsumer, 54 | dlq DLQ, 55 | options *Options, 56 | scope tally.Scope, 57 | logger *zap.Logger, 58 | ) *TopicConsumer { 59 | logger = logger.With( 60 | zap.String("topic", topic.Name), 61 | zap.String("cluster", topic.Cluster), 62 | ) 63 | scope = scope.Tagged(map[string]string{"topic": topic.Name, "cluster": topic.Cluster}) 64 | return &TopicConsumer{ 65 | topic: topic, 66 | msgC: msgC, 67 | partitionConsumerMap: make(map[int32]PartitionConsumer), 68 | dlq: dlq, 69 | saramaConsumer: consumer, 70 | options: options, 71 | scope: scope, 72 | logger: logger.With(zap.String("topic", topic.Name), zap.String("cluster", topic.Cluster)), 73 | } 74 | } 75 | 76 | // Start the DLQ consumer goroutine. 77 | func (c *TopicConsumer) Start() error { 78 | if err := c.dlq.Start(); err != nil { 79 | c.logger.Error("topic consumer start error", zap.Error(err)) 80 | return err 81 | } 82 | c.logger.Info("topic consumer started", zap.Object("topic", c.topic)) 83 | return nil 84 | } 85 | 86 | // Stop shutdown and frees the resource held by this TopicConsumer and stops the batch DLQ producer. 87 | func (c *TopicConsumer) Stop() { 88 | c.shutdown() // close each partition consumer 89 | c.dlq.Stop() // close DLQ, which also closes sarama AsyncProducer 90 | c.logger.Info("topic consumer stopped", zap.Object("topic", c.topic)) 91 | } 92 | 93 | func (c *TopicConsumer) addPartitionConsumer(pc cluster.PartitionConsumer) { 94 | partition := pc.Partition() 95 | old, ok := c.partitionConsumerMap[partition] 96 | if ok { 97 | old.Stop() 98 | delete(c.partitionConsumerMap, partition) 99 | } 100 | c.logger.Debug("topic consumer adding new partition consumer", zap.Int32("partition", partition)) 101 | p := c.topic.PartitionConsumerFactory(c.topic, c.saramaConsumer, pc, c.options, c.msgC, c.dlq, c.scope, c.logger) 102 | c.partitionConsumerMap[partition] = p 103 | p.Start() 104 | c.scope.Gauge(metrics.KafkaPartitionOwnedCount).Update(float64(len(c.partitionConsumerMap))) 105 | } 106 | 107 | func (c *TopicConsumer) shutdown() { 108 | var wg sync.WaitGroup 109 | for _, pc := range c.partitionConsumerMap { 110 | wg.Add(1) 111 | go func(p PartitionConsumer) { 112 | p.Drain(2 * c.options.OffsetCommitInterval) 113 | wg.Done() 114 | }(pc) 115 | } 116 | wg.Wait() 117 | 118 | for k := range c.partitionConsumerMap { 119 | delete(c.partitionConsumerMap, k) 120 | } 121 | } 122 | 123 | // ResetOffset will reset the consumer offset for the specified topic, partition. 124 | func (c *TopicConsumer) ResetOffset(partition int32, offsetRange kafka.OffsetRange) error { 125 | pc, ok := c.partitionConsumerMap[partition] 126 | if !ok { 127 | c.logger.Warn("failed to reset offset for non existent topic-partition", zap.Int32("partition", partition), zap.Object("offsetRange", offsetRange)) 128 | return nil 129 | } 130 | 131 | return pc.ResetOffset(offsetRange) 132 | } 133 | -------------------------------------------------------------------------------- /internal/consumer/types.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package consumer 22 | 23 | import ( 24 | "errors" 25 | 26 | "github.com/Shopify/sarama" 27 | "github.com/bsm/sarama-cluster" 28 | "github.com/golang/protobuf/proto" 29 | "github.com/uber-go/kafka-client/internal/util" 30 | "github.com/uber-go/kafka-client/kafka" 31 | "go.uber.org/zap/zapcore" 32 | ) 33 | 34 | type ( 35 | // SaramaConsumer is an interface for external consumer library (sarama) 36 | SaramaConsumer interface { 37 | Close() error 38 | Errors() <-chan error 39 | Notifications() <-chan *cluster.Notification 40 | Partitions() <-chan cluster.PartitionConsumer 41 | CommitOffsets() error 42 | Messages() <-chan *sarama.ConsumerMessage 43 | HighWaterMarks() map[string]map[int32]int64 44 | MarkOffset(msg *sarama.ConsumerMessage, metadata string) 45 | MarkPartitionOffset(topic string, partition int32, offset int64, metadata string) 46 | ResetPartitionOffset(topic string, partition int32, offset int64, metadata string) 47 | } 48 | 49 | // saramaConsumer is an internal version of SaramaConsumer that implements a close method that can be safely called 50 | // multiple times. 51 | saramaConsumer struct { 52 | SaramaConsumer 53 | lifecycle *util.RunLifecycle 54 | } 55 | 56 | // saramaProducer is an internal version of SaramaConsumer that implements a close method that can be safely called 57 | // multiple times. 58 | saramaProducer struct { 59 | sarama.AsyncProducer 60 | lifecycle *util.RunLifecycle 61 | } 62 | 63 | // saramaClient is an internal version of sarama Client that implements a close method that can be safely called 64 | // multiple times 65 | saramaClient struct { 66 | sarama.Client 67 | lifecycle *util.RunLifecycle 68 | } 69 | 70 | // Constructors wraps multiple Sarama Constructors, which can be used for tests. 71 | Constructors struct { 72 | NewSaramaProducer func(sarama.Client) (sarama.AsyncProducer, error) 73 | NewSaramaConsumer func([]string, string, []string, *cluster.Config) (SaramaConsumer, error) 74 | NewSaramaClient func([]string, *sarama.Config) (sarama.Client, error) 75 | } 76 | 77 | // Topic is an internal wrapper around kafka.ConsumerTopic 78 | Topic struct { 79 | kafka.ConsumerTopic 80 | DLQMetadataDecoder 81 | PartitionConsumerFactory 82 | ConsumerGroupSuffix string 83 | } 84 | 85 | // DLQMetadataDecoder decodes a byte array into DLQMetadata. 86 | DLQMetadataDecoder func([]byte) (*DLQMetadata, error) 87 | ) 88 | 89 | // NoopDLQMetadataDecoder does no decoding and returns a default DLQMetadata object. 90 | func NoopDLQMetadataDecoder(b []byte) (*DLQMetadata, error) { 91 | return newDLQMetadata(), nil 92 | } 93 | 94 | // ProtobufDLQMetadataDecoder uses proto.Unmarshal to decode protobuf encoded binary into the DLQMetadata object. 95 | func ProtobufDLQMetadataDecoder(b []byte) (*DLQMetadata, error) { 96 | dlqMetadata := newDLQMetadata() 97 | if b == nil { 98 | return nil, errors.New("expected to decode non-nil byte array to DLQ metadata") 99 | } 100 | if err := proto.Unmarshal(b, dlqMetadata); err != nil { 101 | return nil, err 102 | } 103 | return dlqMetadata, nil 104 | } 105 | 106 | // MarshalLogObject implements zapcore.ObjectMarshaler for structured logging. 107 | func (t Topic) MarshalLogObject(e zapcore.ObjectEncoder) error { 108 | t.ConsumerTopic.MarshalLogObject(e) 109 | return nil 110 | } 111 | 112 | // NewSaramaConsumer returns a new SaramaConsumer that has a Close method that can be called multiple times. 113 | func NewSaramaConsumer(brokers []string, groupID string, topics []string, config *cluster.Config) (SaramaConsumer, error) { 114 | c, err := cluster.NewConsumer(brokers, groupID, topics, config) 115 | if err != nil { 116 | return nil, err 117 | } 118 | 119 | return newSaramaConsumer(c) 120 | } 121 | 122 | func newSaramaConsumer(c SaramaConsumer) (SaramaConsumer, error) { 123 | sc := saramaConsumer{ 124 | SaramaConsumer: c, 125 | lifecycle: util.NewRunLifecycle("sarama-consumer"), 126 | } 127 | 128 | sc.lifecycle.Start(func() error { return nil }) // must start lifecycle so stop will stop 129 | 130 | return &sc, nil 131 | } 132 | 133 | // Close overwrites the underlying SaramaConsumer Close method with one that can be safely called multiple times. 134 | // 135 | // This close will always return nil error. 136 | func (p *saramaConsumer) Close() error { 137 | p.lifecycle.Stop(func() { 138 | p.SaramaConsumer.Close() 139 | }) 140 | return nil 141 | } 142 | 143 | // NewSaramaProducer returns a new AsyncProducer that has Close method that can be called multiple times. 144 | func NewSaramaProducer(client sarama.Client) (sarama.AsyncProducer, error) { 145 | p, err := sarama.NewAsyncProducerFromClient(client) 146 | if err != nil { 147 | return nil, err 148 | } 149 | 150 | return newSaramaProducer(p) 151 | } 152 | 153 | func newSaramaProducer(p sarama.AsyncProducer) (sarama.AsyncProducer, error) { 154 | sp := saramaProducer{ 155 | AsyncProducer: p, 156 | lifecycle: util.NewRunLifecycle("sarama-producer"), 157 | } 158 | 159 | sp.lifecycle.Start(func() error { return nil }) // must start lifecycle so stop will stop 160 | 161 | return &sp, nil 162 | } 163 | 164 | // Close overwrites the underlying Sarama Close method with one that can be safely called multiple times. 165 | // 166 | // This close will always return nil error. 167 | func (p *saramaProducer) Close() error { 168 | p.lifecycle.Stop(func() { 169 | p.AsyncProducer.Close() 170 | }) 171 | return nil 172 | } 173 | 174 | // NewSaramaClient returns an internal sarama Client, which can be safely closed multiple times. 175 | func NewSaramaClient(brokers []string, config *sarama.Config) (sarama.Client, error) { 176 | sc, err := sarama.NewClient(brokers, config) 177 | if err != nil { 178 | return nil, err 179 | } 180 | 181 | return newSaramaClient(sc) 182 | } 183 | 184 | func newSaramaClient(client sarama.Client) (sarama.Client, error) { 185 | c := &saramaClient{ 186 | Client: client, 187 | lifecycle: util.NewRunLifecycle("sarama-client"), 188 | } 189 | c.lifecycle.Start(func() error { return nil }) // must start lifecycle so stop will stop 190 | return c, nil 191 | } 192 | 193 | // Close overwrites the underlying Sarama Close method with one that can be safely called multiple times. 194 | // 195 | // This close will always return nil error. 196 | func (c *saramaClient) Close() error { 197 | c.lifecycle.Stop(func() { 198 | c.Client.Close() 199 | }) 200 | return nil 201 | } 202 | -------------------------------------------------------------------------------- /internal/consumer/types_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package consumer 22 | 23 | import ( 24 | "sync/atomic" 25 | "testing" 26 | 27 | "github.com/golang/protobuf/proto" 28 | "github.com/stretchr/testify/assert" 29 | ) 30 | 31 | func TestSaramaConsumer(t *testing.T) { 32 | mockC := newMockSaramaConsumer() 33 | c, err := newSaramaConsumer(mockC) 34 | assert.NoError(t, err) 35 | assert.NoError(t, c.Close()) 36 | assert.EqualValues(t, 1, atomic.LoadInt64(&mockC.closed)) 37 | 38 | // Second close should return no error and not increment closed counter 39 | assert.NoError(t, c.Close()) 40 | assert.EqualValues(t, 1, atomic.LoadInt64(&mockC.closed)) 41 | } 42 | 43 | func TestSaramaProducer(t *testing.T) { 44 | mockP := newMockSaramaProducer() 45 | c, err := newSaramaProducer(mockP) 46 | assert.NoError(t, err) 47 | assert.NoError(t, c.Close()) 48 | assert.EqualValues(t, 1, atomic.LoadInt32(&mockP.closed)) 49 | 50 | // Second close should return no error and not increment closed counter 51 | assert.NoError(t, c.Close()) 52 | assert.EqualValues(t, 1, atomic.LoadInt32(&mockP.closed)) 53 | } 54 | 55 | func TestSaramaClient(t *testing.T) { 56 | mock := newMockSaramaClient() 57 | c, err := newSaramaClient(mock) 58 | assert.NoError(t, err) 59 | assert.NoError(t, c.Close()) 60 | assert.EqualValues(t, 1, atomic.LoadInt32(&mock.closed)) 61 | 62 | // Second close should return no error and not increment closed counter 63 | assert.NoError(t, c.Close()) 64 | assert.EqualValues(t, 1, atomic.LoadInt32(&mock.closed)) 65 | } 66 | 67 | func TestProtobufDLQMetadataDecoder(t *testing.T) { 68 | dlqMetadata := newDLQMetadata() 69 | b, err := proto.Marshal(dlqMetadata) 70 | assert.NoError(t, err) 71 | decodedDLQMetadata, err := ProtobufDLQMetadataDecoder(b) 72 | assert.NoError(t, err) 73 | assert.EqualValues(t, dlqMetadata, decodedDLQMetadata) 74 | 75 | _, err = ProtobufDLQMetadataDecoder(nil) 76 | assert.Error(t, err) 77 | 78 | _, err = ProtobufDLQMetadataDecoder([]byte{1, 2, 3, 4}) 79 | assert.Error(t, err) 80 | } 81 | 82 | func TestNoopDLQMetadataDecoder(t *testing.T) { 83 | dlqMetadata := newDLQMetadata() 84 | dlqMetadata.Offset = 100 85 | b, err := proto.Marshal(dlqMetadata) 86 | assert.NoError(t, err) 87 | decodedDLQMetadata, err := NoopDLQMetadataDecoder(b) 88 | assert.NoError(t, err) 89 | assert.EqualValues(t, newDLQMetadata(), decodedDLQMetadata) 90 | } 91 | -------------------------------------------------------------------------------- /internal/list/list.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package list 22 | 23 | import ( 24 | "errors" 25 | "math" 26 | ) 27 | 28 | type ( 29 | // listNode is a single linked list node 30 | listNode struct { 31 | value int64 32 | prev int 33 | next int 34 | } 35 | 36 | // Address refers to the internal address for a list item 37 | Address int 38 | 39 | // IntegerList refers to a linked list of integers 40 | // with fixed maximum capacity 41 | IntegerList struct { 42 | array []listNode 43 | size int 44 | capacity int 45 | head int 46 | freeListHead int 47 | } 48 | ) 49 | 50 | const ( 51 | // Null refers to the null list address 52 | Null = Address(-1) 53 | // beginAddr is the index where the actual data starts 54 | // the first and second entries are *special* - they serve 55 | // as the head of linked list and head of free list respectively 56 | // Having the head as part of the list helps avoid branches (ifs) 57 | // during add/remove operations 58 | beginAddr = 2 59 | ) 60 | 61 | var ( 62 | // ErrEmpty indicates that the list is empty 63 | ErrEmpty = errors.New("list is empty") 64 | // ErrCapacity indicates that the list is out of capacity 65 | ErrCapacity = errors.New("list out of capacity") 66 | // ErrAddrOutOfRange indicates the provided address is out of range 67 | ErrAddrOutOfRange = errors.New("list address out of range") 68 | // ErrAddrMissingValue indicates that the requested address does not contain any value 69 | ErrAddrMissingValue = errors.New("list address missing value") 70 | ) 71 | 72 | // NewIntegerList returns a linked list of integers with a fixed capacity 73 | // this list is tailor made to keep track of out of order message offsets 74 | // while processing kafka messages. Following are the allowed operations on 75 | // this list: 76 | // 77 | // Add(int64): 78 | // Adds the given value to the end of list. Returns an address where this 79 | // value is stored within the list. This address must be provided to remove 80 | // this item from the list later 81 | // 82 | // Remove(listAddr): 83 | // Removes the item with the given address from the list. Returns address out 84 | // of range error if the address is invalid. The provided address is the address 85 | // previously received from an Add operation 86 | // 87 | // PeekHead(): 88 | // Returns the value at the head of the list without removing the item from the list 89 | // This operation is used to identify the same commit level during periodic checkpoints 90 | // 91 | // Note: This list implementation uses math.MinInt64 as a special sentinel value 92 | // 93 | // Implementation Notes: 94 | // The underlying implementation uses a fixed size array where each item in the array is 95 | // a linked list node. In addition to the array, there is also a free list which also 96 | // refers to the same backing array i.e. the implementation uses an intrusive free list 97 | // This implementation choice is chosen for the following reasons: 98 | // - No dynamic memory allocation needed 99 | // - Array based list is very cache friendly 100 | // - With this list, keeping track of out of order offsets is O(1) overhead per message 101 | func NewIntegerList(capacity int) *IntegerList { 102 | list := &IntegerList{ 103 | array: make([]listNode, capacity+beginAddr), 104 | capacity: capacity, 105 | } 106 | list.init() 107 | return list 108 | } 109 | 110 | // Add adds the value to the end of list. Returns the address 111 | // where this value is stored on success 112 | func (list *IntegerList) Add(value int64) (Address, error) { 113 | if list.empty(list.freeListHead) { 114 | return Null, ErrCapacity 115 | } 116 | // remove first node from free list and add it to list 117 | nextAddr := list.array[list.freeListHead].next 118 | list.remove(nextAddr) 119 | list.array[nextAddr].value = value 120 | list.append(list.head, nextAddr) 121 | list.size++ 122 | 123 | return Address(nextAddr), nil 124 | } 125 | 126 | // Get returns the value at the given address 127 | func (list *IntegerList) Get(getAddr Address) (int64, error) { 128 | addr := int(getAddr) 129 | if !list.validAddr(addr) { 130 | return 0, ErrAddrOutOfRange 131 | } 132 | if list.Empty() { 133 | return 0, ErrEmpty 134 | } 135 | value := list.array[addr].value 136 | if value == math.MinInt64 { 137 | return value, ErrAddrMissingValue 138 | } 139 | return value, nil 140 | } 141 | 142 | // Remove removes the item stored at the specified address 143 | // from the list 144 | func (list *IntegerList) Remove(removeAddr Address) error { 145 | addr := int(removeAddr) 146 | if !list.validAddr(addr) { 147 | return ErrAddrOutOfRange 148 | } 149 | list.remove(addr) 150 | list.append(list.freeListHead, addr) 151 | list.size-- 152 | return nil 153 | } 154 | 155 | // PeekHead returns the value at the head of the list if 156 | // the list is not empty. If the list empty, error is returned 157 | func (list *IntegerList) PeekHead() (int64, error) { 158 | if list.Empty() { 159 | return 0, ErrEmpty 160 | } 161 | firstAddr := list.array[list.head].next 162 | return list.array[firstAddr].value, nil 163 | } 164 | 165 | // Size returns the size of the list 166 | func (list *IntegerList) Size() int { 167 | return list.size 168 | } 169 | 170 | // Empty returns true if the list is empty 171 | func (list *IntegerList) Empty() bool { 172 | return list.size == 0 173 | } 174 | 175 | func (list *IntegerList) empty(head int) bool { 176 | return list.array[head].next == head 177 | } 178 | 179 | func (list *IntegerList) append(head int, newAddr int) { 180 | list.array[newAddr].next = head 181 | list.array[newAddr].prev = list.array[head].prev 182 | list.array[list.array[head].prev].next = newAddr 183 | list.array[head].prev = newAddr 184 | } 185 | 186 | func (list *IntegerList) remove(addr int) { 187 | list.array[addr].value = math.MinInt64 188 | list.array[list.array[addr].prev].next = list.array[addr].next 189 | list.array[list.array[addr].next].prev = list.array[addr].prev 190 | } 191 | 192 | // asArray returns the list as array. Only for unit tests 193 | func (list *IntegerList) asArray() []int64 { 194 | var i int 195 | result := make([]int64, list.size) 196 | next := list.array[list.head].next 197 | for next != list.head { 198 | result[i] = list.array[next].value 199 | next, i = list.array[next].next, i+1 200 | } 201 | return result 202 | } 203 | 204 | func (list *IntegerList) validAddr(addr int) bool { 205 | return addr >= beginAddr && addr < len(list.array) 206 | } 207 | 208 | func (list *IntegerList) init() { 209 | list.head = 0 210 | list.freeListHead = 1 211 | list.array[list.head].next = list.head 212 | list.array[list.head].prev = list.head 213 | list.array[list.freeListHead].next = beginAddr 214 | list.array[list.freeListHead].prev = len(list.array) - 1 215 | 216 | for i := beginAddr; i < len(list.array); i++ { 217 | list.array[i].prev = i - 1 218 | list.array[i].next = i + 1 219 | list.array[i].value = math.MinInt64 220 | } 221 | list.array[len(list.array)-1].value = math.MinInt64 222 | list.array[len(list.array)-1].next = list.freeListHead 223 | } 224 | -------------------------------------------------------------------------------- /internal/list/list_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package list 22 | 23 | import ( 24 | "testing" 25 | 26 | "math/rand" 27 | "time" 28 | 29 | "github.com/stretchr/testify/suite" 30 | ) 31 | 32 | type ( 33 | ListTestSuite struct { 34 | suite.Suite 35 | } 36 | 37 | testItem struct { 38 | addr Address 39 | value int64 40 | } 41 | ) 42 | 43 | func TestListTestSuite(t *testing.T) { 44 | suite.Run(t, new(ListTestSuite)) 45 | } 46 | 47 | func (s *ListTestSuite) SetupTest() { 48 | rand.Seed(time.Now().UnixNano()) 49 | } 50 | 51 | func (s *ListTestSuite) TestCapacityOfOne() { 52 | list := NewIntegerList(1) 53 | for i := 0; i < 10; i++ { 54 | s.Equal(0, list.size) 55 | s.True(list.Empty()) 56 | addr := s.testAddNoErr(list, 5) 57 | s.testPeekNoErr(list, 5) 58 | s.testAddErr(list, 6) 59 | s.testRemove(list, addr) 60 | s.testPeekErr(list) 61 | } 62 | } 63 | 64 | func (s *ListTestSuite) TestCapacityOfN() { 65 | cap := 128 66 | list := NewIntegerList(cap) 67 | for i := 0; i < 1; i++ { 68 | items := make([]testItem, cap) 69 | removed := make(map[int64]struct{}) 70 | for c := 0; c < cap; c++ { 71 | item := testItem{value: int64(c)} 72 | item.addr = s.testAddNoErr(list, c) 73 | items[c] = item 74 | } 75 | s.testAddErr(list, cap) 76 | // remove items from list in random order 77 | s.shuffle(items) 78 | for c := 0; c < cap/2; c++ { 79 | s.testRemove(list, items[c].addr) 80 | removed[items[c].value] = struct{}{} 81 | s.testSanity(list, removed) 82 | } 83 | for c := 0; c < cap/2; c++ { 84 | item := testItem{value: int64(cap + c)} 85 | item.addr = s.testAddNoErr(list, cap+c) 86 | items[c] = item 87 | } 88 | s.testAddErr(list, cap) 89 | s.shuffle(items) 90 | for c := 0; c < cap; c++ { 91 | s.testRemove(list, items[c].addr) 92 | removed[items[c].value] = struct{}{} 93 | s.testSanity(list, removed) 94 | } 95 | s.True(list.Empty()) 96 | s.Equal(0, list.Size()) 97 | _, err := list.Get(beginAddr) 98 | s.Equal(ErrEmpty, err) 99 | } 100 | } 101 | 102 | func (s *ListTestSuite) TestAddrOutOfRange() { 103 | list := NewIntegerList(1) 104 | for i := 3; i < 10; i++ { 105 | err := list.Remove(Address(i)) 106 | s.Equal(ErrAddrOutOfRange, err) 107 | _, err = list.Get(Address(i)) 108 | s.Equal(ErrAddrOutOfRange, err) 109 | } 110 | } 111 | 112 | func (s *ListTestSuite) TestAddrMissingValue() { 113 | list := NewIntegerList(2) 114 | addr, err := list.Add(888) 115 | s.NoError(err) 116 | _, err = list.Get(Address(addr + 1)) 117 | s.Equal(ErrAddrMissingValue, err) 118 | } 119 | 120 | func (s *ListTestSuite) testAddNoErr(list *IntegerList, value int) Address { 121 | size := list.Size() 122 | addr, err := list.Add(int64(value)) 123 | s.NoError(err) 124 | s.False(list.Empty()) 125 | s.Equal(size+1, list.Size()) 126 | return addr 127 | } 128 | 129 | func (s *ListTestSuite) testAddErr(list *IntegerList, value int) { 130 | addr, err := list.Add(int64(value)) 131 | s.Equal(ErrCapacity, err) 132 | s.Equal(Null, addr) 133 | } 134 | 135 | func (s *ListTestSuite) testRemove(list *IntegerList, addr Address) { 136 | size := list.Size() 137 | err := list.Remove(addr) 138 | s.NoError(err) 139 | s.Equal(size-1, list.Size()) 140 | if size == 0 { 141 | s.True(list.Empty()) 142 | } 143 | } 144 | 145 | func (s *ListTestSuite) testPeekNoErr(list *IntegerList, expected int) { 146 | val, err := list.PeekHead() 147 | s.NoError(err) 148 | s.Equal(int64(expected), val) 149 | } 150 | 151 | func (s *ListTestSuite) testPeekErr(list *IntegerList) { 152 | _, err := list.PeekHead() 153 | s.Equal(ErrEmpty, err) 154 | s.True(list.Empty()) 155 | } 156 | 157 | // assert list does not contain removed entries 158 | // assert items in the list are in ascending order 159 | func (s *ListTestSuite) testSanity(list *IntegerList, removed map[int64]struct{}) { 160 | array := list.asArray() 161 | if len(array) == 0 { 162 | return 163 | } 164 | 165 | prev := array[0] 166 | _, ok := removed[prev] 167 | s.False(ok) 168 | for i := 1; i < len(array); i++ { 169 | _, ok := removed[array[i]] 170 | s.False(ok) 171 | s.True(array[i] > prev) 172 | prev = array[i] 173 | } 174 | } 175 | 176 | func (s *ListTestSuite) shuffle(items []testItem) { 177 | n := len(items) 178 | for i := n - 1; i >= 0; i-- { 179 | idx := rand.Intn(n) 180 | items[i], items[idx] = items[idx], items[i] 181 | n-- 182 | } 183 | } 184 | -------------------------------------------------------------------------------- /internal/metrics/defs.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package metrics 22 | 23 | // Counters 24 | const ( 25 | KafkaConsumerStarted = "kafka.consumer.started" 26 | KafkaConsumerStopped = "kafka.consumer.stopped" 27 | 28 | KafkaPartitionStarted = "kafka.partition.started" 29 | KafkaPartitionStopped = "kafka.partition.stopped" 30 | KafkaPartitionMessagesIn = "kafka.partition.messages-in" 31 | 32 | KafkaDLQStarted = "kafka.dlq.started" 33 | KafkaDLQStopped = "kafka.dlq.stopped" 34 | KafkaDLQMessagesOut = "kafka.dlq.messages-out" 35 | KafkaDLQErrors = "kafka.dlq.errors" 36 | KafkaDLQMetadataError = "kafka.dlq.metadata-error" 37 | 38 | KafkaPartitionAckMgrDups = "kafka.partition.ackmgr.duplicates" 39 | KafkaPartitionAckMgrSkipped = "kafka.partition.ackmgr.skipped" 40 | KafkaPartitionAckMgrListFull = "kafka.partition.ackmgr.list-full-error" 41 | KafkaPartitionGetAckIDErrors = "kafka.partition.ackmgr.get-ackid-error" 42 | KafkaPartitionAck = "kafka.partition.ack" 43 | KafkaPartitionAckErrors = "kafka.partition.ackmgr.ack-error" 44 | KafkaPartitionNack = "kafka.partition.nack" 45 | KafkaPartitionNackErrors = "kafka.partition.ackmgr.nack-error" 46 | KafkaPartitionRebalance = "kafka.partition.rebalance" 47 | ) 48 | 49 | // Gauges 50 | const ( 51 | KafkaPartitionTimeLag = "kafka.partition.time-lag" 52 | KafkaPartitionOffsetLag = "kafka.partition.offset-lag" 53 | KafkaPartitionOffsetFreshnessLag = "kafka.partition.freshness-lag" 54 | KafkaPartitionReadOffset = "kafka.partition.read-offset" 55 | KafkaPartitionCommitOffset = "kafka.partition.commit-offset" 56 | KafkaPartitionOwnedCount = "kafka.partition.owned.count" 57 | KafkaPartitionOwned = "kafka.partition.owned" 58 | ) 59 | -------------------------------------------------------------------------------- /internal/util/lifecycle.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package util 22 | 23 | import ( 24 | "errors" 25 | "sync" 26 | ) 27 | 28 | // RunLifecycle manages the start/stop lifecycle 29 | // for a runnable implementation 30 | type RunLifecycle struct { 31 | sync.Mutex 32 | name string 33 | started bool 34 | stopped bool 35 | } 36 | 37 | // NewRunLifecycle returns a lifecycle object that can be 38 | // used by a Runnable implementation to make sure the start/stop 39 | // operations are idempotent 40 | func NewRunLifecycle(name string) *RunLifecycle { 41 | return &RunLifecycle{name: name} 42 | } 43 | 44 | // Start executes the given action if and on if lifecycle 45 | // state is previously not started or stopped 46 | func (r *RunLifecycle) Start(action func() error) error { 47 | r.Lock() 48 | defer r.Unlock() 49 | if r.stopped { 50 | return errors.New(r.name + " cannot start a previously stopped runnable") 51 | } 52 | if !r.started { 53 | r.started = true 54 | if err := action(); err != nil { 55 | r.stopped = true 56 | return err 57 | } 58 | } 59 | return nil 60 | } 61 | 62 | // Stop stops the given action if and only if lifecycle 63 | // state is previously started 64 | func (r *RunLifecycle) Stop(action func()) { 65 | r.Lock() 66 | defer r.Unlock() 67 | if !r.started { 68 | r.stopped = true 69 | return 70 | } 71 | if !r.stopped { 72 | action() 73 | r.stopped = true 74 | } 75 | } 76 | 77 | // Status returns the status of this lifecycle as a pair of booleans (started, stopped). 78 | func (r *RunLifecycle) Status() (bool, bool) { 79 | r.Lock() 80 | defer r.Unlock() 81 | return r.started, r.stopped 82 | } 83 | -------------------------------------------------------------------------------- /internal/util/lifecycle_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package util 22 | 23 | import ( 24 | "fmt" 25 | "testing" 26 | 27 | "github.com/stretchr/testify/suite" 28 | ) 29 | 30 | type RunLifecycleTestSuite struct { 31 | suite.Suite 32 | } 33 | 34 | func TestRunLifecycleTestSuite(t *testing.T) { 35 | suite.Run(t, new(RunLifecycleTestSuite)) 36 | } 37 | 38 | func (s *RunLifecycleTestSuite) TestSuccessFlow() { 39 | state := "" 40 | lc := NewRunLifecycle("test") 41 | err := lc.Start(func() error { 42 | state = "started" 43 | return nil 44 | }) 45 | s.NoError(err) 46 | s.Equal("started", state) 47 | lc.Stop(func() { state = "stopped" }) 48 | s.Equal("stopped", state) 49 | } 50 | 51 | func (s *RunLifecycleTestSuite) TestStartError() { 52 | state := "" 53 | lc := NewRunLifecycle("test") 54 | err := lc.Start(func() error { 55 | state = "started" 56 | return fmt.Errorf("test error") 57 | }) 58 | s.Error(err) 59 | s.Equal("started", state) 60 | } 61 | 62 | func (s *RunLifecycleTestSuite) TestMultipleStarts() { 63 | state := "test" 64 | lc := NewRunLifecycle("test") 65 | for i := 0; i < 10; i++ { 66 | err := lc.Start(func() error { 67 | state = state + "_started" 68 | return nil 69 | }) 70 | s.NoError(err) 71 | s.Equal("test_started", state) 72 | } 73 | } 74 | 75 | func (s *RunLifecycleTestSuite) TestMultipleStops() { 76 | state := "" 77 | lc := NewRunLifecycle("test") 78 | lc.Start(func() error { 79 | state = "started" 80 | return nil 81 | }) 82 | for i := 0; i < 10; i++ { 83 | lc.Stop(func() { state = state + "_stopped" }) 84 | s.Equal("started_stopped", state) 85 | } 86 | } 87 | 88 | func (s *RunLifecycleTestSuite) TestStopBeforeStart() { 89 | state := "not_started" 90 | lc := NewRunLifecycle("test") 91 | lc.Stop(func() { state = "stopped" }) 92 | s.Equal("not_started", state) 93 | s.True(lc.stopped) 94 | } 95 | -------------------------------------------------------------------------------- /internal/util/testutil.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package util 22 | 23 | import ( 24 | "fmt" 25 | "sync" 26 | "time" 27 | ) 28 | 29 | // AwaitCondition waits until a given condition is satisfied or times out 30 | func AwaitCondition(cond func() bool, d time.Duration) bool { 31 | ticker := time.NewTicker(time.Millisecond * 50) 32 | defer ticker.Stop() 33 | deadline := time.After(d) 34 | for { 35 | select { 36 | case <-ticker.C: 37 | if cond() { 38 | return true 39 | } 40 | case <-deadline: 41 | return false 42 | } 43 | } 44 | } 45 | 46 | // AwaitWaitGroup does a timed wait on the wait group 47 | func AwaitWaitGroup(wg *sync.WaitGroup, d time.Duration) bool { 48 | ch := make(chan struct{}) 49 | go func() { 50 | wg.Wait() 51 | close(ch) 52 | }() 53 | select { 54 | case <-ch: 55 | return true 56 | case <-time.After(d): 57 | fmt.Println("awaitWaitGroup returns false") 58 | return false 59 | } 60 | } 61 | -------------------------------------------------------------------------------- /kafka/config.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package kafka 22 | 23 | import ( 24 | "crypto/tls" 25 | "fmt" 26 | "time" 27 | 28 | "github.com/Shopify/sarama" 29 | "go.uber.org/zap/zapcore" 30 | ) 31 | 32 | const ( 33 | // OffsetOldest uses sequence number of oldest known message as the current offset 34 | OffsetOldest = sarama.OffsetOldest 35 | // OffsetNewest option uses sequence number of newest message as the current offset 36 | OffsetNewest = sarama.OffsetNewest 37 | ) 38 | 39 | type ( 40 | // Topic contains information for a topic. 41 | // Our topics are uniquely defined by a Topic Name and Cluster pair. 42 | Topic struct { 43 | // Name for the topic 44 | Name string 45 | // Cluster is the logical name of the cluster to find this topic on. 46 | Cluster string 47 | // Delay is msg consumption delay applied on the topic. 48 | Delay time.Duration 49 | } 50 | 51 | // ConsumerTopic contains information for a consumer topic. 52 | ConsumerTopic struct { 53 | Topic 54 | RetryQ Topic 55 | DLQ Topic 56 | MaxRetries int64 // MaxRetries = -1 for infinite retries. 57 | } 58 | 59 | // ConsumerTopicList is a list of consumer topics 60 | ConsumerTopicList []ConsumerTopic 61 | 62 | // ConsumerConfig describes the config for a consumer group 63 | ConsumerConfig struct { 64 | // GroupName identifies your consumer group. Unless your application creates 65 | // multiple consumer groups (in which case it's suggested to have application name as 66 | // prefix of the group name), this should match your application name. 67 | GroupName string 68 | 69 | // TopicList is a list of consumer topics 70 | TopicList ConsumerTopicList 71 | 72 | // OffsetConfig is the offset-handling policy for this consumer group. 73 | Offsets struct { 74 | // Initial specifies the fallback offset configuration on consumer start. 75 | // The consumer will use the offsets persisted from its last run unless \ 76 | // the offsets are too old or too new. 77 | Initial struct { 78 | // Offset is the initial offset to use if there is no previous offset 79 | // committed. Use OffsetNewest for high watermark and OffsetOldest for 80 | // low watermark. Defaults to OffsetOldest. 81 | Offset int64 82 | } 83 | 84 | // Commits a policy for committing consumer offsets to Kafka. 85 | Commits struct { 86 | // Enabled if you want the library to commit offsets on your behalf. 87 | // Defaults to true. 88 | // 89 | // The retry and dlq topic commit will always be committed for you since those topics are abstracted away from you. 90 | Enabled bool 91 | } 92 | } 93 | 94 | // Concurrency determines the number of concurrent messages to process. 95 | // When using the handler based API, this corresponds to the number of concurrent go 96 | // routines handler functions the library will run. Default is 1. 97 | Concurrency int 98 | 99 | 100 | // TLSConfig is the configuration to use for secure connections if 101 | // enabled (not nil) (defaults to disabled, nil). 102 | TLSConfig *tls.Config 103 | } 104 | ) 105 | 106 | // NewConsumerConfig returns ConsumerConfig with sane defaults. 107 | func NewConsumerConfig(groupName string, topicList ConsumerTopicList) *ConsumerConfig { 108 | cfg := new(ConsumerConfig) 109 | cfg.GroupName = groupName 110 | cfg.TopicList = topicList 111 | cfg.Offsets.Initial.Offset = OffsetOldest 112 | cfg.Offsets.Commits.Enabled = true 113 | cfg.Concurrency = 1 114 | return cfg 115 | } 116 | 117 | // MarshalLogObject implements zapcore.ObjectMarshaler for structured logging. 118 | func (c ConsumerConfig) MarshalLogObject(e zapcore.ObjectEncoder) error { 119 | e.AddString("groupName", c.GroupName) 120 | e.AddArray("topicList", c.TopicList) 121 | e.AddObject("offset", zapcore.ObjectMarshalerFunc(func(ee zapcore.ObjectEncoder) error { 122 | ee.AddObject("initial", zapcore.ObjectMarshalerFunc(func(eee zapcore.ObjectEncoder) error { 123 | eee.AddInt64("offset", c.Offsets.Initial.Offset) 124 | return nil 125 | })) 126 | ee.AddObject("commits", zapcore.ObjectMarshalerFunc(func(eee zapcore.ObjectEncoder) error { 127 | eee.AddBool("enabled", c.Offsets.Commits.Enabled) 128 | return nil 129 | })) 130 | return nil 131 | })) 132 | e.AddInt("concurrency", c.Concurrency) 133 | return nil 134 | } 135 | 136 | // MarshalLogArray implements zapcore.ArrayMarshaler for structured logging. 137 | func (c ConsumerTopicList) MarshalLogArray(e zapcore.ArrayEncoder) error { 138 | for _, topic := range c { 139 | e.AppendObject(topic) 140 | } 141 | return nil 142 | } 143 | 144 | // MarshalLogObject implements zapcore.ObjectMarshaler for structured logging. 145 | func (c ConsumerTopic) MarshalLogObject(e zapcore.ObjectEncoder) error { 146 | e.AddObject("defaultQ", c.Topic) 147 | e.AddObject("retryQ", c.RetryQ) 148 | e.AddObject("DLQ", c.DLQ) 149 | e.AddInt64("maxRetries", c.MaxRetries) 150 | return nil 151 | } 152 | 153 | // MarshalLogObject implements zapcore.ObjectMarshaler for structured logging. 154 | func (t Topic) MarshalLogObject(e zapcore.ObjectEncoder) error { 155 | e.AddString("name", t.Name) 156 | e.AddString("cluster", t.Cluster) 157 | e.AddDuration("delay", t.Delay) 158 | return nil 159 | } 160 | 161 | // TopicNames returns the list of topics to consume as a string array. 162 | func (c ConsumerTopicList) TopicNames() []string { 163 | output := make([]string, 0, len(c)) 164 | for _, topic := range c { 165 | output = append(output, topic.Name) 166 | } 167 | return output 168 | } 169 | 170 | // GetConsumerTopicByClusterTopic returns the ConsumerTopic for the cluster, topic pair. 171 | func (c ConsumerTopicList) GetConsumerTopicByClusterTopic(clusterName, topicName string) (ConsumerTopic, error) { 172 | for _, topic := range c { 173 | if topic.Cluster == clusterName && topic.Name == topicName { 174 | return topic, nil 175 | } 176 | } 177 | return ConsumerTopic{}, fmt.Errorf("unable to find TopicConfig with cluster %s and topic %s", clusterName, topicName) 178 | } 179 | 180 | // HashKey converts topic to a string for use as a map key 181 | func (t Topic) HashKey() string { 182 | output := t.Name + t.Cluster 183 | return output 184 | } 185 | 186 | // DLQEnabled returns true if DLQ.Name and DLQ.Cluster are not empty. 187 | func (c ConsumerTopic) DLQEnabled() bool { 188 | if c.DLQ.Cluster != "" && c.DLQ.Name != "" { 189 | return true 190 | } 191 | return false 192 | } 193 | -------------------------------------------------------------------------------- /kafka/config_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package kafka 22 | 23 | import ( 24 | "testing" 25 | 26 | "github.com/stretchr/testify/suite" 27 | ) 28 | 29 | type ( 30 | ConsumerConfigTestSuite struct { 31 | suite.Suite 32 | config ConsumerConfig 33 | topic1 ConsumerTopic 34 | topic2 ConsumerTopic 35 | } 36 | ) 37 | 38 | func TestConsumerConfigTestSuite(t *testing.T) { 39 | suite.Run(t, new(ConsumerConfigTestSuite)) 40 | } 41 | 42 | func (s *ConsumerConfigTestSuite) SetupTest() { 43 | s.topic1 = ConsumerTopic{ 44 | Topic: Topic{ 45 | Name: "topic1", 46 | Cluster: "cluster1", 47 | }, 48 | DLQ: Topic{ 49 | Name: "dlq_topic1", 50 | Cluster: "dlq_cluster1", 51 | }, 52 | } 53 | s.topic2 = ConsumerTopic{ 54 | Topic: Topic{ 55 | Name: "topic2", 56 | Cluster: "cluster2", 57 | }, 58 | DLQ: Topic{ 59 | Name: "dlq_topic2", 60 | Cluster: "dlq_cluster1", 61 | }, 62 | } 63 | s.config = ConsumerConfig{ 64 | GroupName: "groupName", 65 | TopicList: []ConsumerTopic{s.topic1, s.topic2}, 66 | } 67 | } 68 | 69 | func (s *ConsumerConfigTestSuite) TestFilterByClusterTopic() { 70 | topic, err := s.config.TopicList.GetConsumerTopicByClusterTopic("cluster1", "topic1") 71 | s.NoError(err) 72 | s.Equal(s.topic1, topic) 73 | 74 | _, err = s.config.TopicList.GetConsumerTopicByClusterTopic("cluster3", "topic1") 75 | s.Error(err) 76 | _, err = s.config.TopicList.GetConsumerTopicByClusterTopic("cluster1", "topic3") 77 | s.Error(err) 78 | } 79 | 80 | func (s *ConsumerConfigTestSuite) TestHashKey() { 81 | s.Equal(Topic{ 82 | Name: "topic1", 83 | Cluster: "cluster1", 84 | }.HashKey(), s.topic1.HashKey()) 85 | 86 | s.NotEqual(Topic{ 87 | Name: "topic1", 88 | Cluster: "cluster2", 89 | }.HashKey(), s.topic1.HashKey()) 90 | } 91 | -------------------------------------------------------------------------------- /kafka/interfaces.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package kafka 22 | 23 | import ( 24 | "time" 25 | 26 | "go.uber.org/zap/zapcore" 27 | ) 28 | 29 | type ( 30 | // Consumer is the interface for a kafka consumer 31 | Consumer interface { 32 | // Name returns the name of this consumer group. 33 | Name() string 34 | // Topics returns the names of the topics being consumed. 35 | Topics() ConsumerTopicList 36 | // Start starts the consumer 37 | Start() error 38 | // Stop stops the consumer 39 | Stop() 40 | // Closed returns a channel which will be closed after this consumer is completely shutdown 41 | Closed() <-chan struct{} 42 | // Messages return the message channel for this consumer 43 | Messages() <-chan Message 44 | // MergeDLQ consumes the offset ranges for the partitions from the DLQ topic for the specified ConsumerTopic 45 | // Topic should be the __dlq topic name. 46 | MergeDLQ(cluster, group, topic string, partition int32, offsetRange OffsetRange) error 47 | } 48 | 49 | // OffsetRange is a range of offsets 50 | OffsetRange struct { 51 | // LowOffset is the low watermark for this offset range. 52 | // -1 indicates the value is not set. 53 | LowOffset int64 54 | // HighOffset is the high watermark for this offset range. 55 | // -1 indicates the value is not set. 56 | HighOffset int64 57 | } 58 | 59 | // Message is the interface for a Kafka message 60 | Message interface { 61 | zapcore.ObjectMarshaler 62 | 63 | // Key is a mutable reference to the message's key. 64 | Key() []byte 65 | // Value is a mutable reference to the message's value. 66 | Value() []byte 67 | // Topic is the topic from which the message was read. 68 | Topic() string 69 | // Partition is the ID of the partition from which the message was read. 70 | Partition() int32 71 | // Offset is the message's offset. 72 | Offset() int64 73 | // Timestamp returns the timestamp for this message 74 | Timestamp() time.Time 75 | // RetryCount is an incrementing integer denoting the number of times this message has been redelivered. 76 | // The first delivery of the message will be 0, incrementing on each subsequent redelivery. 77 | RetryCount() int64 78 | // Ack marks the message as successfully processed. 79 | Ack() error 80 | // Nack marks the message processing as failed and the message will be retried or sent to DLQ. 81 | Nack() error 82 | // NackToDLQ marks the message processing as failed and sends it immediately to DLQ. 83 | NackToDLQ() error 84 | } 85 | 86 | // NameResolver is an interface that will be used by the consumer library to resolve 87 | // (1) topic to cluster name and (2) cluster name to broker IP addresses. 88 | // Implementations of KafkaNameResolver should be threadsafe. 89 | NameResolver interface { 90 | // ResolveCluster returns a list of IP addresses for the brokers 91 | ResolveIPForCluster(cluster string) ([]string, error) 92 | // ResolveClusterForTopic returns the logical cluster names corresponding to a topic name 93 | // 94 | // It is possible for a topic to exist on multiple clusters in order to 95 | // transparently handle topic migration between clusters. 96 | // TODO (gteo): Remove to simplify API because not needed anymore 97 | ResolveClusterForTopic(topic string) ([]string, error) 98 | } 99 | ) 100 | 101 | // MarshalLogObject implements zapcore.ObjectMarshaler for structured logging. 102 | func (o OffsetRange) MarshalLogObject(e zapcore.ObjectEncoder) error { 103 | e.AddInt64("lowOffset", o.LowOffset) 104 | e.AddInt64("highOffset", o.HighOffset) 105 | return nil 106 | } 107 | 108 | // NewOffsetRange returns a new OffsetRange with the LowOffset of the range as specified. 109 | // First variadic argument is used to set the HighOffset and all other variadic arguments are ignored. 110 | // If no variadic arguments are provided, HighOffset is set to -1 to indicate that it is not set. 111 | func NewOffsetRange(low int64, high ...int64) OffsetRange { 112 | or := OffsetRange{ 113 | LowOffset: low, 114 | HighOffset: -1, 115 | } 116 | if len(high) > 0 { 117 | or.HighOffset = high[0] 118 | } 119 | return or 120 | } 121 | -------------------------------------------------------------------------------- /kafka/resolver.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package kafka 22 | 23 | import ( 24 | "errors" 25 | "sync" 26 | ) 27 | 28 | // staticResolver is an implementation of NameResolver 29 | // that's backed by a static map of clusters to list of brokers 30 | // and a map of topics to cluster 31 | type staticResolver struct { 32 | sync.RWMutex 33 | topicsToCluster map[string][]string 34 | clusterToBrokers map[string][]string 35 | } 36 | 37 | // errNoBrokersForCluster is returned when no brokers can be found for a cluster 38 | var errNoBrokersForCluster = errors.New("no brokers found for cluster") 39 | 40 | // errNoClustersForTopic is returned when no cluster can be found for a topic 41 | var errNoClustersForTopic = errors.New("no cluster found for topic") 42 | 43 | // NewStaticNameResolver returns a instance of NameResolver that relies 44 | // on a static map of topic to list of brokers and map of topics to cluster 45 | func NewStaticNameResolver( 46 | topicsToCluster map[string][]string, 47 | clusterToBrokers map[string][]string, 48 | ) NameResolver { 49 | return &staticResolver{ 50 | topicsToCluster: topicsToCluster, 51 | clusterToBrokers: clusterToBrokers, 52 | } 53 | } 54 | 55 | // ResolveIPForCluster returns list of IP addresses by cluster name by looking up in 56 | // the clusterToBrokers map passed into the NewStaticNameResolver constructor. 57 | func (r *staticResolver) ResolveIPForCluster(cluster string) ([]string, error) { 58 | r.RLock() 59 | defer r.RUnlock() 60 | 61 | if brokers, ok := r.clusterToBrokers[cluster]; ok { 62 | return brokers, nil 63 | } 64 | return nil, errNoBrokersForCluster 65 | } 66 | 67 | // ResolveClusterForTopic resolves the cluster name for a specific topic by looking 68 | // up in the topicsToCluster map passed into the NewStaticNameResolver constructor. 69 | func (r *staticResolver) ResolveClusterForTopic(topic string) ([]string, error) { 70 | r.RLock() 71 | defer r.RUnlock() 72 | 73 | if cluster, ok := r.topicsToCluster[topic]; ok { 74 | return cluster, nil 75 | } 76 | return nil, errNoClustersForTopic 77 | } 78 | -------------------------------------------------------------------------------- /kafka/resolver_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017 Uber Technologies, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | 21 | package kafka 22 | 23 | import ( 24 | "testing" 25 | 26 | "github.com/stretchr/testify/suite" 27 | ) 28 | 29 | type ResolverTestSuite struct { 30 | suite.Suite 31 | clusterBrokerMap map[string][]string 32 | topicClusterMap map[string][]string 33 | } 34 | 35 | func TestClientTestSuite(t *testing.T) { 36 | suite.Run(t, new(ResolverTestSuite)) 37 | } 38 | 39 | func (s *ResolverTestSuite) SetupTest() { 40 | s.clusterBrokerMap = map[string][]string{ 41 | "cluster_az1": {"127.0.0.1", "127.0.0.2", "127.0.0.3"}, 42 | "cluster_az2": {"127.1.1.1"}, 43 | "cluster_az3": {"127.2.2.2", "127.2.2.3"}, 44 | } 45 | s.topicClusterMap = map[string][]string{ 46 | "topic1": {"cluster_az1"}, 47 | } 48 | } 49 | 50 | func (s *ResolverTestSuite) TestResolveIPForCluster() { 51 | clusterBrokers := s.clusterBrokerMap 52 | topicClusterMap := s.topicClusterMap 53 | resolver := NewStaticNameResolver(topicClusterMap, clusterBrokers) 54 | s.NotNil(resolver) 55 | _, err := resolver.ResolveIPForCluster("foobar") 56 | s.Equal(errNoBrokersForCluster, err) 57 | for k, v := range clusterBrokers { 58 | brokers, err := resolver.ResolveIPForCluster(k) 59 | s.NoError(err) 60 | s.Equal(v, brokers) 61 | } 62 | } 63 | 64 | func (s *ResolverTestSuite) TestResolveClusterForTopic() { 65 | clusterBrokers := s.clusterBrokerMap 66 | topicClusterMap := s.topicClusterMap 67 | resolver := NewStaticNameResolver(topicClusterMap, clusterBrokers) 68 | s.NotNil(resolver) 69 | _, err := resolver.ResolveClusterForTopic("foobar") 70 | s.Equal(errNoClustersForTopic, err) 71 | for k, v := range topicClusterMap { 72 | cluster, err := resolver.ResolveClusterForTopic(k) 73 | s.NoError(err) 74 | s.Equal(v, cluster) 75 | } 76 | } 77 | -------------------------------------------------------------------------------- /scripts/cover.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -e 4 | echo "" > cover.out 5 | 6 | for d in $(go list $@); do 7 | go test -coverprofile=profile.out $d 8 | if [ -f profile.out ]; then 9 | cat profile.out >> cover.out 10 | rm profile.out 11 | fi 12 | done -------------------------------------------------------------------------------- /version: -------------------------------------------------------------------------------- 1 | v0.2.3.dev0 2 | --------------------------------------------------------------------------------