From 3cf67afeb1916385c6556a06f1fb84193836ba99 Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Tue, 26 Jan 2016 21:59:37 +0100
Subject: [PATCH] generalize multiline logging to remote logs and controllers

---
 artiq/master/worker_impl.py |  3 +-
 artiq/protocols/logging.py  | 68 +++++++++++--------------------------
 artiq/tools.py              | 24 ++++++++++++-
 3 files changed, 44 insertions(+), 51 deletions(-)

diff --git a/artiq/master/worker_impl.py b/artiq/master/worker_impl.py
index 53381cd25..30d2a8168 100644
--- a/artiq/master/worker_impl.py
+++ b/artiq/master/worker_impl.py
@@ -7,8 +7,7 @@ from collections import OrderedDict
 
 import artiq
 from artiq.protocols import pipe_ipc, pyon
-from artiq.protocols.logging import multiline_log_config
-from artiq.tools import file_import
+from artiq.tools import multiline_log_config, file_import
 from artiq.master.worker_db import DeviceManager, DatasetManager, get_hdf5_output
 from artiq.language.environment import is_experiment
 from artiq.language.core import set_watchdog_factory, TerminationRequested
diff --git a/artiq/protocols/logging.py b/artiq/protocols/logging.py
index 75a92094a..b6ddb9bb6 100644
--- a/artiq/protocols/logging.py
+++ b/artiq/protocols/logging.py
@@ -3,7 +3,7 @@ import logging
 import re
 
 from artiq.protocols.asyncio_server import AsyncioServer
-from artiq.tools import TaskObject
+from artiq.tools import TaskObject, MultilineFormatter
 
 
 logger = logging.getLogger(__name__)
@@ -86,36 +86,16 @@ class LogParser:
             stream, self.source_cb())
 
 
-class MultilineFormatter(logging.Formatter):
-    def __init__(self):
-        logging.Formatter.__init__(
-            self, "%(levelname)s:%(name)s:%(message)s")
-
-    def format(self, record):
-        r = logging.Formatter.format(self, record)
-        linebreaks = r.count("\n")
-        if linebreaks:
-            i = r.index(":")
-            r = r[:i] + "<" + str(linebreaks + 1) + ">" + r[i:]
-        return r
-
-
-def multiline_log_config(level):
-    root_logger = logging.getLogger()
-    root_logger.setLevel(level)
-    handler = logging.StreamHandler()
-    handler.setFormatter(MultilineFormatter())
-    root_logger.addHandler(handler)
-
-
 _init_string = b"ARTIQ logging\n"
 
 
 class Server(AsyncioServer):
     """Remote logging TCP server.
 
-    Takes one log entry per line, in the format:
-        source:levelno:name:message
+    Log entries are in the format:
+        source:levelno<total_lines>:name:message
+        continuation...
+        ...continuation
     """
     async def _handle_connection_cr(self, reader, writer):
         try:
@@ -123,6 +103,9 @@ class Server(AsyncioServer):
             if line != _init_string:
                 return
 
+            source = None
+            parser = LogParser(lambda: source)
+
             while True:
                 line = await reader.readline()
                 if not line:
@@ -132,20 +115,16 @@ class Server(AsyncioServer):
                 except:
                     return
                 line = line[:-1]
-                linesplit = line.split(":", 3)
-                if len(linesplit) != 4:
-                    logger.warning("received improperly formatted message, "
-                                   "dropping connection")
-                    return
-                source, level, name, message = linesplit
-                try:
-                    level = int(level)
-                except:
-                    logger.warning("received improperly formatted level, "
-                                   "dropping connection")
-                    return
-                log_with_name(name, level, message,
-                              extra={"source": source})
+                if parser.multiline_count:
+                    parser.line_input(line)
+                else:
+                    linesplit = line.split(":", maxsplit=1)
+                    if len(linesplit) != 2:
+                        logger.warning("received improperly formatted message, "
+                                       "dropping connection")
+                        return
+                    source, remainder = linesplit
+                    parser.line_input(remainder)
         finally:
             writer.close()
 
@@ -172,19 +151,12 @@ class LogForwarder(logging.Handler, TaskObject):
         logging.Handler.__init__(self, **kwargs)
         self.host = host
         self.port = port
-        self.setFormatter(logging.Formatter(
-            "%(name)s:%(message)s"))
+        self.setFormatter(MultilineFormatter())
         self._queue = asyncio.Queue(queue_size)
         self.reconnect_timer = reconnect_timer
 
     def emit(self, record):
-        message = self.format(record)
-        for part in message.split("\n"):
-            part = "{}:{}:{}".format(record.source, record.levelno, part)
-            try:
-                self._queue.put_nowait(part)
-            except asyncio.QueueFull:
-                break
+        self._queue.put_nowait(record.source + ":" + self.format(record))
 
     async def _do(self):
         while True:
diff --git a/artiq/tools.py b/artiq/tools.py
index addd5d55c..d8abc603a 100644
--- a/artiq/tools.py
+++ b/artiq/tools.py
@@ -133,8 +133,30 @@ def simple_network_args(parser, default_port):
             group.add_argument("--port-" + name, default=default, type=int,
                            help=h)
 
+class MultilineFormatter(logging.Formatter):
+    def __init__(self):
+        logging.Formatter.__init__(
+            self, "%(levelname)s:%(name)s:%(message)s")
+
+    def format(self, record):
+        r = logging.Formatter.format(self, record)
+        linebreaks = r.count("\n")
+        if linebreaks:
+            i = r.index(":")
+            r = r[:i] + "<" + str(linebreaks + 1) + ">" + r[i:]
+        return r
+
+
+def multiline_log_config(level):
+    root_logger = logging.getLogger()
+    root_logger.setLevel(level)
+    handler = logging.StreamHandler()
+    handler.setFormatter(MultilineFormatter())
+    root_logger.addHandler(handler)
+
+
 def init_logger(args):
-    logging.basicConfig(level=logging.WARNING + args.quiet*10 - args.verbose*10)
+    multiline_log_config(level=logging.WARNING + args.quiet*10 - args.verbose*10)
 
 
 def bind_address_from_args(args):