Merge cockroachdb#132881

132881: drtprod: add tpcc_run.sh to workload-scale r=nameisbhaskar a=vidit-bhat

This PR add the `tpcc_run.sh` to all `workload-scale` nodes
Epic: none
Release note: None

Co-authored-by: Vidit Bhat <[email protected]>

pkg/cmd/drtprod/configs/drt_scale.yaml

@@ -8,6 +8,8 @@ environment:
   ROACHPROD_GCE_DEFAULT_PROJECT: cockroach-drt
   CLUSTER: drt-scale
   WORKLOAD_CLUSTER: workload-scale
+  CLUSTER_NODES: 150
+  WORKLOAD_NODES: 9
 
 targets:
   # crdb cluster specs
@@ -21,7 +23,7 @@ targets:
           gce-managed: true
           gce-enable-multiple-stores: true
           gce-zones: "us-central1-a"
-          nodes: 150
+          nodes: $CLUSTER_NODES
           gce-machine-type: n2-standard-16
           local-ssd: true
           gce-local-ssd-count: 4
@@ -49,10 +51,6 @@ targets:
           args: --wal-failover=among-stores
           restart: false
           sql-port: 26257
-        on_rollback:
-          - command: stop
-            args:
-              - $CLUSTER
       - command: run
         args:
           - $CLUSTER
@@ -65,7 +63,7 @@ targets:
         flags:
           clouds: gce
           gce-zones: "us-central1-a"
-          nodes: 9
+          nodes: $WORKLOAD_NODES
           gce-machine-type: n2-standard-8
           os-volume-size: 100
           username: workload
@@ -110,12 +108,12 @@ targets:
       - command: put
         args:
           - $WORKLOAD_CLUSTER
-          - bin/roachprod
+          - artifacts/roachprod
           - roachprod
       - command: put
         args:
           - $WORKLOAD_CLUSTER
-          - bin/roachtest
+          - artifacts/roachtest
           - roachtest-operations
       - command: put
         args:
@@ -136,3 +134,13 @@ targets:
         flags:
           warehouses: 640000
           db: cct_tpcc_big
+      - script: "pkg/cmd/drtprod/scripts/generate_tpcc_run.sh"
+        flags:
+          db: cct_tpcc
+          warehouses: 320000
+          max-rate: 10000
+          workers: 5000
+          conns: 5000
+          duration: 12h
+          ramp: 10m
+          wait: 0

pkg/cmd/drtprod/scripts/generate_tpcc_run.sh

@@ -0,0 +1,58 @@
+#!/bin/bash
+
+# Copyright 2024 The Cockroach Authors.
+#
+# Use of this software is governed by the CockroachDB Software License
+# included in the /LICENSE file.
+
+# This script sets up the tpcc run workload script in the workload nodes
+# The --warehouses flag is passed as argument to this script
+# NOTE - This uses CLUSTER and WORKLOAD_CLUSTER environment variable, if not set the script fails
+
+if [ -z "${CLUSTER}" ]; then
+  echo "environment CLUSTER is not set"
+  exit 1
+fi
+
+if [ -z "${WORKLOAD_CLUSTER}" ]; then
+  echo "environment WORKLOAD_CLUSTER is not set"
+  exit 1
+fi
+
+if [ -z "${WORKLOAD_NODES}" ]; then
+  echo "environment WORKLOAD_NODES is not set"
+  exit 1
+fi
+
+# Prepare PGURLS
+PGURLS=$(roachprod pgurl $CLUSTER | sed s/\'//g)
+
+# Loop through each node
+for NODE in $(seq 1 $WORKLOAD_NODES)
+do
+  # Create the workload script
+  cat <<EOF >/tmp/tpcc_run.sh
+#!/usr/bin/env bash
+
+read -r -a PGURLS_ARR <<< "$PGURLS"
+
+j=0
+while true; do
+    echo ">> Starting tpcc workload"
+    ((j++))
+    LOG=./tpcc_\$j.txt
+    ./cockroach workload run tpcc $@ \
+        --tolerate-errors \
+        --families \
+        "\${PGURLS_ARR[@]}" | tee \$LOG
+    if [ \$? -eq 0 ]; then
+        rm "\$LOG"
+    fi
+    sleep 1
+done
+EOF
+
+  # Upload the script to the workload cluster
+  roachprod put $WORKLOAD_CLUSTER:$NODE /tmp/tpcc_run.sh
+  roachprod ssh $WORKLOAD_CLUSTER:$NODE -- chmod +x tpcc_run.sh
+done