Merge pull request #1585 from solowalker27/web-ui-improvements

Start web_ui later to avoid race adding UI routes

locust/env.py

@@ -143,7 +143,14 @@ def create_worker_runner(self, master_host, master_port):
         )
 
     def create_web_ui(
-        self, host="", port=8089, auth_credentials=None, tls_cert=None, tls_key=None, stats_csv_writer=None
+        self,
+        host="",
+        port=8089,
+        auth_credentials=None,
+        tls_cert=None,
+        tls_key=None,
+        stats_csv_writer=None,
+        delayed_start=False,
     ):
         """
         Creates a :class:`WebUI <locust.web.WebUI>` instance for this Environment and start running the web server
@@ -157,6 +164,8 @@ def create_web_ui(
         :param tls_key: An optional path (str) to a TLS private key. If this is provided the web UI will be
                         served over HTTPS
         :param stats_csv_writer: `StatsCSV <stats_csv.StatsCSV>` instance.
+        :param delayed_start: Whether or not to delay starting web UI until `start()` is called. Delaying web UI start
+                              allows for adding Flask routes or Blueprints before accepting requests, avoiding errors.
         """
         self.web_ui = WebUI(
             self,
@@ -166,6 +175,7 @@ def create_web_ui(
             tls_cert=tls_cert,
             tls_key=tls_key,
             stats_csv_writer=stats_csv_writer,
+            delayed_start=delayed_start,
         )
         return self.web_ui
 

locust/main.py

@@ -298,19 +298,22 @@ def timelimit_stop():
                 tls_cert=options.tls_cert,
                 tls_key=options.tls_key,
                 stats_csv_writer=stats_csv_writer,
+                delayed_start=True,
             )
         except AuthCredentialsError:
             logger.error("Credentials supplied with --web-auth should have the format: username:password")
             sys.exit(1)
-        else:
-            main_greenlet = web_ui.greenlet
     else:
         web_ui = None
 
     # Fire locust init event which can be used by end-users' code to run setup code that
     # need access to the Environment, Runner or WebUI.
     environment.events.init.fire(environment=environment, runner=runner, web_ui=web_ui)
 
+    if web_ui:
+        web_ui.start()
+        main_greenlet = web_ui.greenlet
+
     if options.headless:
         # headless mode
         if options.master:

locust/web.py

@@ -67,7 +67,15 @@ def my_custom_route():
     extending index.html."""
 
     def __init__(
-        self, environment, host, port, auth_credentials=None, tls_cert=None, tls_key=None, stats_csv_writer=None
+        self,
+        environment,
+        host,
+        port,
+        auth_credentials=None,
+        tls_cert=None,
+        tls_key=None,
+        stats_csv_writer=None,
+        delayed_start=False,
     ):
         """
         Create WebUI instance and start running the web server in a separate greenlet (self.greenlet)
@@ -80,6 +88,8 @@ def __init__(
                            Should be supplied in the format: "user:pass".
         tls_cert: A path to a TLS certificate
         tls_key: A path to a TLS private key
+        delayed_start: Whether or not to delay starting web UI until `start()` is called. Delaying web UI start
+                       allows for adding Flask routes or Blueprints before accepting requests, avoiding errors.
         """
         environment.web_ui = self
         self.stats_csv_writer = stats_csv_writer or StatsCSV(environment, stats_module.PERCENTILES_TO_REPORT)
@@ -110,6 +120,8 @@ def __init__(
                 )
         if environment.runner:
             self.update_template_args()
+        if not delayed_start:
+            self.start()
 
         @app.route("/")
         @self.auth_required_if_enabled
@@ -372,10 +384,11 @@ def exceptions_csv():
 
             return _download_csv_response(data.getvalue(), "exceptions")
 
-        self.greenlet = gevent.spawn(self.start)
+    def start(self):
+        self.greenlet = gevent.spawn(self.start_server)
         self.greenlet.link_exception(greenlet_exception_handler)
 
-    def start(self):
+    def start_server(self):
         if self.tls_cert and self.tls_key:
             self.server = pywsgi.WSGIServer(
                 (self.host, self.port), self.app, log=None, keyfile=self.tls_key, certfile=self.tls_cert