Stampede a bunch of ddoc and VDU calls on a CouchDB cluster

stampede_dbs_ddocs_vdpus.py

#!/usr/bin/env python

#
# Use a 3 node cluster
#
# Args examples:
# ./script.py -x 20 -n 100 -w 5 -p 10 -q
# ./script.py -x 100 -n 2 -p 20 -w 1 -t 10
# ./script.py -x 101 -n 2 -p 10 -w 5 -t 2

import copy
import sys
import time
import threading
import os
import argparse
import uuid
import traceback
import random

from multiprocessing.dummy import Pool as ThreadPool
from multiprocessing import Pool

import couchdb

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

DB_NAME = 'db'


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


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


def get_design_doc(key):
    _id = "_design/%s" % uuid.uuid4().hex
    return {
        "_id": _id,
        "views": {
            "v1": {
                 "map": 'function(d){emit("%s",1);}' % key,
                 "reduce": '''function(keys, values, rereduce) {
                    if (rereduce) {
                       return sum(values);
                    } else {
                       return 2 * values.length;
                    }
                 }
                 '''
            }
        },
        "autoupdate": False,
        "validate_doc_update":
           '''
              function(doc, olddoc, ctx) {
                  if (doc.bad) {
                      throw({forbidden:doc._id});
                  }
              }
           ''',
    }


def wait(args):
    sleep = 0
    if args.random_wait > 0.0:
        sleep = random.randrange(args.random_wait * 1000) / 1000.0
    elif args.fixed_wait > 0.0:
        sleep = args.fixed_wait
    if sleep > 0:
        time.sleep(sleep)


def execute_fun(args, pid, tid, i, db):
    db.update([{'x':'y', '_id':uuid.uuid4().hex} for docnum in range(args.num_docs)])
    wait(args)
    #if i % 10 == 0:
    log(" *** process:", pid, "tid", tid, "i:", i)


def thread_worker(args):
    tid = args.tid
    pid = os.getpid()
    url = pick_server(args.urls)
    srv = couchdb.Server(url)
    srv.version()
    dbname = "%s_%s_%s" % (args.dbname, pid, tid)
    if dbname not in srv:
        db = srv.create(dbname)
        time.sleep(5)
    else:
        db = srv[dbname]
    key = '%s_%s' % (pid, tid)
    db.update([get_design_doc(key) for docnum in range(args.design_docs)])
    tries = args.tries
    for i in range(tries):
        try:
            execute_fun(args, pid, tid, i, db)
        except Exception as e:
            log(" >>> Worker exception caught", e)
            traceback.print_exc(file=sys.stderr)
            time.sleep(1)
            continue
            #raise
    return tid


def set_worker_id(args, tid):
    args = copy.deepcopy(args)
    args.tid = tid
    return args

def process_worker(args):
    wcount = args.worker_count
    pool = ThreadPool(wcount)
    worker_args = [set_worker_id(args, i) for i in xrange(wcount)]
    res = pool.map(thread_worker, worker_args)


def main(args):
    if args.urls == []:
        args.urls = DB_URLS
    pool = Pool(processes=args.processes)
    pool_args = [args for pnum in xrange(args.processes)]
    pool.map(process_worker, pool_args)


def _args():
    description = "Do a few crud operations as a stampede"
    p = argparse.ArgumentParser(description = description)
    p.add_argument('-u', '--urls', action="append", default=[], help = "Server URL(s)")
    p.add_argument('-d', '--dbname', default=DB_NAME, help = "DB name")
    p.add_argument('-x', '--design_docs', type=int, default=10)
    p.add_argument('-n', '--num_docs', type=int, default=2000)
    p.add_argument('-w', '--worker-count', type=int, default=1)
    p.add_argument('-t', '--tries', type=int, default=1)
    p.add_argument('-r', '--random-wait', type=float, default=0)
    p.add_argument('-f', '--fixed-wait', type=float, default=0)
    p.add_argument('-p', '--processes', type=int, default=1)
    return p.parse_args()

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