edq.testing.serverrunner

  1import argparse
  2import atexit
  3import logging
  4import signal
  5import subprocess
  6import time
  7import typing
  8
  9import edq.net.settings
 10import edq.util.dirent
 11
 12DEFAULT_SERVER_STARTUP_INITIAL_WAIT_SECS: float = 0.2
 13DEFAULT_STARTUP_WAIT_SECS: float = 10.0
 14SERVER_STOP_WAIT_SECS: float = 5.00
 15
 16DEFAULT_IDENTIFY_MAX_ATTEMPTS: int = 100
 17DEFAULT_IDENTIFY_WAIT_SECS: float = 0.25
 18
 19_logger = logging.getLogger(__name__)
 20
 21class ServerRunner():
 22    """
 23    A class for running an external HTTP server for some sort of larger process (like testing or generating data).
 24    """
 25
 26    def __init__(self,
 27            server: typing.Union[str, None] = None,
 28            server_start_command: typing.Union[str, None] = None,
 29            server_stop_command: typing.Union[str, None] = None,
 30            http_exchanges_out_dir: typing.Union[str, None] = None,
 31            server_output_path: typing.Union[str, None] = None,
 32            startup_initial_wait_secs: float = DEFAULT_SERVER_STARTUP_INITIAL_WAIT_SECS,
 33            startup_wait_secs: typing.Union[float, None] = None,
 34            startup_skip_identify: typing.Union[bool, None] = False,
 35            identify_max_attempts: int = DEFAULT_IDENTIFY_MAX_ATTEMPTS,
 36            identify_wait_secs: float = DEFAULT_IDENTIFY_WAIT_SECS,
 37            skip_restart: bool = False,
 38            **kwargs: typing.Any) -> None:
 39        if (server is None):
 40            raise ValueError('No server specified.')
 41
 42        self.server: str = server
 43        """ The server address to point requests to. """
 44
 45        if (server_start_command is None):
 46            raise ValueError('No command to start the server was specified.')
 47
 48        self.server_start_command: str = server_start_command
 49        """ The server_start_command to run the LMS server. """
 50
 51        self.server_stop_command: typing.Union[str, None] = server_stop_command
 52        """ An optional command to stop the server. """
 53
 54        if (http_exchanges_out_dir is None):
 55            http_exchanges_out_dir = edq.util.dirent.get_temp_dir(prefix = 'edq-serverrunner-http-exchanges-', rm = False)
 56
 57        self.http_exchanges_out_dir: str = http_exchanges_out_dir
 58        """ Where to output the HTTP exchanges. """
 59
 60        if (server_output_path is None):
 61            server_output_path = edq.util.dirent.get_temp_path(prefix = 'edq-serverrunner-server-output-', rm = False) + '.txt'
 62
 63        self.server_output_path: str = server_output_path
 64        """ Where to write server output (stdout and stderr). """
 65
 66        self.startup_initial_wait_secs: float = startup_initial_wait_secs
 67        """ The duration to wait after giving the initial startup command. """
 68
 69        if (startup_wait_secs is None):
 70            startup_wait_secs = DEFAULT_STARTUP_WAIT_SECS
 71
 72        self.startup_wait_secs = startup_wait_secs
 73        """ How long to wait after the server start command is run before making requests to the server. """
 74
 75        if (startup_skip_identify is None):
 76            startup_skip_identify = False
 77
 78        self.startup_skip_identify: bool = startup_skip_identify
 79        """
 80        Whether to skip trying to identify the server after it has been started.
 81        This acts as a way to have a variable wait for the server to start.
 82        When not used, self.startup_wait_secs is the only way to wait for the server to start.
 83        """
 84
 85        self.identify_max_attempts: int = identify_max_attempts
 86        """ The maximum number of times to try an identity check before starting the server. """
 87
 88        self.identify_wait_secs: float = identify_wait_secs
 89        """ The number of seconds each identify request will wait for the server to respond. """
 90
 91        self.skip_restart: bool = skip_restart
 92        """ While set, self.restart() should become a no-op. """
 93
 94        self._old_exchanges_out_dir: typing.Union[str, None] = None
 95        """
 96        The value of edq.net.settings.get_exchanges_out_dir() when start() is called.
 97        The original value may be changed in start(), and will be reset in stop().
 98        """
 99
100        self._process: typing.Union[subprocess.Popen, None] = None
101        """ The server process. """
102
103        self._server_output_file: typing.Union[typing.IO, None] = None
104        """ The file that server output is written to. """
105
106    def start(self) -> None:
107        """ Start the server. """
108
109        if (self._process is not None):
110            return
111
112        # Ensure stop() is called.
113        atexit.register(self.stop)
114
115        # Store and set networking config.
116
117        self._old_exchanges_out_dir = edq.net.settings.get_exchanges_out_dir()
118        edq.net.settings.set_exchanges_out_dir(self.http_exchanges_out_dir)
119
120        # Start the server.
121
122        _logger.info("Writing HTTP exchanges to '%s'.", self.http_exchanges_out_dir)
123        _logger.info("Writing server output to '%s'.", self.server_output_path)
124        _logger.info("Starting the server ('%s') and waiting for it.", self.server)
125
126        self._server_output_file = open(self.server_output_path, 'a', encoding = edq.util.dirent.DEFAULT_ENCODING)  # pylint: disable=consider-using-with
127
128        self._start_server()
129        _logger.info("Server is started up.")
130
131    def _start_server(self) -> None:
132        """ Start the server. """
133
134        if (self._process is not None):
135            return
136
137        self._process = subprocess.Popen(self.server_start_command,  # pylint: disable=consider-using-with
138                shell = True, stdout = self._server_output_file, stderr = subprocess.STDOUT)
139
140        status = None
141        try:
142            # Wait for a short period for the process to start.
143            status = self._process.wait(self.startup_initial_wait_secs)
144        except subprocess.TimeoutExpired:
145            # Good, the server is running.
146            pass
147
148        if (status is not None):
149            hint = f"code: '{status}'"
150            if (status == 125):
151                hint = 'server may already be running'
152
153            raise ValueError(f"Server was unable to start successfully ('{hint}').")
154
155        _logger.info("Completed initial server start wait.")
156
157        # Ping the server to check if it has started.
158        if (not self.startup_skip_identify):
159            for _ in range(self.identify_max_attempts):
160                if (self.identify_server()):
161                    # The server is running and responding, exit early.
162                    return
163
164                time.sleep(self.identify_wait_secs)
165
166        status = None
167        try:
168            # Ensure the server is running cleanly.
169            status = self._process.wait(self.startup_wait_secs)
170        except subprocess.TimeoutExpired:
171            # Good, the server is running.
172            pass
173
174        if (status is not None):
175            raise ValueError(f"Server was unable to start successfully ('code: {status}').")
176
177    def stop(self) -> bool:
178        """
179        Stop the server.
180        Return true if child classes should perform shutdown behavior.
181        """
182
183        if (self._process is None):
184            return False
185
186        # Stop the server.
187        _logger.info('Stopping the server.')
188        self._stop_server()
189
190        # Restore networking config.
191
192        edq.net.settings.set_exchanges_out_dir(self._old_exchanges_out_dir)
193        self._old_exchanges_out_dir = None
194
195        if (self._server_output_file is not None):
196            self._server_output_file.close()
197            self._server_output_file = None
198
199        return True
200
201    def restart(self) -> None:
202        """ Restart the server. """
203
204        if (self.skip_restart):
205            return
206
207        _logger.debug('Restarting the server.')
208        self._stop_server()
209        self._start_server()
210
211    def identify_server(self) ->  bool:
212        """
213        Attempt to identify the target server and return true on a successful attempt.
214        This is used on startup to wait for the server to complete startup.
215
216        Child classes must implement this or set self.startup_skip_identify to true.
217        """
218
219        raise NotImplementedError('identify_server')
220
221    def _stop_server(self) -> typing.Union[int, None]:
222        """ Stop the server process and return the exit status. """
223
224        if (self._process is None):
225            return None
226
227        # Mark the process as dead, so it can be restarted (if need be).
228        current_process = self._process
229        self._process = None
230
231        # Check if the process is already dead.
232        status = current_process.poll()
233        if (status is not None):
234            return status
235
236        # If the user provided a special command, try it.
237        if (self.server_stop_command is not None):
238            subprocess.run(self.server_stop_command,
239                    shell = True, stdout = subprocess.DEVNULL, stderr = subprocess.DEVNULL,
240                    check = False)
241
242        status = current_process.poll()
243        if (status is not None):
244            return status
245
246        # Try to end the server gracefully.
247        try:
248            current_process.send_signal(signal.SIGINT)
249            current_process.wait(SERVER_STOP_WAIT_SECS)
250        except subprocess.TimeoutExpired:
251            pass
252
253        status = current_process.poll()
254        if (status is not None):
255            return status
256
257        # End the server hard.
258        try:
259            current_process.kill()
260            current_process.wait(SERVER_STOP_WAIT_SECS)
261        except subprocess.TimeoutExpired:
262            pass
263
264        status = current_process.poll()
265        if (status is not None):
266            return status
267
268        return None
269
270def modify_parser(parser: argparse.ArgumentParser) -> None:
271    """ Modify the parser to add arguments for running a server. """
272
273    parser.add_argument('server_start_command', metavar = 'RUN_SERVER_COMMAND',
274        action = 'store', type = str,
275        help = 'The command to run the LMS server that will be the target of the data generation commands.')
276
277    group = parser.add_argument_group('server control')
278
279    group.add_argument('--server-output-file', dest = 'server_output_path',
280        action = 'store', type = str, default = None,
281        help = 'Where server output will be written. Defaults to a random temp file.')
282
283    group.add_argument('--server-stop-command', dest = 'server_stop_command',
284        action = 'store', type = str, default = None,
285        help = 'An optional command to stop the server. After this the server will be sent a SIGINT and then a SIGKILL.')
286
287    group.add_argument('--startup-skip-identify', dest = 'startup_skip_identify',
288        action = 'store_true', default = False,
289        help = 'If set, startup will skip trying to identify the server as a means of checking that the server is started.')
290
291    group.add_argument('--startup-wait', dest = 'startup_wait_secs',
292        action = 'store', type = float, default = DEFAULT_STARTUP_WAIT_SECS,
293        help = 'The time to wait between starting the server and sending commands (default: %(default)s).')