112882Sspwilson2@wisc.edu# Copyright (c) 2017 Mark D. Hill and David A. Wood
212882Sspwilson2@wisc.edu# All rights reserved.
312882Sspwilson2@wisc.edu#
412882Sspwilson2@wisc.edu# Redistribution and use in source and binary forms, with or without
512882Sspwilson2@wisc.edu# modification, are permitted provided that the following conditions are
612882Sspwilson2@wisc.edu# met: redistributions of source code must retain the above copyright
712882Sspwilson2@wisc.edu# notice, this list of conditions and the following disclaimer;
812882Sspwilson2@wisc.edu# redistributions in binary form must reproduce the above copyright
912882Sspwilson2@wisc.edu# notice, this list of conditions and the following disclaimer in the
1012882Sspwilson2@wisc.edu# documentation and/or other materials provided with the distribution;
1112882Sspwilson2@wisc.edu# neither the name of the copyright holders nor the names of its
1212882Sspwilson2@wisc.edu# contributors may be used to endorse or promote products derived from
1312882Sspwilson2@wisc.edu# this software without specific prior written permission.
1412882Sspwilson2@wisc.edu#
1512882Sspwilson2@wisc.edu# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
1612882Sspwilson2@wisc.edu# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
1712882Sspwilson2@wisc.edu# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
1812882Sspwilson2@wisc.edu# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
1912882Sspwilson2@wisc.edu# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
2012882Sspwilson2@wisc.edu# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
2112882Sspwilson2@wisc.edu# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
2212882Sspwilson2@wisc.edu# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
2312882Sspwilson2@wisc.edu# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
2412882Sspwilson2@wisc.edu# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
2512882Sspwilson2@wisc.edu# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
2612882Sspwilson2@wisc.edu#
2712882Sspwilson2@wisc.edu# Authors: Sean Wilson
2812882Sspwilson2@wisc.edu
2912882Sspwilson2@wisc.edu'''
3012882Sspwilson2@wisc.eduHandlers for the testlib Log.
3112882Sspwilson2@wisc.edu
3212882Sspwilson2@wisc.edu
3312882Sspwilson2@wisc.edu'''
3412882Sspwilson2@wisc.edufrom __future__ import print_function
3512882Sspwilson2@wisc.edu
3612882Sspwilson2@wisc.eduimport multiprocessing
3712882Sspwilson2@wisc.eduimport os
3812882Sspwilson2@wisc.eduimport Queue
3912882Sspwilson2@wisc.eduimport sys
4012882Sspwilson2@wisc.eduimport threading
4112882Sspwilson2@wisc.eduimport time
4212882Sspwilson2@wisc.eduimport traceback
4312882Sspwilson2@wisc.edu
4412882Sspwilson2@wisc.eduimport helper
4512882Sspwilson2@wisc.eduimport log
4612882Sspwilson2@wisc.eduimport result
4712882Sspwilson2@wisc.eduimport state
4812882Sspwilson2@wisc.eduimport test
4912882Sspwilson2@wisc.eduimport terminal
5012882Sspwilson2@wisc.edu
5112882Sspwilson2@wisc.edufrom config import config, constants
5212882Sspwilson2@wisc.edu
5312882Sspwilson2@wisc.edu
5412882Sspwilson2@wisc.educlass _TestStreamManager(object):
5512882Sspwilson2@wisc.edu    def __init__(self):
5612882Sspwilson2@wisc.edu        self._writers = {}
5712882Sspwilson2@wisc.edu
5812882Sspwilson2@wisc.edu    def open_writer(self, test_result):
5912882Sspwilson2@wisc.edu        if test_result in self._writers:
6012882Sspwilson2@wisc.edu            raise ValueError('Cannot have multiple writters on a single test.')
6112882Sspwilson2@wisc.edu        self._writers[test_result] = _TestStreams(test_result.stdout,
6212882Sspwilson2@wisc.edu                test_result.stderr)
6312882Sspwilson2@wisc.edu
6412882Sspwilson2@wisc.edu    def get_writer(self, test_result):
6512882Sspwilson2@wisc.edu        if test_result not in self._writers:
6612882Sspwilson2@wisc.edu            self.open_writer(test_result)
6712882Sspwilson2@wisc.edu        return self._writers[test_result]
6812882Sspwilson2@wisc.edu
6912882Sspwilson2@wisc.edu    def close_writer(self, test_result):
7012882Sspwilson2@wisc.edu        if test_result in self._writers:
7112882Sspwilson2@wisc.edu            writer = self._writers.pop(test_result)
7212882Sspwilson2@wisc.edu            writer.close()
7312882Sspwilson2@wisc.edu
7412882Sspwilson2@wisc.edu    def close(self):
7512882Sspwilson2@wisc.edu        for writer in self._writers.values():
7612882Sspwilson2@wisc.edu            writer.close()
7712882Sspwilson2@wisc.edu        self._writers.clear()
7812882Sspwilson2@wisc.edu
7912882Sspwilson2@wisc.educlass _TestStreams(object):
8012882Sspwilson2@wisc.edu    def __init__(self, stdout, stderr):
8112882Sspwilson2@wisc.edu        helper.mkdir_p(os.path.dirname(stdout))
8212882Sspwilson2@wisc.edu        helper.mkdir_p(os.path.dirname(stderr))
8312882Sspwilson2@wisc.edu        self.stdout = open(stdout, 'w')
8412882Sspwilson2@wisc.edu        self.stderr = open(stderr, 'w')
8512882Sspwilson2@wisc.edu
8612882Sspwilson2@wisc.edu    def close(self):
8712882Sspwilson2@wisc.edu        self.stdout.close()
8812882Sspwilson2@wisc.edu        self.stderr.close()
8912882Sspwilson2@wisc.edu
9012882Sspwilson2@wisc.educlass ResultHandler(log.Handler):
9112882Sspwilson2@wisc.edu    '''
9212882Sspwilson2@wisc.edu    Log handler which listens for test results and output saving data as
9312882Sspwilson2@wisc.edu    it is reported.
9412882Sspwilson2@wisc.edu
9512882Sspwilson2@wisc.edu    When the handler is closed it writes out test results in the python pickle
9612882Sspwilson2@wisc.edu    format.
9712882Sspwilson2@wisc.edu    '''
9812882Sspwilson2@wisc.edu    def __init__(self, schedule, directory):
9912882Sspwilson2@wisc.edu        '''
10012882Sspwilson2@wisc.edu        :param schedule: The entire schedule as a :class:`LoadedLibrary`
10112882Sspwilson2@wisc.edu            object.
10212882Sspwilson2@wisc.edu
10312882Sspwilson2@wisc.edu        :param directory: Directory to save test stdout/stderr and aggregate
10412882Sspwilson2@wisc.edu            results to.
10512882Sspwilson2@wisc.edu        '''
10612882Sspwilson2@wisc.edu        self.directory = directory
10712882Sspwilson2@wisc.edu        self.internal_results = result.InternalLibraryResults(schedule,
10812882Sspwilson2@wisc.edu                directory)
10912882Sspwilson2@wisc.edu        self.test_stream_manager = _TestStreamManager()
11012882Sspwilson2@wisc.edu        self._closed = False
11112882Sspwilson2@wisc.edu
11212882Sspwilson2@wisc.edu        self.mapping = {
11312882Sspwilson2@wisc.edu            log.LibraryStatus.type_id: self.handle_library_status,
11412882Sspwilson2@wisc.edu
11512882Sspwilson2@wisc.edu            log.SuiteResult.type_id: self.handle_suite_result,
11612882Sspwilson2@wisc.edu            log.TestResult.type_id: self.handle_test_result,
11712882Sspwilson2@wisc.edu
11812882Sspwilson2@wisc.edu            log.TestStderr.type_id: self.handle_stderr,
11912882Sspwilson2@wisc.edu            log.TestStdout.type_id: self.handle_stdout,
12012882Sspwilson2@wisc.edu        }
12112882Sspwilson2@wisc.edu
12212882Sspwilson2@wisc.edu    def handle(self, record):
12312882Sspwilson2@wisc.edu        if not self._closed:
12412882Sspwilson2@wisc.edu            self.mapping.get(record.type_id, lambda _:None)(record)
12512882Sspwilson2@wisc.edu
12612882Sspwilson2@wisc.edu    def handle_library_status(self, record):
12712882Sspwilson2@wisc.edu        if record['status'] in (state.Status.Complete, state.Status.Avoided):
12812882Sspwilson2@wisc.edu            self.test_stream_manager.close()
12912882Sspwilson2@wisc.edu
13012882Sspwilson2@wisc.edu    def handle_suite_result(self, record):
13112882Sspwilson2@wisc.edu        suite_result = self.internal_results.get_suite_result(
13212882Sspwilson2@wisc.edu                    record['metadata'].uid)
13312882Sspwilson2@wisc.edu        suite_result.result = record['result']
13412882Sspwilson2@wisc.edu
13512882Sspwilson2@wisc.edu    def handle_test_result(self, record):
13612882Sspwilson2@wisc.edu        test_result = self._get_test_result(record)
13712882Sspwilson2@wisc.edu        test_result.result = record['result']
13812882Sspwilson2@wisc.edu
13912882Sspwilson2@wisc.edu    def handle_stderr(self, record):
14012882Sspwilson2@wisc.edu        self.test_stream_manager.get_writer(
14112882Sspwilson2@wisc.edu            self._get_test_result(record)
14212882Sspwilson2@wisc.edu        ).stderr.write(record['buffer'])
14312882Sspwilson2@wisc.edu
14412882Sspwilson2@wisc.edu    def handle_stdout(self, record):
14512882Sspwilson2@wisc.edu        self.test_stream_manager.get_writer(
14612882Sspwilson2@wisc.edu            self._get_test_result(record)
14712882Sspwilson2@wisc.edu        ).stdout.write(record['buffer'])
14812882Sspwilson2@wisc.edu
14912882Sspwilson2@wisc.edu    def _get_test_result(self, test_record):
15012882Sspwilson2@wisc.edu        return self.internal_results.get_test_result(
15112882Sspwilson2@wisc.edu                    test_record['metadata'].uid,
15212882Sspwilson2@wisc.edu                    test_record['metadata'].suite_uid)
15312882Sspwilson2@wisc.edu
15412882Sspwilson2@wisc.edu    def _save(self):
15512882Sspwilson2@wisc.edu        #FIXME Hardcoded path name
15612882Sspwilson2@wisc.edu        result.InternalSavedResults.save(
15712882Sspwilson2@wisc.edu            self.internal_results,
15812882Sspwilson2@wisc.edu            os.path.join(self.directory, constants.pickle_filename))
15912882Sspwilson2@wisc.edu        result.JUnitSavedResults.save(
16012882Sspwilson2@wisc.edu            self.internal_results,
16112882Sspwilson2@wisc.edu            os.path.join(self.directory, constants.xml_filename))
16212882Sspwilson2@wisc.edu
16312882Sspwilson2@wisc.edu    def close(self):
16412882Sspwilson2@wisc.edu        if self._closed:
16512882Sspwilson2@wisc.edu            return
16612882Sspwilson2@wisc.edu        self._closed = True
16712882Sspwilson2@wisc.edu        self._save()
16812882Sspwilson2@wisc.edu
16913789Sjason@lowepower.com    def unsuccessful(self):
17013789Sjason@lowepower.com        '''
17113789Sjason@lowepower.com        Performs an or reduce on all of the results.
17213789Sjason@lowepower.com        Returns true if at least one test is unsuccessful, false when all tests
17313789Sjason@lowepower.com        pass
17413789Sjason@lowepower.com        '''
17513789Sjason@lowepower.com        for suite_result in self.internal_results:
17613789Sjason@lowepower.com            if suite_result.unsuccessful:
17713789Sjason@lowepower.com                return True
17813789Sjason@lowepower.com        # If all are successful, then this wasn't "unsuccessful"
17913789Sjason@lowepower.com        return False
18013789Sjason@lowepower.com
18112882Sspwilson2@wisc.edu
18212882Sspwilson2@wisc.edu#TODO Change from a handler to an internal post processor so it can be used
18312882Sspwilson2@wisc.edu# to reprint results
18412882Sspwilson2@wisc.educlass SummaryHandler(log.Handler):
18512882Sspwilson2@wisc.edu    '''
18612882Sspwilson2@wisc.edu    A log handler which listens to the log for test results
18712882Sspwilson2@wisc.edu    and reports the aggregate results when closed.
18812882Sspwilson2@wisc.edu    '''
18912882Sspwilson2@wisc.edu    color = terminal.get_termcap()
19012882Sspwilson2@wisc.edu    reset = color.Normal
19112882Sspwilson2@wisc.edu    colormap = {
19212882Sspwilson2@wisc.edu            state.Result.Errored: color.Red,
19312882Sspwilson2@wisc.edu            state.Result.Failed: color.Red,
19412882Sspwilson2@wisc.edu            state.Result.Passed: color.Green,
19512882Sspwilson2@wisc.edu            state.Result.Skipped: color.Cyan,
19612882Sspwilson2@wisc.edu    }
19712882Sspwilson2@wisc.edu    sep_fmtkey = 'separator'
19812882Sspwilson2@wisc.edu    sep_fmtstr = '{%s}' % sep_fmtkey
19912882Sspwilson2@wisc.edu
20012882Sspwilson2@wisc.edu    def __init__(self):
20112882Sspwilson2@wisc.edu        self.mapping = {
20212882Sspwilson2@wisc.edu            log.TestResult.type_id: self.handle_testresult,
20312882Sspwilson2@wisc.edu            log.LibraryStatus.type_id: self.handle_library_status,
20412882Sspwilson2@wisc.edu        }
20512882Sspwilson2@wisc.edu        self._timer = helper.Timer()
20612882Sspwilson2@wisc.edu        self.results = []
20712882Sspwilson2@wisc.edu
20812882Sspwilson2@wisc.edu    def handle_library_status(self, record):
20912882Sspwilson2@wisc.edu        if record['status'] == state.Status.Building:
21012882Sspwilson2@wisc.edu            self._timer.restart()
21112882Sspwilson2@wisc.edu
21212882Sspwilson2@wisc.edu    def handle_testresult(self, record):
21312882Sspwilson2@wisc.edu        result = record['result'].value
21412882Sspwilson2@wisc.edu        if result in (state.Result.Skipped, state.Result.Failed,
21512882Sspwilson2@wisc.edu                state.Result.Passed, state.Result.Errored):
21612882Sspwilson2@wisc.edu            self.results.append(result)
21712882Sspwilson2@wisc.edu
21812882Sspwilson2@wisc.edu    def handle(self, record):
21912882Sspwilson2@wisc.edu        self.mapping.get(record.type_id, lambda _:None)(record)
22012882Sspwilson2@wisc.edu
22112882Sspwilson2@wisc.edu    def close(self):
22212882Sspwilson2@wisc.edu        print(self._display_summary())
22312882Sspwilson2@wisc.edu
22412882Sspwilson2@wisc.edu    def _display_summary(self):
22512882Sspwilson2@wisc.edu        most_severe_outcome = None
22612882Sspwilson2@wisc.edu        outcome_fmt = ' {count} {outcome}'
22712882Sspwilson2@wisc.edu        strings = []
22812882Sspwilson2@wisc.edu
22912882Sspwilson2@wisc.edu        outcome_count = [0] * len(state.Result.enums)
23012882Sspwilson2@wisc.edu        for result in self.results:
23112882Sspwilson2@wisc.edu            outcome_count[result] += 1
23212882Sspwilson2@wisc.edu
23312882Sspwilson2@wisc.edu        # Iterate over enums so they are in order of severity
23412882Sspwilson2@wisc.edu        for outcome in state.Result.enums:
23512882Sspwilson2@wisc.edu            outcome = getattr(state.Result, outcome)
23612882Sspwilson2@wisc.edu            count  = outcome_count[outcome]
23712882Sspwilson2@wisc.edu            if count:
23812882Sspwilson2@wisc.edu                strings.append(outcome_fmt.format(count=count,
23912882Sspwilson2@wisc.edu                        outcome=state.Result.enums[outcome]))
24012882Sspwilson2@wisc.edu                most_severe_outcome = outcome
24112882Sspwilson2@wisc.edu        string = ','.join(strings)
24212882Sspwilson2@wisc.edu        if most_severe_outcome is None:
24312882Sspwilson2@wisc.edu            string = ' No testing done'
24412882Sspwilson2@wisc.edu            most_severe_outcome = state.Result.Passed
24512882Sspwilson2@wisc.edu        else:
24612882Sspwilson2@wisc.edu            string = ' Results:' + string + ' in {:.2} seconds '.format(
24712882Sspwilson2@wisc.edu                    self._timer.active_time())
24812882Sspwilson2@wisc.edu        string += ' '
24912882Sspwilson2@wisc.edu        return terminal.insert_separator(
25012882Sspwilson2@wisc.edu                string,
25112882Sspwilson2@wisc.edu                color=self.colormap[most_severe_outcome] + self.color.Bold)
25212882Sspwilson2@wisc.edu
25312882Sspwilson2@wisc.educlass TerminalHandler(log.Handler):
25412882Sspwilson2@wisc.edu    color = terminal.get_termcap()
25512882Sspwilson2@wisc.edu    verbosity_mapping = {
25612882Sspwilson2@wisc.edu        log.LogLevel.Warn: color.Yellow,
25712882Sspwilson2@wisc.edu        log.LogLevel.Error: color.Red,
25812882Sspwilson2@wisc.edu    }
25912882Sspwilson2@wisc.edu    default = color.Normal
26012882Sspwilson2@wisc.edu
26112882Sspwilson2@wisc.edu    def __init__(self, verbosity=log.LogLevel.Info, machine_only=False):
26212882Sspwilson2@wisc.edu        self.stream = verbosity >= log.LogLevel.Trace
26312882Sspwilson2@wisc.edu        self.verbosity = verbosity
26412882Sspwilson2@wisc.edu        self.machine_only = machine_only
26512882Sspwilson2@wisc.edu        self.mapping = {
26612882Sspwilson2@wisc.edu            log.TestResult.type_id: self.handle_testresult,
26712882Sspwilson2@wisc.edu            log.SuiteStatus.type_id: self.handle_suitestatus,
26812882Sspwilson2@wisc.edu            log.TestStatus.type_id: self.handle_teststatus,
26912882Sspwilson2@wisc.edu            log.TestStderr.type_id: self.handle_stderr,
27012882Sspwilson2@wisc.edu            log.TestStdout.type_id: self.handle_stdout,
27112882Sspwilson2@wisc.edu            log.TestMessage.type_id: self.handle_testmessage,
27212882Sspwilson2@wisc.edu            log.LibraryMessage.type_id: self.handle_librarymessage,
27312882Sspwilson2@wisc.edu        }
27412882Sspwilson2@wisc.edu
27512882Sspwilson2@wisc.edu    def _display_outcome(self, name, outcome, reason=None):
27612882Sspwilson2@wisc.edu        print(self.color.Bold
27712882Sspwilson2@wisc.edu                 + SummaryHandler.colormap[outcome]
27812882Sspwilson2@wisc.edu                 + name
27912882Sspwilson2@wisc.edu                 + ' '
28012882Sspwilson2@wisc.edu                 + state.Result.enums[outcome]
28112882Sspwilson2@wisc.edu                 + SummaryHandler.reset)
28212882Sspwilson2@wisc.edu
28312882Sspwilson2@wisc.edu        if reason is not None:
28412882Sspwilson2@wisc.edu            log.test_log.info('')
28512882Sspwilson2@wisc.edu            log.test_log.info('Reason:')
28612882Sspwilson2@wisc.edu            log.test_log.info(reason)
28712882Sspwilson2@wisc.edu            log.test_log.info(terminal.separator('-'))
28812882Sspwilson2@wisc.edu
28912882Sspwilson2@wisc.edu    def handle_teststatus(self, record):
29012882Sspwilson2@wisc.edu        if record['status'] == state.Status.Running:
29112882Sspwilson2@wisc.edu            log.test_log.debug('Starting Test Case: %s' %\
29212882Sspwilson2@wisc.edu                    record['metadata'].name)
29312882Sspwilson2@wisc.edu
29412882Sspwilson2@wisc.edu    def handle_testresult(self, record):
29512882Sspwilson2@wisc.edu        self._display_outcome(
29612882Sspwilson2@wisc.edu            'Test: %s'  % record['metadata'].name,
29712882Sspwilson2@wisc.edu            record['result'].value)
29812882Sspwilson2@wisc.edu
29912882Sspwilson2@wisc.edu    def handle_suitestatus(self, record):
30012882Sspwilson2@wisc.edu        if record['status'] == state.Status.Running:
30112882Sspwilson2@wisc.edu              log.test_log.debug('Starting Test Suite: %s ' %\
30212882Sspwilson2@wisc.edu                    record['metadata'].name)
30312882Sspwilson2@wisc.edu
30412882Sspwilson2@wisc.edu    def handle_stderr(self, record):
30512882Sspwilson2@wisc.edu        if self.stream:
30612882Sspwilson2@wisc.edu            print(record.data['buffer'], file=sys.stderr, end='')
30712882Sspwilson2@wisc.edu
30812882Sspwilson2@wisc.edu    def handle_stdout(self, record):
30912882Sspwilson2@wisc.edu        if self.stream:
31012882Sspwilson2@wisc.edu            print(record.data['buffer'], file=sys.stdout, end='')
31112882Sspwilson2@wisc.edu
31212882Sspwilson2@wisc.edu    def handle_testmessage(self, record):
31312882Sspwilson2@wisc.edu        if self.stream:
31412882Sspwilson2@wisc.edu            print(self._colorize(record['message'], record['level']))
31512882Sspwilson2@wisc.edu
31612882Sspwilson2@wisc.edu    def handle_librarymessage(self, record):
31712882Sspwilson2@wisc.edu        if not self.machine_only or record.data.get('machine_readable', False):
31812882Sspwilson2@wisc.edu            print(self._colorize(record['message'], record['level'],
31912882Sspwilson2@wisc.edu                    record['bold']))
32012882Sspwilson2@wisc.edu
32112882Sspwilson2@wisc.edu    def _colorize(self, message, level, bold=False):
32212882Sspwilson2@wisc.edu        return '%s%s%s%s' % (
32312882Sspwilson2@wisc.edu                self.color.Bold if bold else '',
32412882Sspwilson2@wisc.edu                self.verbosity_mapping.get(level, ''),
32512882Sspwilson2@wisc.edu                message,
32612882Sspwilson2@wisc.edu                self.default)
32712882Sspwilson2@wisc.edu
32812882Sspwilson2@wisc.edu    def handle(self, record):
32912882Sspwilson2@wisc.edu        if record.data.get('level', self.verbosity) > self.verbosity:
33012882Sspwilson2@wisc.edu            return
33112882Sspwilson2@wisc.edu        self.mapping.get(record.type_id, lambda _:None)(record)
33212882Sspwilson2@wisc.edu
33312882Sspwilson2@wisc.edu    def set_verbosity(self, verbosity):
33412882Sspwilson2@wisc.edu        self.verbosity = verbosity
33512882Sspwilson2@wisc.edu
33612882Sspwilson2@wisc.edu
33712882Sspwilson2@wisc.educlass PrintHandler(log.Handler):
33812882Sspwilson2@wisc.edu    def __init__(self):
33912882Sspwilson2@wisc.edu        pass
34012882Sspwilson2@wisc.edu
34112882Sspwilson2@wisc.edu    def handle(self, record):
34212882Sspwilson2@wisc.edu        print(str(record).rstrip())
34312882Sspwilson2@wisc.edu
34412882Sspwilson2@wisc.edu    def close(self):
34512882Sspwilson2@wisc.edu        pass
34612882Sspwilson2@wisc.edu
34712882Sspwilson2@wisc.edu
34812882Sspwilson2@wisc.educlass MultiprocessingHandlerWrapper(log.Handler):
34912882Sspwilson2@wisc.edu    '''
35012882Sspwilson2@wisc.edu    A handler class which forwards log records to subhandlers, enabling
35112882Sspwilson2@wisc.edu    logging across multiprocessing python processes.
35212882Sspwilson2@wisc.edu
35312882Sspwilson2@wisc.edu    The 'parent' side of the handler should execute either
35412882Sspwilson2@wisc.edu    :func:`async_process` or :func:`process` to forward
35512882Sspwilson2@wisc.edu    log records to subhandlers.
35612882Sspwilson2@wisc.edu    '''
35712882Sspwilson2@wisc.edu    def __init__(self, *subhandlers):
35812882Sspwilson2@wisc.edu        # Create thread to spin handing recipt of messages
35912882Sspwilson2@wisc.edu        # Create queue to push onto
36012882Sspwilson2@wisc.edu        self.queue = multiprocessing.Queue()
36112882Sspwilson2@wisc.edu        self.queue.cancel_join_thread()
36212882Sspwilson2@wisc.edu        self._shutdown = threading.Event()
36312882Sspwilson2@wisc.edu
36412882Sspwilson2@wisc.edu        # subhandlers should be accessed with the _handler_lock
36512882Sspwilson2@wisc.edu        self._handler_lock = threading.Lock()
36612882Sspwilson2@wisc.edu        self._subhandlers = subhandlers
36712882Sspwilson2@wisc.edu
36812882Sspwilson2@wisc.edu    def add_handler(self, handler):
36912882Sspwilson2@wisc.edu        self._handler_lock.acquire()
37012882Sspwilson2@wisc.edu        self._subhandlers = (handler, ) + self._subhandlers
37112882Sspwilson2@wisc.edu        self._handler_lock.release()
37212882Sspwilson2@wisc.edu
37312882Sspwilson2@wisc.edu    def _with_handlers(self, callback):
37412882Sspwilson2@wisc.edu        exception = None
37512882Sspwilson2@wisc.edu        self._handler_lock.acquire()
37612882Sspwilson2@wisc.edu        for handler in self._subhandlers:
37712882Sspwilson2@wisc.edu            # Prevent deadlock when using this handler by delaying
37812882Sspwilson2@wisc.edu            # exception raise until we get a chance to unlock.
37912882Sspwilson2@wisc.edu            try:
38012882Sspwilson2@wisc.edu                callback(handler)
38112882Sspwilson2@wisc.edu            except Exception as e:
38212882Sspwilson2@wisc.edu                exception = e
38312882Sspwilson2@wisc.edu                break
38412882Sspwilson2@wisc.edu        self._handler_lock.release()
38512882Sspwilson2@wisc.edu
38612882Sspwilson2@wisc.edu        if exception is not None:
38712882Sspwilson2@wisc.edu            raise exception
38812882Sspwilson2@wisc.edu
38912882Sspwilson2@wisc.edu    def async_process(self):
39012882Sspwilson2@wisc.edu        self.thread = threading.Thread(target=self.process)
39112882Sspwilson2@wisc.edu        self.thread.daemon = True
39212882Sspwilson2@wisc.edu        self.thread.start()
39312882Sspwilson2@wisc.edu
39412882Sspwilson2@wisc.edu    def process(self):
39512882Sspwilson2@wisc.edu        while not self._shutdown.is_set():
39612882Sspwilson2@wisc.edu            try:
39712882Sspwilson2@wisc.edu                item = self.queue.get(timeout=0.1)
39812882Sspwilson2@wisc.edu                self._handle(item)
39912882Sspwilson2@wisc.edu            except (KeyboardInterrupt, SystemExit):
40012882Sspwilson2@wisc.edu                raise
40112882Sspwilson2@wisc.edu            except EOFError:
40212882Sspwilson2@wisc.edu                return
40312882Sspwilson2@wisc.edu            except Queue.Empty:
40412882Sspwilson2@wisc.edu                continue
40512882Sspwilson2@wisc.edu
40612882Sspwilson2@wisc.edu    def _drain(self):
40712882Sspwilson2@wisc.edu        while True:
40812882Sspwilson2@wisc.edu            try:
40912882Sspwilson2@wisc.edu                item = self.queue.get(block=False)
41012882Sspwilson2@wisc.edu                self._handle(item)
41112882Sspwilson2@wisc.edu            except (KeyboardInterrupt, SystemExit):
41212882Sspwilson2@wisc.edu                raise
41312882Sspwilson2@wisc.edu            except EOFError:
41412882Sspwilson2@wisc.edu                return
41512882Sspwilson2@wisc.edu            except Queue.Empty:
41612882Sspwilson2@wisc.edu                return
41712882Sspwilson2@wisc.edu
41812882Sspwilson2@wisc.edu    def _handle(self, record):
41912882Sspwilson2@wisc.edu        self._with_handlers(lambda handler: handler.handle(record))
42012882Sspwilson2@wisc.edu
42112882Sspwilson2@wisc.edu    def handle(self, record):
42212882Sspwilson2@wisc.edu        self.queue.put(record)
42312882Sspwilson2@wisc.edu
42412882Sspwilson2@wisc.edu    def _close(self):
42512882Sspwilson2@wisc.edu        if hasattr(self, 'thread'):
42612882Sspwilson2@wisc.edu            self.thread.join()
42712882Sspwilson2@wisc.edu        _wrap(self._drain)
42812882Sspwilson2@wisc.edu        self._with_handlers(lambda handler: _wrap(handler.close))
42912882Sspwilson2@wisc.edu
43012882Sspwilson2@wisc.edu        # NOTE Python2 has an known bug which causes IOErrors to be raised
43112882Sspwilson2@wisc.edu        # if this shutdown doesn't go cleanly on both ends.
43212882Sspwilson2@wisc.edu        # This sleep adds some time for the sender threads on this process to
43312882Sspwilson2@wisc.edu        # finish pickling the object and complete shutdown after the queue is
43412882Sspwilson2@wisc.edu        # closed.
43512882Sspwilson2@wisc.edu        time.sleep(.2)
43612882Sspwilson2@wisc.edu        self.queue.close()
43712882Sspwilson2@wisc.edu        time.sleep(.2)
43812882Sspwilson2@wisc.edu
43912882Sspwilson2@wisc.edu    def close(self):
44012882Sspwilson2@wisc.edu        if not self._shutdown.is_set():
44112882Sspwilson2@wisc.edu            self._shutdown.set()
44212882Sspwilson2@wisc.edu            self._close()
44312882Sspwilson2@wisc.edu
44412882Sspwilson2@wisc.edu
44512882Sspwilson2@wisc.edudef _wrap(callback, *args, **kwargs):
44612882Sspwilson2@wisc.edu    try:
44712882Sspwilson2@wisc.edu        callback(*args, **kwargs)
44812882Sspwilson2@wisc.edu    except:
44912882Sspwilson2@wisc.edu        traceback.print_exc()
450