edq.testing.serverrunner

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