forked from influxdata/telegraf
-
Notifications
You must be signed in to change notification settings - Fork 0
/
kazoo.go
210 lines (174 loc) · 5.08 KB
/
kazoo.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
package kazoo
import (
"encoding/json"
"errors"
"fmt"
"path"
"strconv"
"strings"
"time"
"github.com/samuel/go-zookeeper/zk"
)
var (
FailedToClaimPartition = errors.New("Failed to claim partition for this consumer instance. Do you have a rogue consumer running?")
)
// ParseConnectionString parses a zookeeper connection string in the form of
// host1:2181,host2:2181/chroot and returns the list of servers, and the chroot.
func ParseConnectionString(zookeeper string) (nodes []string, chroot string) {
nodesAndChroot := strings.SplitN(zookeeper, "/", 2)
if len(nodesAndChroot) == 2 {
chroot = fmt.Sprintf("/%s", nodesAndChroot[1])
}
nodes = strings.Split(nodesAndChroot[0], ",")
return
}
// Kazoo interacts with the Kafka metadata in Zookeeper
type Kazoo struct {
conn *zk.Conn
conf *Config
}
// Config holds configuration values f.
type Config struct {
// The chroot the Kafka installation is registerde under. Defaults to "".
Chroot string
// The amount of time the Zookeeper client can be disconnected from the Zookeeper cluster
// before the cluster will get rid of watches and ephemeral nodes. Defaults to 1 second.
Timeout time.Duration
}
// NewConfig instantiates a new Config struct with sane defaults.
func NewConfig() *Config {
return &Config{Timeout: 1 * time.Second}
}
// NewKazoo creates a new connection instance
func NewKazoo(servers []string, conf *Config) (*Kazoo, error) {
if conf == nil {
conf = NewConfig()
}
conn, _, err := zk.Connect(servers, conf.Timeout)
if err != nil {
return nil, err
}
return &Kazoo{conn, conf}, nil
}
// NewKazooFromConnectionString creates a new connection instance
// based on a zookeeer connection string that can include a chroot.
func NewKazooFromConnectionString(connectionString string, conf *Config) (*Kazoo, error) {
if conf == nil {
conf = NewConfig()
}
nodes, chroot := ParseConnectionString(connectionString)
conf.Chroot = chroot
return NewKazoo(nodes, conf)
}
// Brokers returns a map of all the brokers that make part of the
// Kafka cluster that is regeistered in Zookeeper.
func (kz *Kazoo) Brokers() (map[int32]string, error) {
root := fmt.Sprintf("%s/brokers/ids", kz.conf.Chroot)
children, _, err := kz.conn.Children(root)
if err != nil {
return nil, err
}
type brokerEntry struct {
Host string `json:"host"`
Port int `json:"port"`
}
result := make(map[int32]string)
for _, child := range children {
brokerID, err := strconv.ParseInt(child, 10, 32)
if err != nil {
return nil, err
}
value, _, err := kz.conn.Get(path.Join(root, child))
if err != nil {
return nil, err
}
var brokerNode brokerEntry
if err := json.Unmarshal(value, &brokerNode); err != nil {
return nil, err
}
result[int32(brokerID)] = fmt.Sprintf("%s:%d", brokerNode.Host, brokerNode.Port)
}
return result, nil
}
// BrokerList returns a slice of broker addresses that can be used to connect to
// the Kafka cluster, e.g. using `sarama.NewAsyncProducer()`.
func (kz *Kazoo) BrokerList() ([]string, error) {
brokers, err := kz.Brokers()
if err != nil {
return nil, err
}
result := make([]string, 0, len(brokers))
for _, broker := range brokers {
result = append(result, broker)
}
return result, nil
}
// Controller returns what broker is currently acting as controller of the Kafka cluster
func (kz *Kazoo) Controller() (int32, error) {
type controllerEntry struct {
BrokerID int32 `json:"brokerid"`
}
node := fmt.Sprintf("%s/controller", kz.conf.Chroot)
data, _, err := kz.conn.Get(node)
if err != nil {
return -1, err
}
var controllerNode controllerEntry
if err := json.Unmarshal(data, &controllerNode); err != nil {
return -1, err
}
return controllerNode.BrokerID, nil
}
func (kz *Kazoo) Close() error {
kz.conn.Close()
return nil
}
////////////////////////////////////////////////////////////////////////
// Util methods
////////////////////////////////////////////////////////////////////////
// Exists checks existence of a node
func (kz *Kazoo) exists(node string) (ok bool, err error) {
ok, _, err = kz.conn.Exists(node)
return
}
// DeleteAll deletes a node recursively
func (kz *Kazoo) deleteRecursive(node string) (err error) {
children, stat, err := kz.conn.Children(node)
if err == zk.ErrNoNode {
return nil
} else if err != nil {
return
}
for _, child := range children {
if err = kz.deleteRecursive(path.Join(node, child)); err != nil {
return
}
}
return kz.conn.Delete(node, stat.Version)
}
// MkdirAll creates a directory recursively
func (kz *Kazoo) mkdirRecursive(node string) (err error) {
parent := path.Dir(node)
if parent != "/" {
if err = kz.mkdirRecursive(parent); err != nil {
return
}
}
_, err = kz.conn.Create(node, nil, 0, zk.WorldACL(zk.PermAll))
if err == zk.ErrNodeExists {
err = nil
}
return
}
// Create stores a new value at node. Fails if already set.
func (kz *Kazoo) create(node string, value []byte, ephemeral bool) (err error) {
if err = kz.mkdirRecursive(path.Dir(node)); err != nil {
return
}
flags := int32(0)
if ephemeral {
flags = zk.FlagEphemeral
}
_, err = kz.conn.Create(node, value, flags, zk.WorldACL(zk.PermAll))
return
}