#!/usr/bin/env python # Copyright 2014 Cloudera, Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import getpass import optparse import os import signal import sys import toml from cm_api.api_client import ApiResource from supervisor import childutils # ------------------------------------------------------------------------------ class HueServer(object): def __init__(self, address, state): self.address = address self.state = state def __hash__(self): return hash((self.address, self.state)) def __str__(self): return '%s (%s)' % (self.address, self.state) def __eq__(self, other): return isinstance(other, HueServer) and \ self.address == other.address and \ self.state == other.state def __cmp__(self, other): return cmp( (self.address, self.state), (other.address, other.state)) # ------------------------------------------------------------------------------ class MonitorHue(object): def __init__( self, cm_client, stdin=sys.stdin, stdout=sys.stdout, stderr=sys.stderr): self.cm_client = cm_client self.stdin = stdin self.stdout = stdout self.stderr = stderr self.hue_servers = None self.listeners = [] def run_forever(self): self.tick() while True: headers, payload = childutils.listener.wait(self.stdin, self.stdout) # ignore non-tick events. if not headers['eventname'].startswith('TICK'): childutils.listener.ok(self.stdout) continue try: self.tick() finally: childutils.listener.ok(self.stdout) def tick(self): """Update the load balancer if any Hue servers have been added or removed""" hue_servers = self.get_hue_servers() # Don't do anything if we've already processed this set of hue servers. if self.hue_servers == hue_servers: return self.hue_servers = hue_servers print >> sys.stderr, 'updating server list:' for server in sorted(hue_servers): print >> sys.stderr, ' %s' % server for listener in self.listeners: listener.update(hue_servers) def get_hue_servers(self): """Fetch all the known hue servers""" hue_servers = set() for cluster in self.cm_client.get_all_clusters(): for service in cluster.get_all_services(): if service.type == 'HUE': for role in service.get_roles_by_type("HUE_SERVER"): host = self.cm_client.get_host(role.hostRef.hostId) hostname = host.hostname config = role.get_config(view='full') port = config['hue_http_port'] port = (port.value is None and port.default) or port.value address = '%s:%s' % (hostname, port) hue_servers.add(HueServer(address, role.roleState)) return hue_servers # ------------------------------------------------------------------------------ class ConfigListener(object): def __init__(self, rpc, config, process_names): self.rpc = rpc self.config = config self.process_names = process_names self.config_file = self.config['config_file'] with open(self.config['config_template']) as f: self.config_template = f.read() with open(self.config['server_template']) as f: self.server_template = f.read() def update(self, hue_servers): processes = self.rpc.supervisor.getAllProcessInfo() config = self.expand_config_template(hue_servers) with open(self.config_file, 'w') as f: print >> f, config for process in processes: # Ignore down processes. if not process['pid']: continue # Ignore processes we aren't monitoring. if process['name'] not in self.process_names: continue self.reload_process_config(process) def expand_config_template(self, hue_servers): config_servers = [] for index, server in enumerate(hue_servers): config_servers.append(self.expand_server_template(index, server)) return self.config_template % { 'servers': '\n'.join(config_servers) } def expand_server_template(self, index, server): raise NotImplementedError def reload_process(self, process): raise NotImplementedError # ------------------------------------------------------------------------------ class HAProxyListener(ConfigListener): def expand_server_template(self, index, server): """Expand the HAProxy server line""" return self.server_template % { 'index': index, 'address': server.address, } def reload_process_config(self, process): """Let haproxy-wrapper know the config file has changed.""" os.kill(process['pid'], signal.SIGUSR2) # ------------------------------------------------------------------------------ class NginxListener(ConfigListener): def __init__(self, *args, **kwargs): super(NginxListener, self).__init__(*args, **kwargs) self.alias_file = self.config['alias_file'] self.generate_alias_file() def find_static_files(self): paths = ( '/opt/cloudera/parcels/CDH/lib/hue/build/static/', '/usr/lib/hue/build/static/', os.path.join(os.path.dirname(__file__), '../../../build/static/'), ) for path in paths: if os.path.exists(path): return path else: raise Exception('could not find the Hue static files') def generate_alias_file(self): static_files = self.find_static_files() with open(self.alias_file, 'w') as f: print >> f, 'alias %s;' % static_files def expand_server_template(self, index, server): """Expand the Nginx server line""" config_server = self.server_template % { 'address': server.address, } # Let nginx know if this backend is offline. if server.state == 'STOPPED': config_server += ' down' return config_server + ';' def reload_process_config(self, process): """Let nginx know the config file has changed.""" os.kill(process['pid'], signal.SIGHUP) # ------------------------------------------------------------------------------ def main(argv): parser = optparse.OptionParser() parser.add_option('-c', '--config_file', help='config file') parser.add_option('--haproxy', default=[], type='str', action='append', help='monitor this haproxy process') parser.add_option('--nginx', default=[], type='str', action='append', help='monitor this haproxy process') options, args = parser.parse_args(argv[1:]) if options.config_file is None: options.config_file = 'etc/hue-lb.toml' try: with open(options.config_file) as f: config = toml.load(f) except IOError, e: print >> sys.stderr, 'error opening %s: %s' % (options.config_file, e) return 1 try: hostname = config['cloudera-manager']['hostname'] except KeyError: hostname = 'localhost' try: username = config['cloudera-manager']['username'] except KeyError: username = 'admin' try: password = config['cloudera-manager']['password'] except KeyError: password = 'admin' if len(options.haproxy) == 0 and len(options.nginx) == 0: print >> sys.stderr, 'no processes specified' return 1 client = ApiResource(hostname, username=username, password=password) monitor = MonitorHue(client) rpc = childutils.getRPCInterface(os.environ) if len(options.haproxy) != 0: listener = HAProxyListener(rpc, config['haproxy'], options.haproxy) monitor.listeners.append(listener) if len(options.nginx) != 0: listener = NginxListener(rpc, config['nginx'], options.nginx) monitor.listeners.append(listener) monitor.run_forever() return 0 # ------------------------------------------------------------------------------ if __name__ == '__main__': sys.exit(main(sys.argv))