cluster_run

#!/usr/bin/env python
import os, subprocess, sys, signal, atexit, socket, getopt, select, shlex
import platform

base_direct_port = 12000
base_api_port = 9000
base_couch_port = 9500
base_projector_port = 10000

LOGLEVELS = ["debug", "info", "warn", "error", "critical"]

def read_configuration():
    with open("build/cluster_run.configuration") as f:
        def fn(line):
            k, v = line.strip().split('=')
            return k, shlex.split(v)[0]

        return dict(fn(line) for line in f.readlines())

config = read_configuration()
PREFIX = config['prefix']

def setup_path():
    def ebin_search(path):
        dirs = os.walk(path)
        ebins = []

        for d, _, _ in dirs:
            if os.path.basename(d) == "ebin":
                ebins.append(d)

        return ebins

    path = ebin_search(".")
    couchpath = ebin_search("{0}/lib/couchdb/erlang/lib".format(PREFIX))
    couch_plugins = ebin_search("{0}/lib/couchdb/plugins".format(PREFIX))

    if len(couchpath) == 0:
       sys.exit("Couch libs wasn't found.\nCan't handle it")

    return couchpath + path + couch_plugins

def mk_node_couch_config(i):
    try:
        os.mkdir("couch")
    except os.error:
        pass

    with open("couch/n_{0}_conf.ini".format(i), "w") as f:
        f.write("[httpd]\n")
        f.write("port={0}\n".format(base_couch_port + i))
        f.write("[couchdb]\n")
        f.write("database_dir={0}/data/n_{1}/data\n".format(os.getcwd(), i))
        f.write("view_index_dir={0}/data/n_{1}/data\n".format(os.getcwd(), i))
        f.write("max_dbs_open=10000\n")
        f.write("[upr]\n")
        f.write("port={0}\n".format(base_direct_port + i * 2))
        f.write("[dcp]\n")
        f.write("port={0}\n".format(base_direct_port + i * 2))


def couch_configs(i):
    mk_node_couch_config(i)
    return ["{0}/etc/couchdb/default.ini".format(PREFIX),
            "{0}/etc/couchdb/default.d/capi.ini".format(PREFIX),
            "{0}/etc/couchdb/default.d/geocouch.ini".format(PREFIX),
            "couch/n_{0}_conf.ini".format(i)]

def os_specific(args, params):
    """Add os-specific junk to the cluster startup."""
    if platform.system() == 'Darwin':
        import resource
        ## OS X has a pretty tiny default fd limit.  Let's increase it (if it hasn't already been).
        (soft, hard) = resource.getrlimit(resource.RLIMIT_NOFILE)
        if soft < 2048:
            resource.setrlimit(resource.RLIMIT_NOFILE, (2048, 2048))
        ## Also, we build stuff in kind of dumb ways, so fix the path.
        dypath = ("%(pwd)s/../install/lib/memcached:"
                  "%(pwd)s/../install/lib") % {'pwd': os.getcwd()}
        params['env'] = {"DYLD_LIBRARY_PATH": dypath,
                         "ERL_MAX_PORTS": "2048"}
        params['env'].update(os.environ)

ebin_path = None
cluster_extra_args = None
cluster_args_prefix = None

def prepare_start_cluster(extra_args, args_prefix):
    global cluster_args_prefix
    global cluster_extra_args
    global ebin_path

    ebin_path = setup_path()
    cluster_extra_args = extra_args
    cluster_args_prefix = args_prefix

def quote_string_for_erl(s):
    return '"' + s.replace("\\", "\\\\").replace("\"", "\\\"") + '"'

def erlang_args_for_node(i):
    logdir = "logs/n_{0}".format(i)

    args = cluster_args_prefix + ["erl", "+MMmcs" "30",
                                  "+A", "16", "+sbtu",
                                  "+P", "327680", "-pa"] + ebin_path
    args += [
        "-setcookie", "nocookie",
        "-kernel", "inet_dist_listen_min", "21100",
        "inet_dist_listen_max", "21199",
        "error_logger", "false",
        "-sasl", "sasl_error_logger", "false",
        "-couch_ini"] + couch_configs(i)

    datadir = os.path.abspath('data/n_{0}'.format(i))
    tempdir = os.path.abspath('tmp/')
    nodefile = os.path.join(datadir, "nodefile")

    args += [
        "-name", "babysitter_of_n_{0}@127.0.0.1".format(i),
        "-hidden",
        "-ns_server", "config_path", '"etc/static_config.in"',
        "-ns_server", "enable_mlockall", "false",
        "error_logger_mf_dir", quote_string_for_erl(logdir),
        "dont_suppress_stderr_logger", "true",
        "path_config_etcdir", '"priv"',
        "path_config_bindir", quote_string_for_erl(PREFIX+"/bin"),
        "path_config_libdir", quote_string_for_erl(PREFIX+"/lib"),
        "path_config_datadir", quote_string_for_erl(datadir),
        "path_config_tmpdir", quote_string_for_erl(tempdir),
        "rest_port", str(base_api_port + i),
        "query_port", str(base_couch_port - 1 - i),
        "projector_port", str(base_projector_port - 1 - i),
        "ssl_rest_port", str(10000 + base_api_port + i),
        "capi_port", str(base_couch_port + i),
        "ssl_capi_port", str(10000 + base_couch_port + i),
        "memcached_port", str(base_direct_port + i * 2),
        "moxi_port", str(base_direct_port + i * 2 + 1),
        "memcached_dedicated_port", str(base_direct_port - 1 - i * 4),
        "ssl_proxy_downstream_port", str(base_direct_port - 2 - i * 4),
        "ssl_proxy_upstream_port", str(base_direct_port - 3 - i * 4),
        "memcached_ssl_port", str(base_direct_port - 4 - i * 4),
        "nodefile", quote_string_for_erl(nodefile),
        "short_name", quote_string_for_erl('n_{0}'.format(i))
        ] + cluster_extra_args

    return args

def start_cluster(num_nodes, start_index, host, extra_args, args_prefix):
    prepare_start_cluster(extra_args, args_prefix)

    def start_node(i):
        logdir = "logs/n_{0}".format(i)
        try:
            os.makedirs(logdir)
        except:
            pass

        args = erlang_args_for_node(i)

        params = {}

        os_specific(args, params)

        if not params.has_key('env'):
            params['env'] = {}
            params['env'].update(os.environ)
        path = params['env']['PATH']
        path = (PREFIX+"/bin") + os.pathsep + path
        if not params['env'].has_key('ERL_FULLSWEEP_AFTER'):
            params['env']['ERL_FULLSWEEP_AFTER'] = '512'
        params['env']['PATH'] = path

        crash_dump_base = 'erl_crash.dump.n_%d' % i
        params['env']['ERL_CRASH_DUMP_BASE'] = crash_dump_base
        params['env']['ERL_CRASH_DUMP'] = crash_dump_base + '.babysitter'

        params['env']['COUCHBASE_SMALLER_PKEYS'] = '1'

        params['close_fds'] = True
        if platform.system() == "Windows":
            params['close_fds'] = False

        w = None

        if "-noinput" in args:
            (r,w) = os.pipe()

            params['stdin'] = r

            if 'setpgrp' in os.__dict__ and params.get('close_fds'):
                # this puts child out of our process group. So that
                # Ctrl-C doesn't deliver SIGINT to it, leaving us
                # ability to it shutdown carefully or otherwise
                params['preexec_fn'] = os.setpgrp

        pr = subprocess.Popen(args, **params)
        if w != None:
            os.close(r)
        pr.write_side = w
        return pr

    return [start_node(i + start_index) for i in xrange(num_nodes)]

def usage():
    sys.exit("Usage: {0} [--nodes=N] [--dont-rename] [--dont-start] "
             "[--interactive] [--static-cookie] [--start-index=N] "
             "[--static-cookie] [--host=H] [--loglevel=L] "
             "[ns_server args]".format(sys.argv[0]))

def find_primary_addr():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(("8.8.8.8", 56))
        addr, port = s.getsockname()
        return addr
    except socket.error:
        return None
    finally:
        s.close()

def main():
    try:
        optlist, args = getopt.gnu_getopt(sys.argv[1:], "hn:i",
                                          ["help", "start-index=", "nodes=",
                                           "dont-rename", "interactive",
                                           "static-cookie", "dont-start",
                                           "host=", "loglevel=",
                                           "prepend-extras"])
    except getopt.GetoptError, err:
        # print help information and exit:
        print str(err) # will print something like "option -a not recognized"
        usage()
        sys.exit(2)

    dont_rename = False
    dont_start = False
    static_cookie = False
    interactive_shell = False
    start_index = 0
    num_nodes = 1
    prepend_extras = False
    host = "127.0.0.1"
    loglevel = 'debug'

    for o, a in optlist:
        if o in ("--nodes", "-n"):
            num_nodes = int(a)
        elif o in ("--interactive", "-i"):
            interactive_shell = True
        elif o == '--dont-start':
            dont_start = True
        elif o == '--host':
            host = a
        elif o == '--start-index':
            start_index = int(a)
        elif o == '--dont-rename':
            dont_rename = True
        elif o in ("--help", "-h"):
            usage()
            exit(0)
        elif o in("--static-cookie"):
            static_cookie = True
        elif o == '--loglevel':
            loglevel = a
        elif o == "--prepend-extras":
            prepend_extras = True
        else:
            assert False, "unhandled options"

    nodes = []
    terminal_attrs = None

    def kill_nodes(*args):
        for n in nodes:
            if n.write_side != None:
                print("Closing %d\n" % n.write_side)
                # os.write(n.write_side, "shutdown\n") # this line does graceful shutdown versus quick
                os.close(n.write_side)
            else:
                try:
                    n.kill()
                except OSError:
                    pass

        for n in nodes:
            n.wait()

        if terminal_attrs != None:
            termios.tcsetattr(sys.stdin, termios.TCSANOW, terminal_attrs)

    atexit.register(kill_nodes)

    try:
        import termios
        terminal_attrs = termios.tcgetattr(sys.stdin)
    except:
        pass

    extra_args = []
    if not dont_rename:
        primary_addr = find_primary_addr()
        if primary_addr == None:
            print("was unable to detect 'internet' address of this machine."
                  + " node rename will be disabled")
        else:
            extra_args += ["rename_ip", '"' + primary_addr + '"']

    extra_args += args[1:]
    if prepend_extras:
        prepend_args = args[0:]
    else:
        prepend_args = []
        extra_args += args[0:]

    if static_cookie:
        extra_args += ["-ns_server", "dont_reset_cookie", "true"]

    if not interactive_shell:
        extra_args += ["-noinput"]
        if not dont_start:
            extra_args += ["-run", "child_erlang", "child_start", "ns_babysitter_bootstrap"]
    else:
        if not dont_start:
            extra_args += ["-run", "ns_babysitter_bootstrap"]

    if loglevel not in LOGLEVELS:
        print "Valid log levels are the following: %s" % ', '.join(LOGLEVELS)
        sys.exit(1)
    extra_args += ["-ns_server", "loglevel_stderr", loglevel]

    nodes = start_cluster(num_nodes, start_index, host, extra_args, prepend_args)

    for node in nodes:
        node.wait()


if __name__ == '__main__':
    main()