From ad37336794022a86a324cce4c5f72e74f178bec2 Mon Sep 17 00:00:00 2001
From: Aaron Molitor <amolitor@apple.com>
Date: Mon, 29 Nov 2021 12:47:50 -0800
Subject: [PATCH] port #6032 to release-7.0

---
 packaging/docker/run_ycsb.sh | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/packaging/docker/run_ycsb.sh b/packaging/docker/run_ycsb.sh
index 96337cccf0..0f7896201f 100755
--- a/packaging/docker/run_ycsb.sh
+++ b/packaging/docker/run_ycsb.sh
@@ -2,12 +2,22 @@
 set -Eeuo pipefail
 
 namespace=$(cat /var/run/secrets/kubernetes.io/serviceaccount/namespace)
-
-echo "WAITING FOR ALL PODS TO COME UP"
-while [[ $(kubectl get pods -n ${namespace} -l name=ycsb,run=${RUN_ID} --field-selector=status.phase=Running | grep -cv NAME) -lt ${NUM_PODS} ]]; do
-    sleep 0.1
-done
-echo "ALL PODS ARE UP"
+POD_NUM=$(echo $POD_NAME | cut -d - -f3)
+KEY="ycsb_load_${POD_NUM}_of_${NUM_PODS}_complete"
+CLI=$(ls /var/dynamic-conf/bin/*/fdbcli | head -n1)
+if [ ${MODE} != "load" ]; then
+    echo "WAITING FOR ALL PODS TO COME UP"
+    while [[ $(kubectl get pods -n ${namespace} -l name=ycsb,run=${RUN_ID} --field-selector=status.phase=Running | grep -cv NAME) -lt ${NUM_PODS} ]]; do
+        sleep 0.1
+    done
+    echo "ALL PODS ARE UP"
+else
+    if ${CLI} --exec "get ${KEY}" | grep is ;
+    then
+        # load already completed
+        exit 0
+    fi
+fi;
 
 echo "RUNNING YCSB"
 ./bin/ycsb.sh ${MODE} foundationdb -s -P workloads/${WORKLOAD} ${YCSB_ARGS}
@@ -15,4 +25,9 @@ echo "YCSB FINISHED"
 
 echo "COPYING HISTOGRAMS TO S3"
 aws s3 sync --sse aws:kms --exclude "*" --include "histogram.*" /tmp s3://${BUCKET}/ycsb_histograms/${namespace}/${POD_NAME}
-echo "COPYING HISTOGRAMS TO S3 FINISHED"
\ No newline at end of file
+echo "COPYING HISTOGRAMS TO S3 FINISHED"
+
+if [ ${MODE} == "load" ]; then
+    ${CLI} --exec "writemode on; set ${KEY} 1"
+    echo "WROTE LOAD COMPLETION KEY"
+fi