1# Lint as: python2, python3
2# The source code is from following Python documentation:
3# https://docs.python.org/2/howto/logging-cookbook.html#network-logging
4
5# Classes in this file are used to create a simple TCP socket-based logging
6# receiver. The receiver listens to default logging port (9020) and save log to
7# any given log configuration, e.g., a local file. Once the receiver is running,
8# client can add a logging handler to write log to the receiver with following
9# sample code:
10# socketHandler = logging.handlers.SocketHandler('localhost',
11#         logging.handlers.DEFAULT_TCP_LOGGING_PORT)
12# logging.getLogger().addHandler(socketHandler)
13
14from __future__ import absolute_import
15from __future__ import division
16from __future__ import print_function
17
18import ctypes
19import pickle
20import logging
21import multiprocessing
22import select
23import six.moves.socketserver
24import struct
25import time
26
27import common
28from autotest_lib.client.common_lib import utils
29
30class LogRecordStreamHandler(six.moves.socketserver.StreamRequestHandler):
31    """Handler for a streaming logging request.
32
33    This basically logs the record using whatever logging policy is
34    configured locally.
35    """
36
37    def handle(self):
38        """
39        Handle multiple requests - each expected to be a 4-byte length,
40        followed by the LogRecord in pickle format. Logs the record
41        according to whatever policy is configured locally.
42        """
43        while True:
44            chunk = self.connection.recv(4)
45            if len(chunk) < 4:
46                return
47            slen = struct.unpack('>L', chunk)[0]
48            chunk = self.connection.recv(slen)
49            while len(chunk) < slen:
50                chunk = chunk + self.connection.recv(slen - len(chunk))
51            obj = self.unpickle(chunk)
52            record = logging.makeLogRecord(obj)
53            self.handle_log_record(record)
54
55
56    def unpickle(self, data):
57        """Unpickle data received.
58
59        @param data: Received data.
60        @returns: unpickled data.
61        """
62        return pickle.loads(data)
63
64
65    def handle_log_record(self, record):
66        """Process log record.
67
68        @param record: log record.
69        """
70        # if a name is specified, we use the named logger rather than the one
71        # implied by the record.
72        if self.server.logname is not None:
73            name = self.server.logname
74        else:
75            name = record.name
76        logger = logging.getLogger(name)
77        # N.B. EVERY record gets logged. This is because Logger.handle
78        # is normally called AFTER logger-level filtering. If you want
79        # to do filtering, do it at the client end to save wasting
80        # cycles and network bandwidth!
81        logger.handle(record)
82
83
84class LogRecordSocketReceiver(six.moves.socketserver.ThreadingTCPServer):
85    """Simple TCP socket-based logging receiver.
86    """
87
88    allow_reuse_address = 1
89
90    def __init__(self, host='localhost', port=None,
91                 handler=LogRecordStreamHandler):
92        if not port:
93            port = utils.get_unused_port()
94        six.moves.socketserver.ThreadingTCPServer.__init__(self, (host, port), handler)
95        self.abort = 0
96        self.timeout = 1
97        self.logname = None
98        self.port = port
99
100
101    def serve_until_stopped(self):
102        """Run the socket receiver until aborted."""
103        print('Log Record Socket Receiver is started.')
104        abort = 0
105        while not abort:
106            rd, wr, ex = select.select([self.socket.fileno()], [], [],
107                                       self.timeout)
108            if rd:
109                self.handle_request()
110            abort = self.abort
111        print('Log Record Socket Receiver is stopped.')
112
113
114class LogSocketServer:
115    """A wrapper class to start and stop a TCP server for logging."""
116
117    process = None
118    port = None
119
120    @staticmethod
121    def start(**kwargs):
122        """Start Log Record Socket Receiver in a new process.
123
124        @param kwargs: log configuration, e.g., format, filename.
125
126        @raise Exception: if TCP server is already running.
127        """
128        if LogSocketServer.process:
129            raise Exception('Log Record Socket Receiver is already running.')
130        server_started = multiprocessing.Value(ctypes.c_bool, False)
131        port = multiprocessing.Value(ctypes.c_int, 0)
132        LogSocketServer.process = multiprocessing.Process(
133                target=LogSocketServer._start_server,
134                args=(server_started, port),
135                kwargs=kwargs)
136        LogSocketServer.process.start()
137        while not server_started.value:
138            time.sleep(0.1)
139        LogSocketServer.port = port.value
140        print('Log Record Socket Server is started at port %d.' % port.value)
141
142
143    @staticmethod
144    def _start_server(server_started, port, **kwargs):
145        """Start the TCP server to receive log.
146
147        @param server_started: True if socket log server is started.
148        @param port: Port used by socket log server.
149        @param kwargs: log configuration, e.g., format, filename.
150        """
151        # Clear all existing log handlers.
152        logging.getLogger().handlers = []
153        if not kwargs:
154            logging.basicConfig(
155                    format='%(asctime)s - %(levelname)s - %(message)s')
156        else:
157            logging.basicConfig(**kwargs)
158
159        tcp_server = LogRecordSocketReceiver()
160        print('Starting TCP server...')
161        server_started.value = True
162        port.value = tcp_server.port
163        tcp_server.serve_until_stopped()
164
165
166    @staticmethod
167    def stop():
168        """Stop Log Record Socket Receiver.
169        """
170        if LogSocketServer.process:
171            LogSocketServer.process.terminate()
172            LogSocketServer.process = None
173            LogSocketServer.port = None
174