/
whirlwind.py
executable file
·170 lines (124 loc) · 4.63 KB
/
whirlwind.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
#!/usr/bin/env python
from argparse import ArgumentParser
import logging
import os
import random
import sys
import time
from jujuclient import Environment
import yaml
log = logging.getLogger('whirlwind')
class Worker(object):
def __init__(self, config):
self.config = self.prepare_config(config)
self.env = None
def start(self):
self.env = self.connect_state_server()
if not self.env:
return False
while True:
try:
if self.perform_change():
time.sleep(self.config.change_interval)
else:
time.sleep(self.config.change_retry_interval)
except KeyboardInterrupt:
return
except:
log.exception('Error occured during proactive change')
return
def connect_state_server(self):
while True:
try:
log.debug('Connecting to state server for environment: %s', self.config.environment)
return Environment.connect(self.config.environment)
except KeyboardInterrupt:
return
except:
log.exception('Could not connect to state server')
log.debug('Retrying connection in 30 seconds')
time.sleep(30)
def perform_change(self):
services = self.config.services.keys()
random.shuffle(services)
for service in services:
if self.try_change_service(service):
return True
log.debug('Service not ready for change: %s', service)
log.warn('No services ready for proactive change')
def try_change_service(self, service):
started_units = self.fetch_units(service)
if not started_units:
return
log.debug('Performing proactive change for service: %s', service)
log.debug('Adding unit to service')
self.env.add_unit(service)
while True:
if len(self.fetch_units(service)) > len(started_units):
break
log.debug('Waiting 30 seconds for unit to be added')
time.sleep(30)
unit_to_remove = random.choice(started_units.keys())
log.debug('Removing unit: %s', unit_to_remove)
self.env.remove_units([unit_to_remove])
while True:
if unit_to_remove not in self.fetch_units(service, started=False):
break
log.debug('Waiting 15 seconds for unit to be removed')
time.sleep(15)
if self.config.remove_machines:
machine = started_units[unit_to_remove]['Machine']
log.debug('Removing machine: %s', machine)
self.env.destroy_machines([machine])
log.debug('Proactive change successful')
return True
def fetch_units(self, service, started=True):
status = self.env.status()
if service not in status['Services']:
return
if not started:
return status['Services'][service]['Units']
return dict((k, v) for k, v in status['Services'][service]['Units'].iteritems() if v['AgentState'] == 'started')
def prepare_config(self, config):
config = normalize_structure(config)
for k, v in config.services.items():
config.services[k] = normalize_structure(v)
return config
def normalize_structure(val, recursive=False):
if isinstance(val, dict):
s = Structure()
for k, v in val.iteritems():
if recursive:
v = normalize_config(v, True)
setattr(s, k, v)
return s
elif recursive and isinstance(val, list):
return [normalize_structure(v, True) for v in val]
return val
class Structure(object):
pass
def main():
parser = ArgumentParser()
parser.add_argument('-v', '--verbose', action='store_true')
parser.add_argument('-c', '--config-file', default='config.yml')
args = parser.parse_args()
configure_logger(args.verbose)
if not os.path.isfile(args.config_file):
log.error('Could not open config file: %s', args.config_file)
sys.exit(1)
with open(args.config_file, 'r') as fp:
config = yaml.load(fp)
worker = Worker(config)
if not worker.start():
sys.exit(1)
def configure_logger(verbose):
formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s', '%Y-%m-%d %H:%M:%S')
handler = logging.StreamHandler(sys.stdout)
handler.setFormatter(formatter)
log.addHandler(handler)
if verbose:
log.setLevel(logging.DEBUG)
else:
log.setLevel(logging.WARN)
if __name__ == '__main__':
main()