set -euo pipefail

export ROACHPROD_GCE_DEFAULT_PROJECT=cockroach-drt

export ROACHPROD_DNS="drt.crdb.io"

if [ "$#" -lt 1 ]; then

echo

echo "drtprod runs roachprod commands against DRT clusters"

echo

echo "Usage: drtprod <command>"

echo

echo "Commands include:"

echo " push-hosts: write the ips and pgurl files for a cluster to a node/cluster"

echo " dns: update/create DNS entries in drt.crdb.io for a cluster"

echo " create: a wrapper for the 'roachprod' with predefined specs for named clusters"

echo " *: any other command is passed to roachprod, potentialy with flags addded"

echo

roachprod | sed -e 's/roachprod/drtprod/g'

exit 0

case $1 in

"start")

cluster="$2"

# This is to avoid the command to have --restart and --restart=false both in case of drt-large

if [[ $cluster != "drt-large" && "$*" != *"--restart"* ]]; then

# implied for long-lived DRT clusters; avoid on init w/ --restart=false.

shift

set -- start "--enable-fluent-sink" "--restart" "$@"

fi

if [[ "$*" != *"--secure"* ]]; then

shift

set -- start "--enable-fluent-sink" "--secure" "$@"

fi

if [[ "$*" != *"--sql-port 26257"* ]]; then

shift

set -- start "--enable-fluent-sink" "--sql-port" "26257" "$@"

fi

case $cluster in

"drt-large")

shift

set -- start "--binary" "./cockroach" --args=--log="file-defaults: {dir: 'logs', max-group-size: 1GiB}" --store-count=16 --restart=false --args="--wal-failover=among-stores" "$@"

roachprod run $cluster -- "sudo systemctl unmask cron.service ; sudo systemctl enable cron.service ; echo \"crontab -l ; echo '@reboot sleep 100 && ~/cockroach.sh' | crontab -\" > t.sh ; sh t.sh ; rm t.sh"

;;

"drt-chaos")

shift

set -- start "--binary" "./cockroach" --args=--log="file-defaults: {dir: 'logs', max-group-size: 1GiB}" --store-count=4 --args="--wal-failover=among-stores" "$@"

;;

"drt-ldr1"|"drt-ldr2")

shift

set -- start "--binary" "./cockroach" --args=--log="file-defaults: {dir: 'logs', max-group-size: 1GiB}" --restart=false "$@"

;;

*)

;;

esac

;;

"side-eye")

shift

cluster="$1"

key="$(gcloud secrets versions access latest --secret='side-eye-key')"

if [[ -z "${key}" ]]; then

echo "no key"

exit 1

fi

echo "installing side-eye on ${cluster}..."

roachprod ssh $cluster -- "curl https://sh.side-eye.io/ | SIDE_EYE_API_TOKEN=${key} SIDE_EYE_ENVIRONMENT=${cluster} sh"

exit 0

;;

"dns")

if [ "$#" -lt 2 ]; then

echo "usage: drtprod $1 <cluster> [create]"

exit 1

fi

# roachprod only manages DNS in ephemeral, so we just do this ourselves.

# These are very low-churn clusters so this is fine being manual and in a

# wrapper.

shift

cluster=$1

roachprod adminurl $cluster --ips --insecure=false |

awk '{printf "%04d\t%s\n", NR, $0}' | # prepend the padded node IDs.

sed -e 's,https://\(.*\):26258/,\1,g' | # remove the HTTPS part.

while read node ip; do

host="${cluster}-${node}.drt.crdb.io."

gcloud dns --project=cockroach-shared record-sets ${2:-update} "${host}" --rrdatas="${ip}" \

--type="A" --zone="drt" --ttl=60

done

exit 0

;;

"datadog")

shift

cluster="${1}"

dd_api_key="$(gcloud --project=cockroach-drt secrets versions access latest --secret datadog-api-key)"

if [ -z "${dd_api_key}" ]; then

echo "Missing Datadog API key!"

exit 1

fi

dd_site="us5.datadoghq.com"

case $cluster in

"drt-chaos")

roachprod ssh ${cluster} -- "sudo mkdir -p /etc/fluent-bit && sudo tee /etc/fluent-bit/config-override.yaml > /dev/null << EOF

;;

"workload-chaos")

roachprod ssh ${cluster} -- "sudo mkdir -p /etc/otelcol-contrib && sudo tee /etc/otelcol-contrib/config-override.yaml > /dev/null << EOF

- "localhost:2112"

- "localhost:2113"

- "localhost:2114"

;;

"workload-large")

roachprod ssh ${cluster} -- "sudo mkdir -p /etc/otelcol-contrib && sudo tee /etc/otelcol-contrib/config-override.yaml > /dev/null << EOF

- "localhost:2112"

;;

esac

# TODO: Move this into roachtest_ops.sh per

# https://github.com/cockroachdb/cockroach/pull/126530#discussion_r1668760594

roachprod ssh ${cluster} -- "sudo tee /etc/profile.d/99-datadog.sh > /dev/null << EOF

roachprod opentelemetry-start ${cluster} \

--datadog-api-key "${dd_api_key}" \

--datadog-tags 'service:drt-cockroachdb,team:drt'

roachprod fluent-bit-start ${cluster} \

--datadog-api-key "${dd_api_key}" \

--datadog-service drt-cockroachdb \

--datadog-tags 'service:drt-cockroachdb,team:drt'

echo

echo "Updated ${cluster} configuration to send telemetry data to Datadog."

echo

echo "If this was the first time this script was run against ${cluster} then"

echo "CockroachDB must be restarted to reload its logging configuration."

echo

exit 0

;;

"ua-failover")

if [ "$#" -lt 2 ]; then

echo "usage: $0 failover {drt-ua1,drt-ua2}"

echo

echo "fails the ua workload over to the specified ua cluster, setting the other as a standby"

exit 1

fi

shift

cluster=$1

shift

case "${cluster}" in

"drt-ua1")

standby=drt-ua2

pgurls=pgurls.ua1.txt

;;

"drt-ua2")

standby=drt-ua1

pgurls=pgurls.ua2.txt

;;

*)

echo "unknown cluster: ${cluster}"

exit 1

;;

esac

echo "stopping workload..."

$0 ssh workload-ua:1 -- "sudo systemctl stop tpcc"

echo "cutting over..."

$0 sql $cluster:1 --cluster=system -- -e "ALTER VIRTUAL CLUSTER main COMPLETE REPLICATION TO LATEST"

$0 sql $cluster:1 --cluster=system -- -e "SELECT status FROM [SHOW JOBS WHEN COMPLETE SELECT id FROM system.jobs WHERE job_type = 'REPLICATION STREAM INGESTION'];"

echo "cutover complete; restarting workload..."

$0 sql $cluster:1 --cluster=system -- -e "ALTER TENANT MAIN START SERVICE SHARED"

$0 ssh workload-ua:1 -- "ln -sf $pgurls pgurls.txt; sudo systemctl start tpcc"

echo "restarted workload; configuing new standby..."

$0 sql $standby:1 --cluster=system -- -e "ALTER VIRTUAL CLUSTER main STOP SERVICE"

$0 sql $standby:1 --cluster=system -- -e "ALTER VIRTUAL CLUSTER main START REPLICATION OF main ON 'external://$cluster'"

echo "started replication in other direction."

exit 0

;;

"create")

if [ "$#" -lt 2 ]; then

echo "usage: drtprod create {drt-large,drt-chaos,drt-ua1,drt-ua2,workload-large,workload-chaos}"

exit 1

fi

case "${2}" in

"drt-large")

roachprod create drt-large \

--clouds gce \

--gce-managed \

--gce-enable-multiple-stores \

--gce-zones "northamerica-northeast2-a:2,northamerica-northeast2-b:2,northamerica-northeast2-c:1,us-east5-a:2,us-east5-b:2,us-east5-c:1,us-east1-b:2,us-east1-c:2,us-east1-d:1" \

--gce-use-spot \

--nodes 15 \

--gce-machine-type n2-standard-16 \

--local-ssd=true \

--gce-local-ssd-count 16 \

--os-volume-size 100 \

--username drt \

--lifetime 8760h

# setup dns

$0 dns drt-large create

;;

"workload-large")

roachprod create workload-large \

--clouds gce \

--gce-zones "northamerica-northeast2-a,us-east5-a,us-east1-b" \

--nodes 3 \

--gce-machine-type n2-standard-4 \

--os-volume-size 100 \

--username workload \

--lifetime 8760h

;;

"drt-chaos")

roachprod create drt-chaos \

--clouds gce \

--gce-managed \

--gce-zones "us-east1-d,us-east1-b,us-east1-c" \

--nodes 6 \

--gce-machine-type n2-standard-16 \

--gce-local-ssd-count="4" \

--gce-enable-multiple-stores \

--local-ssd=true \

--username drt \

--lifetime="8760h" \

--gce-image "ubuntu-2204-jammy-v20240319"

# setup dns

$0 dns drt-chaos create

;;

"workload-chaos")

roachprod create workload-chaos \

--clouds gce \

--gce-zone "us-east1-c" \

--nodes 1 \

--gce-machine-type n2-standard-8 \

--os-volume-size 100 \

--username workload \

--lifetime 8760h

;;

"ua")

roachprod create drt-ua1 \

--clouds="gce" \

--gce-zones="us-east1-c" \

--nodes="5" \

--gce-machine-type="n2-standard-16" \

--local-ssd="true" \

--gce-local-ssd-count="8" \

--username="drt" \

--lifetime="8760h"

roachprod create drt-ua2 \

--clouds="gce" \

--gce-zones="us-east1-c" \

--nodes="5" \

--gce-machine-type="n2-standard-16" \

--local-ssd="true" \

--gce-local-ssd-count="8" \

--username="drt" \

--lifetime="8760h"

# setup dns

$0 dns drt-ua1 create || $0 dns drt-ua1 || (echo "dns setup failed -- run dns manually"; true)

$0 dns drt-ua2 create || $0 dns drt-ua2 || (echo "dns setup failed -- run dns manually"; true)

# start cockroach

$0 stage drt-ua1 cockroach

$0 stage drt-ua2 cockroach

$0 start drt-ua1 --restart=false

$0 start drt-ua2 --restart=false

# tell the clusters about eachother.

ua2="$(roachprod ssh drt-ua2:1 -- ./cockroach encode-uri --inline --ca-cert ./certs/ca.crt --key ./certs/client.root.key --cert ./certs/client.root.crt {pgurl:1:system} | sed 's,postgresql://postgres://,postgres://,')"

ua1="$(roachprod ssh drt-ua1:1 -- ./cockroach encode-uri --inline --ca-cert ./certs/ca.crt --key ./certs/client.root.key --cert ./certs/client.root.crt {pgurl:1:system} | sed 's,postgresql://postgres://,postgres://,')"

$0 sql drt-ua1:1 --cluster=system -- -e "CREATE EXTERNAL CONNECTION 'drt-ua2' AS '${ua2}'"

$0 sql drt-ua2:1 --cluster=system -- -e "CREATE EXTERNAL CONNECTION 'drt-ua1' AS '${ua1}'"

# enable rangefeeds.

$0 sql drt-ua1:1 --cluster=system -- -e "SET CLUSTER SETTING kv.rangefeed.enabled = true"

$0 sql drt-ua2:1 --cluster=system -- -e "SET CLUSTER SETTING kv.rangefeed.enabled = true"

# enable weak transaction isolation levels.

$0 sql drt-ua1:1 --cluster=system -- -e "SET CLUSTER SETTING sql.txn.read_committed_isolation.enabled = true"

$0 sql drt-ua2:1 --cluster=system -- -e "SET CLUSTER SETTING sql.txn.read_committed_isolation.enabled = true"

$0 sql drt-ua1:1 --cluster=system -- -e "SET CLUSTER SETTING sql.txn.snapshot_isolation.enabled = true"

$0 sql drt-ua2:1 --cluster=system -- -e "SET CLUSTER SETTING sql.txn.snapshot_isolation.enabled = true"

# Setup the main tenant, on ua1 initilly.

$0 sql drt-ua1:1 --cluster=system -- -e "CREATE VIRTUAL CLUSTER main"

$0 sql drt-ua1:1 --cluster=system -- -e "ALTER VIRTUAL CLUSTER main START SERVICE SHARED"

$0 sql drt-ua1:1 --cluster=main --auth-mode=root -- -e "CREATE USER roachprod WITH PASSWORD 'cockroachdb'"

$0 sql drt-ua1:1 --cluster=main --auth-mode=root -- -e "GRANT admin TO roachprod"

# setup the second to follow it.

$0 sql drt-ua2:1 --cluster=system -- -e "CREATE VIRTUAL CLUSTER main FROM REPLICATION OF main ON 'external://drt-ua1'"

# point traffic at the main tenant.

$0 sql drt-ua1:1 --cluster=system -- -e "SET CLUSTER SETTING server.controller.default_target_cluster = 'main';"

$0 sql drt-ua2:1 --cluster=system -- -e "SET CLUSTER SETTING server.controller.default_target_cluster = 'main';"

echo

echo "Completed setup of ua1 and ua2, with virtual cluster main on ua1 replicating to ua2."

echo

echo "To restore the 150k warehouse fixture, run:"

echo " restore database tpcc from latest in 'gs://cockroach-fixtures-us-east1/backups/tpc-c/v24.1/db/warehouses=150k/?AUTH=implicit';'"

echo

echo "Then run 'drtprod create workload-ua' after the restore completes to start the workload."

echo

;;

"workload-ua")

set -x

roachprod create workload-ua \

--clouds="gce" \

--gce-zones="us-east1-c" \

--nodes="1" \

--gce-machine-type="n2-standard-8" \

--os-volume-size 100 \

--username workload \

--lifetime 8760h \

--gce-image "ubuntu-2204-jammy-v20240319"

# setup dns

$0 dns workload-ua create || $0 dns workload-ua

# push cockroach to the worker to run its built-in workloads

roachprod stage workload-ua:1 cockroach

roachprod stage workload-ua:1 workload

# Push certs for both clusters to the worker.

roachprod get drt-ua1:1 certs certs.ua1

roachprod put workload-ua:1 certs.ua1

roachprod get drt-ua2:1 certs certs.ua2

roachprod put workload-ua:1 certs.ua2

roachprod ssh workload-ua:1 -- "chmod 0600 certs.ua1/* certs.ua2/*"

rm -rf certs.ua1 certs.ua2

# Setup pgurls for worklooads to use.

roachprod ssh workload-ua:1 -- "

# Setup the tpcc workload runner as a script invoked by systemd.

roachprod ssh workload-ua:1 -- 'cat - > tpcc << EOF

roachprod ssh workload-ua:1 -- "chmod +x ./tpcc;

roachprod ssh workload-ua -- "sudo systemctl daemon-reload; sudo systemctl enable tpcc.service; sudo systemctl start tpcc.service"

;;

"ldr")

roachprod create drt-ldr1 \

--clouds="gce" \

--gce-zones="us-east1-c" \

--nodes="5" \

--gce-machine-type="n2-standard-16" \

--local-ssd="false" \

--gce-pd-volume-size="500" \

--gce-pd-volume-type="pd-balanced" \

--username="drt" \

--lifetime="8760h"

roachprod create drt-ldr2 \

--clouds="gce" \

--gce-zones="us-east1-c" \

--nodes="5" \

--gce-machine-type="n2-standard-16" \

--local-ssd="false" \

--gce-pd-volume-size="500" \

--gce-pd-volume-type="pd-balanced" \

--username="drt" \

--lifetime="8760h"

# setup dns

$0 dns drt-ldr1 create || $0 dns drt-ldr1 || (echo "dns setup failed -- run dns manually"; true)

$0 dns drt-ldr2 create || $0 dns drt-ldr2 || (echo "dns setup failed -- run dns manually"; true)

# start cockroach

$0 stage drt-ldr1 cockroach

$0 stage drt-ldr2 cockroach

$0 start drt-ldr1 --restart=false

$0 start drt-ldr2 --restart=false

$0 sql drt-ldr1:1 -- -e "SET CLUSTER SETTING kv.rangefeed.enabled = true"

$0 sql drt-ldr1:1 -- -e "SET CLUSTER SETTING sql.ttl.default_delete_batch_size = 1000"

$0 sql drt-ldr1:1 -- -e "SET CLUSTER SETTING sql.ttl.default_delete_rate_limit = 2000"

$0 sql drt-ldr1:1 -- -e "SET CLUSTER SETTING sql.ttl.default_select_batch_size = 5000"

$0 sql drt-ldr1:1 -- -e "SET CLUSTER SETTING sql.txn.read_committed_isolation.enabled = true"

$0 sql drt-ldr1:1 -- -e "SET CLUSTER SETTING sql.txn.snapshot_isolation.enabled = true"

$0 sql drt-ldr2:1 -- -e "SET CLUSTER SETTING kv.rangefeed.enabled = true"

$0 sql drt-ldr2:1 -- -e "SET CLUSTER SETTING sql.ttl.default_delete_batch_size = 1000"

$0 sql drt-ldr2:1 -- -e "SET CLUSTER SETTING sql.ttl.default_delete_rate_limit = 2000"

$0 sql drt-ldr2:1 -- -e "SET CLUSTER SETTING sql.ttl.default_select_batch_size = 5000"

$0 sql drt-ldr2:1 -- -e "SET CLUSTER SETTING sql.txn.read_committed_isolation.enabled = true"

$0 sql drt-ldr2:1 -- -e "SET CLUSTER SETTING sql.txn.snapshot_isolation.enabled = true"

# import the workload

$0 sql drt-ldr1:1 -- -e "CREATE DATABASE ycsb"

$0 sql drt-ldr1:1 -- -e "ALTER DATABASE ycsb CONFIGURE ZONE USING gc.ttlseconds = 600"

$0 sql drt-ldr2:1 -- -e "CREATE DATABASE ycsb"

$0 sql drt-ldr2:1 -- -e "ALTER DATABASE ycsb CONFIGURE ZONE USING gc.ttlseconds = 600"

$0 ssh drt-ldr1:1 "./cockroach workload init ycsb --workload=A --insert-count=1000 --families=false {pgurl:1}"

$0 ssh drt-ldr2:1 "./cockroach workload init ycsb --workload=A --insert-count=1000 --families=false --insert-start=4511686018427387904 {pgurl:1}"

# tell the clusters about eachother.

ldr2="$(roachprod ssh drt-ldr2:1 -- ./cockroach encode-uri --inline --ca-cert ./certs/ca.crt --key ./certs/client.root.key --cert ./certs/client.root.crt {pgurl:1} | sed 's,postgresql://postgres://,postgres://,' | sed 's/roachprod:cockroachdb/root/' | sed 's/defaultdb/ycsb/')"

ldr1="$(roachprod ssh drt-ldr1:1 -- ./cockroach encode-uri --inline --ca-cert ./certs/ca.crt --key ./certs/client.root.key --cert ./certs/client.root.crt {pgurl:1} | sed 's,postgresql://postgres://,postgres://,' | sed 's/roachprod:cockroachdb/root/' | sed 's/defaultdb/ycsb/')"

$0 sql drt-ldr1:1 -- -e "CREATE EXTERNAL CONNECTION 'drt-ldr2' AS '${ldr2}'"

$0 sql drt-ldr2:1 -- -e "CREATE EXTERNAL CONNECTION 'drt-ldr1' AS '${ldr1}'"

$0 sql drt-ldr1:1 -- -e "ALTER TABLE ycsb.public.usertable ADD COLUMN expired_at TIMESTAMPTZ NOT NULL DEFAULT now() + '30 minutes';"

# set row level ttl

$0 sql drt-ldr1:1 -- -e "ALTER TABLE ycsb.public.usertable SET (ttl_expiration_expression = 'expired_at', ttl_job_cron = '*/30 * * * *');"

$0 sql drt-ldr2:1 -- -e "ALTER TABLE ycsb.public.usertable ADD COLUMN expired_at TIMESTAMPTZ NOT NULL DEFAULT now() + '30 minutes';"

$0 sql drt-ldr2:1 -- -e "ALTER TABLE ycsb.public.usertable SET (ttl_expiration_expression = 'expired_at', ttl_job_cron = '*/30 * * * *');"

roachprod sql drt-ldr1:1 -- -e "CREATE LOGICAL REPLICATION STREAM FROM TABLE usertable ON 'external://drt-ldr2' INTO TABLE ycsb.public.usertable;"

roachprod sql drt-ldr2:1 -- -e "CREATE LOGICAL REPLICATION STREAM FROM TABLE usertable ON 'external://drt-ldr1' INTO TABLE ycsb.public.usertable;"

echo

echo "Completed setup of ldr1 and ldr2, with logical data replication between them on table ycsb.usertable."

echo

echo "Run 'drtprod create workload-ldr' to create the workload node and start the workload."

echo

;;

"workload-ldr")

set -x

roachprod create workload-ldr \

--clouds="gce" \

--gce-zones="us-east1-c" \

--nodes="1" \

--gce-machine-type="n2-standard-16" \

--os-volume-size 100 \

--username workload \

--lifetime 8760h \

# setup dns

$0 dns workload-ldr create || $0 dns workload-ldr

# push cockroach to the worker to run its built-in workloads

roachprod stage workload-ldr:1 cockroach

roachprod stage workload-ldr:1 workload

# Push certs for both clusters to the worker.

roachprod get drt-ldr1:1 certs certs.ldr1

roachprod put workload-ldr:1 certs.ldr1

roachprod get drt-ldr2:1 certs certs.ldr2

roachprod put workload-ldr:1 certs.ldr2

roachprod ssh workload-ldr:1 -- "chmod 0600 certs.ldr1/* certs.ldr2/*"

rm -rf certs.ldr1 certs.ldr2

# Setup pgurls for worklooads to use.

roachprod ssh workload-ldr:1 -- "

# Setup the ycsb workload runner as a script invoked by systemd.

# The runner cycles through all lettered ycsb workloads in a

# round-robin fashion, plus an insert-only workload.

roachprod ssh workload-ldr:1 -- 'cat - > ycsb_run_ldr1.sh << EOF

--workload='custom' --read-freq=0.2 --read-modify-write-freq 0.1 --insert-freq 0.3 --scan-freq 0.1 --update-freq 0.2 --delete-freq 0.1 \\

roachprod ssh workload-ldr:1 -- 'cat - > ycsb_run_ldr2.sh << EOF

--workload='custom' --read-freq=0.2 --read-modify-write-freq 0.1 --insert-freq 0.3 --scan-freq 0.1 --update-freq 0.2 --delete-freq 0.1 \\

roachprod ssh workload-ldr:1 -- "chmod +x ./ycsb_run_ldr1.sh;

roachprod ssh workload-ldr:1 -- "chmod +x ./ycsb_run_ldr2.sh;

roachprod ssh workload-ldr -- "sudo systemctl daemon-reload; sudo systemctl enable ycsb_ldr1.service; sudo systemctl start ycsb_ldr1.service"

roachprod ssh workload-ldr -- "sudo systemctl enable ycsb_ldr2.service; sudo systemctl start ycsb_ldr2.service"

;;

*)

echo

echo "ATN: If $2 is intended to be long-lived, please define it by name in drtprod instead."

echo

exec roachprod "$@"

;;

esac

exit 0

roachprod "$@"