import math
import psycopg2
import random
import select
import time

# To use this test you need to create 150 test databases like so:
# for i in $(seq -w 0 149) ; do createdb --cluster 11/main test_${i} ; done
# For CPU usage it assumes postgresql is running under systemd
PORT=5432
NUM_DATABASES = 150
NUM_PER_DATABASE = 2
NOTIFY_PER_SECOND = 100

def get_usage():
    with open("/sys/fs/cgroup/cpu/system.slice/system-postgresql.slice/cpuacct.usage") as f:
        return int(f.read())

def perc(d, p):
        pos = (len(d)-1)*p

        pos_low = d[int(math.floor(pos))] * (1-math.modf(pos)[0])
        pos_high = d[int(math.ceil(pos))] * math.modf(pos)[0]
        return pos_low+pos_high

CONNS = {}
REV_CONNS = {}
poller = select.poll()

for i in range(NUM_DATABASES * NUM_PER_DATABASE):
    CONNS[i] = psycopg2.connect("dbname=test_%03d port=%s" % (int(i/NUM_PER_DATABASE), PORT))
    CONNS[i].set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
    REV_CONNS[CONNS[i].fileno()] = i
    cur = CONNS[i].cursor()
    cur.execute("LISTEN changes")

    poller.register(CONNS[i], select.POLLIN)

counter = 0
next_time = time.time()
last_update = int(next_time / 5) * 5

stat_sent = 0
stat_recv = 0
stat_delays = []
stat_usage = get_usage()

while True:
    t = time.time()
    while t > next_time:
        r = random.randrange(NUM_DATABASES * NUM_PER_DATABASE)
        CONNS[r].cursor().execute("NOTIFY changes, '%s'" % counter)

#        print "%.2f NOTIFY on %d" % (t, r)
        counter += 1
        stat_sent += 1

        next_time += 1.0/NOTIFY_PER_SECOND
        t = time.time()

        # Process self-notify
        while CONNS[r].notifies:
            stat_recv += 1
            notify = CONNS[r].notifies.pop(0)
#            print "%.2f Got NOTIFY on %d:" % (t, r), notify.pid, notify.channel, notify.payload
            stat_delays.append(float(counter - int(notify.payload))/NOTIFY_PER_SECOND)

    while t < next_time:
        r = poller.poll(1000*(next_time-t))

        for fd, _event in r:
            c = CONNS[REV_CONNS[fd]]
            c.poll()
            while c.notifies:
                stat_recv += 1
                notify = c.notifies.pop(0)
#                print "%.2f Got NOTIFY on %d:" % (t, REV_CONNS[fd]), notify.pid, notify.channel, notify.payload
                stat_delays.append(float(counter - int(notify.payload))/NOTIFY_PER_SECOND)
        t = time.time()

    if t-last_update > 5:
        stat_delays.sort()
        new_usage = get_usage()
        pct_cpu = (new_usage-stat_usage)/(t-last_update)/1000000000
        print "%d Sent: %d, Recv: %d, Delays: Min: %.2f, Max: %.2f, Avg: %.2f [%.2f/%.2f/%.2f/%.2f/%.2f], %.2f%%" % (last_update,
            stat_sent, stat_recv, min(stat_delays), max(stat_delays), sum(stat_delays)/len(stat_delays),
            perc(stat_delays, 0.05), perc(stat_delays, 0.25), perc(stat_delays, 0.5), perc(stat_delays, 0.75), perc(stat_delays, 0.95),
            pct_cpu*100)
        stat_sent = 0
        stat_recv = 0
        stat_delays = []

        last_update += 5
        stat_usage = new_usage