Enable private repos - merge conflict fixed

Dockerfile

@@ -1,7 +1,8 @@
 FROM ubuntu:14.04
 
 RUN apt-get update
-RUN apt-get install -y wget make gcc binutils python-pip python-dev libssl-dev libffi-dev bash
+
+RUN apt-get install -y wget make gcc binutils python-pip python-dev libssl-dev libffi-dev
 
 WORKDIR /root
 

main.py

@@ -1,81 +1,190 @@
-#!/usr/bin/python
+#!/usr/bin/env python
 
-import etcd
-from jinja2 import Environment, PackageLoader
-import os
+from collections import defaultdict
+from distutils.spawn import find_executable
 from subprocess import call
+import os
 import signal
 import sys
 import time
 
+from jinja2 import Environment, PackageLoader
+import etcd
+
+
 env = Environment(loader=PackageLoader('haproxy', 'templates'))
-POLL_TIMEOUT=5
+POLL_TIMEOUT = 5
+HAPROXY_CONFIG = '/etc/haproxy.cfg'
+HAPROXY_PID = '/var/run/haproxy.pid'
+
 
 signal.signal(signal.SIGCHLD, signal.SIG_IGN)
 
 def get_etcd_addr():
- if "ETCD_HOST" not in os.environ:
- print "ETCD_HOST not set"
- sys.exit(1)
+ """
+ Determine the host and port that etcd should be available on using the
+ `ETCD_HOST` environment variable..
+
+ :returns:
+ A 2-tuple with the hostname/IP and the numeric TCP port at which etcd
+ can be reached.
+
+ :raises SystemExit:
+ If the `ETCD_HOST` environment variable is not defined or is empty.
+
+ """
 
- etcd_host = os.environ["ETCD_HOST"]
+ etcd_host = os.environ.get("ETCD_HOST", None)
  if not etcd_host:
- print "ETCD_HOST not set"
+ print("ETCD_HOST not set")
  sys.exit(1)
 
- port = 4001
- host = etcd_host
+ host, port = etcd_host, 4001
+ if ":" in host:
+ host, port = host.split(":")
 
- if ":" in etcd_host:
- host, port = etcd_host.split(":")
+ return host, int(port)
 
- return host, port
 
-def get_services():
+def get_haproxy_path():
+ """
+ Return the absolute path to the `haproxy` executable.
 
- host, port = get_etcd_addr()
- client = etcd.Client(host=host, port=int(port))
- backends = client.read('/backends', recursive = True)
- services = {}
+ :raises SystemExit:
+ If haproxy cannot be found on the PATH.
 
- for i in backends.children:
+ """
+
+ path = find_executable('haproxy')
+ if not path:
+ print('haproxy was not found on your PATH, and it must be installed '
+ 'to use this script')
+ sys.exit(1)
+
+ return path
+
+
+def get_services(client):
+ """
+ Find all services which have been published to etcd and have exposed a
+ port.
 
- if i.key[1:].count("/") != 2:
+ :param etcd.Client client:
+ A handle to an etcd server.
+
+ :returns:
+ A dictionary of dictionaries keyed on service name. The inner
+ dictionary includes the TCP port that the service uses, along with a
+ list of IP:port values that refer to containers which have exposed the
+ service port (thereby acting as backend services).
+
+ """
+
+ # TODO: handle severed connection, etc
+ backends = client.read('/backends', recursive=True)
+ services = defaultdict(lambda: {
+ 'port': None,
+ 'backends': []
+ })
+
+ for i in backends.children:
+ if i.key[1:].count("/") < 2:
  continue
 
- ignore, service, container = i.key[1:].split("/")
- endpoints = services.setdefault(service, dict(port="", backends=[]))
+ ignore, service, container = i.key[1:].rsplit("/", 2)
+ endpoints = services[service]
  if container == "port":
  endpoints["port"] = i.value
  continue
+
  endpoints["backends"].append(dict(name=container, addr=i.value))
+
+ # filter out services with no "port" value in etcd
+ for svc, data in tuple(services.items()):
+ if data['port'] is None:
+ services.pop(svc)
+
+ for service in services:
+ services[service]["backends"].sort()
+
  return services
 
+
 def generate_config(services):
+ """
+ Generate a configuration file for haproxy and save it to disk.
+
+ It is expected that the results of :py:func:`get_services` will be passed
+ to this function.
+
+ :param dict services:
+ A dictionary of dictionaries, keyed on service name.
+
+ """
+
  template = env.get_template('haproxy.cfg.tmpl')
- with open("/etc/haproxy.cfg", "w") as f:
+ with open(HAPROXY_CONFIG, "w") as f:
  f.write(template.render(services=services))
 
-if __name__ == "__main__":
+
+def restart_haproxy():
+ """
+ Restart haproxy.
+
+ :returns:
+ ``True`` when haproxy appears to have restarted successfully, ``False``
+ otherwise.
+
+ """
+
+ path = get_haproxy_path()
+ cmd = '{haproxy} -f {cfg} -p {pid} -sf $(cat {pid})'.format(
+ haproxy=path,
+ cfg=HAPROXY_CONFIG,
+ pid=HAPROXY_PID,
+ )
+
+ # TODO: shell=True is yucky... read in PID rather than using the cat
+ return call(cmd, shell=True) == 0
+
+
+def main():
+ """
+ Periodically poll etcd for the list of available services running in docker
+ containers. When a new service becomes available or a service disappears,
+ update the configuration for haproxy and restart it.
+
+ """
+
+ # check for haproxy and etcd config before getting into the real code
+ get_haproxy_path()
+ host, port = get_etcd_addr()
+
+ client = None
  current_services = {}
+
  while True:
- try:
- services = get_services()
+ if client is None:
+ # TODO: connection error handling
+ client = etcd.Client(host=host, port=port)
+
+ services = get_services(client)
 
- if not services or services == current_services:
- time.sleep(POLL_TIMEOUT)
- continue
+ if services != current_services:
+ print("config changed. reload haproxy")
 
- print "config changed. reload haproxy"
  generate_config(services)
- ret = call(["./reload-haproxy.sh"])
- if ret != 0:
- print "reloading haproxy returned: ", ret
- time.sleep(POLL_TIMEOUT)
- continue
- current_services = services
 
- except Exception, e:
- print "Error:", e
+ if restart_haproxy():
+ current_services = services
+ else:
+ print("failed to restart haproxy!")
 
- time.sleep(POLL_TIMEOUT)
+ time.sleep(POLL_TIMEOUT)
+
+
+if __name__ == "__main__":
+ try:
+ main()
+ except KeyboardInterrupt:
+ pass