[NL Load Testing] Support various instance and fix a bug when specify…

… total requests (#4267)

tools/nl/loadtest/README.md

@@ -1,10 +1,12 @@
+# NodeJS Server Load Testing
+
 ## NL SVG Latency tester
 
 A few tens of queries against the NodeJS NL server (`nodejs/query`) in series and tracks times.
 
 Run it as:
 
-```
+```bash
 ./run_latencies.sh
 ```
 
@@ -14,14 +16,19 @@ Runs specified number of queries against the NodeJS NL server (`nodejs/query`),
 
 Run it as:
 
-```
-./run_load.sh --total_requests=<total queries to send> --parallel_requests=<number of queries to send at a time> --apikey=<apikey for bard instance>
+```bash
+./run_load.sh \
+    --instance=<bard|dev> \
+    --total_requests=<total queries to send> \
+    --parallel_requests=<number of queries to send at a time> \
+    --apikey=<apikey for bard instance>
 ```
 
 Notes about the arguments:
+
 - If total requests is not set, this tool will run through all the queries in queryset.csv once
 - If parallel requests is not set, this tool will run requests one after another in series
 - With the current setup, 20000 total requests and 150 parallel requests will run for about 7 mins and achieve a qps of around 45
 - For the apikey, you can go to the [Bard Apigee Console](https://pantheon.corp.google.com/apigee/apps?mods=-monitoring_api_staging&project=datcom-bard) and either:
-    1. Use an existing key
-    2. Create your own key by [creating an "API Product", "App", and "Developer"](https://cloud.google.com/apigee/docs/api-platform/security/api-keys).
+  1. Use an existing key
+  2. Create your own key by [creating an "API Product", "App", and "Developer"](https://cloud.google.com/apigee/docs/api-platform/security/api-keys).

tools/nl/loadtest/load.py

@@ -22,14 +22,15 @@
 
 FLAGS = flags.FLAGS
 
+flags.DEFINE_string('instance', 'bard', 'The instance name of the website')
 flags.DEFINE_integer('parallel_requests', 1,
                      'Number of requests to run in parallel')
 flags.DEFINE_integer('total_requests', None, 'Total number of requests to run')
 flags.DEFINE_string('apikey', '', "apikey to use when making query requests.")
 
 _OUTPUT_FILE = 'load_results.json'
 _QUERY_FILE = 'queryset.csv'
-_URL = 'https://bard.datacommons.org/nodejs/query?apikey={apikey}&q={query}'
+_URL = 'https://{instance}.datacommons.org/nodejs/query?apikey={apikey}&q={query}'
 _RESULT_KEY_CODE = 'code'
 _RESULT_KEY_ERROR_CONTENT = 'errorContent'
 _RESULT_KEY_NL_TIME = 'nlTime'
@@ -43,8 +44,8 @@ def _run_query(query):
   # Make the API request
   print(f'Running query: {query}')
   try:
-    response = session.get(_URL.format(apikey=FLAGS.apikey, query=query),
-                           timeout=None)
+    url = _URL.format(instance=FLAGS.instance, apikey=FLAGS.apikey, query=query)
+    response = session.get(url, timeout=None)
     if response.status_code == 200:
       debug_timing = response.json().get('debug', {}).get('timing', {})
       return {
@@ -129,7 +130,7 @@ def run_load(total_requests, parallel_requests):
   # If number of total requests is specified, generate a list of queries that is
   # total requests long
   if total_requests:
-    while len(query_list_to_run) < total_requests or len(query_list):
+    while len(query_list_to_run) < total_requests and len(query_list):
       num_missing = total_requests - len(query_list_to_run)
       query_list_to_run += query_list[0:min(num_missing, len(query_list))]
   else: