forked from chutium/saiki-buku
-
Notifications
You must be signed in to change notification settings - Fork 0
/
rebalance_partitions.py
263 lines (214 loc) · 9.87 KB
/
rebalance_partitions.py
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
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
from kazoo.client import KazooClient, KazooState, NodeExistsError, NoNodeError
from time import sleep
from collections import Counter
import logging
import json
import os
from kazoo.protocol.states import EventType
class NotEnoughBrokersException(Exception):
def __init__(self):
logging.warning("NotEnoughBrokersException")
import sys
sys.stdout.flush()
def state_listener(state):
if state == KazooState.LOST:
# Register somewhere that the session was lost
pass
elif state == KazooState.SUSPENDED:
# Handle being disconnected from Zookeeper
pass
else:
# Handle being connected/reconnected to Zookeeper
pass
def readout_brokerids(zk):
try:
return zk.get_children("/brokers/ids")
except NoNodeError:
logging.error("there are no brokers registrated in this zookeeper cluster, therefore exiting")
exit(1)
def readout_topics(zk):
try:
return zk.get_children("/brokers/topics")
except NoNodeError:
logging.info("there are no topics registrated in this zookeeper cluster")
def readout_topic_details(zk, topic):
try:
return json.loads(zk.get("/brokers/topics/" + topic)[0].decode('utf-8'))
except NoNodeError:
logging.info("no information for topic %s existing", topic)
def readout_partitions(zk, topic):
try:
return zk.get_children("/brokers/topics/" + topic + "/partitions")
except NoNodeError:
logging.info("there are no partitions for topic %s in this zookeeper cluster", topic)
def check_for_broken_partitions(zk_dict):
result = {}
for topic in zk_dict['topics']:
logging.debug("checking topic: %s", topic['name'])
for partition, brokers in topic['partitions'].items():
logging.debug("checking partition: %s", partition)
for part_broker_id in brokers:
logging.debug("checking if this broker is still existing: %s", part_broker_id)
if str(part_broker_id) not in zk_dict['broker']:
if topic['name'] not in result:
result[topic['name']] = {}
result[topic['name']][partition] = part_broker_id
logging.debug("no, it isn't ...")
break
logging.debug("yes, it is ...")
return result
def get_own_ip():
import requests
return requests.get('http://169.254.169.254/latest/dynamic/instance-identity/document').json()['privateIp']
def update_broker_weigths(weights, brokers):
for i, broker in enumerate(reversed(brokers)):
if broker in weights:
weights[broker] += 2 ** (i + i)
def get_broker_weights(zk_dict, target_brokers="all", ignore_existing=False):
if target_brokers == "all":
target_brokers = zk_dict['broker']
weights = {int(broker): 0 for broker in target_brokers}
if not ignore_existing:
for topic in zk_dict['topics']:
for brokers in topic['partitions'].values():
update_broker_weigths(weights, brokers)
return weights
def generate_json(zk_dict, topics_to_reassign="all", target_brokers="all"):
ignore_existing = False
new_broker_ass = False
if set(zk_dict['broker']) - set(target_brokers) != set() and target_brokers != "all":
avail_brokers_init = target_brokers
new_broker_ass = True
logging.debug("new broker assignment: " + str(avail_brokers_init))
else:
avail_brokers_init = zk_dict['broker']
if topics_to_reassign == "all" or new_broker_ass is True:
# logging.info("reassigning all topics")
topics_to_reassign = {}
for topic in zk_dict['topics']:
topics_to_reassign[topic['name']] = {}
for partition in topic['partitions']:
topics_to_reassign[topic['name']][partition] = 0
ignore_existing = True
logging.debug("topics_to_reassign:")
logging.debug(topics_to_reassign)
tmp_topic_dict = {t['name']: t['partitions'] for t in zk_dict['topics']}
# this is needed for creating new topics (currently in pemetaan)
for topic_to_reassign, value in topics_to_reassign.items():
if topic_to_reassign not in tmp_topic_dict:
tmp_topic_dict[topic_to_reassign] = {}
for partition, brokers in value.items():
tmp_topic_dict[topic_to_reassign][str(partition)] = brokers
if len(topics_to_reassign) > 0:
logging.info("topics_to_reassign found, generating new assignment pattern")
logging.info("reading out broker id's")
final_result = {'version': 1, 'partitions': []}
logging.info("generating now ")
weights = get_broker_weights(zk_dict, avail_brokers_init, ignore_existing)
for topic, partitions in topics_to_reassign.items():
counter = Counter(map(len, tmp_topic_dict[topic].values()))
replication_factor = sorted(counter.items(), key=lambda e: e[1])[0][0]
logging.debug("Available Brokers: %s", len(avail_brokers_init))
logging.debug("Replication Factor: %s", replication_factor)
if len(avail_brokers_init) < replication_factor:
raise NotEnoughBrokersException
for partition in partitions:
logging.debug("finding new brokers for topic: %s, partition: %s", topic, partition)
broker_list = [b for b, w in sorted(weights.items(), key=lambda v: v[1])][:replication_factor]
final_result['partitions'].append({'topic': topic,
'partition': int(partition),
'replicas': broker_list})
update_broker_weigths(weights, broker_list)
return final_result
else:
return {}
def write_json_to_zk(zk, final_result):
watcher_event = zk.handler.event_object()
def reassign_partitions_watcher(event):
if event.type == EventType.DELETED:
watcher_event.set()
else:
if zk.exists('/admin/reassign_partitions', reassign_partitions_watcher) is None:
watcher_event.set()
logging.info("writing reassigned partitions in ZK")
count_steps_left = len(final_result['partitions'])
for step in final_result['partitions']:
if count_steps_left % 20 == 0 or count_steps_left == len(final_result['partitions']):
logging.info("steps left: %s", count_steps_left)
logging.info("trying to write zk node for repairing %s", step)
done = False
while not done:
try:
zk.create("/admin/reassign_partitions",
json.dumps({'version': 1, 'partitions': [step]}, separators=(',', ':')).encode('utf-8'))
done = True
logging.info("done")
count_steps_left -= 1
watcher_event.clear()
if zk.exists('/admin/reassign_partitions', reassign_partitions_watcher) is not None:
watcher_event.wait(300)
except NodeExistsError:
try:
watcher_event.clear()
check = zk.get("/admin/reassign_partitions", reassign_partitions_watcher)
if check[0] == b'{"version": 1, "partitions": []}':
zk.delete("/admin/reassign_partitions", recursive=True)
continue
for timeout_count in range(0, 6):
# only output message every 10mins
logging.info("there seems to be a reassigning already taking place: %s",
check[0].decode('utf-8'))
logging.info("waiting ...")
watcher_event.wait(300)
if watcher_event.isSet():
break
except NoNodeError:
pass
# logging.info("NoNodeError")
if done is False:
logging.warning("Reassignment was not successfull due to timeout issues of the previous reassignment")
break
def get_zk_dict(zk):
result = {'topics': [], 'broker': readout_brokerids(zk)}
for topic in readout_topics(zk):
result['topics'].append({'name': topic, 'partitions': readout_topic_details(zk, topic)['partitions']})
return result
def connect_to_zk():
zookeeper_connect_string = os.getenv('ZOOKEEPER_CONN_STRING')
zk = KazooClient(hosts=zookeeper_connect_string)
zk.start()
zk.add_listener(state_listener)
logging.info("connected to Zookeeper")
return zk
def run():
import wait_for_kafka_startup
logging.info("waiting for kafka to start up")
if os.getenv('WAIT_FOR_KAFKA') != 'no':
wait_for_kafka_startup.run(get_own_ip())
else:
sleep(10)
logging.info("kafka port is open, continuing")
zk = connect_to_zk()
zk_dict = get_zk_dict(zk)
logging.info("checking for broken topics")
result = generate_json(zk_dict, topics_to_reassign=check_for_broken_partitions(zk_dict))
if result != {}:
logging.info("there are %s partitions to repair", len(result['partitions']))
logging.debug(result)
if os.getenv('WRITE_TO_JSON') != 'no':
logging.info("writing to ZooKeeper ...")
write_json_to_zk(zk, result)
else:
logging.info("no broken topics found, no JSON generated")
if any(weight == 0 for weight in get_broker_weights(zk_dict).values()):
logging.info("there are unused brokers, reassigning all topics ...")
result = generate_json(zk_dict)
if result != {}:
logging.info("JSON generated")
logging.debug(result)
if os.getenv('WRITE_TO_JSON') != 'no':
write_json_to_zk(zk, result)
else:
logging.info("no unused Broker found, no JSON generated")
zk.stop()
logging.info("exiting")