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