diff options
Diffstat (limited to 'tools/tasks.py')
| -rw-r--r-- | tools/tasks.py | 354 |
1 files changed, 354 insertions, 0 deletions
diff --git a/tools/tasks.py b/tools/tasks.py new file mode 100644 index 00000000..f8f11d4e --- /dev/null +++ b/tools/tasks.py @@ -0,0 +1,354 @@ +# Copyright 2015 Intel Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Task management helper functions and classes. +""" + +import select +import subprocess +import logging +import pexpect +import threading +import sys +import os +import locale + +from conf import settings + + +CMD_PREFIX = 'cmd : ' +_MY_ENCODING = locale.getdefaultlocale()[1] + +def _get_stdout(): + """Get stdout value for ``subprocess`` calls. + """ + stdout = None + + if settings.getValue('VERBOSITY') != 'debug': + stdout = open(os.devnull, 'wb') + + return stdout + + +def run_task(cmd, logger, msg=None, check_error=False): + """Run task, report errors and log overall status. + + Run given task using ``subprocess.Popen``. Log the commands + used and any errors generated. Prints stdout to screen if + in verbose mode and returns it regardless. Prints stderr to + screen always. + + :param cmd: Exact command to be executed + :param logger: Logger to write details to + :param msg: Message to be shown to user + :param check_error: Throw exception on error + + :returns: (stdout, stderr) + """ + def handle_error(exception): + """Handle errors by logging and optionally raising an exception. + """ + logger.error( + 'Unable to execute %(cmd)s. Exception: %(exception)s', + {'cmd': ' '.join(cmd), 'exception': exception}) + if check_error: + raise exception + + stdout = [] + stderr = [] + + if msg: + logger.info(msg) + + logger.debug('%s%s', CMD_PREFIX, ' '.join(cmd)) + + try: + proc = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=0) + + while True: + reads = [proc.stdout.fileno(), proc.stderr.fileno()] + ret = select.select(reads, [], []) + + for file_d in ret[0]: + if file_d == proc.stdout.fileno(): + line = proc.stdout.readline() + if settings.getValue('VERBOSITY') == 'debug': + sys.stdout.write(line.decode(_MY_ENCODING)) + stdout.append(line) + if file_d == proc.stderr.fileno(): + line = proc.stderr.readline() + sys.stderr.write(line.decode(_MY_ENCODING)) + stderr.append(line) + + if proc.poll() is not None: + break + except OSError as ex: + handle_error(ex) + else: + if proc.returncode: + ex = subprocess.CalledProcessError(proc.returncode, cmd, stderr) + handle_error(ex) + + return ('\n'.join(sout.decode(_MY_ENCODING).strip() for sout in stdout), + ('\n'.join(sout.decode(_MY_ENCODING).strip() for sout in stderr))) + +def run_background_task(cmd, logger, msg): + """Run task in background and log when started. + + Run given task using ``subprocess.Popen``. Log the command + used. Print stdout to screen if in verbose mode. Prints stderr + to screen always. + + :param cmd: Exact command to be executed + :param logger: Logger to write details to + :param msg: Message to be shown to user + + :returns: Process PID + """ + logger.info(msg) + logger.debug('%s%s', CMD_PREFIX, ' '.join(cmd)) + + proc = subprocess.Popen(cmd, stdout=_get_stdout(), bufsize=0) + + return proc.pid + + +def run_interactive_task(cmd, logger, msg): + """Run a task interactively and log when started. + + Run given task using ``pexpect.spawn``. Log the command used. + Performs neither validation of the process - if the process + successfully started or is still running - nor killing of the + process. The user must do both. + + :param cmd: Exact command to be executed + :param logger: Logger to write details to + :param msg: Message to be shown to user + + :returns: ``pexpect.child`` object + """ + logger.info(msg) + logger.debug('%s%s', CMD_PREFIX, cmd) + child = pexpect.spawnu(cmd) + + if settings.getValue('VERBOSITY') == 'debug': + child.logfile_read = sys.stdout + + return child + + +class Process(object): + """Control an instance of a long-running process. + + This is basically a context-manager wrapper around the + ``run_interactive_task`` function above (with some extra helper + functions). + """ + _cmd = None + _child = None + _logfile = None + _logger = logging.getLogger(__name__) + _expect = None + _timeout = -1 + _proc_name = 'unnamed process' + _relinquish_thread = None + + # context manager + + def __enter__(self): + """Start process instance using context manager. + """ + self.start() + return self + + def __exit__(self, type_, value, traceback): + """Shutdown process instance. + """ + self.kill() + + # startup/shutdown + + def start(self): + """Start process instance. + """ + self._start_process() + if self._timeout > 0: + self._expect_process() + + def _start_process(self): + """Start process instance. + """ + cmd = ' '.join(settings.getValue('SHELL_CMD') + + ['"%s"' % ' '.join(self._cmd)]) + + self._child = run_interactive_task(cmd, self._logger, + 'Starting %s...' % self._proc_name) + self._child.logfile = open(self._logfile, 'w') + + def expect(self, msg, timeout=None): + """Expect string from process. + + Expect string and die if not received. + + :param msg: String to expect. + :param timeout: Time to wait for string. + + :returns: None + """ + self._expect_process(msg, timeout) + + def _expect_process(self, msg=None, timeout=None): + """Expect string from process. + """ + if not msg: + msg = self._expect + if not timeout: + timeout = self._timeout + + # we use exceptions rather than catching conditions in ``expect`` list + # as we want to fail catastrophically after handling; there is likely + # little we can do from within the scripts to fix issues such as + # hugepages not being mounted + try: + self._child.expect([msg], timeout=timeout) + except pexpect.EOF as exc: + self._logger.critical( + 'An error occurred. Please check the logs (%s) for more' + ' information. Exiting...', self._logfile) + raise exc + except pexpect.TIMEOUT as exc: + self._logger.critical( + 'Failed to execute in \'%d\' seconds. Please check the logs' + ' (%s) for more information. Exiting...', + timeout, self._logfile) + self.kill() + raise exc + except (Exception, KeyboardInterrupt) as exc: + self._logger.critical('General exception raised. Exiting...') + self.kill() + raise exc + + def kill(self): + """Kill process instance if it is alive. + """ + if self._child and self._child.isalive(): + run_task(['sudo', 'kill', '-2', str(self._child.pid)], + self._logger) + + if self.is_relinquished(): + self._relinquish_thread.join() + + self._logger.info( + 'Log available at %s', self._logfile) + + def is_relinquished(self): + """Returns True if process is relinquished. + + If relinquished the process is no longer controllable and can + only be killed. + + :returns: True if process is relinquished, else False. + """ + return self._relinquish_thread + + def is_running(self): + """Returns True if process is running. + + :returns: True if process is running, else False + """ + return self._child is not None + + def _affinitize_pid(self, core, pid): + """Affinitize a process with ``pid`` to ``core``. + + :param core: Core to affinitize process to. + :param pid: Process ID to affinitize. + + :returns: None + """ + run_task(['sudo', 'taskset', '-c', '-p', str(core), + str(pid)], + self._logger) + + def affinitize(self, core): + """Affinitize process to a specific ``core``. + + :param core: Core to affinitize process to. + + :returns: None + """ + self._logger.info('Affinitizing process') + + if self._child and self._child.isalive(): + self._affinitize_pid(core, self._child.pid) + + class ContinueReadPrintLoop(threading.Thread): + """Thread to read output from child and log. + + Taken from: https://github.com/pexpect/pexpect/issues/90 + """ + def __init__(self, child): + self.child = child + threading.Thread.__init__(self) + + def run(self): + while True: + try: + self.child.read_nonblocking() + except (pexpect.EOF, pexpect.TIMEOUT): + break + + def relinquish(self): + """Relinquish control of process. + + Give up control of application in order to ensure logging + continues for the application. After relinquishing control it + will no longer be possible to :func:`expect` anything. + + This works around an issue described here: + + https://github.com/pexpect/pexpect/issues/90 + + It is hoped that future versions of pexpect will avoid this + issue. + """ + self._relinquish_thread = self.ContinueReadPrintLoop(self._child) + self._relinquish_thread.start() + + +class CustomProcess(Process): + """An sample implementation of ``Process``. + + This is essentially a more detailed version of the + ``run_interactive_task`` function that checks for process execution + and kills the process (assuming use of the context manager). + """ + def __init__(self, cmd, timeout, logfile, expect, name): + """Initialise process state. + + :param cmd: Command to execute. + :param timeout: Time to wait for ``expect``. + :param logfile: Path to logfile. + :param expect: String to expect indicating startup. This is a + regex and should be escaped as such. + :param name: Name of process to use in logs. + + :returns: None + """ + self._cmd = cmd + self._logfile = logfile + self._expect = expect + self._proc_name = name + self._timeout = timeout |