path: root/dev/run

#!/usr/bin/env python
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.

import atexit
import contextlib as ctx
import glob
import httplib
import optparse as op
import os
import re
import select
import subprocess as sp
import sys
import time
import traceback
import urllib


USAGE = "%prog [options] [command to run...]"
DEV_PATH = os.path.dirname(os.path.abspath(__file__))
COUCHDB = os.path.dirname(DEV_PATH)

N = 3
PROCESSES = []


def init_log_dir():
    logdir = os.path.join(DEV_PATH, "logs")
    if not os.path.exists(logdir):
        os.makedirs(logdir)


def init_beams():
    # Including this for people that forget to run
    # make dev.
    for fname in glob.glob(os.path.join(DEV_PATH, "*.erl")):
        cmd = [
            "erlc",
            "-o", DEV_PATH + os.sep,
            fname
        ]
        sp.check_call(cmd)


def hack_default_ini(opts, node, args, contents):
    # Replace log file
    logfile = os.path.join(DEV_PATH, "logs", "%s.log" % node)
    repl = "file = %s" % logfile
    contents = re.sub("(?m)^file.*$", repl, contents)

    # Replace couchjs command
    couchjs = os.path.join(COUCHDB, "src", "couch", "priv", "couchjs")
    mainjs = os.path.join(COUCHDB, "share", "server", "main.js")
    coffeejs = os.path.join(COUCHDB, "share", "server", "main-coffee.js")

    repl = "javascript = %s %s" % (couchjs, mainjs)
    contents = re.sub("(?m)^javascript.*$", repl, contents)

    repl = "coffeescript = %s %s" % (couchjs, coffeejs)
    contents = re.sub("(?m)^coffeescript.*$", repl, contents)

    return contents


def hack_local_ini(opts, node, args, contents):
    if opts.admin is None:
        return contents
    usr, pwd = opts.admin.split(":", 1)
    return contents + "\n[admins]\n%s = %s" % (usr, pwd)


def write_config(opts, node, args):
    etc_src = os.path.join(COUCHDB, "rel", "overlay", "etc")
    etc_tgt = os.path.join(DEV_PATH, "lib", node, "etc")
    if not os.path.exists(etc_tgt):
        os.makedirs(etc_tgt)

    etc_files = glob.glob(os.path.join(etc_src, "*"))
    for fname in etc_files:
        base = os.path.basename(fname)
        tgt = os.path.join(etc_tgt, base)
        with open(fname) as handle:
            contents = handle.read()
        for key in args:
            contents = re.sub("{{%s}}" % key, args[key], contents)
        if base == "default.ini":
            contents = hack_default_ini(opts, node, args, contents)
        elif base == "local.ini":
            contents = hack_local_ini(opts, node, args, contents)
        with open(tgt, "w") as handle:
            handle.write(contents)


def write_configs(opts):
    datadir = os.path.join(DEV_PATH, "data")
    if not os.path.exists(datadir):
        os.makedirs(datadir)
    for i in range(1,4):
        node = "node%d" % i
        args = {
            "prefix": COUCHDB,
            "package_author_name": "The Apache Software Foundation",
            "data_dir": os.path.join(DEV_PATH, "lib", node, "data"),
            "view_index_dir": os.path.join(DEV_PATH, "lib", node, "data"),
            "node_name": "-name %s@127.0.0.1" % node,
            "cluster_port": str((10000 * i) + 5984),
            "backend_port" : str((10000 * i) + 5986)
        }
        if not os.path.exists(args["data_dir"]):
            os.makedirs(args["data_dir"])
        write_config(opts, node, args)


def all_nodes_alive(n):
    for i in range(1, n+1):
        url = "http://127.0.0.1:{0}/".format(local_port(i))
        while True:
            try:
                with ctx.closing(urllib.urlopen(url)) as resp:
                    pass
            except IOError:
                time.sleep(0.25)
                continue
            break
    return True


def local_port(n):
    return 10000 * n + 5986


def node_port(n):
    return 10000 * n + 5984


def boot_node(node):
    apps = os.path.join(COUCHDB, "src")
    env = os.environ.copy()
    env["ERL_LIBS"] = os.pathsep.join([apps])
    cmd = [
        "erl",
        "-args_file", os.path.join(DEV_PATH, "lib", node, "etc", "vm.args"),
        "-config", os.path.join(COUCHDB, "rel", "files", "sys"),
        "-couch_ini",
            os.path.join(DEV_PATH, "lib", node, "etc", "default.ini"),
            os.path.join(DEV_PATH, "lib", node, "etc", "local.ini"),
        "-reltool_config", os.path.join(COUCHDB, "rel", "reltool.config"),
        "-parent_pid", str(os.getpid()),
        "-pa", DEV_PATH,
        "-pa", os.path.join(COUCHDB, "src", "*"),
        "-s", "boot_node"
    ]
    logfname = os.path.join(DEV_PATH, "logs", "%s.log" % node)
    log = open(logfname, "w")
    return sp.Popen(
            cmd,
            stdin=sp.PIPE,
            stdout=log,
            stderr=sp.STDOUT,
            env=env)


def connect_nodes(host, port):
    global N
    for i in range(1, N+1):
        body = "{}"
        conn = httplib.HTTPConnection(host, port)
        conn.request("PUT", "/nodes/node%d@127.0.0.1" % i, body)
        resp = conn.getresponse()
        if resp.status not in (200, 201, 202, 409):
            print resp.reason
            exit(1)


def kill_processes():
    global PROCESSES
    for p in PROCESSES:
        if p.returncode is None:
            p.kill()


def boot_nodes():
    global N, PROCESSES
    for i in range(1, N+1):
        p = boot_node("node%d" % i)
        PROCESSES.append(p)

    for i in range(30):
        if all_nodes_alive(N):
            break
        time.sleep(1)


def reboot_nodes():
    kill_processes()
    boot_nodes()


def run_command(cmd):
    p = sp.Popen(cmd, shell=True, stdout=sp.PIPE, stderr=sys.stderr)
    while True:
        line = p.stdout.readline()
        if not line:
            break
        try:
            eval(line)
        except:
            traceback.print_exc()
            exit(1)
    p.wait()
    exit(p.returncode)


def wait_for_procs():
    global PROCESSES
    while True:
        for p in PROCESSES:
            if p.returncode is not None:
                exit(1)
        time.sleep(2)


def options():
    return [
        op.make_option("-a", "--admin", metavar="USER:PASS", default=None,
            help="Add an admin account to the development cluster")
    ]


def main():
    parser = op.OptionParser(usage=USAGE, option_list=options())
    opts, args = parser.parse_args()

    init_log_dir()
    init_beams()
    write_configs(opts)

    atexit.register(kill_processes)

    boot_nodes()
    connect_nodes("127.0.0.1", 15986)

    if len(args):
        run_command(" ".join(args))
    else:
        wait_for_procs()


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass