cli: BST_FORCE_START_METHOD only sets if necessary

Allow situations where the start method is already set, this enables us
to use this in testing situations.

Also, print a diagnostic if it's already set to something we didn't
want.

Now this block got more complex, split out into a new function.

Now we're using this string a lot, extract it to a variable, to make
sure we're spelling it correctly everywhere.

1 files changed, 39 insertions, 10 deletions

diff --git a/src/buildstream/_frontend/cli.py b/src/buildstream/_frontend/cli.py
index 60527e698..42d557615 100644
--- a/src/buildstream/_frontend/cli.py
+++ b/src/buildstream/_frontend/cli.py
@@ -202,6 +202,41 @@ def validate_output_streams():
             sys.exit(-1)
 
 
+def handle_bst_force_start_method_env():
+    bst_force_start_method_str = "BST_FORCE_START_METHOD"
+    if bst_force_start_method_str in os.environ:
+        start_method = os.environ[bst_force_start_method_str]
+        existing_start_method = multiprocessing.get_start_method(allow_none=True)
+        if existing_start_method is None:
+            multiprocessing.set_start_method(start_method)
+            print(
+                bst_force_start_method_str + ": multiprocessing start method forced to:",
+                start_method,
+                file=sys.stderr,
+                flush=True,
+            )
+        elif existing_start_method == start_method:
+            # Note that when testing, we run the buildstream entrypoint
+            # multiple times in the same executable, so guard against that
+            # here.
+            print(
+                bst_force_start_method_str + ": multiprocessing start method already set to:",
+                existing_start_method,
+                file=sys.stderr,
+                flush=True,
+            )
+        else:
+            print(
+                bst_force_start_method_str + ": cannot set multiprocessing start method to:",
+                start_method,
+                ", already set to:",
+                existing_start_method,
+                file=sys.stderr,
+                flush=True,
+            )
+            sys.exit(-1)
+
+
 def override_main(self, args=None, prog_name=None, complete_var=None,
                   standalone_mode=True, **extra):
 
@@ -229,16 +264,10 @@ def override_main(self, args=None, prog_name=None, complete_var=None,
     validate_output_streams()
 
     # We can only set the global multiprocessing start method once; for that
-    # reason we're advised to do it inside the entrypoint, where it is easy to
-    # ensure the code path is only followed once.
-    if 'BST_FORCE_START_METHOD' in os.environ:
-        multiprocessing.set_start_method(os.environ['BST_FORCE_START_METHOD'])
-        print(
-            "BST_FORCE_START_METHOD: multiprocessing start method forced to:",
-            os.environ['BST_FORCE_START_METHOD'],
-            file=sys.stderr,
-            flush=True,
-        )
+    # reason we're advised to do it inside the entrypoint, where it's more
+    # likely that you can ensure the code path is only followed once. In the
+    # case of testing, our tests preceed our entrypoint, so we do our best.
+    handle_bst_force_start_method_env()
 
     original_main(self, args=args, prog_name=prog_name, complete_var=None,
                   standalone_mode=standalone_mode, **extra)