Add side-band-64k support to ReceivePackHandler.

To do this, we add a BufferedPktLineWriter class that buffers writes
independently of pkt-lines. This class is used to wrap the pkt-lines of
the status report inside a longer pkt-line written to the sideband.

Change-Id: Ie5d9294e66d4174e2fe111a8424651a70869a629

NEWS

@@ -35,6 +35,8 @@
   * ObjectStore.iter_tree_contents can optionally yield tree objects as well.
     (Dave Borowitz).
 
+  * Add side-band-64k support to ReceivePackHandler. (Dave Borowitz)
+
   * Change server capabilities methods to classmethods. (Dave Borowitz)
 
   * Tweak server handler injection. (Dave Borowitz)

dulwich/protocol.py

@@ -318,3 +318,46 @@ def ack_type(capabilities):
     elif 'multi_ack' in capabilities:
         return MULTI_ACK
     return SINGLE_ACK
+
+
+class BufferedPktLineWriter(object):
+    """Writer that wraps its data in pkt-lines and has an independent buffer.
+
+    Consecutive calls to write() wrap the data in a pkt-line and then buffers it
+    until enough lines have been written such that their total length (including
+    length prefix) reach the buffer size.
+    """
+
+    def __init__(self, write, bufsize=65515):
+        """Initialize the BufferedPktLineWriter.
+
+        :param write: A write callback for the underlying writer.
+        :param bufsize: The internal buffer size, including length prefixes.
+        """
+        self._write = write
+        self._bufsize = bufsize
+        self._wbuf = StringIO()
+        self._buflen = 0
+
+    def write(self, data):
+        """Write data, wrapping it in a pkt-line."""
+        line = pkt_line(data)
+        line_len = len(line)
+        over = self._buflen + line_len - self._bufsize
+        if over >= 0:
+            start = line_len - over
+            self._wbuf.write(line[:start])
+            self.flush()
+        else:
+            start = 0
+        saved = line[start:]
+        self._wbuf.write(saved)
+        self._buflen += len(saved)
+
+    def flush(self):
+        """Flush all data from the buffer."""
+        data = self._wbuf.getvalue()
+        if data:
+            self._write(data)
+        self._len = 0
+        self._wbuf = StringIO()

dulwich/server.py

@@ -57,6 +57,7 @@ from dulwich.protocol import (
     ack_type,
     extract_capabilities,
     extract_want_line_capabilities,
+    BufferedPktLineWriter,
     )
 from dulwich.repo import (
     Repo,
@@ -577,7 +578,7 @@ class ReceivePackHandler(Handler):
 
     @classmethod
     def capabilities(cls):
-        return ("report-status", "delete-refs")
+        return ("report-status", "delete-refs", "side-band-64k")
 
     def _apply_pack(self, refs):
         f, commit = self.repo.object_store.add_thin_pack()
@@ -622,14 +623,27 @@ class ReceivePackHandler(Handler):
         return status
 
     def _report_status(self, status):
+        if self.has_capability('side-band-64k'):
+            writer = BufferedPktLineWriter(
+              lambda d: self.proto.write_sideband(1, d))
+            write = writer.write
+
+            def flush():
+                writer.flush()
+                self.proto.write_pkt_line(None)
+        else:
+            write = self.proto.write_pkt_line
+            flush = lambda: None
+
         for name, msg in status:
             if name == 'unpack':
-                self.proto.write_pkt_line('unpack %s\n' % msg)
+                write('unpack %s\n' % msg)
             elif msg == 'ok':
-                self.proto.write_pkt_line('ok %s\n' % name)
+                write('ok %s\n' % name)
             else:
-                self.proto.write_pkt_line('ng %s %s\n' % (name, msg))
-        self.proto.write_pkt_line(None)
+                write('ng %s %s\n' % (name, msg))
+        write(None)
+        flush()
 
     def handle(self):
         refs = self.repo.get_refs().items()

dulwich/tests/test_protocol.py

@@ -30,6 +30,7 @@ from dulwich.protocol import (
     SINGLE_ACK,
     MULTI_ACK,
     MULTI_ACK_DETAILED,
+    BufferedPktLineWriter,
     )
 from dulwich.tests import TestCase
 
@@ -192,3 +193,57 @@ class CapabilitiesTestCase(TestCase):
         self.assertEquals(MULTI_ACK_DETAILED,
                           ack_type(['foo', 'bar', 'multi_ack',
                                     'multi_ack_detailed']))
+
+
+class BufferedPktLineWriterTests(TestCase):
+
+    def setUp(self):
+        self._output = StringIO()
+        self._writer = BufferedPktLineWriter(self._output.write, bufsize=16)
+
+    def assertOutputEquals(self, expected):
+        self.assertEquals(expected, self._output.getvalue())
+
+    def _truncate(self):
+        self._output.seek(0)
+        self._output.truncate()
+
+    def test_write(self):
+        self._writer.write('foo')
+        self.assertOutputEquals('')
+        self._writer.flush()
+        self.assertOutputEquals('0007foo')
+
+    def test_write_none(self):
+        self._writer.write(None)
+        self.assertOutputEquals('')
+        self._writer.flush()
+        self.assertOutputEquals('0000')
+
+    def test_flush_empty(self):
+        self._writer.flush()
+        self.assertOutputEquals('')
+
+    def test_write_multiple(self):
+        self._writer.write('foo')
+        self._writer.write('bar')
+        self.assertOutputEquals('')
+        self._writer.flush()
+        self.assertOutputEquals('0007foo0007bar')
+
+    def test_write_across_boundary(self):
+        self._writer.write('foo')
+        self._writer.write('barbaz')
+        self.assertOutputEquals('0007foo000abarba')
+        self._truncate()
+        self._writer.flush()
+        self.assertOutputEquals('z')
+
+    def test_write_to_boundary(self):
+        self._writer.write('foo')
+        self._writer.write('barba')
+        self.assertOutputEquals('0007foo0009barba')
+        self._truncate()
+        self._writer.write('z')
+        self._writer.flush()
+        self.assertOutputEquals('0005z')