A rotating log

src/calibre/__init__.py

@@ -166,14 +166,15 @@ def prints(*args, **kwargs):
     Has the same signature as the print function from Python 3, except for the
     additional keyword argument safe_encode, which if set to True will cause the
     function to use repr when encoding fails.
+
+    Returns the number of bytes written.
     '''
     file = kwargs.get('file', sys.stdout)
-    sep  = kwargs.get('sep', ' ')
-    end  = kwargs.get('end', '\n')
-    enc = preferred_encoding
+    sep  = bytes(kwargs.get('sep', ' '))
+    end  = bytes(kwargs.get('end', '\n'))
+    enc = 'utf-8' if 'CALIBRE_WORKER' in os.environ else preferred_encoding
     safe_encode = kwargs.get('safe_encode', False)
-    if 'CALIBRE_WORKER' in os.environ:
-        enc = 'utf-8'
+    count = 0
     for i, arg in enumerate(args):
         if isinstance(arg, unicode):
             if iswindows:
@@ -181,8 +182,10 @@ def prints(*args, **kwargs):
                 cs = Detect(file)
                 if cs.is_console:
                     cs.write_unicode_text(arg)
+                    count += len(arg)
                     if i != len(args)-1:
-                        file.write(bytes(sep))
+                        file.write(sep)
+                        count += len(sep)
                     continue
             try:
                 arg = arg.encode(enc)
@@ -211,12 +214,18 @@ def prints(*args, **kwargs):
 
         try:
             file.write(arg)
+            count += len(arg)
         except:
             import repr as reprlib
-            file.write(reprlib.repr(arg))
+            arg = reprlib.repr(arg)
+            file.write(arg)
+            count += len(arg)
         if i != len(args)-1:
-            file.write(bytes(sep))
-    file.write(bytes(end))
+            file.write(sep)
+            count += len(sep)
+    file.write(end)
+    count += len(sep)
+    return count
 
 class CommandLineError(Exception):
     pass

src/calibre/srv/tests/loop.py

@@ -8,6 +8,7 @@ __copyright__ = '2015, Kovid Goyal <kovid at kovidgoyal.net>'
 
 import httplib, ssl, os, socket, time
 from unittest import skipIf
+from glob import glob
 
 try:
     from calibre.utils.certgen import create_server_cert
@@ -21,6 +22,34 @@ from calibre.ptempfile import TemporaryDirectory
 
 class LoopTest(BaseTest):
 
+    def test_log_rotation(self):
+        'Test log rotation'
+        from calibre.srv.utils import RotatingLog
+        from calibre.ptempfile import TemporaryDirectory
+        with TemporaryDirectory() as tdir:
+            fname = os.path.join(tdir, 'log')
+            l = RotatingLog(fname, max_size=100)
+
+            def history():
+                return {int(x.rpartition('.')[-1]) for x in glob(fname + '.*')}
+
+            def log_size():
+                ssize = l.outputs[0].stream.tell()
+                self.ae(ssize, l.outputs[0].current_pos)
+                self.ae(ssize, os.path.getsize(fname))
+                return ssize
+
+            self.ae(log_size(), 0)
+            l('a' * 99)
+            self.ae(log_size(), 100)
+            l('b'), l('c')
+            self.ae(log_size(), 2)
+            self.ae(history(), {1})
+            for i in 'abcdefg':
+                l(i * 101)
+            self.assertLessEqual(log_size(), 100)
+            self.ae(history(), {1,2,3,4,5})
+
     def test_workers(self):
         ' Test worker semantics '
         with TestServer(lambda data:(data.path[0] + data.read()), worker_count=3) as server:
@@ -42,6 +71,7 @@ class LoopTest(BaseTest):
             self.ae(1, sum(int(w.is_alive()) for w in pool.workers))
 
     def test_ring_buffer(self):
+        'Test the ring buffer used for reads'
         class FakeSocket(object):
             def __init__(self, data):
                 self.data = data

src/calibre/srv/utils.py

@@ -13,8 +13,11 @@ import repr as reprlib
 from email.utils import formatdate
 from operator import itemgetter
 
+from calibre import prints
 from calibre.constants import iswindows
+from calibre.utils.filenames import atomic_rename
 from calibre.utils.socket_inheritance import set_socket_inherit
+from calibre.utils.logging import ThreadSafeLog
 
 HTTP1  = 'HTTP/1.0'
 HTTP11 = 'HTTP/1.1'
@@ -175,6 +178,49 @@ def eintr_retry_call(func, *args, **kwargs):
                 continue
             raise
 
+class RotatingStream(object):
+
+    def __init__(self, filename, max_size=None, history=5):
+        self.filename, self.history, self.max_size = filename, history, max_size
+        self.set_output()
+
+    def set_output(self):
+        self.stream = lopen(self.filename, 'ab', 1)  # line buffered
+        try:
+            self.current_pos = self.stream.tell()
+        except EnvironmentError:
+            # Happens if filename is /dev/stdout for example
+            self.current_pos = 0
+            self.max_size = None
+
+    def flush(self):
+        self.stream.flush()
+
+    def prints(self, level, *args, **kwargs):
+        kwargs['safe_encode'] = True
+        kwargs['file'] = self.stream
+        self.current_pos += prints(*args, **kwargs)
+        self.rollover()
+
+    def rollover(self):
+        if self.max_size is None or self.current_pos <= self.max_size:
+            return
+        self.stream.close()
+        for i in xrange(self.history - 1, 0, -1):
+            try:
+                atomic_rename('%s.%d' % (self.filename, i), '%s.%d' % (self.filename, i+1))
+            except EnvironmentError as e:
+                if e.errno != errno.ENOENT:  # the source of the rename does not exist
+                    raise
+        atomic_rename(self.filename, '%s.%d' % (self.filename, 1))
+        self.set_output()
+
+class RotatingLog(ThreadSafeLog):
+
+    def __init__(self, filename, max_size=None, history=5):
+        ThreadSafeLog.__init__(self)
+        self.outputs = [RotatingStream(filename, max_size, history)]
+
 class HandleInterrupt(object):  # {{{
 
     # On windows socket functions like accept(), recv(), send() are not