mt_test.py

text/x-python

Download

#!/usr/bin/python

import sys, pg, threading, random, time, datetime

num_threads = 10
num_loops = 1000

lock = threading.RLock()
total_nontrans = 0
total_commits = 0
total_rollbacks = 0
total_inserts = 0
total_updates = 0
total_deletes = 0
total_violations = 0
total_deadlocks = 0
total_unknown_errors = 0
total_errors = 0
total_duplicates = 0

def open():
    return pg.DB("pgdevel", "localhost", -1,
                 "-c client_min_messages=WARNING", None, "pgdevel", "")

def find_duplicates(db):
    result = db.query("SELECT max(c)-1 AS dups FROM "+\
                      "(SELECT count(*) AS c FROM foo GROUP BY a) AS foo")
    result = result.dictresult()[0]["dups"]

    if result == None: return 0
    return result

def setup():
    db = open()
    db.query("DROP TABLE IF EXISTS foo")
#    db.query("CREATE TABLE foo(a int UNIQUE)")
#    db.query("CREATE TABLE foo(a int UNIQUE DEFERRABLE INITIALLY DEFERRED)")
    db.query("CREATE TABLE foo(a int)")
    db.query("ALTER TABLE foo ADD CONSTRAINT foo_u EXCLUSION"+\
             " USING btree (a CHECK WITH =) DEFERRABLE INITIALLY DEFERRED")
    db.close()

def do_crud(db):
    global total_nontrans, total_commits, total_rollbacks
    global total_inserts, total_updates, total_deletes
    global total_violations, total_deadlocks, total_unknown_errors
    global total_errors, total_duplicates

    inserts = 0
    updates = 0
    deletes = 0

    do_trans = random.random() > 0.2
    do_commit = random.random() > 0.2
    do_loop = True

    duplicates = find_duplicates(db)
    lock.acquire()
    total_duplicates += duplicates
    if duplicates > 0: print "1> FOUND DUPLICATES"
    lock.release()
    if total_duplicates > 0: sys.exit(1)

    try:
        if do_trans:
            db.query("BEGIN")

        while do_loop:
            if random.random() > 0.5:
                val = int(random.random()*100)
                db.query("INSERT INTO foo VALUES("+str(val)+")")
                inserts += 1
            if random.random() > 0.5:
                val1 = int(random.random()*100)
                val2 = int(random.random()*100)
                db.query("UPDATE foo SET a="+str(val2)+" WHERE a="+str(val1))
                updates += 1
            if random.random() > 0.5:
                val = int(random.random()*100)
                db.query("DELETE FROM foo WHERE a="+str(val))
                deletes += 1
            if random.random() > 0.5:
                do_loop = False

        if do_trans:
            if do_commit:
                db.query("COMMIT")
            else:
                db.query("ROLLBACK")
                inserts = 0
                updates = 0
                deletes = 0

        duplicates = find_duplicates(db)

        lock.acquire()

        if do_trans:
            if do_commit: total_commits += 1
            else: total_rollbacks += 1
        else:
            total_nontrans += 1

        total_inserts += inserts
        total_updates += updates
        total_deletes += deletes

        total_duplicates += duplicates
        if duplicates > 0: print "2> FOUND DUPLICATES"

        lock.release()
        if total_duplicates > 0: sys.exit(1)

    except pg.ProgrammingError, detail:
        if do_trans:
            db.query("ROLLBACK")
        duplicates = find_duplicates(db)

        lock.acquire()

        if str(detail).find("operator exclusion constraint violation detected") != -1:
            total_violations += 1
        elif str(detail).find("deadlock detected") != -1:
            total_deadlocks += 1
        else:
            print detail
            total_unknown_errors += 1
        total_errors += 1

        total_duplicates += duplicates
        if duplicates > 0: print "3> FOUND DUPLICATES"

        lock.release()
        if total_duplicates > 0: sys.exit(1)

def do_cruds():
    db = open()

    for ii in range(num_loops):
        do_crud(db)

    db.close()

class Inserter(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)

    def run(self):
        do_cruds()

if __name__ == "__main__":
    start_time = time.time()

    if len(sys.argv) > 1:
        num_threads = int(sys.argv[1])
    if len(sys.argv) > 2:
        num_loops = int(sys.argv[2])

    print "Testing with %d threads, %d loops..." % (num_threads, num_loops)

    setup()

    inserters = []
    for ii in range(num_threads):
        inserter = Inserter()
        inserter.start()
        inserters.append(inserter)

    for inserter in inserters:
        inserter.join()

    db = open()

    result = db.query("SELECT count(*) FROM foo;")
    count = result.dictresult()[0]["count"]

    duplicates = find_duplicates(db)
    total_duplicates += duplicates
    if duplicates > 0: print "4> FOUND DUPLICATES"

    db.close()

    print ""
    print "Ops outside of a transaction: %d" % total_nontrans
    print "Committed transactions: %d" % total_commits
    print "Rolled back transactions: %d" % total_rollbacks
    print "  Unique key violations: %d" % total_violations
    print "  Deadlocks: %d" % total_deadlocks
    print "  Unexpected errors: %d" % total_unknown_errors
    print "Failed transactions: %d" % total_errors
    print ""
    print "Num inserts: %d" % total_inserts
    print "Num updates: %d" % total_updates
    print "Num deletes: %d" % total_deletes
    print ""
    print "Number of rows: %d" % count
    print "Number of duplicates: %d" % total_duplicates

    elapsed_time = time.time() - start_time
    print "Elapsed time: %s" % datetime.timedelta(seconds=elapsed_time)

    print ""
    if total_duplicates > 0 or total_unknown_errors > 0:
        print "FAILED (something dodgy happened)"
    else:
        print "PASSED (everything looks OK)"