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, WEXITSTATUS
from popen2 import Popen3
from signal import SIGTERM
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()
    exitcode = WEXITSTATUS(w)
    if exitcode != 0:
        print>>stderr
        print>>stderr, "monotone process exited with status %d" % exitcode
        print>>stderr, "cmd: %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 rev in graph:
            print rev
            print pars
        assert rev not in graph
        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(len(frontier)))
    assert len(frontier)
    monotone(db, "automate 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(num):
        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
        monotone(db, 'serve --pid-file=%s %s' % (PID_FILE, TEST_BRANCH))

def pull(db, branch):

    monotone(db, 'pull %S' % TEST_BRANCH)

def serve_pull(serve_db, pull_db):

    pid = serve(serve_db, TEST_BRANCH)

    pull(pull_db, TEST_BRANCH)

    os.kill(pid, SIGTERM)

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, frontier, revs_per_ns)

        serve_pull(SERVE_DB, PULL_DB)

        num_to_go -= revs_per_ns

if __name__ == '__main__':
    main()