thehub/testing/test-parallel

#!/usr/bin/env python2
# Based on https://github.com/google/gtest-parallel
# Copyright 2013 Google Inc. All rights reserved.
#
# Adaptations for running Bitcoin's tests (using runner based on
# Boost Test Library instead of Google Test):
# Copyright (c) 2017 freetrader <freetrader@tuta.io> (a Bitcoin developer)
# Copyright (C) 2017 Tom Zander <tom@flowee.org>
#
# 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.

import cPickle
import errno
import getpass
import gzip
import json
import multiprocessing
import optparse
import os
import re
import shutil
import signal
import subprocess
import sys
import tempfile
import thread
import threading
import time
import zlib

# An object that catches SIGINT sent to the Python process and notices
# if processes passed to wait() die by SIGINT (we need to look for
# both of those cases, because pressing Ctrl+C can result in either
# the main process or one of the subprocesses getting the signal).
#
# Before a SIGINT is seen, wait(p) will simply call p.wait() and
# return the result. Once a SIGINT has been seen (in the main process
# or a subprocess, including the one the current call is waiting for),
# wait(p) will call p.terminate() and raise ProcessWasInterrupted.


class SigintHandler(object):
    class ProcessWasInterrupted(Exception):
        pass
    sigint_returncodes = {-signal.SIGINT,  # Unix
                          -1073741510,     # Windows
                          }

    def __init__(self):
        self.__lock = threading.Lock()
        self.__processes = set()
        self.__got_sigint = False
        signal.signal(signal.SIGINT, lambda signal_num,
                      frame: self.interrupt())

    def __on_sigint(self):
        self.__got_sigint = True
        while self.__processes:
            try:
                self.__processes.pop().terminate()
            except OSError:
                pass

    def interrupt(self):
        with self.__lock:
            self.__on_sigint()

    def got_sigint(self):
        with self.__lock:
            return self.__got_sigint

    def wait(self, p):
        with self.__lock:
            if self.__got_sigint:
                p.terminate()
            self.__processes.add(p)
        code = p.wait()
        with self.__lock:
            self.__processes.discard(p)
            if code in self.sigint_returncodes:
                self.__on_sigint()
            if self.__got_sigint:
                raise self.ProcessWasInterrupted
        return code


sigint_handler = SigintHandler()

# Return the width of the terminal, or None if it couldn't be
# determined (e.g. because we're not being run interactively).


def term_width(out):
    if not out.isatty():
        return None
    try:
        p = subprocess.Popen(["stty", "size"],
                             stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        (out, err) = p.communicate()
        if p.returncode != 0 or err:
            return None
        return int(out.split()[1])
    except (IndexError, OSError, ValueError):
        return None

# Output transient and permanent lines of text. If several transient
# lines are written in sequence, the new will overwrite the old. We
# use this to ensure that lots of unimportant info (tests passing)
# won't drown out important info (tests failing).


class Outputter(object):
    def __init__(self, out_file):
        self.__out_file = out_file
        self.__previous_line_was_transient = False
        # Line width, or None if not a tty.
        self.__width = term_width(out_file)

    def transient_line(self, msg):
        if self.__width is None:
            self.__out_file.write(msg + "\n")
        else:
            self.__out_file.write(
                "\r" + msg[:self.__width].ljust(self.__width))
            self.__previous_line_was_transient = True

    def flush_transient_output(self):
        if self.__previous_line_was_transient:
            self.__out_file.write("\n")
            self.__previous_line_was_transient = False

    def permanent_line(self, msg):
        self.flush_transient_output()
        self.__out_file.write(msg + "\n")


stdout_lock = threading.Lock()


class FilterFormat:
    if sys.stdout.isatty():
        # stdout needs to be unbuffered since the output is interactive.
        sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', 0)

    out = Outputter(sys.stdout)
    total_tests = 0
    finished_tests = 0

    tests = {}
    outputs = {}
    passed = []
    failed = []
    started = set()

    def move_to(self, destination_dir, test_ids):
        destination_dir = os.path.join(self.output_dir, destination_dir)
        os.makedirs(destination_dir)
        for test_id in test_ids:
            log_file = self.outputs[test_id]
            test_name = os.path.basename(log_file)
            shutil.move(log_file, os.path.join(destination_dir, test_name))

    def print_tests(self, message, test_ids):
        self.out.permanent_line("%s (%s/%s):" %
                                (message, len(test_ids), self.total_tests))
        test_ids = sorted(test_ids, key=lambda test_id: self.tests[test_id])
        for test_id in test_ids:
            self.out.permanent_line(" %s: %s" % self.tests[test_id])

    def print_test_status(self):
        lines = list()
        for test_id in sorted(self.started):
            (binary, test) = self.tests[test_id]
            lines.append(test)

        if (len(lines) == 0):
            self.out.transient_line("Finished %d test-suites" % self.total_tests)
        else:
            self.out.transient_line("[%d/%d] [%s]"
                                    % (self.finished_tests, self.total_tests, " ".join(lines)))

    def log(self, job_id, command, args=tuple(), file_name=None):
        with stdout_lock:
            if command == "TEST":
                self.tests[job_id] = args
            elif command == "START":
                self.started.add(job_id)
                self.outputs[job_id] = file_name
                self.print_test_status()
            elif command == "EXIT":
                self.started.remove(job_id)
                self.finished_tests += 1
                (exit_code, time_ms) = args
                if exit_code == 0 or (exit_code == 200 and not options.fail_empty_suites):
                    self.passed.append(job_id)
                else:
                    self.failed.append(job_id)
                    (binary, test) = self.tests[job_id]
                    self.out.transient_line("-+==[%d/%d] Test failed; `%s', Loglines follow:"
                            % (self.finished_tests, self.total_tests, test))
                    with open(file_name) as f:
                        for line in f.readlines():
                            self.out.permanent_line(" | %s" % line.rstrip())
                    self.out.permanent_line(
                        " +- %s returned/aborted with exit code %d (%d ms)\n"
                        % (test, exit_code, time_ms))
                self.print_test_status()
            elif command == "TESTCNT":
                self.total_tests = args[0]
                self.out.transient_line(
                    "[0/%d] Running tests..." % self.total_tests)

    def end(self):
        if self.passed:
            self.move_to('passed', self.passed)
        if self.failed:
            self.print_tests("FAILED TESTS", self.failed)
            self.move_to('failed', self.failed)
        if self.started:
            self.print_tests("INTERRUPTED TESTS", self.started)
            self.move_to('interrupted', self.started)
        self.out.flush_transient_output()


class RawFormat:
    def log(self, job_id, command, args=tuple(), file_name=None):
        stdout_lock.acquire()
        line = "%d: %s %s" % (
            job_id, command, ' '.join(str(arg) for arg in args))
        sys.stdout.write(line + "\n")
        if command == "EXIT":
            with open(file_name) as f:
                for line in f:
                    sys.stdout.write("%d > %s" % (job_id, line))
        sys.stdout.flush()
        stdout_lock.release()

    def end(self):
        pass


class CollectTestResults(object):
    def __init__(self, json_dump_filepath):
        self.test_results_lock = threading.Lock()
        self.json_dump_file = open(json_dump_filepath, 'w')
        self.test_results = {
            "interrupted": False,
            "path_delimiter": ".",
            # Third version of the file format. See the link in the flag description
            # for details.
            "version": 3,
            "seconds_since_epoch": int(time.time()),
            "num_failures_by_type": {
                "PASS": 0,
                "FAIL": 0,
            },
            "tests": {},
        }

    def log(self, test, result):
        with self.test_results_lock:
            self.test_results['num_failures_by_type'][result['actual']] += 1
            results = self.test_results['tests']
            for name in test.split('.'):
                results = results.setdefault(name, {})
            results.update(result)

    def dump_to_file_and_close(self):
        json.dump(self.test_results, self.json_dump_file)
        self.json_dump_file.close()


class IgnoreTestResults(object):
    def log(self, test, result):
        pass

    def dump_to_file_and_close(self):
        pass


class DummyTimer(object):
    def start(self):
        pass

    def cancel(self):
        pass

# Record of test runtimes. Has built-in locking.


class TestTimes(object):
    def __init__(self, save_file):
        "Create new object seeded with saved test times from the given file."
        self.__times = {}  # (test binary, test name) -> runtime in ms

        # Protects calls to record_test_time(); other calls are not
        # expected to be made concurrently.
        self.__lock = threading.Lock()

        try:
            with gzip.GzipFile(save_file, "rb") as f:
                times = cPickle.load(f)
        except (EOFError, IOError, cPickle.UnpicklingError, zlib.error):
            # File doesn't exist, isn't readable, is malformed---whatever.
            # Just ignore it.
            return

        # Discard saved times if the format isn't right.
        if type(times) is not dict:
            return
        for ((test_binary, test_name), runtime) in times.items():
            if (type(test_binary) is not str or type(test_name) not in (str, unicode)
                    or type(runtime) not in {int, long, type(None)}):
                return

        self.__times = times

    def get_test_time(self, binary, testname):
        """Return the last duration for the given test as an integer number of
        milliseconds, or None if the test failed or if there's no record for it."""
        return self.__times.get((binary, testname), None)

    def record_test_time(self, binary, testname, runtime_ms):
        """Record that the given test ran in the specified number of
        milliseconds. If the test failed, runtime_ms should be None."""
        with self.__lock:
            self.__times[(binary, testname)] = runtime_ms

    def write_to_file(self, save_file):
        "Write all the times to file."
        try:
            with open(save_file, "wb") as f:
                with gzip.GzipFile("", "wb", 9, f) as gzf:
                    cPickle.dump(self.__times, gzf, cPickle.HIGHEST_PROTOCOL)
        except IOError:
            pass  # ignore errors---saving the times isn't that important


# Remove additional arguments (anything after --).
additional_args = []

for i in range(len(sys.argv)):
    if sys.argv[i] == '--':
        additional_args = sys.argv[i + 1:]
        sys.argv = sys.argv[:i]
        break

parser = optparse.OptionParser(
    usage='usage: %prog [options] binary [binary ...] -- [additional args]')

parser.add_option('-d', '--output_dir', type='string',
                  default=os.path.join(
                      tempfile.gettempdir(), "gtest-parallel-bitcoin.%s" %
                      getpass.getuser()),
                  help='output directory for test logs')
parser.add_option('-r', '--repeat', type='int', default=1,
                  help='repeat tests')
parser.add_option('--fail_empty_suites', action='store_true', default=True,
                  help='fail empty test suites')
parser.add_option('--failed', action='store_true', default=False,
                  help='run only failed and new tests')
parser.add_option('-w', '--workers', type='int',
                  default=multiprocessing.cpu_count(),
                  help='number of workers to spawn')
parser.add_option('--format', type='string', default='filter',
                  help='output format (raw,filter)')
parser.add_option('--print_test_times', action='store_true', default=False,
                  help='list the run time of each test at the end of execution')
parser.add_option('--shard_count', type='int', default=1,
                  help='total number of shards (for sharding test execution '
                       'between multiple machines)')
parser.add_option('--shard_index', type='int', default=0,
                  help='zero-indexed number identifying this shard (for '
                       'sharding test execution between multiple machines)')
parser.add_option('--dump_json_test_results', type='string', default=None,
                  help='Saves the results of the tests as a JSON machine-'
                       'readable file. The format of the file is specified at '
                       'https://www.chromium.org/developers/the-json-test-results-format')
parser.add_option('--timeout', type='int', default=None,
                  help='Interrupt all remaining processes after the given '
                       'time (in seconds).')

(options, binaries) = parser.parse_args()

if binaries == []:
    parser.print_usage()
    sys.exit(1)

logger = RawFormat()
if options.format == 'raw':
    pass
elif options.format == 'filter':
    logger = FilterFormat()
    logger.output_dir = options.output_dir
else:
    parser.error("Unknown output format: " + options.format)

if options.shard_count < 1:
    parser.error("Invalid number of shards: %d. Must be at least 1." %
                 options.shard_count)
if not (0 <= options.shard_index < options.shard_count):
    parser.error("Invalid shard index: %d. Must be between 0 and %d "
                 "(less than the number of shards)." %
                 (options.shard_index, options.shard_count - 1))

timeout = (DummyTimer() if options.timeout is None
           else threading.Timer(options.timeout, sigint_handler.interrupt))

test_results = (IgnoreTestResults() if options.dump_json_test_results is None
                else CollectTestResults(options.dump_json_test_results))

# Find tests.
save_file = os.path.join(os.path.expanduser(
    "~"), ".gtest-parallel-bitcoin-times")
times = TestTimes(save_file)
tests = []
for test_binary in binaries:
    command = [test_binary]

    list_command = list(command)
    list_command += ['--list_content']

    try:
        proc = subprocess.Popen(list_command,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.STDOUT)
        proc_output = ""
        while True:
                # Read line from stdout, break if EOF reached, append line to
                # output
            line = proc.stdout.readline()
            line = line.decode()
            if (line == ""):
                break
            if (line.find("*") == -1):
                continue
            proc_output += line
        proc_output = proc_output.split('\n')
        test_list = [x.replace('*', '')
                     for x in proc_output if x and not x[0].isspace()]

    except OSError as e:
        sys.exit("%s: %s" % (test_binary, str(e)))

    command += additional_args

    for entry in test_list:
        test = entry
        tests.append((times.get_test_time(test_binary, test),
                      test_binary, test, command))

tests = tests[options.shard_index::options.shard_count]

if options.failed:
    # The first element of each entry is the runtime of the most recent
    # run if it was successful, or None if the test is new or the most
    # recent run failed.
    tests = [x for x in tests if x[0] is None]
    if not tests:
        print("No new or failed tests.")
        sys.exit(0)

# Sort tests by falling runtime (with None, which is what we get for
# new and failing tests, being considered larger than any real
# runtime).
tests.sort(reverse=True, key=lambda x: ((1 if x[0] is None else 0), x))

# Repeat tests (-r flag).
tests = [(test, i + 1) for test in tests for i in range(options.repeat)]
test_lock = threading.Lock()
job_id = 0
logger.log(-1, "TESTCNT", (len(tests),))

exit_code = 0

# Remove files from old test runs.
if os.path.isdir(options.output_dir):
    shutil.rmtree(options.output_dir)
# Create directory for test log output.
try:
    os.makedirs(options.output_dir)
except OSError as e:
    # Ignore errors if this directory already exists.
    if e.errno != errno.EEXIST or not os.path.isdir(options.output_dir):
        raise e

# Run the specified job. Return the elapsed time in milliseconds if
# the job succeeds, or None if the job fails. (This ensures that
# failing tests will run first the next time.)


def run_job((command, job_id, test, test_index)):
    begin = time.time()

    test_name = re.sub('[^A-Za-z0-9]', '_', test) + \
        '-' + str(test_index) + '.log'
    test_file = os.path.join(options.output_dir, test_name)
    logger.log(job_id, "START", file_name=test_file)
    with open(test_file, 'w') as log:
        sub = subprocess.Popen(command + ['--run_test=' + test],
                               stdout=log,
                               stderr=log)
        try:
            code = sigint_handler.wait(sub)
        except sigint_handler.ProcessWasInterrupted:
            thread.exit()
        runtime_ms = int(1000 * (time.time() - begin))

    test_results.log(test, {
        "expected": "PASS",
        "actual": "PASS" if code == 0 else "FAIL",
        "time": runtime_ms,
    })
    logger.log(job_id, "EXIT", (code, runtime_ms), file_name=test_file)
    if code == 0 or (code == 200 and not options.fail_empty_suites):
        return runtime_ms
    global exit_code
    exit_code = code
    return None


def worker():
    global job_id
    while True:
        job = None
        test_lock.acquire()
        if job_id < len(tests):
            ((_, test_binary, test, command), test_index) = tests[job_id]
            logger.log(job_id, "TEST", (test_binary, test))
            job = (command, job_id, test, test_index)
        job_id += 1
        test_lock.release()
        if job is None:
            return
        times.record_test_time(test_binary, test, run_job(job))


def start_daemon(func):
    t = threading.Thread(target=func)
    t.daemon = True
    t.start()
    return t


try:
    timeout.start()
    workers = [start_daemon(worker) for i in range(options.workers)]
    [t.join() for t in workers]
finally:
    timeout.cancel()

logger.end()
times.write_to_file(save_file)
if options.print_test_times:
    tt = [(test_binary, test) for ((_, test_binary, test, _), __)
          in tests if times.get_test_time(test_binary, test) is not None]
    tts = sorted((times.get_test_time(test_binary, test), test_binary, test)
                 for (test_binary, test) in tt)
    for (time_ms, test_binary, test) in tts:
        print "%8s %s" % ("%dms" % time_ms, test)

test_results.dump_to_file_and_close()
sys.exit(-signal.SIGINT if sigint_handler.got_sigint() else exit_code)