tools/tasks.py

   1 # Copyright 2015 Intel Corporation.
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at
   6 #
   7 #   http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14
  15 """Task management helper functions and classes.
  16 """
  17
  18 import select
  19 import subprocess
  20 import logging
  21 import pexpect
  22 import threading
  23 import sys
  24 import os
  25 import locale
  26 import time
  27
  28 from conf import settings
  29 from tools import systeminfo
  30
  31
  32 CMD_PREFIX = 'cmd : '
  33
  34 def _get_stdout():
  35     """Get stdout value for ``subprocess`` calls.
  36     """
  37     stdout = None
  38
  39     if settings.getValue('VERBOSITY') != 'debug':
  40         stdout = open(os.devnull, 'wb')
  41
  42     return stdout
  43
  44
  45 def run_task(cmd, logger, msg=None, check_error=False):
  46     """Run task, report errors and log overall status.
  47
  48     Run given task using ``subprocess.Popen``. Log the commands
  49     used and any errors generated. Prints stdout to screen if
  50     in verbose mode and returns it regardless. Prints stderr to
  51     screen always.
  52
  53     :param cmd: Exact command to be executed
  54     :param logger: Logger to write details to
  55     :param msg: Message to be shown to user
  56     :param check_error: Throw exception on error
  57
  58     :returns: (stdout, stderr)
  59     """
  60     def handle_error(exception):
  61         """Handle errors by logging and optionally raising an exception.
  62         """
  63         logger.error(
  64             'Unable to execute %(cmd)s. Exception: %(exception)s',
  65             {'cmd': ' '.join(cmd), 'exception': exception})
  66         if check_error:
  67             raise exception
  68
  69     stdout = []
  70     stderr = []
  71     my_encoding = locale.getdefaultlocale()[1]
  72
  73     if msg:
  74         logger.info(msg)
  75
  76     logger.debug('%s%s', CMD_PREFIX, ' '.join(cmd))
  77
  78     try:
  79         proc = subprocess.Popen(map(os.path.expanduser, cmd),
  80                                 stdout=subprocess.PIPE,
  81                                 stderr=subprocess.PIPE, bufsize=0)
  82
  83         while True:
  84             reads = [proc.stdout.fileno(), proc.stderr.fileno()]
  85             ret = select.select(reads, [], [])
  86
  87             for file_d in ret[0]:
  88                 if file_d == proc.stdout.fileno():
  89                     line = proc.stdout.readline()
  90                     if settings.getValue('VERBOSITY') == 'debug':
  91                         sys.stdout.write(line.decode(my_encoding))
  92                     stdout.append(line)
  93                 if file_d == proc.stderr.fileno():
  94                     line = proc.stderr.readline()
  95                     sys.stderr.write(line.decode(my_encoding))
  96                     stderr.append(line)
  97
  98             if proc.poll() is not None:
  99                 break
 100     except OSError as ex:
 101         handle_error(ex)
 102     else:
 103         if proc.returncode:
 104             ex = subprocess.CalledProcessError(proc.returncode, cmd, stderr)
 105             handle_error(ex)
 106
 107     return ('\n'.join(sout.decode(my_encoding).strip() for sout in stdout),
 108             ('\n'.join(sout.decode(my_encoding).strip() for sout in stderr)))
 109
 110 def run_background_task(cmd, logger, msg):
 111     """Run task in background and log when started.
 112
 113     Run given task using ``subprocess.Popen``. Log the command
 114     used. Print stdout to screen if in verbose mode. Prints stderr
 115     to screen always.
 116
 117     :param cmd: Exact command to be executed
 118     :param logger: Logger to write details to
 119     :param msg: Message to be shown to user
 120
 121     :returns: Process PID
 122     """
 123     logger.info(msg)
 124     logger.debug('%s%s', CMD_PREFIX, ' '.join(cmd))
 125
 126     proc = subprocess.Popen(map(os.path.expanduser, cmd), stdout=_get_stdout(), bufsize=0)
 127
 128     return proc.pid
 129
 130
 131 def run_interactive_task(cmd, logger, msg):
 132     """Run a task interactively and log when started.
 133
 134     Run given task using ``pexpect.spawn``. Log the command used.
 135     Performs neither validation of the process - if the process
 136     successfully started or is still running - nor killing of the
 137     process. The user must do both.
 138
 139     :param cmd: Exact command to be executed
 140     :param logger: Logger to write details to
 141     :param msg: Message to be shown to user
 142
 143     :returns: ``pexpect.child`` object
 144     """
 145     logger.info(msg)
 146     logger.debug('%s%s', CMD_PREFIX, cmd)
 147     child = pexpect.spawnu(cmd)
 148
 149     if settings.getValue('VERBOSITY') == 'debug':
 150         child.logfile_read = sys.stdout
 151
 152     return child
 153
 154 def terminate_task_subtree(pid, signal='-15', sleep=10, logger=None):
 155     """Terminate given process and all its children
 156
 157     Function will sent given signal to the process. In case
 158     that process will not terminate within given sleep interval
 159     and signal was not SIGKILL, then process will be killed by SIGKILL.
 160     After that function will check if all children of the process
 161     are terminated and if not the same terminating procedure is applied
 162     on any living child (only one level of children is considered).
 163
 164     :param pid: Process ID to terminate
 165     :param signal: Signal to be sent to the process
 166     :param sleep: Maximum delay in seconds after signal is sent
 167     :param logger: Logger to write details to
 168     """
 169     try:
 170         output = subprocess.check_output("pgrep -P " + str(pid), shell=True).decode().rstrip('\n')
 171     except subprocess.CalledProcessError:
 172         output = ""
 173
 174     terminate_task(pid, signal, sleep, logger)
 175
 176     # just for case children were kept alive
 177     children = output.split('\n')
 178     for child in children:
 179         terminate_task(child, signal, sleep, logger)
 180
 181 def terminate_task(pid, signal='-15', sleep=10, logger=None):
 182     """Terminate process with given pid
 183
 184     Function will sent given signal to the process. In case
 185     that process will not terminate within given sleep interval
 186     and signal was not SIGKILL, then process will be killed by SIGKILL.
 187
 188     :param pid: Process ID to terminate
 189     :param signal: Signal to be sent to the process
 190     :param sleep: Maximum delay in seconds after signal is sent
 191     :param logger: Logger to write details to
 192     """
 193     if systeminfo.pid_isalive(pid):
 194         run_task(['sudo', 'kill', signal, str(pid)], logger)
 195         logger.debug('Wait for process %s to terminate after signal %s', pid, signal)
 196         for dummy in range(sleep):
 197             time.sleep(1)
 198             if not systeminfo.pid_isalive(pid):
 199                 break
 200
 201         if signal.lstrip('-').upper() not in ('9', 'KILL', 'SIGKILL') and systeminfo.pid_isalive(pid):
 202             terminate_task(pid, '-9', sleep, logger)
 203
 204 class Process(object):
 205     """Control an instance of a long-running process.
 206
 207     This is basically a context-manager wrapper around the
 208     ``run_interactive_task`` function above (with some extra helper
 209     functions).
 210     """
 211     _cmd = None
 212     _child = None
 213     _logfile = None
 214     _logger = logging.getLogger(__name__)
 215     _expect = None
 216     _timeout = -1
 217     _proc_name = 'unnamed process'
 218     _relinquish_thread = None
 219
 220     # context manager
 221
 222     def __enter__(self):
 223         """Start process instance using context manager.
 224         """
 225         self.start()
 226         return self
 227
 228     def __exit__(self, type_, value, traceback):
 229         """Shutdown process instance.
 230         """
 231         self.kill()
 232
 233     # startup/shutdown
 234
 235     def start(self):
 236         """Start process instance.
 237         """
 238         self._start_process()
 239         if self._timeout > 0:
 240             self._expect_process()
 241
 242     def _start_process(self):
 243         """Start process instance.
 244         """
 245         cmd = ' '.join(settings.getValue('SHELL_CMD') +
 246                        ['"%s"' % ' '.join(self._cmd)])
 247
 248         self._child = run_interactive_task(cmd, self._logger,
 249                                            'Starting %s...' % self._proc_name)
 250         self._child.logfile = open(self._logfile, 'w')
 251
 252     def expect(self, msg, timeout=None):
 253         """Expect string from process.
 254
 255         Expect string and die if not received.
 256
 257         :param msg: String to expect.
 258         :param timeout: Time to wait for string.
 259
 260         :returns: None
 261         """
 262         self._expect_process(msg, timeout)
 263
 264     def _expect_process(self, msg=None, timeout=None):
 265         """Expect string from process.
 266         """
 267         if not msg:
 268             msg = self._expect
 269         if not timeout:
 270             timeout = self._timeout
 271
 272         # we use exceptions rather than catching conditions in ``expect`` list
 273         # as we want to fail catastrophically after handling; there is likely
 274         # little we can do from within the scripts to fix issues such as
 275         # hugepages not being mounted
 276         try:
 277             self._child.expect([msg], timeout=timeout)
 278         except pexpect.EOF as exc:
 279             self._logger.critical(
 280                 'An error occurred. Please check the logs (%s) for more'
 281                 ' information. Exiting...', self._logfile)
 282             raise exc
 283         except pexpect.TIMEOUT as exc:
 284             self._logger.critical(
 285                 'Failed to execute in \'%d\' seconds. Please check the logs'
 286                 ' (%s) for more information. Exiting...',
 287                 timeout, self._logfile)
 288             self.kill()
 289             raise exc
 290         except (Exception, KeyboardInterrupt) as exc:
 291             self._logger.critical('General exception raised. Exiting...')
 292             self.kill()
 293             raise exc
 294
 295     def kill(self, signal='-15', sleep=10):
 296         """Kill process instance if it is alive.
 297
 298         :param signal: signal to be sent to the process
 299         :param sleep: delay in seconds after signal is sent
 300         """
 301         if self.is_running():
 302             terminate_task_subtree(self._child.pid, signal, sleep, self._logger)
 303
 304             if self.is_relinquished():
 305                 self._relinquish_thread.join()
 306
 307         self._logger.info(
 308             'Log available at %s', self._logfile)
 309
 310     def is_relinquished(self):
 311         """Returns True if process is relinquished.
 312
 313         If relinquished the process is no longer controllable and can
 314         only be killed.
 315
 316         :returns: True if process is relinquished, else False.
 317         """
 318         return self._relinquish_thread
 319
 320     def is_running(self):
 321         """Returns True if process is running.
 322
 323         :returns: True if process is running, else False
 324         """
 325         return self._child and self._child.isalive()
 326
 327     def _affinitize_pid(self, core, pid):
 328         """Affinitize a process with ``pid`` to ``core``.
 329
 330         :param core: Core to affinitize process to.
 331         :param pid: Process ID to affinitize.
 332
 333         :returns: None
 334         """
 335         run_task(['sudo', 'taskset', '-c', '-p', str(core),
 336                   str(pid)],
 337                  self._logger)
 338
 339     def affinitize(self, core):
 340         """Affinitize process to a specific ``core``.
 341
 342         :param core: Core to affinitize process to.
 343
 344         :returns: None
 345         """
 346         self._logger.info('Affinitizing process')
 347
 348         if self.is_running():
 349             self._affinitize_pid(core, self._child.pid)
 350
 351     class ContinueReadPrintLoop(threading.Thread):
 352         """Thread to read output from child and log.
 353
 354         Taken from: https://github.com/pexpect/pexpect/issues/90
 355         """
 356         def __init__(self, child):
 357             self.child = child
 358             threading.Thread.__init__(self)
 359
 360         def run(self):
 361             while True:
 362                 try:
 363                     self.child.read_nonblocking()
 364                 except (pexpect.EOF, pexpect.TIMEOUT):
 365                     break
 366
 367     def relinquish(self):
 368         """Relinquish control of process.
 369
 370         Give up control of application in order to ensure logging
 371         continues for the application. After relinquishing control it
 372         will no longer be possible to :func:`expect` anything.
 373
 374         This works around an issue described here:
 375
 376             https://github.com/pexpect/pexpect/issues/90
 377
 378         It is hoped that future versions of pexpect will avoid this
 379         issue.
 380         """
 381         self._relinquish_thread = self.ContinueReadPrintLoop(self._child)
 382         self._relinquish_thread.start()
 383
 384
 385 class CustomProcess(Process):
 386     """An sample implementation of ``Process``.
 387
 388     This is essentially a more detailed version of the
 389     ``run_interactive_task`` function that checks for process execution
 390     and kills the process (assuming use of the context manager).
 391     """
 392     def __init__(self, cmd, timeout, logfile, expect, name):
 393         """Initialise process state.
 394
 395         :param cmd: Command to execute.
 396         :param timeout: Time to wait for ``expect``.
 397         :param logfile: Path to logfile.
 398         :param expect: String to expect indicating startup. This is a
 399             regex and should be escaped as such.
 400         :param name: Name of process to use in logs.
 401
 402         :returns: None
 403         """
 404         self._cmd = cmd
 405         self._logfile = logfile
 406         self._expect = expect
 407         self._proc_name = name
 408         self._timeout = timeout