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                     while True:
  90                         line = proc.stdout.readline()
  91                         if not line:
  92                             break
  93                         if settings.getValue('VERBOSITY') == 'debug':
  94                             sys.stdout.write(line.decode(my_encoding))
  95                         stdout.append(line)
  96                 if file_d == proc.stderr.fileno():
  97                     while True:
  98                         line = proc.stderr.readline()
  99                         if not line:
 100                             break
 101                         sys.stderr.write(line.decode(my_encoding))
 102                         stderr.append(line)
 103
 104             if proc.poll() is not None:
 105                 break
 106
 107     except OSError as ex:
 108         handle_error(ex)
 109     else:
 110         if proc.returncode:
 111             ex = subprocess.CalledProcessError(proc.returncode, cmd, stderr)
 112             handle_error(ex)
 113
 114     return ('\n'.join(sout.decode(my_encoding).strip() for sout in stdout),
 115             ('\n'.join(sout.decode(my_encoding).strip() for sout in stderr)))
 116
 117 def run_background_task(cmd, logger, msg):
 118     """Run task in background and log when started.
 119
 120     Run given task using ``subprocess.Popen``. Log the command
 121     used. Print stdout to screen if in verbose mode. Prints stderr
 122     to screen always.
 123
 124     :param cmd: Exact command to be executed
 125     :param logger: Logger to write details to
 126     :param msg: Message to be shown to user
 127
 128     :returns: Process PID
 129     """
 130     logger.info(msg)
 131     logger.debug('%s%s', CMD_PREFIX, ' '.join(cmd))
 132
 133     proc = subprocess.Popen(map(os.path.expanduser, cmd), stdout=_get_stdout(), bufsize=0)
 134
 135     return proc.pid
 136
 137
 138 def run_interactive_task(cmd, logger, msg):
 139     """Run a task interactively and log when started.
 140
 141     Run given task using ``pexpect.spawn``. Log the command used.
 142     Performs neither validation of the process - if the process
 143     successfully started or is still running - nor killing of the
 144     process. The user must do both.
 145
 146     :param cmd: Exact command to be executed
 147     :param logger: Logger to write details to
 148     :param msg: Message to be shown to user
 149
 150     :returns: ``pexpect.child`` object
 151     """
 152     logger.info(msg)
 153     logger.debug('%s%s', CMD_PREFIX, cmd)
 154     child = pexpect.spawnu(cmd)
 155
 156     if settings.getValue('VERBOSITY') == 'debug':
 157         child.logfile_read = sys.stdout
 158
 159     return child
 160
 161 def terminate_task_subtree(pid, signal='-15', sleep=10, logger=None):
 162     """Terminate given process and all its children
 163
 164     Function will sent given signal to the process. In case
 165     that process will not terminate within given sleep interval
 166     and signal was not SIGKILL, then process will be killed by SIGKILL.
 167     After that function will check if all children of the process
 168     are terminated and if not the same terminating procedure is applied
 169     on any living child (only one level of children is considered).
 170
 171     :param pid: Process ID to terminate
 172     :param signal: Signal to be sent to the process
 173     :param sleep: Maximum delay in seconds after signal is sent
 174     :param logger: Logger to write details to
 175     """
 176     try:
 177         output = subprocess.check_output("pgrep -P " + str(pid), shell=True).decode().rstrip('\n')
 178     except subprocess.CalledProcessError:
 179         output = ""
 180
 181     terminate_task(pid, signal, sleep, logger)
 182
 183     # just for case children were kept alive
 184     children = output.split('\n')
 185     for child in children:
 186         terminate_task(child, signal, sleep, logger)
 187
 188 def terminate_task(pid, signal='-15', sleep=10, logger=None):
 189     """Terminate process with given pid
 190
 191     Function will sent given signal to the process. In case
 192     that process will not terminate within given sleep interval
 193     and signal was not SIGKILL, then process will be killed by SIGKILL.
 194
 195     :param pid: Process ID to terminate
 196     :param signal: Signal to be sent to the process
 197     :param sleep: Maximum delay in seconds after signal is sent
 198     :param logger: Logger to write details to
 199     """
 200     if systeminfo.pid_isalive(pid):
 201         run_task(['sudo', 'kill', signal, str(pid)], logger)
 202         logger.debug('Wait for process %s to terminate after signal %s', pid, signal)
 203         for dummy in range(sleep):
 204             time.sleep(1)
 205             if not systeminfo.pid_isalive(pid):
 206                 break
 207
 208         if signal.lstrip('-').upper() not in ('9', 'KILL', 'SIGKILL') and systeminfo.pid_isalive(pid):
 209             terminate_task(pid, '-9', sleep, logger)
 210
 211 class Process(object):
 212     """Control an instance of a long-running process.
 213
 214     This is basically a context-manager wrapper around the
 215     ``run_interactive_task`` function above (with some extra helper
 216     functions).
 217     """
 218     _cmd = None
 219     _child = None
 220     _logfile = None
 221     _logger = logging.getLogger(__name__)
 222     _expect = None
 223     _timeout = -1
 224     _proc_name = 'unnamed process'
 225     _relinquish_thread = None
 226
 227     # context manager
 228
 229     def __enter__(self):
 230         """Start process instance using context manager.
 231         """
 232         self.start()
 233         return self
 234
 235     def __exit__(self, type_, value, traceback):
 236         """Shutdown process instance.
 237         """
 238         self.kill()
 239
 240     # startup/shutdown
 241
 242     def start(self):
 243         """Start process instance.
 244         """
 245         self._start_process()
 246         if self._timeout > 0:
 247             self._expect_process()
 248
 249     def _start_process(self):
 250         """Start process instance.
 251         """
 252         cmd = ' '.join(settings.getValue('SHELL_CMD') +
 253                        ['"%s"' % ' '.join(self._cmd)])
 254
 255         self._child = run_interactive_task(cmd, self._logger,
 256                                            'Starting %s...' % self._proc_name)
 257         self._child.logfile = open(self._logfile, 'w')
 258
 259     def expect(self, msg, timeout=None):
 260         """Expect string from process.
 261
 262         Expect string and die if not received.
 263
 264         :param msg: String to expect.
 265         :param timeout: Time to wait for string.
 266
 267         :returns: None
 268         """
 269         self._expect_process(msg, timeout)
 270
 271     def _expect_process(self, msg=None, timeout=None):
 272         """Expect string from process.
 273         """
 274         if not msg:
 275             msg = self._expect
 276         if not timeout:
 277             timeout = self._timeout
 278
 279         # we use exceptions rather than catching conditions in ``expect`` list
 280         # as we want to fail catastrophically after handling; there is likely
 281         # little we can do from within the scripts to fix issues such as
 282         # hugepages not being mounted
 283         try:
 284             self._child.expect([msg], timeout=timeout)
 285         except pexpect.EOF as exc:
 286             self._logger.critical(
 287                 'An error occurred. Please check the logs (%s) for more'
 288                 ' information. Exiting...', self._logfile)
 289             raise exc
 290         except pexpect.TIMEOUT as exc:
 291             self._logger.critical(
 292                 'Failed to execute in \'%d\' seconds. Please check the logs'
 293                 ' (%s) for more information. Exiting...',
 294                 timeout, self._logfile)
 295             self.kill()
 296             raise exc
 297         except (Exception, KeyboardInterrupt) as exc:
 298             self._logger.critical('General exception raised. Exiting...')
 299             self.kill()
 300             raise exc
 301
 302     def kill(self, signal='-15', sleep=10):
 303         """Kill process instance if it is alive.
 304
 305         :param signal: signal to be sent to the process
 306         :param sleep: delay in seconds after signal is sent
 307         """
 308         if self.is_running():
 309             terminate_task_subtree(self._child.pid, signal, sleep, self._logger)
 310
 311             if self.is_relinquished():
 312                 self._relinquish_thread.join()
 313
 314         self._logger.info(
 315             'Log available at %s', self._logfile)
 316
 317     def is_relinquished(self):
 318         """Returns True if process is relinquished.
 319
 320         If relinquished the process is no longer controllable and can
 321         only be killed.
 322
 323         :returns: True if process is relinquished, else False.
 324         """
 325         return self._relinquish_thread
 326
 327     def is_running(self):
 328         """Returns True if process is running.
 329
 330         :returns: True if process is running, else False
 331         """
 332         return self._child and self._child.isalive()
 333
 334     def _affinitize_pid(self, core, pid):
 335         """Affinitize a process with ``pid`` to ``core``.
 336
 337         :param core: Core to affinitize process to.
 338         :param pid: Process ID to affinitize.
 339
 340         :returns: None
 341         """
 342         run_task(['sudo', 'taskset', '-c', '-p', str(core),
 343                   str(pid)],
 344                  self._logger)
 345
 346     def affinitize(self, core):
 347         """Affinitize process to a specific ``core``.
 348
 349         :param core: Core to affinitize process to.
 350
 351         :returns: None
 352         """
 353         self._logger.info('Affinitizing process')
 354
 355         if self.is_running():
 356             self._affinitize_pid(core, self._child.pid)
 357
 358     class ContinueReadPrintLoop(threading.Thread):
 359         """Thread to read output from child and log.
 360
 361         Taken from: https://github.com/pexpect/pexpect/issues/90
 362         """
 363         def __init__(self, child):
 364             self.child = child
 365             threading.Thread.__init__(self)
 366
 367         def run(self):
 368             while True:
 369                 try:
 370                     self.child.read_nonblocking()
 371                 except (pexpect.EOF, pexpect.TIMEOUT):
 372                     break
 373
 374     def relinquish(self):
 375         """Relinquish control of process.
 376
 377         Give up control of application in order to ensure logging
 378         continues for the application. After relinquishing control it
 379         will no longer be possible to :func:`expect` anything.
 380
 381         This works around an issue described here:
 382
 383             https://github.com/pexpect/pexpect/issues/90
 384
 385         It is hoped that future versions of pexpect will avoid this
 386         issue.
 387         """
 388         self._relinquish_thread = self.ContinueReadPrintLoop(self._child)
 389         self._relinquish_thread.start()
 390
 391
 392 class CustomProcess(Process):
 393     """An sample implementation of ``Process``.
 394
 395     This is essentially a more detailed version of the
 396     ``run_interactive_task`` function that checks for process execution
 397     and kills the process (assuming use of the context manager).
 398     """
 399     def __init__(self, cmd, timeout, logfile, expect, name):
 400         """Initialise process state.
 401
 402         :param cmd: Command to execute.
 403         :param timeout: Time to wait for ``expect``.
 404         :param logfile: Path to logfile.
 405         :param expect: String to expect indicating startup. This is a
 406             regex and should be escaped as such.
 407         :param name: Name of process to use in logs.
 408
 409         :returns: None
 410         """
 411         self._cmd = cmd
 412         self._logfile = logfile
 413         self._expect = expect
 414         self._proc_name = name
 415         self._timeout = timeout