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):
 245         """Kill process instance if it is alive.
 246         """
 247         if self._child and self._child.isalive():
 248             run_task(['sudo', 'kill', '-15', str(self._child.pid)],
 249                      self._logger)
 250             self._logger.debug('Wait for process to terminate')
 251             time.sleep(2)
 252
 253             if self.is_relinquished():
 254                 self._relinquish_thread.join()
 255
 256         self._logger.info(
 257             'Log available at %s', self._logfile)
 258
 259     def is_relinquished(self):
 260         """Returns True if process is relinquished.
 261
 262         If relinquished the process is no longer controllable and can
 263         only be killed.
 264
 265         :returns: True if process is relinquished, else False.
 266         """
 267         return self._relinquish_thread
 268
 269     def is_running(self):
 270         """Returns True if process is running.
 271
 272         :returns: True if process is running, else False
 273         """
 274         return self._child is not None
 275
 276     def _affinitize_pid(self, core, pid):
 277         """Affinitize a process with ``pid`` to ``core``.
 278
 279         :param core: Core to affinitize process to.
 280         :param pid: Process ID to affinitize.
 281
 282         :returns: None
 283         """
 284         run_task(['sudo', 'taskset', '-c', '-p', str(core),
 285                   str(pid)],
 286                  self._logger)
 287
 288     def affinitize(self, core):
 289         """Affinitize process to a specific ``core``.
 290
 291         :param core: Core to affinitize process to.
 292
 293         :returns: None
 294         """
 295         self._logger.info('Affinitizing process')
 296
 297         if self._child and self._child.isalive():
 298             self._affinitize_pid(core, self._child.pid)
 299
 300     class ContinueReadPrintLoop(threading.Thread):
 301         """Thread to read output from child and log.
 302
 303         Taken from: https://github.com/pexpect/pexpect/issues/90
 304         """
 305         def __init__(self, child):
 306             self.child = child
 307             threading.Thread.__init__(self)
 308
 309         def run(self):
 310             while True:
 311                 try:
 312                     self.child.read_nonblocking()
 313                 except (pexpect.EOF, pexpect.TIMEOUT):
 314                     break
 315
 316     def relinquish(self):
 317         """Relinquish control of process.
 318
 319         Give up control of application in order to ensure logging
 320         continues for the application. After relinquishing control it
 321         will no longer be possible to :func:`expect` anything.
 322
 323         This works around an issue described here:
 324
 325             https://github.com/pexpect/pexpect/issues/90
 326
 327         It is hoped that future versions of pexpect will avoid this
 328         issue.
 329         """
 330         self._relinquish_thread = self.ContinueReadPrintLoop(self._child)
 331         self._relinquish_thread.start()
 332
 333
 334 class CustomProcess(Process):
 335     """An sample implementation of ``Process``.
 336
 337     This is essentially a more detailed version of the
 338     ``run_interactive_task`` function that checks for process execution
 339     and kills the process (assuming use of the context manager).
 340     """
 341     def __init__(self, cmd, timeout, logfile, expect, name):
 342         """Initialise process state.
 343
 344         :param cmd: Command to execute.
 345         :param timeout: Time to wait for ``expect``.
 346         :param logfile: Path to logfile.
 347         :param expect: String to expect indicating startup. This is a
 348             regex and should be escaped as such.
 349         :param name: Name of process to use in logs.
 350
 351         :returns: None
 352         """
 353         self._cmd = cmd
 354         self._logfile = logfile
 355         self._expect = expect
 356         self._proc_name = name
 357         self._timeout = timeout