#include<time.h>

#include<stdio.h>

#include<stdarg.h>

#include<stdlib.h>

#include<inttypes.h>

#include<sys/time.h>

#include<pthread.h>

#include<string>

#include<vector>

#include<pqxx/connection>

#include<pqxx/transaction>

#include<pqxx/nontransaction>

#include<pqxx/pipeline>

template<classT>

{

T* ptr;

unique_ptr(T* p =NULL) : ptr(p) {}

~unique_ptr() {delete ptr; }

T&operator*() {return *ptr; }

T*operator->() {return ptr; }

voidoperator=(T* p) { ptr = p; }

voidoperator=(unique_ptr& other) {

ptr = other.ptr;

other.ptr =NULL;

}

};

typedefvoid* (*thread_proc_t)(void*);

typedefuint32_txid_t;

{

pthread_t t;

size_t proceeded;

size_t aborts;

int id;

voidstart(int tid,thread_proc_t proc) {

id = tid;

proceeded =0;

aborts =0;

pthread_create(&t,NULL, proc,this);

}

voidwait() {

pthread_join(t,NULL);

}

};

{

int nReaders;

int nWriters;

int nIterations;

int nAccounts;

charconst* isolationLevel;

vector<string> connections;

config() {

nReaders =1;

nWriters =10;

nIterations =1000;

nAccounts =1000;

isolationLevel ="read committed";

}

};

config cfg;

bool running;

#defineUSEC1000000

statictime_tgetCurrentTime()

{

structtimeval tv;

gettimeofday(&tv,NULL);

return (time_t)tv.tv_sec*USEC + tv.tv_usec;

}

voidexec(transaction_base& txn,charconst* sql, ...)

{

va_list args;

va_start(args, sql);

char buf[1024];

vsprintf(buf, sql, args);

va_end(args);

txn.exec(buf);

}

xid_texecQuery( transaction_base& txn,charconst* sql, ...)

{

va_list args;

va_start(args, sql);

char buf[1024];

vsprintf(buf, sql, args);

va_end(args);

result r = txn.exec(buf);

return r[0][0].as(xid_t());

}

void*reader(void* arg)

{

thread& t = *(thread*)arg;

vector< unique_ptr<connection> >conns(cfg.connections.size());

for (size_t i =0; i < conns.size(); i++) {

conns[i] =newconnection(cfg.connections[i]);

}

int64_t prevSum =0;

while (running) {

xid_t xid =0;

for (size_t i =0; i < conns.size(); i++) {

worktxn(*conns[i]);

if (i ==0) {

xid =execQuery(txn,"select dtm_begin_transaction()");

}else {

exec(txn,"select dtm_join_transaction(%u)", xid);

}

txn.commit();

}

vector< unique_ptr<nontransaction> >txns(conns.size());

vector< unique_ptr<pipeline> >pipes(conns.size());

vector<pipeline::query_id>results(conns.size());

for (size_t i =0; i < conns.size(); i++) {

txns[i] =newnontransaction(*conns[i]);

pipes[i] =newpipeline(*txns[i]);

results[i] = pipes[i]->insert("select sum(v) from t");

}

int64_t sum =0;

for (size_t i =0; i < conns.size(); i++) {

pipes[i]->complete();

result r = pipes[i]->retrieve(results[i]);

sum += r[0][0].as(int64_t());

}

if (sum != prevSum) {

printf("Total=%ld xid=%u\n", sum, xid);

prevSum = sum;

}

t.proceeded +=1;

}

returnNULL;

}

void*writer(void* arg)

{

thread& t = *(thread*)arg;

vector< unique_ptr<connection> >conns(cfg.connections.size());

for (size_t i =0; i < conns.size(); i++) {

conns[i] =newconnection(cfg.connections[i]);

}

for (int i =0; i < cfg.nIterations; i++)

{

int srcCon, dstCon;

int srcAcc = (random() % ((cfg.nAccounts-cfg.nWriters)/cfg.nWriters))*cfg.nWriters + t.id;

int dstAcc = (random() % ((cfg.nAccounts-cfg.nWriters)/cfg.nWriters))*cfg.nWriters + t.id;

do {

srcCon =random() % cfg.connections.size();

dstCon =random() % cfg.connections.size();

}while (srcCon == dstCon);

nontransactionsrcTx(*conns[srcCon]);

nontransactiondstTx(*conns[dstCon]);

xid_t xid =execQuery(srcTx,"select dtm_begin_transaction()");

exec(dstTx,"select dtm_join_transaction(%u)", xid);

exec(srcTx,"begin transaction isolation level %s", cfg.isolationLevel);

exec(dstTx,"begin transaction isolation level %s", cfg.isolationLevel);

try {

exec(srcTx,"update t set v = v - 1 where u=%d", srcAcc);

exec(dstTx,"update t set v = v + 1 where u=%d", dstAcc);

}catch (pqxx_exceptionconst& x) {

exec(srcTx,"rollback");

exec(srcTx,"rollback");

t.aborts +=1;

i -=1;

continue;

}

pipelinesrcPipe(srcTx);

pipelinedstPipe(dstTx);

srcPipe.insert("commit transaction");

dstPipe.insert("commit transaction");

srcPipe.complete();

dstPipe.complete();

t.proceeded +=1;

}

returnNULL;

}

voidinitializeDatabase()

{

for (size_t i =0; i < cfg.connections.size(); i++) {

connectionconn(cfg.connections[i]);

worktxn(conn);

exec(txn,"drop extension if exists pg_dtm");

exec(txn,"create extension pg_dtm");

exec(txn,"drop table if exists t");

exec(txn,"create table t(u int primary key, v int)");

exec(txn,"insert into t (select generate_series(0,%d), %d)", cfg.nAccounts-1,0);

txn.commit();

}

}

intmain (int argc,char* argv[])

{

bool initialize =false;

for (int i =1; i < argc; i++) {

if (argv[i][0] =='-') {

switch (argv[i][1]) {

case'r':

cfg.nReaders =atoi(argv[++i]);

continue;

case'w':

cfg.nWriters =atoi(argv[++i]);

continue;

case'a':

cfg.nAccounts =atoi(argv[++i]);

continue;

case'n':

cfg.nIterations =atoi(argv[++i]);

continue;

case'c':

cfg.connections.push_back(string(argv[++i]));

continue;

case'l':

cfg.isolationLevel = argv[++i];

continue;

case'i':

initialize =true;

continue;

}

}

printf("Options:\n"

"\t-i\tinitialize datanase\n");

return1;

}

if (initialize) {

initializeDatabase();

}

time_t start =getCurrentTime();

running =true;

vector<thread>readers(cfg.nReaders);

vector<thread>writers(cfg.nWriters);

size_t nReads =0;

size_t nWrites =0;

size_t nAborts =0;

for (int i =0; i < cfg.nReaders; i++) {

readers[i].start(i, reader);

}

for (int i =0; i < cfg.nWriters; i++) {

writers[i].start(i, writer);

}

for (int i =0; i < cfg.nWriters; i++) {

writers[i].wait();

nWrites += writers[i].proceeded;

nAborts += writers[i].aborts;

}

running =false;

for (int i =0; i < cfg.nReaders; i++) {

readers[i].wait();

nReads += readers[i].proceeded;

}

time_t elapsed =getCurrentTime() - start;

printf("TPS(updates)=%f, TPS(selects)=%f, aborts=%ld\n", (double)(nWrites*USEC)/elapsed, (double)(nReads*USEC)/elapsed, nAborts);

return0;

}

CXX=g++

CXXFLAGS=-g -Wall -O2 -pthread

all: dtmbench

dtmbench: dtmbench.cpp

$(CXX)$(CXXFLAGS) -o dtmbench dtmbench.cpp -lpqxx

clean:

rm -f dtmbench

./dtmbench \

-c"dbname=postgres host=localhost user=knizhnik port=5432 sslmode=disable" \

-c"dbname=postgres host=localhost user=knizhnik port=5433 sslmode=disable" \

-c"dbname=postgres host=localhost user=knizhnik port=5434 sslmode=disable" \

-n 1000 -a 1000 -w 10 -r 1$*