Merge pull request #104 from gcv/segment-serialization

Add support for segment serialization

capnp/includes/schema_cpp.pxd

@@ -667,6 +667,8 @@ cdef extern from "capnp/message.h" namespace " ::capnp":
         DynamicStruct_Builder initRootDynamicStruct'initRoot< ::capnp::DynamicStruct>'(StructSchema)
         void setRootDynamicStruct'setRoot< ::capnp::DynamicStruct::Reader>'(DynamicStruct.Reader)
 
+        ConstWordArrayArrayPtr getSegmentsForOutput'getSegmentsForOutput'()
+
         AnyPointer.Builder getRootAnyPointer'getRoot< ::capnp::AnyPointer>'()
 
         DynamicOrphan newOrphan'getOrphanage().newOrphan'(StructSchema)
@@ -691,6 +693,10 @@ cdef extern from "capnp/message.h" namespace " ::capnp":
         MallocMessageBuilder()
         MallocMessageBuilder(int)
 
+    cdef cppclass SegmentArrayMessageReader(MessageReader):
+        SegmentArrayMessageReader(ConstWordArrayArrayPtr array) except +reraise_kj_exception
+        SegmentArrayMessageReader(ConstWordArrayArrayPtr array, ReaderOptions) except +reraise_kj_exception
+
     cdef cppclass FlatMessageBuilder(MessageBuilder):
         FlatMessageBuilder(WordArrayPtr array)
         FlatMessageBuilder(WordArrayPtr array, ReaderOptions)
@@ -714,6 +720,16 @@ cdef extern from "kj/common.h" namespace " ::kj":
         ByteArrayPtr(byte *, size_t size)
         size_t size()
         byte& operator[](size_t index)
+    cdef cppclass ConstWordArrayPtr " ::kj::ArrayPtr< const ::capnp::word>":
+        ConstWordArrayPtr()
+        ConstWordArrayPtr(word *, size_t size)
+        size_t size()
+        const word* begin()
+    cdef cppclass ConstWordArrayArrayPtr " ::kj::ArrayPtr< const ::kj::ArrayPtr< const ::capnp::word>>":
+        ConstWordArrayArrayPtr()
+        ConstWordArrayArrayPtr(ConstWordArrayPtr*, size_t size)
+        size_t size()
+        ConstWordArrayPtr& operator[](size_t index)
 
 cdef extern from "kj/array.h" namespace " ::kj":
     # Cython can't handle Array[word] as a function argument

capnp/lib/capnp.pxd

@@ -53,6 +53,7 @@ cdef class _DynamicStructBuilder:
 
     cdef _check_write(self)
     cpdef to_bytes(_DynamicStructBuilder self) except +reraise_kj_exception
+    cpdef to_segments(_DynamicStructBuilder self) except +reraise_kj_exception
     cpdef _to_bytes_packed_helper(_DynamicStructBuilder self, word_count) except +reraise_kj_exception
     cpdef to_bytes_packed(_DynamicStructBuilder self) except +reraise_kj_exception
 

capnp/lib/capnp.pyx

@@ -1192,6 +1192,20 @@ cdef class _DynamicStructBuilder:
         self._is_written = True
         return ret
 
+    cpdef to_segments(_DynamicStructBuilder self) except +reraise_kj_exception:
+        """Returns the struct's containing message as a Python list of Python bytes objects.
+
+        This avoids making copies.
+
+        NB: This is not currently supported on PyPy.
+
+        :rtype: list
+        """
+        self._check_write()
+        cdef _MessageBuilder builder = self._parent
+        segments = builder.get_segments_for_output()
+        return segments
+
     cpdef _to_bytes_packed_helper(_DynamicStructBuilder self, word_count) except +reraise_kj_exception:
         cdef _MessageBuilder builder = self._parent
         array = helpers.messageToPackedBytes(deref(builder.thisptr), word_count)
@@ -2989,6 +3003,17 @@ class _StructModule(object):
         else:
             message = _FlatArrayMessageReader(buf, traversal_limit_in_words, nesting_limit)
             return message.get_root(self.schema)
+    def from_segments(self, segments, traversal_limit_in_words = None, nesting_limit = None):
+        """Returns a Reader for a list of segment bytes.
+
+        This avoids making copies.
+
+        NB: This is not currently supported on PyPy.
+
+        :rtype: list
+        """
+        message = _SegmentArrayMessageReader(segments, traversal_limit_in_words, nesting_limit)
+        return message.get_root(self.schema)
     def from_bytes_packed(self, buf, traversal_limit_in_words = None, nesting_limit = None):
         """Returns a Reader for the packed object in buf.
 
@@ -3285,6 +3310,18 @@ cdef class _MessageBuilder:
             self.thisptr.setRootDynamicStruct((<_DynamicStructReader>value).thisptr)
             return self.get_root(value.schema)
 
+    cpdef get_segments_for_output(self) except +reraise_kj_exception:
+        segments = self.thisptr.getSegmentsForOutput()
+        res = []
+        cdef const char* ptr
+        cdef bytes segment_bytes
+        for i in range(0, segments.size()):
+            segment = segments[i]
+            ptr = <const char *> segment.begin()
+            segment_bytes = ptr[:8*segment.size()]
+            res.append(segment_bytes)
+        return res
+
     cpdef new_orphan(self, schema) except +reraise_kj_exception:
         """A method for instantiating Cap'n Proto orphans
 
@@ -3738,6 +3775,44 @@ cdef class _FlatArrayMessageReader(_MessageReader):
         del self.thisptr
 
 
+@cython.internal
+cdef class _SegmentArrayMessageReader(_MessageReader):
+
+    cdef object _objects_to_pin
+    cdef schema_cpp.ConstWordArrayPtr* _seg_ptrs
+
+    def __init__(self, segments, traversal_limit_in_words = None, nesting_limit = None):
+        cdef schema_cpp.ReaderOptions opts
+        if traversal_limit_in_words is not None:
+            opts.traversalLimitInWords = traversal_limit_in_words
+        if nesting_limit is not None:
+            opts.nestingLimit = nesting_limit
+        # take a Python array of bytes and constructs a ConstWordArrayArrayPtr
+        num_segments = len(segments)
+        cdef const void* ptr
+        cdef Py_ssize_t segment_size
+        cdef schema_cpp.ConstWordArrayPtr seg_ptr
+        self._seg_ptrs = <schema_cpp.ConstWordArrayPtr*>malloc(num_segments * sizeof(schema_cpp.ConstWordArrayPtr))
+        self._objects_to_pin = []
+        for i in range(0, num_segments):
+            PyObject_AsReadBuffer(segments[i], &ptr, &segment_size)
+            if (<uintptr_t>ptr) % 8 != 0:
+                aligned = _AlignedBuffer(segments[i])
+                ptr = aligned.buf
+                self._objects_to_pin.append(aligned)
+            else:
+                self._objects_to_pin.append(segments[i])
+            seg_ptr = schema_cpp.ConstWordArrayPtr(<schema_cpp.word*>ptr, segment_size//8)
+            self._seg_ptrs[i] = seg_ptr
+        self.thisptr = new schema_cpp.SegmentArrayMessageReader(
+            schema_cpp.ConstWordArrayArrayPtr(self._seg_ptrs, num_segments),
+            opts)
+
+    def __dealloc__(self):
+        free(self._seg_ptrs)
+        del self.thisptr
+
+
 @cython.internal
 cdef class _FlatMessageBuilder(_MessageBuilder):
     cdef object _object_to_pin

docs/quickstart.rst

@@ -304,6 +304,29 @@ There are also packed versions::
 
     alice2 = addressbook_capnp.Person.from_bytes_packed(alice.to_bytes_packed())
 
+
+Byte Segments
+~~~~~~~~~~~~~
+
+.. note:: This feature is not supported in PyPy at the moment, pending investigation.
+
+Cap'n Proto supports a serialization mode which minimizes object copies. In the C++ interface, ``capnp::MessageBuilder::getSegmentsForOutput()`` returns an array of pointers to segments of the message's content without copying. ``capnp::SegmentArrayMessageReader`` performs the reverse operation, i.e., takes an array of pointers to segments and uses the underlying data, again without copying. This produces a different wire serialization format from ``to_bytes()`` serialization, which uses ``capnp::messageToFlatArray()`` and ``capnp::FlatArrayMessageReader`` (both of which use segments internally, but write them in an incompatible way).
+
+For compatibility on the Python side, use the ``to_segments()`` and ``from_segments()`` functions::
+
+    segments = alice.to_segments()
+
+This returns a list of segments, each a byte buffer. Each segment can be, e.g., turned into a ZeroMQ message frame. The list of segments can also be turned back into an object::
+
+    alice = addressbook_capnp.Person.from_segments(segments)
+
+For more information, please refer to the following links:
+
+- `Advice on minimizing copies from Cap'n Proto <https://stackoverflow.com/questions/28149139/serializing-mutable-state-and-sending-it-asynchronously-over-the-network-with-ne/28156323#28156323>`_ (from the author of Cap'n Proto)
+- `Advice on using Cap'n Proto over ZeroMQ <https://stackoverflow.com/questions/32041315/how-to-send-capn-proto-message-over-zmq/32042234#32042234>`_ (from the author of Cap'n Proto)
+- `Discussion about sending and reassembling Cap'n Proto message segments in C++ <https://groups.google.com/forum/#!topic/capnproto/ClDjGbO7egA>`_ (from the Cap'n Proto mailing list; includes sample code)
+
+
 RPC
 ----------
 

test/test_serialization.py

@@ -42,6 +42,14 @@ def test_roundtrip_bytes(all_types):
     msg = all_types.TestAllTypes.from_bytes(message_bytes)
     test_regression.check_all_types(msg)
 
+@pytest.mark.skipif(platform.python_implementation() == 'PyPy', reason="TODO: Investigate why this works on CPython but fails on PyPy.")
+def test_roundtrip_segments(all_types):
+    msg = all_types.TestAllTypes.new_message()
+    test_regression.init_all_types(msg)
+    segments = msg.to_segments()
+    msg = all_types.TestAllTypes.from_segments(segments)
+    test_regression.check_all_types(msg)
+
 @pytest.mark.skipif(sys.version_info[0] < 3, reason="mmap doesn't implement the buffer interface under python 2.")
 def test_roundtrip_bytes_mmap(all_types):
     msg = all_types.TestAllTypes.new_message()