skytools.skylog: initial implementation of UdpTNetStringsHandler

2 files changed, 170 insertions, 2 deletions

diff --git a/python/skytools/skylog.py b/python/skytools/skylog.py
index 0279d7d0..e3f63086 100644
--- a/python/skytools/skylog.py
+++ b/python/skytools/skylog.py
@@ -1,11 +1,22 @@
 """Our log handlers for Python's logging package.
 """
 
-import os, time, socket
-import logging, logging.handlers
+import logging
+import logging.handlers
+import os
+import socket
+import time
 
 import skytools
 
+# use fast implementation if available, otherwise fall back to reference one
+try:
+    import tnetstring as tnetstrings
+    tnetstrings.parse = tnetstrings.pop
+except ImportError:
+    import skytools.tnetstrings as tnetstrings
+    tnetstrings.dumps = tnetstrings.dump
+
 __all__ = ['getLogger']
 
 # add TRACE level
@@ -64,6 +75,7 @@ class EasyRotatingFileHandler(logging.handlers.RotatingFileHandler):
         fn = os.path.expanduser(filename)
         logging.handlers.RotatingFileHandler.__init__(self, fn, maxBytes=maxBytes, backupCount=backupCount)
 
+
 # send JSON message over UDP
 class UdpLogServerHandler(logging.handlers.DatagramHandler):
     """Sends log records over UDP to logserver in JSON format."""
@@ -114,6 +126,45 @@ class UdpLogServerHandler(logging.handlers.DatagramHandler):
         sock.sendto(s, (self.host, self.port))
         sock.close()
 
+
+# send TNetStrings message over UDP
+class UdpTNetStringsHandler(logging.handlers.DatagramHandler):
+    """ Sends log records in TNetStrings format over UDP. """
+
+    # LogRecord fields to send
+    send_fields = [
+        'created', 'exc_text', 'levelname', 'levelno', 'message', 'msecs', 'name',
+        #'hostaddr', 'hostname', 'job_name', 'service_name'
+    ]
+
+    def __init__(self, host, port):
+        """ Initializes the handler with a specific host address and port.
+        """
+        logging.handlers.DatagramHandler.__init__(self, host, port)
+        self._udp_reset = 0
+
+    def makePickle(self, record):
+        """ Create message in TNetStrings format.
+        """
+        msg = {}
+        self.format(record) # render 'message' attribute and others
+        for k in self.send_fields:
+            msg[k] = record.__dict__[k]
+        tnetstr = tnetstrings.dumps(msg)
+        return tnetstr
+
+    def send(self, s):
+        """ Cache socket for a moment, then recreate it.
+        """
+        now = time.time()
+        if now - 1 > self._udp_reset:
+            if self.sock:
+                self.sock.close()
+            self.sock = self.makeSocket()
+            self._udp_reset = now
+        self.sock.sendto(s, (self.host, self.port))
+
+
 class LogDBHandler(logging.handlers.SocketHandler):
     """Sends log records into PostgreSQL server.
 
@@ -234,6 +285,7 @@ class LogDBHandler(logging.handlers.SocketHandler):
             query = "select * from log.add(%s, %s, %s)"
             logcur.execute(query, [type, service, msg])
 
+
 # fix unicode bug in SysLogHandler
 class SysLogHandler(logging.handlers.SysLogHandler):
     """Fixes unicode bug in logging.handlers.SysLogHandler."""
@@ -301,6 +353,7 @@ class SysLogHostnameHandler(SysLogHandler):
                                msg)
         return msg
 
+
 try:
     from logging import LoggerAdapter
 except ImportError:
diff --git a/python/skytools/tnetstrings.py b/python/skytools/tnetstrings.py
new file mode 100644
index 00000000..afacc09e
--- /dev/null
+++ b/python/skytools/tnetstrings.py
@@ -0,0 +1,115 @@
+# Note this implementation is more strict than necessary to demonstrate
+# minimum restrictions on types allowed in dictionaries.
+
+def dump(data):
+    if type(data) is long or type(data) is int:
+        out = str(data)
+        return '%d:%s#' % (len(out), out)
+    elif type(data) is float:
+        out = '%f' % data
+        return '%d:%s^' % (len(out), out)
+    elif type(data) is str:
+        return '%d:' % len(data) + data + ',' 
+    elif type(data) is dict:
+        return dump_dict(data)
+    elif type(data) is list:
+        return dump_list(data)
+    elif data == None:
+        return '0:~'
+    elif type(data) is bool:
+        out = repr(data).lower()
+        return '%d:%s!' % (len(out), out)
+    else:
+        assert False, "Can't serialize stuff that's %s." % type(data)
+
+
+def parse(data):
+    payload, payload_type, remain = parse_payload(data)
+
+    if payload_type == '#':
+        value = int(payload)
+    elif payload_type == '}':
+        value = parse_dict(payload)
+    elif payload_type == ']':
+        value = parse_list(payload)
+    elif payload_type == '!':
+        value = payload == 'true'
+    elif payload_type == '^':
+        value = float(payload)
+    elif payload_type == '~':
+        assert len(payload) == 0, "Payload must be 0 length for null."
+        value = None
+    elif payload_type == ',':
+        value = payload
+    else:
+        assert False, "Invalid payload type: %r" % payload_type
+
+    return value, remain
+
+def parse_payload(data):
+    assert data, "Invalid data to parse, it's empty."
+    length, extra = data.split(':', 1)
+    length = int(length)
+
+    payload, extra = extra[:length], extra[length:]
+    assert extra, "No payload type: %r, %r" % (payload, extra)
+    payload_type, remain = extra[0], extra[1:]
+
+    assert len(payload) == length, "Data is wrong length %d vs %d" % (length, len(payload))
+    return payload, payload_type, remain
+
+def parse_list(data):
+    if len(data) == 0: return []
+
+    result = []
+    value, extra = parse(data)
+    result.append(value)
+
+    while extra:
+        value, extra = parse(extra)
+        result.append(value)
+
+    return result
+
+def parse_pair(data):
+    key, extra = parse(data)
+    assert extra, "Unbalanced dictionary store."
+    value, extra = parse(extra)
+
+    return key, value, extra
+
+def parse_dict(data):
+    if len(data) == 0: return {}
+
+    key, value, extra = parse_pair(data)
+    assert type(key) is str, "Keys can only be strings."
+
+    result = {key: value}
+
+    while extra:
+        key, value, extra = parse_pair(extra)
+        result[key] = value
+  
+    return result
+    
+
+
+def dump_dict(data):
+    result = []
+    for k,v in data.items():
+        result.append(dump(str(k)))
+        result.append(dump(v))
+
+    payload = ''.join(result)
+    return '%d:' % len(payload) + payload + '}'
+
+
+def dump_list(data):
+    result = []
+    for i in data:
+        result.append(dump(i))
+
+    payload = ''.join(result)
+    return '%d:' % len(payload) + payload + ']'
+
+