#!/usr/bin/env python import psycopg2 import sys import time import os from psycopg2.extensions import ( ISOLATION_LEVEL_READ_COMMITTED, ISOLATION_LEVEL_SERIALIZABLE, STATUS_BEGIN, STATUS_READY) import signal import subprocess import random ######## CONFIGURATION ########## debug = False serializable = False ################################# conn = psycopg2.connect("") if serializable: conn.set_isolation_level(ISOLATION_LEVEL_SERIALIZABLE) else: conn.set_isolation_level(ISOLATION_LEVEL_READ_COMMITTED) children = [] def main_slave(slave_number, master_pid): print "SLAVE[%02s] Spawned" % slave_number n = 0 # Acquire the first finish lock curs = conn.cursor() if debug: print "SLAVE[%02s] LS(1,0)" % slave_number curs.execute("SELECT pg_advisory_lock_shared(1,0);") # Register our successful start to the controller curs.execute("INSERT INTO control (slave) VALUES (%s);" % slave_number) conn.commit(); print "SLAVE[%02s] Started" % slave_number # We enter this loop holding the finish lock for n, and no other locks. while True: # Aqcuire the finish lock for the next record pre-emptively if debug: print "SLAVE[%02s] LS(1,%s)" % (slave_number, (n+1)) curs.execute("SELECT pg_advisory_lock_shared(1,%s);" % (n+1)); # Then wait on the start lock for our current record. if debug: print "SLAVE[%02s] LS(0,%s)" % (slave_number, n) curs.execute("SELECT pg_advisory_lock_shared(0,%s);" % n) # The controller let us off the leash. Race! try: curs.execute("INSERT INTO x (a, slave) SELECT %s, %s WHERE NOT EXISTS (SELECT 1 FROM x WHERE a = %s);" % (n,slave_number,n)) conn.commit(); except psycopg2.IntegrityError, e: print "SLAVE[%02s]: Integrity error inserting %s!" % (slave_number, n) print e # SIGUSR1 the master os.kill( int(master_pid), signal.SIGUSR1) # and terminate. exit(2) # Release the start lock now we're done with it. if debug: print "SLAVE[%02s] US(0,%s)" % (slave_number, n) curs.execute("SELECT pg_advisory_unlock_shared(0,%s);" % n); # Done with n. Tell the controller about our success so it can go ahead # and unlock the next record. if debug: print "SLAVE[%02s] US(1,%s)" % (slave_number, n) curs.execute("SELECT pg_advisory_unlock_shared(1,%s);" % n) n += 1 def master_start_slaves(nslaves): for n in range(0, nslaves): child = subprocess.Popen([sys.argv[0], "slave", str(n), str(os.getpid())], shell=False, stdin=None, stdout=None, stderr=None, close_fds=True) children.append(child) def master_wait_for_slaves(curs, nslaves): while True: curs.execute("SELECT count(slave) from control;") ret = curs.fetchone() if ret is not None and ret[0] == nslaves: break; else: conn.rollback() time.sleep(0.01) def master_num_alive_slaves(): alive = 0 for child in children: child.poll() if child.returncode is None: alive += 1 return alive def master_sigusr1_handler(signum, frame): """We get SIGUSR1 from a slave when the test fails. We should terminate our children and then ourselves.""" print "MASTER: Test failed, terminating children" master_killall_slaves() print "MASTER: Test failed, terminating" sys.exit(2) def master_killall_slaves(): for child in children: child.poll() if child.returncode is None: try: os.kill(child.pid, signal.SIGTERM) except: pass time.sleep(0.1) if master_num_alive_slaves(): time.sleep(1) for child in children: child.poll() if child.returncode is None: try: os.kill(child.pid, signal.SIGKILL) except: pass def main_master(): """Test controller. This keeps the inserters in lockstep by forcing them to wait on an advisory lock for the id they're about to try to insert.""" nslaves = 2 # We use 0,n advisory locks for start locks. These locks # stop the slaves starting to insert "n" until the lock is # released by the controller. # # 1,n locks are finish locks. They're acquired by the slaves # as share locks before trying to get the finish lock, and only # released once the INSERT is done. The controller must acquire # this lock exclusively before proceeding to unlock the next # start lock. signal.signal(signal.SIGUSR1, master_sigusr1_handler) curs = conn.cursor() curs.execute("TRUNCATE TABLE x;") curs.execute("TRUNCATE TABLE control;") conn.commit() n = 0; # Take the first start lock so the slaves don't race off as soon # as they're started. if debug: print "MASTER: L(0,%s)" % n curs.execute("SELECT pg_advisory_lock(0,%s);" % n) # Start slaves print "MASTER: Spawning slaves..." master_start_slaves(nslaves); print "MASTER: slaves spawned, waiting for them to report in." # Wait until they've all reported in master_wait_for_slaves(curs, nslaves); print "MASTER: Slaves have reported in. Starting test." # and start them going. State on loop entry is start lock # for "n" is held, no other locks held. while True: # Lock the record they're going to try to insert # once we let them loose on this one. This keeps # them from running ahead. if debug: print "MASTER: L(0,%s)" % (n+1) curs.execute("SELECT pg_advisory_lock(0,%s);" % (n+1)) # Unlock the current record, letting them try to # insert. if debug: print "MASTER: U(0,%s): Slaves inserting %s" % (n,n) curs.execute("SELECT pg_advisory_unlock(0,%s);" % n) # Attempt to acquire the finish lock on the record we just # let the slaves insert. We'll only be able to obtain it once # all slaves have released their finish locks. if debug: print "MASTER: L(1,%s)" % n curs.execute("SELECT pg_advisory_lock(1,%s)" % n); # Release it again, since we know they're all done. if debug: print "MASTER: U(1,%s): Slaves finished inserting %s" % (n,n) curs.execute("SELECT pg_advisory_unlock(1,%s)" % n); # and get ready for the next run. n += 1 sys.stdout.write(" %s" % n); sys.stdout.flush() if not master_num_alive_slaves() == nslaves: print "One or more slaves appear to have died. Terminating the test" master_killall_slaves() sys.exit(3) if len(sys.argv) > 1: if sys.argv[1] == "slave": main_slave(sys.argv[2], sys.argv[3]) elif sys.argv[1] == "master": main_master() else: print "Usage: %s master|slave" % sys.argv[0] sys.exit(1); else: main_master()