## setup the schema

# find out YCSB install
if [ -z "$YCSB" ]; then
  if [ -d ~/bin/ycsb-0.14.0 ]; then
    export YCSB=~/bin/ycsb-0.14.0
  elif [ -d ~/ycsb-0.14.0 ]; then
    export YCSB=~/ycsb-0.14.0
  else
    echo "Cannot find YCSB install"
    return 1
  fi  
fi

## variations on the tests
export verbose="-s"           # display stat every second
export operationcount=50000  # stop after this many operations
export threadperinstance=1    # thread per YCSB client
export insertstart=0          # key start from user0
export insertcount=50000      # each YCSB client inserts this many keys 
export requestdistribution="uniform"  # uniform zipfian
#requestdistribution="zipfian"
export insertorder="hashed"           # ordered hashed
#insertorder="hashed"
export workloads="a"

# roundrobin list of hosts:port and return host port
# rr 0 "$HOSTS"
rr() {
  local host port
  read -r host port <<< `echo "$2" | awk '{i=n % NF + 1; split($i,a,":"); print a[1] " " a[2]}' n=$1`
  port=${port:-26257}
  echo "$host $port"
}

# setup CRBD_NODES environment variable from cockroach node status
# default
#  -h 127.0.0.1:26257  
crdb_nodes() {
  local OPTIND opt host port
  host=127.0.0.1
  port=26257
  while getopts "h:" opt; do
    case "${opt}" in
      h)
        host=`echo "${OPTARG}" | awk -F: '{print $1}'`
        port=`echo "${OPTARG}" | awk -F: '{print $2}'`
        port=${port:-26257}
        ;;
      esac
  done
  shift $((OPTIND-1))

  # DB nodes where the YCSB clients will also run
  export CRDB_NODES=`cockroach node status --host $host --port $port --insecure --format tsv | tail -n +2 | awk '{print $2}' | tr "\n" " "`
}

# create/overview a default YCSB config
ycsb_config() {
  if [ ! -d "$YCSB/cockroachdb" ]; then
    mkdir $YCSB/cockroachdb
  fi

  if [ ! -f "$YCSB/lib/postgresql-42.2.4.jar" ]; then
    curl --location https://jdbc.postgresql.org/download/postgresql-42.2.4.jar -o $YCSB/lib/postgresql-42.2.4.jar
  fi

  cat >$YCSB/cockroachdb/db.properties<<EOF
db.driver=org.postgresql.Driver
db.url=jdbc:postgresql://127.0.0.1:26257/ycsb
db.user=root
db.passwd=
EOF
}

# HOSTS=host:port host:port
# create schema
# use the first node from $CRDB_NODES to create the schema
ycsb_schema() {
  local port
  local host

  # DB nodes where the YCSB clients will also run
  if [ -z "$CRDB_NODES" ]; then 
    crdb_nodes
  fi

  export SECURE="--insecure"

  read -r host port <<<$(rr 0 "$CRDB_NODES") 
  host="--host $host"
  port="--port $port"

  cockroach sql $SECURE $host $port <<EOF
  create database if not exists ycsb;
  use ycsb;
  CREATE TABLE usertable(YCSB_KEY VARCHAR PRIMARY KEY,
    FIELD0 VARCHAR, FIELD1 VARCHAR,
    FIELD2 VARCHAR, FIELD3 VARCHAR,
    FIELD4 VARCHAR, FIELD5 VARCHAR,
    FIELD6 VARCHAR, FIELD7 VARCHAR,
    FIELD8 VARCHAR, FIELD9 VARCHAR);
  alter table usertable split at select left(concat('user', generate_series(1,99)::string, '0'),6);
  alter table usertable scatter;
EOF
}

# truncate and split ranges
# split user10,user20 ...
# each truncate hurt performance for now. shutdown, wipe out database, restart for now
ycsb_truncate() {
  local port
  local host

  # DB nodes where the YCSB clients will also run
  if [ -z "$CRDB_NODES" ]; then 
    echo "CRDB_NODES not defined.  Run crdb_nodes first"
    return 1
  fi

  read -r host port <<<$(rr 0 "$HOSTS") 
  host="--host $host"
  port="--port $port"

  cockroach sql $SECURE $host $port <<EOF
  truncate ycsb.usertable;
  alter table ycsb.usertable split at select left(concat('user', generate_series(1,99)::string, '0'),6);
  alter table ycsb.usertable scatter;
EOF
}

# load dataset for number of YCSB Clients $instances
# usually number of $HOSTS and $instances are the the same
# when there are more $instances than $HOSTS then round robin the $HOSTS
# $1=batchsize
ycsb_load() {
  local port
  local host
  local n
  local i
  local ts
  local dburl
  local batch
  local parallel
  local OPTIND opt bs
  bs=0
  ts=`date "+%y%m%d-%H%M"`
  while getopts "pb:t:" opt; do
    case "${opt}" in
      p)
        parallel="1"
        ;;
      b)
        bs=${OPTARG}
        batch="-p jdbc.batchupdateapi=true -p db.batchsize=$bs"
        ;;
      t)
        ts=${OPTARG}
        ;;
      esac
  done
  shift $((OPTIND-1))

  w=${w:-a}
  r=${r:-uniform}
  o=${o:-hashed}
 
  export instances=`echo "$HOSTS" | awk '{print NF}'`

  n=$insertstart
  for i in `seq $instances`; do

    echo load $w $r $o $i $n

    # host:port
    read -r host port <<<$(rr $i "$HOSTS") 
    host=$host
    port=${port:-26257}
    dburl="jdbc:postgresql://$host:$port/ycsb"

    if [ "${parallel}" ]; then
      $YCSB/bin/ycsb load jdbc $verbose -P $YCSB/workloads/workload${w} -threads 1 -P $YCSB/cockroachdb/db.properties -p db.url=$dburl -p insertstart=${n} -p insertcount=$insertcount -p insertorder=${o} -p recordcount=0 $batch > ${ts}_ycsb_load_${instances}_${w}_${r}_${o}_${i}_${threadperinstance}_${n}_${bs}.log &
    ((n=$n + $insertcount))
    else
      $YCSB/bin/ycsb load jdbc $verbose -P $YCSB/workloads/workload${w} -threads 1 -P $YCSB/cockroachdb/db.properties -p db.url=$dburl -p insertstart=${n} -p insertcount=$insertcount -p insertorder=${o} -p recordcount=0 $batch > ${ts}_ycsb_load_${instances}_${w}_${r}_${o}_${i}_${threadperinstance}_${n}_${bs}.log 
    fi
    ((n=$n + $insertcount))
  done
}

#  
# -p {instance numner} runnin in parallel.  
#    $instances indicate how many other instances are there.  -p indicate which instance this is
#
ycsb_load_local() {
  local hostname host port jdbcopt
  local n insertlen
  local i
  local ts
  local dburl
  local batch
  local parallel
  local OPTIND opt bs
  bs=0
  ts=`date "+%y%m%d-%H%M"`


  while getopts "b:h:jt:" opt; do
    case "${opt}" in
      b)
        bs=${OPTARG}
        batch="-p jdbc.batchupdateapi=true -p db.batchsize=$bs"
        ;;
      j)
        jdbcopt="reWriteBatchedInserts=true"
        ;;
      h)
        hostname=${OPTARG}
        ;;        
      t)
        ts=${OPTARG}
        ;;
      esac
  done
  shift $((OPTIND-1))

  w=${w:-a}
  r=${r:-uniform}
  o=${o:-hashed}

# host:port
  hostport=$(cockroach node status --insecure --format tsv | grep "$(hostname):" | head -n 1 | awk '{print $1 ":" $2}')
  i=$(echo $hostport | awk -F: '{print $1}')
  host=$(echo $hostport | awk -F: '{print $2}')
  port=$(echo $hostport | awk -F: '{print $3}')

  echo $host
  echo $port

  insertlen=$insertcount
  ((n = ($i - 1) * $insertcount))

  echo load workload=$w request=$r order=$o instance=$i rec=$n

  if [ -z "${jdbcopt}" ]; then
    dburl="jdbc:postgresql://${host}:${port}/ycsb"
  else
    dburl="jdbc:postgresql://${host}:${port}/ycsb?${jdbcopt}"
  fi    


  $YCSB/bin/ycsb load jdbc $verbose -P $YCSB/workloads/workload${w} -threads ${threadperinstance} -P $YCSB/cockroachdb/db.properties -p db.url=$dburl -p insertstart=${n} -p insertcount=${insertlen} -p insertorder=${o} -p recordcount=0 $batch > ${ts}_ycsb_load_${instances}_${w}_${r}_${o}_${i}_${threadperinstance}_${n}_${bs}.log 
}


# w=workload (a,b,...)
# r=requestdistriution (uniform, zipfian)
# o=insertorder (ordered, hashed)
# i=instance (YCSB client instance)
# n=range start
ycsb_run() {
  local port
  local host
  local n
  local i
  local ts
  local batch
  local OPTIND opt bs
  bs=0
  ts=`date "+%y%m%d-%H%M"`
  while getopts "b:t:" opt; do
    case "${opt}" in
      b)
        bs=${OPTARG}
        batch="-p jdbc.batchupdateapi=true -p db.batchsize=$bs"
        ;;
      t)
        ts=${OPTARG}
        ;;
      esac
  done
  shift $((OPTIND-1))

  w=${w:-a}
  r=${r:-uniform}
  o=${o:-hashed}

  n=$insertstart
  for i in `seq $instances`; do
    echo run $w $r $o $i $n

    # host:port
    read -r host port <<<$(rr $i "$HOSTS") 
    host=$host
    port=${port:-26257}
    dburl="jdbc:postgresql://$host:$port/ycsb"

    $YCSB/bin/ycsb run jdbc  $verbose -P $YCSB/workloads/workload${w} -threads ${threadperinstance} -P $YCSB/cockroachdb/db.properties -p operationcount=${operationcount} -p db.url=$dburl -p insertstart=${n} -p insertcount=$insertcount -p insertorder=${o} -p recordcount=0 -p requestdistribution=${r} ${batch} > ${ts}_ycsb_run_${instances}_${w}_${r}_${o}_${i}_${threadperinstance}_${n}_${bs}.log &
    ((n=$n + $insertcount))
  done
}

# run against the entire rnage by starting from $insertstart and ending at $instances * $insertcount
# probably may not run well with zifian distribution.                   
# $1 = batchsize                    database node to target
# $2 = hostname                     database node to target
# $3 = port                         database port on the node to target
#
ycsb_run_local() {

  local port
  local host dburl
  local n
  local i
  local ts
  local batch
  local bs
  local OPTIND opt bs
  bs=0
  ts=`date "+%y%m%d-%H%M"`
  while getopts "b:t:" opt; do
    case "${opt}" in
      b)
        bs=${OPTARG}
        batch="-p jdbc.batchupdateapi=true -p db.batchsize=$bs"
        ;;
      t)
        ts=${OPTARG}
        ;;
      esac
  done
  shift $((OPTIND-1))

  w=${w:-a}
  r=${r:-uniform}
  o=${o:-hashed}

# host:port
  hostport=$(cockroach node status --insecure --format tsv | grep "$(hostname):" | head -n 1 | awk '{print $1 ":" $2}')
  i=$(echo $hostport | awk -F: '{print $1}')
  host=$(echo $hostport | awk -F: '{print $2}')
  port=$(echo $hostport | awk -F: '{print $3}')

  echo $host
  echo $port

  insertlen=$insertcount
  ((n = ($i - 1) * $insertcount))

  echo run workload=$w request=$r order=$o instance=$i rec=$n

  dburl="jdbc:postgresql://${host}:${port}/ycsb"

  $YCSB/bin/ycsb run jdbc  $verbose -P $YCSB/workloads/workload${w} -threads ${threadperinstance} -P $YCSB/cockroachdb/db.properties -p operationcount=${operationcount} -p db.url=$dburl -p insertstart=${n} -p insertcount=$insertlen -p insertorder=${o} -p recordcount=0 -p requestdistribution=${r} ${batch} > ${ts}_ycsb_run_${instances}_${w}_${r}_${o}_${i}_${threadperinstance}_${n}_${bs}.log 
}