Implement dynamic display port selection for xvfb process

Author: diningPhilosopher64

jupyter_matlab_proxy/app_state.py

@@ -303,7 +303,6 @@ def reserve_matlab_port(self):
             # try-except.
             # s.bind(("", 0))
             # self.matlab_port = s.getsockname()[1]
-
             for port in mw.range_matlab_connector_ports():
                 try:
                     s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
@@ -358,22 +357,11 @@ async def start_matlab(self, restart=False):
         except FileNotFoundError:
             pass
 
-        # The presence of xvfb_ready_file indicates if Xvfb is ready to receive
-        # connections, but this could be leftover from a terminated Xvfb, so ensure it
-        # is cleaned up before starting Xvfb
-        display_num = self.settings["matlab_display"].replace(":", "")
-        xvfb_ready_file = Path(tempfile.gettempdir()) / ".X11-unix" / f"X{display_num}"
-        try:
-            xvfb_ready_file.unlink()
-        except FileNotFoundError:
-            pass
-
         # Configure the environment MATLAB needs to start
         matlab_env = os.environ.copy()
         matlab_env["MW_CRASH_MODE"] = "native"
         matlab_env["MATLAB_WORKER_CONFIG_ENABLE_LOCAL_PARCLUSTER"] = "true"
         matlab_env["PCT_ENABLED"] = "true"
-        matlab_env["DISPLAY"] = self.settings["matlab_display"]
         matlab_env["HTTP_MATLAB_CLIENT_GATEWAY_PUBLIC_PORT"] = "1"
         matlab_env["MW_CONNECTOR_SECURE_PORT"] = str(self.matlab_port)
         matlab_env["MW_DOCROOT"] = str(
@@ -410,18 +398,21 @@ async def start_matlab(self, restart=False):
         # matlab_env["CONNECTOR_CONFIGURABLE_WARMUP_TASKS"] = "warmup_hgweb"
         # matlab_env["CONNECTOR_WARMUP"] = "true"
 
-        logger.debug(f"Starting Xvfb on display {self.settings['matlab_display']}")
-        xvfb = await asyncio.create_subprocess_exec(
-            *self.settings["xvfb_cmd"], env=matlab_env
-        )
+        # Start Xvfb process
+        create_xvfb_cmd = self.settings["create_xvfb_cmd"]
+        xvfb_cmd, dpipe = create_xvfb_cmd()
+
+        xvfb, display_port = await mw.create_xvfb_process(xvfb_cmd, dpipe, matlab_env)
+
+        # Update settings and matlab_env dict
+        self.settings["matlab_display"] = ":" + str(display_port)
         self.processes["xvfb"] = xvfb
-        logger.debug(f"Started Xvfb (PID={xvfb.pid})")
 
-        # Wait for Xvfb to be ready
-        while not xvfb_ready_file.exists():
-            logger.debug(f"Waiting for XVFB")
-            await asyncio.sleep(0.1)
+        matlab_env["DISPLAY"] = self.settings["matlab_display"]
+
+        logger.debug(f"Started Xvfb with PID={xvfb.pid} on DISPLAY={display_port}")
 
+        # Start MATLAB Process
         logger.info(f"Starting MATLAB on port {self.matlab_port}")
         master, slave = pty.openpty()
         matlab = await asyncio.create_subprocess_exec(

jupyter_matlab_proxy/devel.py

@@ -3,10 +3,7 @@
 # Development specific functions
 import asyncio, aiohttp
 from aiohttp import web
-import socket
-import time
-import os
-import sys
+import socket, time, os, sys
 
 desktop_html = b"""
 <h1>Fake MATLAB Web Desktop</h1>
@@ -39,17 +36,6 @@ def wait_for_port(port):
         break
 
 
-def xvfb(args):
-    "Writes ready file after 1 second then sleeps forever in a loose loop"
-    time.sleep(1)
-    print(f"Creating {str(args.ready_file)}")
-    args.ready_file.parent.mkdir(parents=True, exist_ok=True)
-    args.ready_file.touch()
-
-    while True:
-        time.sleep(60)
-
-
 async def web_handler(request):
     return web.Response(content_type="text/html", body=desktop_html)
 
@@ -147,9 +133,6 @@ def matlab(args):
 
     parser = argparse.ArgumentParser()
     subparsers = parser.add_subparsers(dest="cmd", required=True)
-    xvfb_parser = subparsers.add_parser("xvfb")
-    xvfb_parser.add_argument("--ready-file", type=Path)
-    xvfb_parser.set_defaults(func=xvfb)
     matlab_parser = subparsers.add_parser("matlab")
     matlab_parser.add_argument(
         "--ready-file",

jupyter_matlab_proxy/settings.py

@@ -4,7 +4,7 @@
 from pathlib import Path
 import tempfile
 import xml.etree.ElementTree as ET
-import uuid
+import uuid, socket
 import shutil
 
 
@@ -49,14 +49,7 @@ def get_dev_settings():
             "--ready-file",
             str(matlab_ready_file),
         ],
-        "xvfb_cmd": [
-            "python",
-            "-u",
-            str(devel_file),
-            "xvfb",
-            "--ready-file",
-            str(Path(tempfile.gettempdir()) / ".X11-unix" / "X1"),
-        ],
+        "create_xvfb_cmd": create_xvfb_cmd,
         "matlab_ready_file": matlab_ready_file,
         "base_url": os.environ.get("BASE_URL", ""),
         "app_port": os.environ.get("APP_PORT", 8000),
@@ -99,6 +92,12 @@ def get(dev=False):
             matlab_cmd[4:4] = ready_delay
             settings["matlab_cmd"] = matlab_cmd
 
+            # Randomly picks an available port and updates the value of settings['app_port'] .
+            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            s.bind(("", 0))
+            settings["app_port"] = s.getsockname()[1]
+            s.close()
+
             # Set NLM Connection string. Server will start using this connection string for licensing
             settings["nlm_conn_str"] = "abc@nlm"
 
@@ -122,32 +121,49 @@ def get(dev=False):
                 "-r",
                 f"try; run('{matlab_startup_file}'); catch; end;",
             ],
-            "xvfb_cmd": [
-                "Xvfb",
-                ":1",
-                "-screen",
-                "0",
-                "1600x1200x24",
-                "-dpi",
-                "100",
-                # "-ac",
-                "-extension",
-                "RANDR",
-                # "+render",
-                # "-noreset",
-            ],
+            "create_xvfb_cmd": create_xvfb_cmd,
             "matlab_ready_file": Path(tempfile.gettempdir()) / "connector.securePort",
             "base_url": os.environ["BASE_URL"],
             "app_port": os.environ["APP_PORT"],
             "host_interface": os.environ.get("APP_HOST"),
             "mwapikey": str(uuid.uuid4()),
             "matlab_protocol": "https",
-            # TODO: Uncomment this ?
-            # "matlab_display": ":1",
             "nlm_conn_str": os.environ.get("MLM_LICENSE_FILE"),
             "matlab_config_file": Path.home() / ".matlab" / "proxy_app_config.json",
             "ws_env": ws_env,
             "mwa_api_endpoint": f"https://login{ws_env_suffix}.mathworks.com/authenticationws/service/v4",
             "mhlm_api_endpoint": f"https://licensing{ws_env_suffix}.mathworks.com/mls/service/v1/entitlement/list",
             "mwa_login": f"https://login{ws_env_suffix}.mathworks.com",
         }
+
+
+def create_xvfb_cmd():
+    """Creates the Xvfb command with a write descriptor.
+
+    Returns:
+        List: Containing 2 lists.
+
+    The second List contains a read and a write descriptor.
+    The first List is the command to launch Xvfb process with the same write descriptor(from the first list) embedded in the command.
+    """
+    dpipe = os.pipe()
+    os.set_inheritable(dpipe[1], True)
+
+    xvfb_cmd = [
+        "Xvfb",
+        "-displayfd",
+        # Write descriptor
+        str(dpipe[1]),
+        "-screen",
+        "0",
+        "1600x1200x24",
+        "-dpi",
+        "100",
+        # "-ac",
+        "-extension",
+        "RANDR",
+        # "+render",
+        # "-noreset",
+    ]
+
+    return xvfb_cmd, dpipe

jupyter_matlab_proxy/util/mw.py

@@ -1,7 +1,7 @@
 # Copyright 2020-2021 The MathWorks, Inc.
 
 import xml.etree.ElementTree as ET
-import aiohttp
+import aiohttp, os, asyncio, select, logging
 from .exceptions import (
     OnlineLicensingError,
     EntitlementError,
@@ -11,6 +11,7 @@
 
 
 LICENSING_URL = "https://github.com/mathworks/jupyter-matlab-proxy/blob/main/MATLAB_Licensing_Info.md"
+logger = logging.getLogger("MATLABProxyApp")
 
 
 async def fetch_entitlements(mhlm_api_endpoint, access_token, matlab_release):
@@ -179,3 +180,63 @@ def parse_other_error(logs):
         "MATLAB returned an unexpected error. For more details, see the log below.",
         logs=logs,
     )
+
+
+async def create_xvfb_process(xvfb_cmd, pipe, matlab_env={}):
+    """Creates the Xvfb process.
+
+    The Xvfb process is run with '-displayfd' flag set. This makes Xvfb choose an available
+    display number and write it into the provided write descriptor. ie: pipe[1]
+
+    We read this display number from the read descriptor.
+    For this, we have 2 things to consider:
+        1. How many bytes to read from the read descriptor. This is handled by the variable number_of_bytes.
+
+        2. For how long to read from the read descriptor. We wait for atmost 10 seconds for Xvfb to write
+           the display number using the 'select' package.
+
+
+    Args:
+        xvfb_cmd (List): A list containing the command to run the Xvfb process
+        pipe (List): A list containing a pair of file descriptor.
+        matlab_env (Dict): A Dict containing environment variables within which the Xvfb process is created.
+
+    Returns:
+        [List]: Containing the Xvfb process object, and display number on which Xvfb process has started.
+    """
+
+    # Creates subprocess asynchronously with environment variables defined in matlab_env
+    # Pipe errors, if any, to the process object instead of stdout.
+    xvfb = await asyncio.create_subprocess_exec(
+        *xvfb_cmd, close_fds=False, env=matlab_env, stderr=asyncio.subprocess.PIPE
+    )
+
+    read_descriptor, write_descriptor = pipe
+    number_of_bytes = 200
+
+    logger.debug("Waiting for XVFB process to initialize and provide Display Number")
+    # Waits upto 10 seconds for the read_descriptor to be ready.
+    ready_descriptors, _, _ = select.select([read_descriptor], [], [], 10)
+
+    # If read_descriptor is in ready_descriptors, read from it.
+    if read_descriptor in ready_descriptors:
+        logger.debug("Reading display number from the read descriptor.")
+        line = os.read(read_descriptor, number_of_bytes).decode("utf-8")
+        # Xvfb process writes the display number and adds a new line character ('\n') at the end. Removing it with .strip()
+        display_port = line.strip()
+
+    else:
+        # Check for errors and raise exception.
+        error = ""
+        while not xvfb.stderr.at_eof():
+            line = await xvfb.stderr.readline()
+            error += line.decode("utf-8")
+
+        await xvfb.wait()
+        raise Exception(f"Unable to start the Xvfb process: \n {error}")
+
+    # Close the read and write descriptors.
+    os.close(read_descriptor)
+    os.close(write_descriptor)
+
+    return xvfb, display_port

tests/test_app.py

@@ -1,6 +1,5 @@
 # Copyright 2021 The MathWorks, Inc.
 
-from jupyter_matlab_proxy.devel import xvfb
 import pytest, asyncio, aiohttp, os, json, psutil, socket, subprocess, time, requests
 from unittest.mock import patch
 from aiohttp import web
@@ -138,74 +137,32 @@ def test_marshal_error(actual_error, expected_error):
 
 
 class FakeServer:
-    """Context Manager class which returns a web server wrapped by aiohttp_client
-    for sending HTTP requests during testing.
-
-    executes the remove_zombie_matlab_process() method before starting the server
-    and after shutting it down so as to clear out any
+    """Context Manager class which returns a web server wrapped in aiohttp_client pytest fixture
+    for testing.
     """
 
     def __init__(self, loop, aiohttp_client):
         self.loop = loop
         self.aiohttp_client = aiohttp_client
-        self.pretest_dev_processes = None
-        self.posttest_dev_processes = None
-        self.zombie_dev_processes = None
 
     def __enter__(self):
-
-        self.pretest_dev_processes = set(self.gather_running_dev_processes())
-
-        # self.patcher = patch(
-        #     "jupyter_matlab_proxy.app_state.AppState.reserve_matlab_port",
-        #     new=MockServerPort.mock_reserve_port,
-        # )
-
-        # self.patcher.start()
-
         self.server = app.create_app()
         self.runner = web.AppRunner(self.server)
-
         self.loop.run_until_complete(self.runner.setup())
 
         self.site = web.TCPSite(
             self.runner,
             host=self.server["settings"]["host_interface"],
             port=self.server["settings"]["app_port"],
         )
-
         self.loop.run_until_complete(self.site.start())
+
         return self.loop.run_until_complete(self.aiohttp_client(self.server))
 
     def __exit__(self, exc_type, exc_value, exc_traceback):
-
         self.loop.run_until_complete(self.runner.shutdown())
         self.loop.run_until_complete(self.runner.cleanup())
 
-        # self.patcher.stop()
-
-        self.posttest_dev_processes = set(self.gather_running_dev_processes())
-        self.zombie_dev_processes = (
-            self.posttest_dev_processes - self.pretest_dev_processes
-        )
-
-        for process in self.zombie_dev_processes:
-            process.terminate()
-
-        gone, alive = psutil.wait_procs(self.zombie_dev_processes)
-
-        for process in alive:
-            process.kill()
-
-    def gather_running_dev_processes(self):
-        running_dev_processes = []
-        for process in psutil.process_iter(["pid", "name"]):
-            cmd = process.cmdline()
-            if len(cmd) > 3 and "devel.py" in cmd[2]:
-                running_dev_processes.append(process)
-
-        return running_dev_processes
-
 
 @pytest.fixture(name="test_server")
 def test_server_fixture(
@@ -255,11 +212,6 @@ async def test_start_matlab_route(test_server):
         assert resp.status == 200
 
         resp_json = json.loads(await resp.text())
-        # print(test_server.app["state"].settings)
-        # print(resp)
-        # print(resp_json)
-        # break
-
         if resp_json["matlab"]["status"] == "up":
             break
         else:
@@ -336,7 +288,8 @@ async def test_matlab_proxy_404(proxy_payload, test_server):
 
     headers = {"content-type": "application/json"}
 
-    # Request a non-existing html file. Request gets proxied to app.matlab_view() which should raise HTTPNotFound() exception ie. return HTTP status code 404
+    # Request a non-existing html file.
+    # Request gets proxied to app.matlab_view() which should raise HTTPNotFound() exception ie. return HTTP status code 404
     resp = await test_server.post(
         "./1234.html", data=json.dumps(proxy_payload), headers=headers
     )