feat: Load test fixups (#584)

Some minor tweaks to the load tests to make them easier to run locally.

These include:
* normalizing the `--config` option across binaries (this is not used in
production)
* Change the load test to use a hash of the payload instead of the
payload itself (freeing memory and making checks faster)
* Make the `WORKER_COUNT` an environment parameter

Closes SYNC-4063

.circleci/config.yml

@@ -295,7 +295,6 @@ jobs:
             docker push "${DOCKERHUB_LOAD_TEST_REPO}:latest"
 
 workflows:
-  version: 2
 
   build-test-deploy:
     jobs:

autoconnect/src/main.rs

@@ -26,14 +26,12 @@ Usage: autoconnect [options]
 
 Options:
     -h, --help                          Show this message.
-    --config-connection=CONFIGFILE      Connection configuration file path.
-    --config-shared=CONFIGFILE          Common configuration file path.
+    --config=CONFIGFILE                 Connection configuration file path.
 ";
 
 #[derive(Debug, Deserialize)]
 struct Args {
-    flag_config_connection: Option<String>,
-    flag_config_shared: Option<String>,
+    flag_config: Option<String>,
 }
 
 #[actix_web::main]
@@ -43,10 +41,7 @@ async fn main() -> Result<()> {
         .and_then(|d| d.deserialize())
         .unwrap_or_else(|e| e.exit());
     let mut filenames = Vec::new();
-    if let Some(shared_filename) = args.flag_config_shared {
-        filenames.push(shared_filename);
-    }
-    if let Some(config_filename) = args.flag_config_connection {
+    if let Some(config_filename) = args.flag_config {
         filenames.push(config_filename);
     }
     let settings =

autoendpoint/src/settings.rs

@@ -96,9 +96,8 @@ impl Settings {
 
         let built = config.build()?;
 
-        built
-            .try_deserialize::<Self>()
-            .map_err(|error| match error {
+        built.try_deserialize::<Self>().map_err(|error| {
+            match error {
                 // Configuration errors are not very sysop friendly, Try to make them
                 // a bit more 3AM useful.
                 ConfigError::Message(error_msg) => {
@@ -115,7 +114,8 @@ impl Settings {
                     error!("Configuration error: Other: {:?}", &error);
                     error
                 }
-            })
+            }
+        })
     }
 
     /// Convert a string like `[item1,item2]` into a iterator over `item1` and `item2`.

tests/load/locustfiles/load.py

@@ -4,6 +4,7 @@
 
 """Load test shape module."""
 import math
+import os
 from typing import Type
 
 import numpy
@@ -51,9 +52,13 @@ class AutopushLoadTestShape(LoadTestShape):
     the WORKERS_COUNT and USERS_PER_WORKER values must be changed respectively.
     """
 
-    MAX_RUN_TIME: int = 600  # 10 minutes
-    WORKER_COUNT: int = 150  # Must match value defined in setup_k8s.sh
-    USERS_PER_WORKER: int = 1000  # Number of users supported on a c3d-standard-4 hosted worker
+    MAX_RUN_TIME: int = int(os.environ.get("MAX_RUN_TIME", 600))  # 10 minutes
+    WORKER_COUNT: int = int(
+        os.environ.get("WORKER_COUNT", 150)
+    )  # Must match value defined in setup_k8s.sh
+    USERS_PER_WORKER: int = int(
+        os.environ.get("USERS_PER_WORKER", 1000)
+    )  # Number of users supported on a worker running on a n1-standard-2
     MAX_USERS: int = WORKER_COUNT * USERS_PER_WORKER
     trend: QuadraticTrend
     user_classes: list[Type[User]] = [AutopushUser]
@@ -77,7 +82,7 @@ def tick(self) -> TickTuple | None:
 
         None: Instruction to stop the load test
         """
-        run_time: int = self.get_run_time()
+        run_time: int = int(self.get_run_time())
         if run_time > self.MAX_RUN_TIME:
             return None
 

tests/load/locustfiles/locustfile.py

@@ -4,13 +4,13 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-import base64
 import json
 import logging
 import random
 import string
 import time
 import uuid
+from hashlib import sha1
 from json import JSONDecodeError
 from logging import Logger
 from typing import Any, TypeAlias
@@ -73,9 +73,9 @@ def __init__(self, environment) -> None:
         super().__init__(environment)
         self.channels: dict[str, str] = {}
         self.hello_record: HelloRecord | None = None
-        self.notification_records: list[NotificationRecord] = []
-        self.register_records: list[RegisterRecord] = []
-        self.unregister_records: list[RegisterRecord] = []
+        self.notification_records: dict[bytes, NotificationRecord] = {}
+        self.register_records: dict[str, RegisterRecord] = {}
+        self.unregister_records: dict[str, RegisterRecord] = {}
         self.uaid: str = ""
         self.ws: WebSocketApp | None = None
         self.ws_greenlet: Greenlet | None = None
@@ -232,7 +232,7 @@ def post_notification(self, endpoint_url: str) -> None:
         )
 
         record = NotificationRecord(send_time=time.perf_counter(), data=data)
-        self.notification_records.append(record)
+        self.notification_records[sha1(data.encode()).digest()] = record  # nosec
 
         with self.client.post(
             url=endpoint_url,
@@ -281,26 +281,19 @@ def recv(self, data: str) -> Message | None:
                 case "notification":
                     message = NotificationMessage(**message_dict)
                     message_data: str = message.data
-                    decode_data: str = base64.urlsafe_b64decode(message_data + "===").decode(
-                        "utf8"
-                    )
-                    record = next(
-                        (r for r in self.notification_records if r.data == decode_data), None
+                    # scan through the notification records to see
+                    # if this matches a record we sent.
+                    record = self.notification_records.get(
+                        sha1(message_data.encode()).digest(), None  # nosec
                     )
                 case "register":
                     message = RegisterMessage(**message_dict)
                     register_chid: str = message.channelID
-                    record = next(
-                        (r for r in self.register_records if r.channel_id == register_chid),
-                        None,
-                    )
+                    record = self.register_records.get(register_chid, None)
                 case "unregister":
                     message = UnregisterMessage(**message_dict)
                     unregister_chid: str = message.channelID
-                    record = next(
-                        (r for r in self.unregister_records if r.channel_id == unregister_chid),
-                        None,
-                    )
+                    record = self.unregister_records.get(unregister_chid)
                 case _:
                     exception = f"Unexpected data was received. Data: {data}"
 
@@ -390,7 +383,7 @@ def send_register(self, ws: WebSocket, channel_id: str):
         message_type: str = "register"
         data: dict[str, Any] = dict(messageType=message_type, channelID=channel_id)
         record = RegisterRecord(send_time=time.perf_counter(), channel_id=channel_id)
-        self.register_records.append(record)
+        self.register_records[channel_id] = record
         self.send(ws, message_type, data)
 
     def send_unregister(self, ws: WebSocketApp, channel_id: str) -> None:
@@ -409,7 +402,7 @@ def send_unregister(self, ws: WebSocketApp, channel_id: str) -> None:
         message_type: str = "unregister"
         data: dict[str, Any] = dict(messageType=message_type, channelID=channel_id)
         record = RegisterRecord(send_time=time.perf_counter(), channel_id=channel_id)
-        self.unregister_records.append(record)
+        self.unregister_records[channel_id] = record
         self.send(ws, message_type, data)
 
     def send(self, ws: WebSocket | WebSocketApp, message_type: str, data: dict[str, Any]) -> None:

tests/load/setup_k8s.sh

@@ -10,7 +10,7 @@ CLUSTER='autopush-locust-load-test'
 TARGET='https://updates-autopush.stage.mozaws.net'
 SCOPE='https://www.googleapis.com/auth/cloud-platform'
 REGION='us-central1'
-WORKER_COUNT=150
+WORKER_COUNT=${WORKER_COUNT:-150}
 MACHINE_TYPE='c3d-standard-4' # 4 CPUs + 16GB Memory
 BOLD=$(tput bold)
 NORM=$(tput sgr0)