From 9dadd56987bd395538b5b826fdd0ee0b42231429 Mon Sep 17 00:00:00 2001
From: Alexey Melnikov <a_melnikov@lesta.group>
Date: Mon, 23 Sep 2024 10:43:27 +0300
Subject: [PATCH] Added explicitly set timescale.

---
 locust_swarm/swarm.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/locust_swarm/swarm.py b/locust_swarm/swarm.py
index b36f62f..8ddbfa8 100644
--- a/locust_swarm/swarm.py
+++ b/locust_swarm/swarm.py
@@ -193,6 +193,12 @@ def parse(self, stream):
     default=False,
     help="Upload locust-plugins to load gens (useful if you are developing locust-plugins)",
 )
+parser.add_argument(
+    "--timescale",
+    action="store_true",
+    default=False,
+    help="Explicitly enable timescale plugin on master/workers nodes",
+)
 
 parser.add_argument(
     "--version",
@@ -373,6 +379,12 @@ def start_worker_process(server, port):
     if args.test_env:
         extra_env.append("LOCUST_TEST_ENV=" + args.test_env)
 
+    timescale_args = []
+    if args.timescale:
+        timescale_args = [
+            "--timescale",
+        ]
+
     cmd = " ".join([
         "ssh",
         "-q",
@@ -388,6 +400,7 @@ def start_worker_process(server, port):
         "--master-port",
         str(port),
         *master_parameters,
+        *timescale_args,
         "--headless",
         "--expect-workers-max-wait",
         "30",
@@ -520,6 +533,12 @@ def get_available_servers_and_lock_them():
     if args.playwright:
         extra_env.append("LOCUST_PLAYWRIGHT=1")
 
+    timescale_args = []
+    if args.timescale:
+        timescale_args = [
+            "--timescale",
+        ]
+
     master_command = [
         *ssh_command,
         *extra_env,
@@ -532,6 +551,7 @@ def get_available_servers_and_lock_them():
         str(worker_process_count),
         "--expect-workers-max-wait",
         "60",
+        *timescale_args,
         "--headless",
         "-f",
         locustfile,