lots_of_session_queries.py

import os
import time
import argparse
import multiprocessing

import psutil
import pymysql

parser = argparse.ArgumentParser()
parser.add_argument(
    "procedure",
    type=str,
    help="name of stored procedure to call"
)
parser.add_argument(
    "--num-processes",
    type=int,
    default=10,
    help="number of concurrent processes to run"
)
parser.add_argument(
    "--num-iterations",
    type=int,
    default=1000,
    help="number of iterations to run per process"
)

args = parser.parse_args()
print "Trying", args.num_processes, "concurrent runs of", args.num_iterations, "calls to", args.procedure


def lots_of_session_checks(*_):
    conn = pymysql.connect(
        host='127.0.0.1',
        user='root',
        db='fxa',
    )
    # This calls the specified procedure with a randomly-selected tokenId,
    # by having the caller pass in a random string and taking the first
    # tokenId greater than that string.
    q = 'call {}((select tokenId from sessionTokens where tokenId > 0x{} order by tokenId limit 1))'
    try:
        for _ in xrange(args.num_iterations):
            with conn.cursor() as cur:
                cur.execute(q.format(args.procedure, os.urandom(32).encode('hex')))
                cur.fetchall()
    finally:
        conn.close()


for proc in psutil.process_iter():
    if proc.name() == "mysqld":
        mysql_process = proc
        break
else:
    raise RuntimeError("Unable to find 'mysqld' process for monitoring")


pool = multiprocessing.Pool(args.num_processes)
try:
    mysql_cpu = 0
    start_time = time.time()
    result = pool.map_async(lots_of_session_checks, range(args.num_processes))
    while not result.ready():
        mysql_cpu = max(mysql_cpu, mysql_process.cpu_percent(0.1))
    print "Time elapsed:  {:.2f}s".format(time.time() - start_time)
    print "MySQL CPU use: {}%".format(mysql_cpu)
finally:
    pool.terminate()
    pool.join()