Below is the file 'on_the_fly.py' from this revision. You can also download the file.

#!/usr/bin/env python
# vim:et:ts=4:

from sys import argv, exit, stdin, stdout, stderr
import random
from random import choice, randint
from shutil import copyfile
from os import environ, fork, WEXITSTATUS, WIFSIGNALED, WIFEXITED, WTERMSIG
import os
from popen2 import Popen3
from signal import SIGTERM, SIGKILL
from time import sleep

TEST_BRANCH = 'benchmark-test-branch'
SERVE_DB = 'serve.db'
PULL_DB = 'pull.db'
PID_FILE = 'serve.pid'
SERVE_WAIT=4
KEY="shootout@example.com"

def usage():
    print>>stderr, "Usage: %s <db> <revs per netsync> <mto-head binary> <testing binary> [optional seed]" % argv[0]
    print>>stderr

def pairwise(l):
    front = l[:-1]
    back = l[1:]
    return zip(front,back)

def set_random():
    if len(argv) == 6:
        seed = int(argv[5])
    else:
        seed = randint(10000, 1000000)

    print "Random seed is %d" % seed

    random.seed(seed)

def monotone(db, args):

    """ args is either a list of arguments, or a string to split with whitespace
    """

    if isinstance(args, basestring):
        args = args.split()

    # hack
    if db == SERVE_DB:
        MONOTONE = SERVE_MTO
    elif db == PULL_DB:
        MONOTONE = PULL_MTO
    else:
        assert 8 == 0

    cmd = [MONOTONE, '-d', db, '--dump=mkpristine.dump', '-k', KEY, '--rcfile=test_hooks.lua'] + args
    print ' '.join(cmd)
    p = Popen3(cmd)
    stdout = p.fromchild
    lines = stdout.readlines()
    w = p.wait()
    exit_code = WEXITSTATUS(w)
    if WIFSIGNALED(w):
        exit_signal = WTERMSIG(w)
        if exit_signal != SIGTERM and exit_signal != SIGKILL:
            print>>stderr
            print>>stderr, "monotone process exited with signal %d" % exit_signal
            print>>stderr, "Command was: %s" % ' '.join(cmd)
            exit(5)
    elif WIFEXITED(w) and exit_code != 0:
        print>>stderr
        print>>stderr, "monotone process exited with status %d" % exit_code
        print>>stderr, "Command was: %s" % ' '.join(cmd)
        exit(3)

    return lines

def get_num_revs(db):
    m = monotone(db, ['auto', 'select', ''])
    return len( m )

def remove_all_branch_certs(db):

    monotone(db, ['db', 'execute', "delete from revision_certs where name = 'branch'"])

def get_rev_graph(db):

    """
    returns (graph, roots)
    where graph is a { par: [children...] } map,
    roots is the set of parent-less revs.
    """

    m = monotone(db, "auto graph")

    graph = {}
    roots = []

    lines = [ l.strip().split() for l in m ]
    for t in lines:
        rev = t[0]
        pars = t[1:]
        if not len(pars):
            roots.append(rev)

        for p in pars:
            graph.setdefault(p, []).append(rev)

    return (graph, roots)

def remove_one_rev(db):

    m = monotone(db, ['auto', 'leaves'])
    leaves = map(lambda l: l.strip(), m)

    if len(leaves) < 1:
        print>>stderr, "Ran out of revs to remove"
        exit(2)

    to_kill = choice(leaves)
    monotone(db, ['db', 'kill_rev_locally', to_kill])

def cert_leaf_rev(db, graph, frontier):

    rev = frontier.pop(randint(0,len(frontier)-1))
    assert len(frontier)
    monotone(db, "cert %s branch %s" % (rev, TEST_BRANCH))
    if rev in graph:
        frontier += graph[rev]
        del graph[rev]

def cert_leaf_revs(db, graph, frontier, revs_per_ns):

    for i in range(revs_per_ns):
        cert_leaf_rev(db, graph, frontier)

def serve(db, branch):

    # hacky, should work. we need to sleep a bit for the client anyway.
    if fork():
        # parent
        sleep(SERVE_WAIT)
        pid = int(file(PID_FILE).read())
        return pid
    else:
        # child
        try:
            os.remove(PID_FILE)
        except OSError:
            pass

        try:
            monotone(db, 'serve --pid-file=%s %s' % (PID_FILE, TEST_BRANCH))
        except SystemExit, e:
            # try and cleanup
            try:
                pid = int(file(PID_FILE).read())
                os.kill(pid, SIGKILL)
                os.remove(PID_FILE)
            except OSError:
                pass
        exit(9)

def pull(db, branch):

    monotone(db, 'pull localhost %s' % TEST_BRANCH)

def serve_pull(serve_db, pull_db):

    pid = serve(serve_db, TEST_BRANCH)

    pull(pull_db, TEST_BRANCH)

    os.kill(pid, SIGKILL)

def fresh_db(db):
    monotone(db, "db init")

def main():
    global SERVE_MTO, PULL_MTO
    if len(argv) != 5 and len(argv) != 6:
        usage()
        exit(1)

    db, revs_per_ns, SERVE_MTO, PULL_MTO = argv[1:5]

    revs_per_ns = int(revs_per_ns)

    set_random()

    copyfile(db, SERVE_DB)

    fresh_db(PULL_DB)

    remove_all_branch_certs(SERVE_DB)

    num_to_go = get_num_revs(SERVE_DB)

    graph, roots = get_rev_graph(SERVE_DB)

    while num_to_go > 0:

        print '\r%d' % num_to_go

        cert_leaf_revs(SERVE_DB, graph, roots, revs_per_ns)

        serve_pull(SERVE_DB, PULL_DB)

        num_to_go -= revs_per_ns

if __name__ == '__main__':
    main()