reworked lit remote.py

executor.py

import os

import tracing

from lit.util import executeCommand  # pylint: disable=import-error


class Executor(object):
    def run(self, exe_path, cmd, local_cwd, env=None):
        """Execute a command.

        Args:
            exe_path: str:   Local path to the executable to be run
            cmd: [str]:      subprocess.call style command
            local_cwd: str:  Local path to the working directory
            env: {str: str}: Environment variables to execute under
        Returns:
            out, err, exitCode
        """
        raise NotImplementedError


class LocalExecutor(Executor):
    def __init__(self):
        super(LocalExecutor, self).__init__(self)

    def run(self, exe_path, cmd=None, work_dir='.', env=None):
        cmd = cmd or [exe_path]
        env_cmd = []
        if env:
            env_cmd += ['env']
            env_cmd += ['%s=%s' % (k, v) for k, v in env.items()]
        if work_dir == '.':
            work_dir = os.getcwd()
        return executeCommand(env_cmd + cmd, cwd=work_dir)


class PrefixExecutor(Executor):
    """Prefix an executor with some other command wrapper.

    Most useful for setting ulimits on commands, or running an emulator like
    qemu and valgrind.
    """
    def __init__(self, commandPrefix, chain):
        super(PrefixExecutor, self).__init__(self)

        self.commandPrefix = commandPrefix
        self.chain = chain

    def run(self, exe_path, cmd=None, work_dir='.', env=None):
        cmd = cmd or [exe_path]
        return self.chain.run(self.commandPrefix + cmd, work_dir, env=env)


class PostfixExecutor(Executor):
    """Postfix an executor with some args."""
    def __init__(self, commandPostfix, chain):
        super(PostfixExecutor, self).__init__(self)

        self.commandPostfix = commandPostfix
        self.chain = chain

    def run(self, exe_path, cmd=None, work_dir='.', env=None):
        cmd = cmd or [exe_path]
        return self.chain.run(cmd + self.commandPostfix, work_dir, env=env)



class TimeoutExecutor(PrefixExecutor):
    """Execute another action under a timeout.

    Deprecated. http://reviews.llvm.org/D6584 adds timeouts to LIT.
    """
    def __init__(self, duration, chain):
        super(TimeoutExecutor, self).__init__(
            self, ['timeout', duration], chain)


class SSHExecutor(Executor):
    def __init__(self, host, username=None):
        super(SSHExecutor, self).__init__(self)

        self.user_prefix = username + '@' if username else ''
        self.host = host
        self.scp_command = 'scp'
        self.ssh_command = 'ssh'
        self.local_executor = LocalExecutor()

        # Maps from local->remote to make sure we consistently return
        # the same name for a given local version of the same file/dir
        self.filemap = {}
        self.dirmap = {}

        # TODO(jroelofs): switch this on some -super-verbose-debug config flag
        if False:
            self.local_executor = tracing.trace_object(
                self.local_executor, log_calls=True, log_results=True,
                label='ssh_local')

    def remote_temp_dir(self):
        return self._remote_temp(is_dir=True)

    def remote_temp_file(self):
        return self._remote_temp(is_dir=False)

    def _remote_temp(self, is_dir):
        # TODO: detect what the target system is, and use the correct
        # mktemp command for it. (linux and darwin differ here, and I'm
        # sure windows has another way to do it)

        # Not sure how to do suffix on osx yet
        dir_arg = '-d' if is_dir else ''
        cmd = 'mktemp -q {} /tmp/libcxx.XXXXXXXXXX'.format(dir_arg)
        temp_path, err, exitCode = self.__execute_command_remote([cmd])
        temp_path = temp_path.strip()
        if exitCode != 0:
            raise RuntimeError(err)
        self.remote_temps.append(temp_path)
        return temp_path

    def remote_from_local_dir(self, local_dir):
        if local_dir not in self.dirmap:
            remote_dir = self.remote_temp_dir()
            self.dirmap[local_dir] = remote_dir
        return self.dirmap[local_dir]

    def remote_from_local_file(self, local_file):
        if not local_file in self.filemap:
            remote_file = self.remote_temp_file()
            self.filemap[local_file] = remote_file
        return self.filemap[local_file]

    def copy_in(self, local_srcs, remote_dsts):
        scp = self.scp_command
        remote = self.host
        remote = self.user_prefix + remote

        # This could be wrapped up in a tar->scp->untar for performance
        # if there are lots of files to be copied/moved
        for src, dst in zip(local_srcs, remote_dsts):
            cmd = [scp, '-p', src, remote + ':' + dst]
            self.local_executor.run(cmd)

    def delete_remote(self, remote):
        try:
            self.__execute_command_remote(['rm', '-rf', remote])
        except OSError:
            # TODO: Log failure to delete?
            pass

    def run(self, exe_path, cmd=None, work_dir='.', env=None):
        try:
            target_exe_path = self.remote_from_local_file(exe_path)
            target_cwd = self.remote_from_local_dir(work_dir)
            if cmd:
                # Replace exe_path with target_exe_path.
                cmd = [c if c != exe_path else target_exe_path for c in cmd]
            else:
                cmd = [target_exe_path]
            self.copy_in([exe_path], [target_exe_path])
            return self.__execute_command_remote(cmd, target_cwd, env)
        except:
            raise
        finally:
            self.delete_remote(target_exe_path)
            self.delete_remote(target_cwd)

    def __execute_command_remote(self, cmd, remote_work_dir='.', env=None):
        remote = self.user_prefix + self.host
        ssh_cmd = [self.ssh_command, '-oBatchMode=yes', remote]
        if env:
            env_cmd = ['env'] + ['%s=%s' % (k, v) for k, v in env.items()]
        else:
            env_cmd = []
        remote_cmd = ' '.join(env_cmd + cmd)
        if remote_work_dir != '.':
            remote_cmd = 'cd ' + remote_work_dir + ' && ' + remote_cmd
        return self.local_executor.run(ssh_cmd + remote_cmd)

tracing.py

import inspect


def trace_function(function, log_calls, log_results, label=''):
    def wrapper(*args, **kwargs):
        kwarg_strs = ['{}={}'.format(k, v) for (k, v) in kwargs]
        arg_str = ', '.join(args + kwarg_strs)
        call_str = '{}({})'.format(function.func_name, arg_str)

        # Perform the call itself, logging before, after, and anything thrown.
        try:
            if log_calls:
                print '{}: Calling {}'.format(label, call_str)
            res = function(*args, **kwargs)
            if log_results:
                print '{}: {} -> {}'.format(label, call_str, res)
            return res
        except Exception as ex:
            if log_results:
                print '{}: {} raised {}'.format(label, call_str, type(ex))
            raise ex

    return wrapper


def trace_object(obj, log_calls, log_results, label=''):
    for name, member in inspect.getmembers(obj):
        if inspect.ismethod(member):
            # Skip meta-functions, decorate everything else
            if not member.func_name.startswith('__'):
                setattr(obj, name, trace_function(member, log_calls,
                                                  log_results, label))
    return obj