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