-
Notifications
You must be signed in to change notification settings - Fork 0
/
k8sbalance.py
executable file
·114 lines (83 loc) · 3.12 KB
/
k8sbalance.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
#!/usr/bin/env python3
import sys
SERVICE_LABEL = "spreadspace.org/onion-service"
INSTANCE_ANNOT = "spreadspace.org/onion-instance"
REFRESH_INTERVAL = 120
SECRETS_PATH = "/var/run/secrets/spreadspace.org/onionbalance"
ONIONBALANCE_CONFIG = "/tmp/onionbalance.yml"
ONIONBALANCE_CONTROL = "/tmp/onionbalance.control"
def get_onion_mapping(client, namespace):
pods = client.list_namespaced_pod(namespace, label_selector=SERVICE_LABEL)
mapping = {}
for pod in pods.items:
if not pod.metadata.annotations:
continue
if INSTANCE_ANNOT not in pod.metadata.annotations:
continue
service = pod.metadata.labels[SERVICE_LABEL]
instance = pod.metadata.annotations[INSTANCE_ANNOT]
if service not in mapping:
mapping[service] = set()
mapping[service].add(instance)
return mapping
def onionbalance_config(mapping):
import json
import os.path
return json.dumps(
{
"STATUS_SOCKET_LOCATION": ONIONBALANCE_CONTROL,
"REFRESH_INTERVAL": REFRESH_INTERVAL,
"services": [
{
"key": os.path.join(SECRETS_PATH, address),
"instances": [{"address": s} for s in instances],
}
for address, instances in mapping.items()
],
}
)
def start_onionbalance(mapping):
from subprocess import Popen
with open(ONIONBALANCE_CONFIG, "w") as fd:
fd.write(onionbalance_config(mapping))
return Popen(["onionbalance", "-c", ONIONBALANCE_CONFIG])
def kill(process):
from subprocess import TimeoutExpired
print("Sending SIGTERM to onionbalance")
process.terminate()
try:
process.wait(timeout=5)
except TimeoutExpired:
print("Onionbalance failed to terminate within 5s")
process.kill()
process.wait(timeout=60)
def log_changes(oldmap, newmap, output=sys.stderr):
output.write("Updating onionbalance config:\n")
for host in set(itertools.chain(newmap.keys(), oldmap.keys())):
if host in newmap and host in oldmap and newmap[host] == oldmap[host]:
continue
output.write(" %s\n" % host)
output.write(" Adding: %s\n" % (newmap[host] - oldmap[host]))
output.write(" Removing: %s\n" % (oldmap[host] - newmap[host]))
output.write(" Keeping: %s\n" % (oldmap[host] & newmap[host]))
output.flush()
if __name__ == "__main__":
import itertools
import os
from kubernetes import client, config, watch
NAMESPACE = os.environ["POD_NAMESPACE"]
config.incluster_config.load_incluster_config()
v1 = client.CoreV1Api()
onionmap = get_onion_mapping(v1, NAMESPACE)
onionbalance = start_onionbalance(onionmap)
stream = watch.Watch().stream(
v1.list_namespaced_pod, NAMESPACE, label_selector=SERVICE_LABEL
)
for _ in stream:
newmap = get_onion_mapping(v1, NAMESPACE)
if newmap == onionmap:
continue
log_changes(onionmap, newmap)
onionmap = newmap
kill(onionbalance)
onionbalance = start_onionbalance(onionmap)