Marathon service discovery: writing HAProxy and hosts file

gistfile1.py

#!/usr/bin/env python

import urllib2
import json
import random
import subprocess

slave_ip = "10.0.10.1"
marathons = ["10.0.10.1", "10.0.10.2"]
domain = ".example.com"

haproxy_header = """
global
  daemon
  log 127.0.0.1 local0
  log 127.0.0.1 local1 notice
  maxconn 4096

defaults
  log         global
  retries     3
  maxconn     2000
  contimeout  5000
  clitimeout  50000
  srvtimeout  50000

listen stats
  bind 127.0.0.1:9090
  balance
  mode http
  stats enable
  stats auth admin:admin

"""
backend_template = "backend {app_name}\n  balance roundrobin\n  mode tcp\n"
backend_server_template = "  server {instance_id} {host}:{port} check\n"
frontend_template = "frontend http\n  bind *:80\n  mode tcp\n"
frontend_host_template = "  acl host_{app_name} hdr(host) -i {app_name}{domain}\n  use_backend {app_name} if host_{app_name}"

hosts_comment = "# Added by service discovery"
hosts_template = "127.0.0.1\t{app_name}{domain} {comment}"


class Application:
    def __init__(self, app_name):
        self.app_name = app_name
        self.backends = {}

    def add_backend(self, host, ports):
        d = self.backends.get(host, [])
        d.extend(ports)
        self.backends[host] = d

    def sanitized_app_name(self):
        return self.app_name.replace("/", "")

    def output_server_backend(self, instance_id, host, port):
        if host == slave_ip:
            host = ""
        return backend_server_template.format(instance_id=instance_id, host=host, port=port)

    def output_backend(self):
        i = 0
        result = backend_template.format(app_name = self.sanitized_app_name())
        for host,ports in self.backends.iteritems():
            for port in ports:
                app_name = "%s-%02d" % (self.sanitized_app_name(), i)
                result += self.output_server_backend(app_name, host, port)
                i += 1
        return result

    def output_frontend(self):
        return frontend_host_template.format(app_name = self.sanitized_app_name(), domain = domain)

    def output_hosts_entry(self):
        return hosts_template.format(app_name = self.sanitized_app_name(), domain = domain, comment = hosts_comment)


class Applications:
    def __init__(self):
        self.applications = {}

        # Pre-load with marathon config
        marathon = Application("marathon")
        for m in marathons:
            host = m.split(":")[0]
            ports = [m.split(":")[1]]
            marathon.add_backend(host, ports)
        self.applications["marathon"] = marathon

    def load_from_tasks(self, tasks):
        for t in tasks:
            app_name  = t["appId"]
            app = self.applications.get(app_name, Application(app_name))
            app.add_backend(t["host"], t["ports"])
            self.applications[app_name] = app

    def generate_backends(self):
        result = []
        for app in self.applications.itervalues():
            result.append(app.output_backend())
        return "\n".join(result) + "\n"

    def generate_frontends(self):
        result = [frontend_template]
        for app in self.applications.itervalues():
            result.append(app.output_frontend())
        return "\n".join(result) + "\n"

    def generate_haproxy_config(self):
        result = haproxy_header
        result += self.generate_backends()
        result += self.generate_frontends()
        return result

    def generate_hosts_file_entries(self):
        result = []
        for app in self.applications.itervalues():
            result.append(app.output_hosts_entry())
        return "\n".join(result) + "\n"


def get_json():
    marathon = random.choice(marathons)
    request = urllib2.urlopen("http://%s/v2/tasks" % marathon)
    payload = json.load(request)
    request.close()
    return payload

def reload_haproxy_if_started_else_start():
    process = subprocess.Popen(["/usr/sbin/service", "haproxy", "status"]).wait()
    command = "reload" if process.returncode == 0 else "start"
    subprocess.Popen(["/usr/sbin/service", "haproxy", command]).wait()

def write_haproxy_cfg_and_reload(haproxy_cfg):
    f = open("/etc/haproxy/haproxy.cfg", "w")
    f.write(haproxy_cfg)
    f.close()
    reload_haproxy_if_started_else_start()

def generate_new_hosts_file(hosts_entries):
    result = []
    f = open("/etc/hosts", "r")
    for line in f.readlines():
        line_ = line.strip()
        if not line_.endswith(hosts_comment):
            result.append(line)
    result.append(hosts_entries)
    f.close()
    return "".join(result)

def write_hosts_file(hosts):
    f = open("/etc/hosts", "w")
    f.write(hosts)
    f.close()

def main():
    tasks = get_json()["tasks"]
    applications = Applications()
    applications.load_from_tasks(tasks)
    haproxy_cfg = applications.generate_haproxy_config()
    hosts_entries = applications.generate_hosts_file_entries()
    hosts_file = generate_new_hosts_file(hosts_entries)
    write_hosts_file(hosts_file)
    write_haproxy_cfg_and_reload(haproxy_cfg)

if __name__ == '__main__':
    main()