tools/tasks.py

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