2 files changed, 14 insertions, 2 deletions
diff --git a/tools/gcp/stress_test/run_client.py b/tools/gcp/stress_test/run_client.py
index f1084f6caa..2004bf6db4 100755
--- a/tools/gcp/stress_test/run_client.py
+++ b/tools/gcp/stress_test/run_client.py
@@ -109,6 +109,11 @@ def run_client():
   dataset_id = env['DATASET_ID']
   summary_table_id = env['SUMMARY_TABLE_ID']
   qps_table_id = env['QPS_TABLE_ID']
+  # The following parameter is to inform us whether the stress client runs
+  # forever until forcefully stopped or will it naturally stop after sometime.
+  # This way, we know that the stress client process should not terminate (even
+  # if it does with a success exit code) and flag the termination as a failure
+  will_run_forever = env.get('WILL_RUN_FOREVER', '1')
 
   bq_helper = BigQueryHelper(run_id, image_type, pod_name, project_id,
                              dataset_id, summary_table_id, qps_table_id)
@@ -146,11 +151,12 @@ def run_client():
   while True:
     # Check if stress_client is still running. If so, collect metrics and upload
     # to BigQuery status table
+    # If stress_p.poll() is not None, it means that the stress client terminated
     if stress_p.poll() is not None:
       end_time = datetime.datetime.now().isoformat()
       event_type = EventType.SUCCESS
       details = 'End time: %s' % end_time
-      if stress_p.returncode != 0:
+      if will_run_forever == '1' or stress_p.returncode != 0:
         event_type = EventType.FAILURE
         details = 'Return code = %d. End time: %s' % (stress_p.returncode,
                                                       end_time)
diff --git a/tools/gcp/stress_test/run_server.py b/tools/gcp/stress_test/run_server.py
index cff6fc8c72..a666ae2900 100755
--- a/tools/gcp/stress_test/run_server.py
+++ b/tools/gcp/stress_test/run_server.py
@@ -74,6 +74,11 @@ def run_server():
   dataset_id = env['DATASET_ID']
   summary_table_id = env['SUMMARY_TABLE_ID']
   qps_table_id = env['QPS_TABLE_ID']
+  # The following parameter is to inform us whether the server runs forever
+  # until forcefully stopped or will it naturally stop after sometime.
+  # This way, we know that the process should not terminate (even if it does
+  # with a success exit code) and flag any termination as a failure.
+  will_run_forever = env.get('WILL_RUN_FOREVER', '1')
 
   logfile_name = env.get('LOGFILE_NAME')
 
@@ -112,7 +117,8 @@ def run_server():
                               stderr=subprocess.STDOUT)
 
   returncode = stress_p.wait()
-  if returncode != 0:
+
+  if will_run_forever == '1' or returncode != 0:
     end_time = datetime.datetime.now().isoformat()
     event_type = EventType.FAILURE
     details = 'Returncode: %d; End time: %s' % (returncode, end_time)