//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);
}