durable.pike

//Test PostgreSQL transaction durability

int pool=100,nthrd=80,testlength=60; //TODO: Support command-line parameters for these
mapping(int:int) possible=([]);
mapping(int:int) confirmed=([]);

object getdb()
{
    object db=Sql.Sql("pgsql://USERNAME:PASSWORD@HOST/DATABASE",(["reconnect":0]));
    db->master_sql->settimeout(10);
    return db;
}

void transact()
{
    object db=getdb();
    mixed ex=catch
    {
        while (testlength)
        {
            db->big_query("begin");
            array(int) cur=({random(pool)+1,0,time()});
            db->big_query("update public.durability set value=value+1 where id=:id",(["id":cur[0]]));
            db->big_query("update public.durability set value=value+1 where id=0");
            possible[cur[0]]++; possible[0]++;
            db->big_query("commit");
            confirmed[cur[0]]++; confirmed[0]++;
        }
    };
    //Currently, I don't actually much care about exceptions. Should probably do something with them.
    catch {db->master_sql->close(); destruct(db);}; //Ignore errors on closing connection
}

int main()
{
    object db=getdb(); //Initialize in autocommit mode
    db->big_query("create table if not exists public.durability (id integer primary key,value integer not null default 0)");
    db->big_query("truncate public.durability");
    db->big_query("insert into public.durability (id) select generate_series(0,:pool)",(["pool":pool]));
    db->master_sql->close(); destruct(db);
    signal(2,lambda() {error("Ctrl-C!");});
    int start=time();
    write("Durability test: thrds %d, pool %d, time %d, start %s",nthrd,pool,testlength,ctime(start));
    array threads=({ }); while (nthrd--) threads+=({Thread.Thread(transact)});
    catch {while (has_value(threads->status(),0))
    {
        float tm=time(start);
        werror("%d/%fs=%dtps...\r",confirmed[0],tm,(int)(confirmed[0]/tm));
        if (tm>testlength) break;
        sleep(1);
    }};
    signal(2); //Reset Ctrl-C handler
    testlength=0; threads->kill();
    write("Completed %d transactions in %fs, %f tps.\n",confirmed[0],time(start),confirmed[0]/time(start));
    write("Reconnecting to db..."); while (catch {db=getdb();} || !db) {sleep(1); write(".");}
    write("\nReconnected, checking.\n");
    int sum=0;
    foreach (db->typed_query("select * from public.durability order by id"),mapping row)
    {
        if (row->id) sum+=row->value; else sum-=row->value;
        if (row->value<confirmed[row->id] || row->value>possible[row->id]) write("Fail: Row %d [%d] s/be %d-%d\n",row->id,row->value,confirmed[row->id],possible[row->id]);
    }
    if (sum) write("*** ERROR: Inconsistent data [sum is wrong by %d] ***\n",sum);
    write("Verification complete.\n");
    destruct(db);
}