From d07f09f133f9495f309150c2c74a4520757f97ce Mon Sep 17 00:00:00 2001
From: Shailendra Patel <shailendra.patel@cockroachlabs.com>
Date: Wed, 6 Nov 2024 14:55:27 +0530
Subject: [PATCH] drtprod: modify tpcc run script.

Modifying generate_tpcc_run.sh script to equally distribute pgurls of cluster
nodes among workload nodes.

Epic: none

Release note: None
---
 pkg/cmd/drtprod/configs/drt_scale.yaml       |  4 +-
 pkg/cmd/drtprod/scripts/generate_tpcc_run.sh | 40 +++++++++++++-------
 2 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/pkg/cmd/drtprod/configs/drt_scale.yaml b/pkg/cmd/drtprod/configs/drt_scale.yaml
index 90ba8d048b24..bf31f742c0d2 100644
--- a/pkg/cmd/drtprod/configs/drt_scale.yaml
+++ b/pkg/cmd/drtprod/configs/drt_scale.yaml
@@ -162,8 +162,8 @@ targets:
         flags:
           db: cct_tpcc
           warehouses: 400000
-          max-rate: 10000
-          workers: 5000
+          max-rate: 1000
+          workers: 400000
           conns: 5000
           duration: 12h
           ramp: 10m
diff --git a/pkg/cmd/drtprod/scripts/generate_tpcc_run.sh b/pkg/cmd/drtprod/scripts/generate_tpcc_run.sh
index d1290035d80a..ab5437c0eaa2 100755
--- a/pkg/cmd/drtprod/scripts/generate_tpcc_run.sh
+++ b/pkg/cmd/drtprod/scripts/generate_tpcc_run.sh
@@ -39,20 +39,32 @@ if [ -z "${WORKLOAD_NODES}" ]; then
   exit 1
 fi
 
+if [ -z "${CLUSTER_NODES}" ]; then
+  echo "environment CLUSTER_NODES is not set"
+  exit 1
+fi
+
 absolute_path=$(roachprod run "${WORKLOAD_CLUSTER}":1 -- "realpath ./cockroach")
 pwd=$(roachprod run "${WORKLOAD_CLUSTER}":1 -- "dirname ${absolute_path}")
-# Prepare PGURLS
-PGURLS=$(roachprod pgurl $CLUSTER | sed s/\'//g)
 
-# Loop through each node
-for NODE in $(seq 1 $WORKLOAD_NODES)
-do
+# Calculate the number of PGURLS each workload node should get
+PGURL_PER_NODE=$((CLUSTER_NODES / WORKLOAD_NODES))
+REMAINDER_NODE=$((CLUSTER_NODES % WORKLOAD_NODES))
+
+# Distribute the PGURLS among the workload nodes
+for ((NODE=0; NODE<WORKLOAD_NODES; NODE++)); do
+  START_OFFSET=$((NODE * PGURL_PER_NODE + (NODE < REMAINDER_NODE ? NODE : REMAINDER_NODE) + 1))
+  END_OFFSET=$((START_OFFSET + PGURL_PER_NODE + (NODE < REMAINDER_NODE ? 1 : 0) - 1))
+
+  # Print or use the PGURLS for the current workload node
+  echo "pgurl for Nodes ${START_OFFSET}:${END_OFFSET}"
+  PGURLS=$(roachprod pgurl $CLUSTER:$START_OFFSET-$END_OFFSET)
+
+
   # Create the workload script
   cat <<EOF >/tmp/tpcc_run_${suffix}.sh
 #!/usr/bin/env bash
 
-read -r -a PGURLS_ARR <<< "$PGURLS"
-
 j=0
 while true; do
     echo ">> Starting tpcc workload"
@@ -61,7 +73,7 @@ while true; do
     ./cockroach workload run tpcc $@ \
         --tolerate-errors \
         --families \
-        "\${PGURLS_ARR[@]}" | tee \$LOG
+         $PGURLS  | tee \$LOG
     if [ \$? -eq 0 ]; then
         rm "\$LOG"
     fi
@@ -69,12 +81,12 @@ while true; do
 done
 EOF
 
-  # Upload the script to the workload cluster
-  roachprod put $WORKLOAD_CLUSTER:$NODE /tmp/tpcc_run_${suffix}.sh
-  roachprod ssh $WORKLOAD_CLUSTER:$NODE -- "chmod +x tpcc_run_${suffix}.sh"
+#   Upload the script to the workload cluster
+  roachprod put $WORKLOAD_CLUSTER:$((NODE + 1)) /tmp/tpcc_run_${suffix}.sh
+  roachprod ssh $WORKLOAD_CLUSTER:$((NODE + 1)) -- "chmod +x tpcc_run_${suffix}.sh"
   if [ "$execute_script" = "true" ]; then
-    roachprod run "${WORKLOAD_CLUSTER}":1 -- "sudo systemd-run --unit tpcc_run_${suffix} --same-dir --uid \$(id -u) --gid \$(id -g) bash ${pwd}/tpcc_run_${suffix}.sh"
+      roachprod run "${WORKLOAD_CLUSTER}":1 -- "sudo systemd-run --unit tpcc_run_${suffix} --same-dir --uid \$(id -u) --gid \$(id -g) bash ${pwd}/tpcc_run_${suffix}.sh"
   else
-    echo "Run --> roachprod run "${WORKLOAD_CLUSTER}":1 -- \"sudo systemd-run --unit tpcc_run_${suffix} --same-dir --uid \\\$(id -u) --gid \\\$(id -g) bash ${pwd}/tpcc_run_${suffix}.sh\""
-  fi
+    echo "Run --> roachprod run "${WORKLOAD_CLUSTER}":$((NODE + 1)) -- \"sudo systemd-run --unit tpcc_run_${suffix} --same-dir --uid \\\$(id -u) --gid \\\$(id -g) bash ${pwd}/tpcc_run_${suffix}.sh\""
+    fi
 done