chronos_docker

#!/usr/bin/env python
import itertools
import json
import os
import re
import signal
import subprocess
import sys
import threading
import time
import traceback
import shlex

import mesos
import mesos_pb2

import socket


# Syslog host information
SYSLOG_HOST = "localhost"
SYSLOG_PORT = 1122
SYSLOG_APPNAME = "Chronos_docker"


class MySysLogger():
    def __init__(self, app_name=SYSLOG_APPNAME, host=SYSLOG_HOST, port=SYSLOG_PORT):
        self.app_name = app_name
        self.host = host
        self.port = port
        self.hostname = socket.gethostname()

    def __getattr__(self, facility):
        def wrapper_func(message):
            localtime = time.localtime()
            timestamp = time.strftime('%b %d %H:%M:%S', localtime)
            sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            data = "%s %s %s %s %s" % (str(timestamp), self.hostname, self.app_name, facility.upper(), message)
            sock.sendto(data, (self.host, self.port))
            sock.close()

        return wrapper_func


def main():
    def handler(signum, _):
        log.info('Exiting due to signal: ' + str(signum))
        exit(-signum)

    signal.signal(signal.SIGINT, handler)
    signal.signal(signal.SIGTERM, handler)
    signal.signal(signal.SIGABRT, handler)
    signal.signal(signal.SIGPIPE, handler)
    signal.signal(signal.SIGSEGV, handler)
    driver = mesos.MesosExecutorDriver(DockerExecutor(sys.argv[1:]))
    log.info('Ready to serve!')
    log.info('Args in handler function: %s' % sys.argv)
    exit(0 if driver.run() == mesos_pb2.DRIVER_STOPPED else 1)

# This script wraps one and only one container.
cid = None
cidfile = None


def read_cid():
    global cid
    if cidfile is not None:
        try:
            with open(cidfile) as f:
                cid = f.read().strip()
        except IOError:
            pass


log = MySysLogger(SYSLOG_APPNAME)


class DockerExecutor(mesos.Executor):
    def __init__(self, args):
        self.args = args
        self.task = None
        self.driver = None
        self.runner_thread = None
        self.shutdown_thread = None
        self.data = {}
        self.env = {}
        log.info('Args Provided __INIT__: %s' % args)

    def run(self):
        global proc
        exitcode = 2
        finalstate = mesos_pb2.TASK_FAILED
        log.info("Task from Chronos: %s" % self.task_json)
        try:
            img = self.task_json['container']
            args = self.task_json['cmd']
            params = self.task_json['params']
            cmd = run_with_settings(img, args, params)
            self.send_state(mesos_pb2.TASK_RUNNING)
            proc = subprocess.call(cmd)
            log.info('Container exited with code: %d' % proc)
            if proc == 0:
                finalstate = mesos_pb2.TASK_FINISHED
            else:
                if self.shutdown_thread:
                    finalstate = mesos_pb2.TASK_KILLED
                else:
                    finalstate = mesos_pb2.TASK_FAILED
        except Exception, e:
            log_exc()
        finally:
            self.send_state(finalstate)
            exit(exitcode)

    def send_state(self, state):
        try:
            update = mesos_pb2.TaskStatus()
            update.task_id.value = self.task.task_id.value
            update.state = state
            self.driver.sendStatusUpdate(update)
        except Exception, e:
            log_exc()

    #### Mesos Executor API methods ####

    def registered(self, driver, executorInfo, frameworkInfo, slaveInfo):
        log.info("Registered with Mesos slave")

    def reregistered(driver, slaveInfo):
        log.info("Reregistered with Mesos slave. More info: %s" % slaveInfo)

    def disconnected(driver):
        log.warning('Disconnected from Mesos slave')

    def launchTask(self, driver, task):
        if self.task is not None:
            log.error('Executor was reused but this executor is not reuseable. Old Task: %s' % self.task)
            exit(2)
        self.task = task
        self.driver = driver
        self.task_data = self.task.data.replace("|", "")
        try:
            self.task_json = json.loads(self.task_data)
        except Exception, e:
            log.info("JSON Task Data: %s" % self.task_data)
            log.info("JSON Exception: %s" % e)
        log.info('Task is: %s' % task.task_id.value)
        try:
            self.run_thread = threading.Thread(target=self.run)
            self.run_thread.daemon = True
            self.run_thread.start()
        except Exception, e:
            log_exc()
            self.send_state(mesos_pb2.TASK_FAILED)
            exit(2)

    def killTask(self, driver, task_id):
        if self.task.task_id.value == task_id.value:
            log.info('Asked to shutdown managed task %s' % task_id.value)
            self.cleanup()
        else:
            log.info('Asked to shutdown unknown task %s' % task_id.value)

    def cleanup(self):
        if self.shutdown_thread is None:
            self.shutdown_thread = threading.Thread(target=cleanup_container)
            self.shutdown_thread.daemon = True
            self.shutdown_thread.start()

    def allocated_ports(self):
        range_resources = [_.ranges.range for _ in self.task.resources
                           if _.name == 'ports']
        ranges = itertools.chain(*range_resources)
        # NB: Casting long() to int() so there's no trailing 'L' in later
        #     stringifications. Ports should only ever be shorts, anyways.
        ports = [range(int(_.begin), int(_.end) + 1) for _ in ranges]
        return list(itertools.chain(*ports))

    def shutdown(self, driver):
        self.cleanup()


# Handles signals, passed as negative numbers, and ensures worker process is
# cleaned up if it exists.
#
# This function shows up in many places but because it's final statement is a
# call to os._exit() we can be sure it is only ever called one time.
def exit(returncode):
    try:
        cleanup_container()
    except Exception, e:
        log_exc()
    finally:
        os._exit(((-returncode) + 128) if returncode < 0 else returncode)


def log_exc():
    for line in traceback.format_exc().splitlines():
        log.error(line)


def json_pp(thing):
    s = json.dumps(thing, indent=2, separators=(',', ': '), sort_keys=True)
    data_lines = s.splitlines()[1:-1]
    return "{ " + '\n'.join([data_lines[0][2:]] + data_lines[1:]) + " }"


def ensure_image(f):
    def f_(image, *args, **kwargs):
        pull_once(image)
        return f(image, *args, **kwargs)

    return f_


def try_cid(f):
    def f_(*args, **kwargs):
        if cid is None:
            read_cid()
        return f(*args, **kwargs)

    return f_


cleaning_up_already = False # Hackish lock.


@try_cid
def cleanup_container():
    global cid
    global cidfile
    global cleaning_up_already
    if cid is not None and not cleaning_up_already:
        cleaning_up_already = True
        log.info('Cleaning up container %s' % cid)
        subprocess.check_call(['docker', 'stop', '-t=2', cid])
        subprocess.check_call(['rm', '-f', cidfile])
        cid = None
        cidfile = None


@ensure_image
def run_with_settings(image, args, params):
    global cidfile
    cidfile = '/tmp/docker_cid.' + os.urandom(8).encode('hex')
    cmd = ['run', '-cidfile', cidfile]
    argv = ['docker'] + cmd + [params] + [image] + [args]
    new_cmd = ' '.join(str(arg) for arg in argv)
    new_args = shlex.split(new_cmd)
    log.info("New Args: %s" % new_args)
    log.info('Docker Command:  ' + ' '.join(str(arg) for arg in argv))
    return new_args


@ensure_image
def inner_ports(image):
    text = subprocess.check_output(['docker', 'inspect', image])
    parsed = json.loads(text)[0]
    config = None
    if 'Config' in parsed:
        config = parsed['Config']
    if 'config' in parsed and config is None:
        config = parsed['config']
    if config:
        exposed = config.get('ExposedPorts', {})
        if exposed and isinstance(exposed, dict):
            return [int(k.split('/')[0]) for k in exposed.keys()]
        specs = config.get('PortSpecs', [])
        if specs and isinstance(specs, list):
            return [int(v.split(':')[-1]) for v in specs]
    return []  # If all else fails...


def pull(image):
    subprocess.check_call(['docker', 'pull', image])
    refresh_docker_image_info(image)


def pull_once(image):
    if not image_info(image):
        pull(image)


def image_info(image):
    if image in images:
        return images[image]
    else:
        return refresh_docker_image_info(image)


images = {}


def refresh_docker_image_info(image):
    delim = re.compile('  +')
    text = subprocess.check_output(['docker', 'images', image])
    records = [delim.split(line) for line in text.splitlines()]
    for record in records:
        if record[0] == image:
            images[image] = image
            return record


if __name__ == '__main__':
    main()