Source

Dad / dad / web / plugins / manager.py

Full commit
import cherrypy
import os
import sys
import shlex
import re

import yaml

from cherrypy.process.plugins import SimplePlugin
from contextlib import contextmanager
from subprocess import Popen, PIPE, STDOUT

from pprint import pprint, pformat

import pdb


class Sandbox(object):
    def __init__(self, name, instance):
        self.name = name
        self.instance = instance

    def path(self):
        base_dir = cherrypy.config['dad.sandboxes']
        sbox = os.path.join(base_dir, self.name)
        if not os.path.isdir(sbox):
            os.makedirs(sbox)
        return sbox
    
    def operation(self, cmd, out=None, err=None):
        '''
        Perform some operation... How's that for vague!

        An operation on a sandbox is when you agree to an API on how
        to do things to your production system. For example, if you
        install packages via RPM, then you might have an install
        operation that calls yum to install packages in the sandbox
        directory using a command line flag.

        Therefore, the only thing that we provide as a sandbox is the
        root directory. You have to do the rest. We'll call the
        command and give you the output, but past that, you're on your
        own.
        '''

        out = out or sys.stdout
        err = err or STDOUT
        proc = Popen(cmd.split(), cwd=self.cwd(), stdout=out, stderr=err)
        return proc
        

class AppManagerPlugin(SimplePlugin):
    '''
    Manage the configuration details for each app.
    '''
    def __init__(self, bus, conf):
        SimplePlugin.__init__(self, bus)
        self.conf = conf
        self.bus.subscribe('dad.service.add', self.add_service)
        self.bus.subscribe('dad.get_service_conf', self.get_conf)
        self.bus.subscribe('dad.reload_conf', self.reload_conf)

    def get_conf(self):
        return self.conf

    def start(self):
        if self.conf:
            self.reload_conf(self.conf)

    def add_service(self, new_service):
        for name, conf in new_service.items():
            self.conf[name] = conf
            return name
        return False

    def reload_conf(self, new_conf):
        '''
        Read our dad/lib/config.Config object for applications.
        '''
        self.conf = new_conf
        # pdb.set_trace()

        self.bus.log('Reloading the config')
        for name, app in self.conf.iteritems():
            self.bus.log('Adding and starting "%s" with %i instances' % (name, app.get('instances', 1)))
            self.bus.publish('dad.action.start', name, app)
        

class ProcessManagerPlugin(SimplePlugin):
    '''
    Manage the sandboxes and processes.
    '''
    start_port = 6000
    env_var_re = re.compile('^\$[\w\d\_]+')
    
    def __init__(self, bus):
        SimplePlugin.__init__(self, bus)
        self.conf = None
        self.processes = {}
        self.bus.subscribe('dad.action.start', self.start_app)
        self.bus.subscribe('dad.action.stop', self.stop_app)
        self.bus.subscribe('dad.action.restart', self.restart_app)        
        self.bus.subscribe('dad.apps_status', self.status)

    def port(self):
        if not self.processes:
            return self.start_port
        return self.start_port + ((len(self.processes) - 1) * 100)

    def update_path(self, env, sbox):
        path = [sbox.path() + '/bin']
        if env.get('PATH'):
            path.append(env['PATH'])
        path.append(os.environ['PATH'])
        env['PATH'] = ':'.join(path)

    def replace_env_vars(self, token, env):
        def repl(m):
            key = m.group(0)[1:]
            return env.get(key, os.environ.get(key, ''))
        return self.env_var_re.sub(repl, token)

    def stop(self):
        for name, procs in self.processes.items():
            for port, proc in procs.items():
                self.bus.log('Stopping: %s %s %i' % (name, port, proc.pid))
                proc.terminate()
                proc.wait()
                self.bus.log('Stopped %s %s %i' % (name, port, proc.pid))
                
    def status(self):
        info = {}
        for name, proc in self.processes.iteritems():
            info[name] = {
                'pid': proc.pid,
            }
        return info

    def restart_app(self, name, conf, **kw):
        pass
    
    def stop_app(self, name, conf, **kw):
        results = {}
        for instance, proc in self.processes.get(name, {}).items():
            if proc.poll():
                results[instance] = {'returncode': None}
            else:
                proc.terminate()
                proc.wait()
                results[instance] = {'returncode': proc.poll()}

        # nothing was running
        if not results:
            results = dict((instance, {'returncode': None})
                           for instance in range(conf.get('instances', 1)))
        return results
    
    def start_app(self, name, conf):
        cmd = shlex.split(conf['command'].encode('utf-8'))

        for instance in xrange(conf.get('instances', 1)):

            if conf.get('PORT'):
                port = conf['PORT']
            else:
                # generate a port
                port = self.port() + instance

            # create a sandbox with a place to run
            sbox = Sandbox(name, instance)

            # add our generated port to the environ
            environ = {
                'HOST': os.environ.get('HOST', '127.0.0.1'),
                'PORT': str(port),
                'PWD': sbox.path(),
            }
            # update it with the config
            environ.update(conf.get('env', {}))

            # provide a path
            self.update_path(environ, sbox)

            # check for replacements in the shell command
            repl = []

            for index, token in enumerate(cmd):
                if token.startswith('$'):
                    value = self.replace_env_vars(token, environ)
                    if value:
                        cmd[index] = value

            # update our environ with any replacements from the dad environ
            environ = dict((k, self.replace_env_vars(v, environ))
                           for k, v in environ.items())

            self.bus.log('Starting process: %s %s %s' %
                         (cmd, environ['HOST'], environ['PORT']))
            # start the process
            proc = Popen(cmd,
                         cwd=sbox.path(),
                         stdout=sys.stdout,
                         stderr=STDOUT,
                         env=environ)
            if not self.processes.get(name):
                self.processes[name] = {port: proc}
            else:
                self.processes[name][port] = proc

        return {'started': name, 'num_instances': conf.get('instances', 1)}