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