memcache_load.py

import multiprocessing
import sys


from multiprocessing import JoinableQueue
from multiprocessing import Pipe
from multiprocessing import Queue
from multiprocessing import Process
from pylibmc import Client
from time import sleep
from time import time


def run(conn, success, fail, running_q, func):
    running_q.put(True)
    
    mc = Client(['127.0.0.1:11211'], binary=True)
    mc.behaviors = {"tcp_nodelay": True}
    
    s = 0
    f = 0
    ops = 1
    key = 'test_key'
    val = 'test_val'
    
    multi_size = 256
    multi_keys = [key + str(x) for x in xrange(0,multi_size)]
    multi_map = dict([(k, k) for k in multi_keys])
    
    if func == 'get':
        mc.set(key, val)
    
    if func == 'get_multi':
        mc.set_multi(multi_map)
    
    if func == 'get_multi' or func == 'set_multi':
        ops = multi_size
    
    if func == 'incr':
        mc.set(key, 0)
    
    if func == 'decr':
        mc.set(key, 10 ** 9)
    
    funcs = {
        'get': lambda: mc.get(key) == val,
        'get_multi': lambda: mc.get_multi(multi_keys) == multi_map,
        'set_multi': lambda: mc.set_multi(multi_map) == [],
        'set': lambda: mc.set(key, val),
        'incr': lambda: mc.incr(key) != False,
        'decr': lambda: mc.decr(key) != False, 
    }
    
    while not conn.poll():
        if funcs[func]():
            s += ops
        else:
            f += ops
    
    success.put(s)
    fail.put(f)
    print '%s is dying %d %d' % (multiprocessing.current_process(), s, f)
    running_q.get()
    running_q.task_done()


if __name__ == '__main__':
    func = sys.argv[1]
    concurrency = int(sys.argv[2])
    timeout = int(sys.argv[3])
    
    success = Queue()
    fail = Queue()
    running_q = JoinableQueue()
    
    pipes = []
    processes = []
    for x in xrange(0,concurrency):
        pipe_out, pipe_in = Pipe()
        p = Process(target=run, args=(pipe_out, success, fail, running_q, func))
        p.daemon = True
        pipes.append(pipe_in)
        processes.append(p)
    
    st = time()
    for p in processes:
        p.start()
    
    sleep(timeout - (time() - st))
    
    print 'sending kill msgs'
    for p in pipes:
        p.send(True)
    
    running_q.join()
    dt = time() - st
    
    succ = 0
    while not success.empty():
        succ += success.get()
    
    fails = 0
    while not fail.empty():
        fails += fail.get()
    
    print '%d succ, %d succ/sec, %d fails' % (succ, float(succ) / dt, fails)