From 1de2c7dd54b359d2a74c941b7d4c59d7eb5b58b1 Mon Sep 17 00:00:00 2001 From: Lars Holmberg Date: Sun, 11 Feb 2024 19:11:00 +0100 Subject: [PATCH 1/5] Store the locustfile in a temp dir and delete it after shutdown when downloading from master ("--locustfile -"), in the same fashion as when downloading from a URL. --- locust/argument_parser.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/locust/argument_parser.py b/locust/argument_parser.py index 7ffc7edbb4..f27040c142 100644 --- a/locust/argument_parser.py +++ b/locust/argument_parser.py @@ -246,11 +246,16 @@ def wait_for_reply(): sys.exit(1) filename = msg.data["filename"] - with open(filename, "w") as local_file: - local_file.write(msg.data["contents"]) + with open(os.path.join(tempfile.gettempdir(), filename), "w") as locustfile: + locustfile.write(msg.data["contents"]) + + def exit_handler(): + os.remove(locustfile.name) + + atexit.register(exit_handler) tempclient.close() - return filename + return locustfile.name def parse_locustfile_option(args=None) -> list[str]: From 8746fa885851e850dc17dc0cd909cef4c253ad9b Mon Sep 17 00:00:00 2001 From: Lars Holmberg Date: Sun, 11 Feb 2024 20:09:18 +0100 Subject: [PATCH 2/5] Send filename instead of full path from master when worker requests locustfile. --- locust/runners.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/locust/runners.py b/locust/runners.py index e03d1edf5b..702db4fe38 100644 --- a/locust/runners.py +++ b/locust/runners.py @@ -1064,7 +1064,7 @@ def client_listener(self) -> NoReturn: self.send_message( "locustfile", client_id=client_id, - data={"filename": filename, "contents": file_contents}, + data={"filename": os.path.basename(filename), "contents": file_contents}, ) continue elif msg.type == "client_stopped": From f4db9fbb533d8e9c9d2f7845f97532f877e37f70 Mon Sep 17 00:00:00 2001 From: Lars Holmberg Date: Sun, 11 Feb 2024 20:10:57 +0100 Subject: [PATCH 3/5] Unit test: Ensure no "Traceback" from workers. --- locust/test/test_main.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/locust/test/test_main.py b/locust/test/test_main.py index a02bd594b9..9981d65f7a 100644 --- a/locust/test/test_main.py +++ b/locust/test/test_main.py @@ -1721,12 +1721,15 @@ def t(self): text=True, ) stdout = proc.communicate()[0] - proc_worker2.communicate() - proc_worker.communicate() + stdout_worker = proc_worker.communicate()[0] + stdout_worker2 = proc_worker2.communicate()[0] self.assertIn('All users spawned: {"User1": 1} (1 total users)', stdout) self.assertIn("Locustfile contents changed on disk after first worker requested locustfile", stdout) self.assertIn("Shutting down (exit code 0)", stdout) + self.assertNotIn("Traceback", stdout) + self.assertNotIn("Traceback", stdout_worker) + self.assertNotIn("Traceback", stdout_worker2) self.assertEqual(0, proc.returncode) self.assertEqual(0, proc_worker.returncode) From db9f0771ef50713de0fd7cb88375e4f5fe3d16d6 Mon Sep 17 00:00:00 2001 From: Lars Holmberg Date: Sun, 11 Feb 2024 20:11:48 +0100 Subject: [PATCH 4/5] Rename download_file_from url to download_locustfile_from_url, more closely matching download_locustfile_from_master --- locust/argument_parser.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/locust/argument_parser.py b/locust/argument_parser.py index f27040c142..4094428f21 100644 --- a/locust/argument_parser.py +++ b/locust/argument_parser.py @@ -156,7 +156,7 @@ def is_url(url: str) -> bool: return False -def download_file_from_url(url: str) -> str: +def download_locustfile_from_url(url: str) -> str: try: response = requests.get(url) except requests.exceptions.RequestException as e: @@ -317,7 +317,7 @@ def parse_locustfile_option(args=None) -> list[str]: # Comma separated string to list locustfile_as_list = [ - download_file_from_url(f) if is_url(f.strip()) else f.strip() for f in options.locustfile.split(",") + download_locustfile_from_url(f) if is_url(f.strip()) else f.strip() for f in options.locustfile.split(",") ] # Checking if the locustfile is a single file, multiple files or a directory From 78d57ed071cd66d2af06b47319eefd4167acb8a7 Mon Sep 17 00:00:00 2001 From: Lars Holmberg Date: Sun, 11 Feb 2024 20:23:58 +0100 Subject: [PATCH 5/5] Catch the exception when deleting downloaded locustfiles if the file has already been deleted. --- locust/argument_parser.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/locust/argument_parser.py b/locust/argument_parser.py index 4094428f21..2bad65a2b4 100644 --- a/locust/argument_parser.py +++ b/locust/argument_parser.py @@ -167,7 +167,10 @@ def download_locustfile_from_url(url: str) -> str: locustfile.write(response.text) def exit_handler(): - os.remove(locustfile.name) + try: + os.remove(locustfile.name) + except FileNotFoundError: + pass # this is normal when multiple workers are running on the same machine atexit.register(exit_handler) return locustfile.name @@ -250,7 +253,10 @@ def wait_for_reply(): locustfile.write(msg.data["contents"]) def exit_handler(): - os.remove(locustfile.name) + try: + os.remove(locustfile.name) + except FileNotFoundError: + pass # this is normal when multiple workers are running on the same machine atexit.register(exit_handler)