tpcc-cons.diff

diff --git a/scripts/gen/tpcc.sh b/scripts/gen/tpcc.sh
index c99761f..1424b42 100644
--- a/scripts/gen/tpcc.sh
+++ b/scripts/gen/tpcc.sh
@@ -10,6 +10,7 @@ f_active_per_core=125
 f_skip_load=''
 f_load_args=''
 f_duration="30m"
+f_nodes=4
 
 function usage() {
   echo "$1
@@ -22,11 +23,12 @@ Usage: $0 [-f] [-w] [-s server] [pgurl,...]
   -L: extra args for load
   -a: number of active warehouses per core
   -d: duration; default 30m
+  -n: nodes: number of nodes in the current cluster
 "
   exit 1
 }
 
-while getopts 'fwsW:A:d:a:L:' flag; do
+while getopts 'fwsW:A:d:a:L:n:' flag; do
   case "${flag}" in
     f) f_force='true' ;;
     w) f_wait='true' ;;
@@ -36,6 +38,7 @@ while getopts 'fwsW:A:d:a:L:' flag; do
     s) f_skip_load='true' ;;
     d) f_duration="${OPTARG}" ;;
     L) f_load_args="${OPTARG}" ;;
+    n) f_nodes="${OPTARG}" ;;
     *) usage "";;
   esac
 done
@@ -92,10 +95,12 @@ then
   echo "done loading"
 fi
 
+num_vcpu=$(cat /proc/cpuinfo | grep processor | wc -l)
+
 if (( f_active == 0 ))
 then
   # Scale active warehouse count by f_active_per_core * number of CPUs.
-  f_active=$(( f_active_per_core * $(cat /proc/cpuinfo | grep processor | wc -l) ))
+  f_active=$(( f_active_per_core * $num_vcpu ))
   if (( f_active > f_warehouses ))
   then
     f_active=0
@@ -104,5 +109,5 @@ fi
 
 report="${logdir}/tpcc-results-$f_active.txt"
 ./cockroach workload run tpcc \
-  --warehouses="$f_warehouses" --active-warehouses="$f_active" --ramp=5m --duration="$f_duration" --tolerate-errors --wait=0 \
+  --warehouses="$f_warehouses" --active-warehouses="$f_active"  --cons=$((num_vcpu * f_nodes * 4))  --ramp=5m --duration="$f_duration" --tolerate-errors --wait=0 \
   "${pgurls[@]}" > "$report"