#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>

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*);

typedefint64_tcsn_t;

{

pthread_t t;

size_t proceeded;

int id;

voidstart(int tid,thread_proc_t proc) {

id = tid;

proceeded =0;

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

}

voidwait() {

pthread_join(t,NULL);

}

};

{

int nReaders;

int nWriters;

int nIterations;

int nAccounts;

vector<string> connections;

config() {

nReaders =1;

nWriters =10;

nIterations =1000;

nAccounts =1000;

}

};

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

}

int64_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(int64_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) {

csn_t snapshot;

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

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

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

}

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

if (i ==0) {

snapshot =execQuery(*txns[i],"select dtm_extend()");

}else {

snapshot =execQuery(*txns[i],"select dtm_access(%ld)", snapshot);

}

}

int64_t sum =0;

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

sum +=execQuery(*txns[i],"select sum(v) from t");

}

if (sum != prevSum) {

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

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++)

{

char gtid[32];

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;

sprintf(gtid,"%d.%d", t.id, i);

do {

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

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

}while (srcCon == dstCon);

nontransactionsrcTx(*conns[srcCon]);

nontransactiondstTx(*conns[dstCon]);

exec(srcTx,"begin transaction");

exec(dstTx,"begin transaction");

csn_t snapshot =execQuery(srcTx,"select dtm_extend('%s')", gtid);

snapshot =execQuery(dstTx,"select dtm_access(%ld, '%s')", snapshot, gtid);

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

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

exec(srcTx,"prepare transaction '%s'", gtid);

exec(dstTx,"prepare transaction '%s'", gtid);

exec(srcTx,"select dtm_begin_prepare('%s')", gtid);

exec(dstTx,"select dtm_begin_prepare('%s')", gtid);

csn_t csn =execQuery(srcTx,"select dtm_prepare('%s', 0)", gtid);

csn =execQuery(dstTx,"select dtm_prepare('%s', %ld)", gtid, csn);

exec(srcTx,"select dtm_end_prepare('%s', %ld)", gtid, csn);

exec(dstTx,"select dtm_end_prepare('%s', %ld)", gtid, csn);

exec(srcTx,"commit prepared '%s'", gtid);

exec(dstTx,"commit prepared '%s'", gtid);

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,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'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;

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;

}

running =false;

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

readers[i].wait();

nReads += writers[i].proceeded;

}

time_t elapsed =getCurrentTime() - start;

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

return0;

}

CXX=g++

CXXFLAGS=-g -Wall -O2 -pthread

all: dtmbench

dtmbench: dtmbench.cpp

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

clean:

rm -f dtmbench

)

typeConnStrings []string

Writersstruct {

"Show progress and other stuff for mortals")

flag.BoolVar(&cfg.Parallel,"p",false,

"Use parallel execs")

flag.StringVar(&cfg.Profile,"cpuprofile","","write cpu profile to file")

repread:=flag.Bool("l",false,

"Use 'repeatable read' isolation level instead of 'read committed'")

flag.Parse()

fmt.Println("ERROR: This test needs at leas two connections")

os.Exit(1)

}

runtime.GOMAXPROCS(100)

ifcfg.Profile!="" {

f,err:=os.Create(cfg.Profile)

iferr!=nil {

fmt.Println("Failed to create profile file")

os.Exit(1)

}

pprof.StartCPUProfile(f)

deferpprof.StopCPUProfile()

}

go run*.go \

-C"dbname=postgres user=knizhnik port=5432 sslmode=disable" \

-C"dbname=postgres user=knizhnik port=5433 sslmode=disable" \

-C"dbname=postgres user=knizhnik port=5434 sslmode=disable" \

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