stampede_dbcreate.py

#!/usr/bin/env python

import sys, couchdb, time, os, threading, uuid, argparse, random, traceback
from multiprocessing.dummy import Pool as ThreadPool

DB_URLS = [
    'http://adm:pass@localhost:15984',
    'http://adm:pass@localhost:25984',
    'http://adm:pass@localhost:35984'
]

DB_NAME = 'zig'

def log(*args):
    msg = " ".join(["%s" % a for a in args])
    sys.stdout.write(msg + '\n')
    sys.stdout.flush()

def pick_server(URLS):
    if isinstance(URLS, list):
        return random.choice(URLS)
    return URLS

def worker(args):
    wid = uuid.uuid4().hex[:4]
    url = pick_server(args.urls)
    srv = couchdb.Server(url)
    event = args.event
    dbname = args.dbname
    tries = args.tries
    event.wait()
    for i in xrange(tries):
        try:
            srv.create(dbname)
        except Exception as e:
            log(wid, "error creating", e)
        try:
            srv.delete(dbname)
        except Exception as e:
            log(wid, "error deleting", e)
    return wid

def main(args):
    wcount = args.worker_count
    pool = ThreadPool(wcount)
    event = threading.Event()
    try:
        event.clear()
        args.event = event
        worker_args = [args for wid in xrange(wcount)]
        asyncres = pool.map_async(worker, worker_args)
        time.sleep(1)
        raw_input("Press any key when ready:")
        event.set()
        log(" >>> Send event <<< ")
        log(" >>> Got", len(asyncres.get(300)), "results <<<")
    except (KeyboardInterrupt, SystemExit):
        log(" >>> Closing pool")
        pool.join()
        pool.close()

def _args():
    description = "Do a few crud operations as a stampede"
    p = argparse.ArgumentParser(description = description)
    p.add_argument('-u', '--urls', action="append", default=DB_URLS, help = "Server URL(s)")
    p.add_argument('-d', '--dbname', default=DB_NAME, help = "DB name")
    p.add_argument('-w', '--worker-count', type=int, default=1)
    p.add_argument('-t', '--tries', type=int, default=1)
    return p.parse_args()

if __name__=='__main__':
    args = _args()
    main(_args())