Add to_segment_views for zero-copy serialization (#405)

Co-authored-by: bigtailfox <leoherz.liu@gmail.com>

capnp/lib/capnp.pxd

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

capnp/lib/capnp.pyx

@@ -19,6 +19,7 @@ from cpython.buffer cimport PyBUF_SIMPLE, PyBUF_WRITABLE, PyBUF_WRITE, PyBUF_REA
 from cpython.memoryview cimport PyMemoryView_FromMemory, PyMemoryView_FromBuffer
 from cpython.bytes cimport PyBytes_FromStringAndSize
 from cpython.exc cimport PyErr_Clear
+from cpython.pyport cimport PY_SSIZE_T_MAX
 from cython.operator cimport dereference as deref
 from libc.stdlib cimport malloc, free
 from libc.string cimport memcpy
@@ -1173,6 +1174,70 @@ cdef class _MessageSize:
         self.word_count = word_count
         self.cap_count = cap_count
 
+
+@cython.internal
+cdef class _SegmentView:
+    cdef object _builder
+    cdef const char* _ptr
+    cdef Py_ssize_t _size
+
+    cdef _init(self, object builder, const char* ptr, Py_ssize_t size):
+        self._builder = builder
+        self._ptr = ptr
+        self._size = size
+        return self
+
+    def __getbuffer__(self, Py_buffer *buffer, int flags):
+        if PyBuffer_FillInfo(buffer, self, <void*>self._ptr, self._size, 1, flags) < 0:
+            raise BufferError("Failed to create segment buffer view")
+
+    def __releasebuffer__(self, Py_buffer *buffer):
+        pass
+
+    def __len__(self):
+        return self._size
+
+    def __repr__(self):
+        return '<capnp segment view size=%d>' % self._size
+
+
+@cython.internal
+cdef class _SegmentViews:
+    cdef object _builder
+    cdef list _views
+
+    cdef _init(self, _MessageBuilder builder):
+        cdef schema_cpp.ConstWordArrayArrayPtr segments = builder.thisptr.getSegmentsForOutput()
+        cdef size_t i
+        cdef size_t word_count
+        cdef Py_ssize_t byte_count
+
+        self._builder = builder
+        self._views = []
+        for i in range(0, segments.size()):
+            word_count = segments[i].size()
+            if word_count > <size_t>(PY_SSIZE_T_MAX // 8):
+                raise OverflowError("segment is too large to expose as a Python buffer")
+            byte_count = <Py_ssize_t>(8 * word_count)
+            self._views.append(_SegmentView()._init(
+                builder,
+                <const char*>segments[i].begin(),
+                byte_count))
+        return self
+
+    def __getitem__(self, index):
+        return self._views[index]
+
+    def __iter__(self):
+        return iter(self._views)
+
+    def __len__(self):
+        return len(self._views)
+
+    def __repr__(self):
+        return '<capnp segment views count=%d>' % len(self)
+
+
 if getattr(_sys, 'subversion', [''])[0] == 'PyPy':
     from pickle_helper import _struct_reducer
 else:
@@ -1451,7 +1516,8 @@ cdef class _DynamicStructBuilder:
     cpdef to_segments(_DynamicStructBuilder self):
         """Returns the struct's containing message as a Python list of Python bytes objects.
 
-        This avoids making copies.
+        This copies each output segment into a Python-owned bytes object. Use
+        to_segment_views() for zero-copy, read-only borrowed segment views.
 
         NB: This is not currently supported on PyPy.
 
@@ -1462,6 +1528,18 @@ cdef class _DynamicStructBuilder:
         segments = builder.get_segments_for_output()
         return segments
 
+    cpdef to_segment_views(_DynamicStructBuilder self):
+        """Returns the struct's containing message as zero-copy, read-only segment views.
+
+        The returned views borrow memory from the message builder. Do not mutate, reset, or reuse
+        the builder while the views are still in use.
+
+        :rtype: sequence
+        """
+        self._check_write()
+        cdef _MessageBuilder builder = self._parent
+        return _SegmentViews()._init(builder)
+
     cpdef _to_bytes_packed_helper(_DynamicStructBuilder self, word_count):
         cdef _MessageBuilder builder = self._parent
         array = helpers.messageToPackedBytes(deref(builder.thisptr), word_count)

docs/quickstart.rst

@@ -338,10 +338,18 @@ For compatibility on the Python side, use the ``to_segments()`` and ``from_segme
 
     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::
+This returns a list of copied, Python-owned ``bytes`` objects. 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 high-throughput code that can safely consume borrowed buffers, ``to_segment_views()`` exposes the same output segments without copying them into Python ``bytes`` objects::
+
+    segment_views = alice.to_segment_views()
+    for segment in segment_views:
+        transport.send(segment)
+
+Each segment view supports the Python buffer protocol and is read-only. The returned views borrow memory from the message builder's arena, so do not mutate, reset, or reuse the builder while any segment view is still in use. If you need data that remains independent of the builder lifetime, use ``to_segments()`` instead.
+
 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)

test/test_serialization.py

@@ -1,6 +1,7 @@
 import warnings
 from contextlib import contextmanager
 
+import gc
 import pytest
 import capnp
 import os
@@ -62,6 +63,93 @@ def test_roundtrip_segments(all_types):
     test_regression.check_all_types(msg)
 
 
+@pytest.mark.skipif(
+    platform.python_implementation() == "PyPy",
+    reason="TODO: Investigate segmented serialization support on PyPy.",
+)
+def test_segment_views_are_read_only_buffers(all_types):
+    msg = all_types.TestAllTypes.new_message()
+    test_regression.init_all_types(msg)
+
+    segments = msg.to_segments()
+    segment_views = msg.to_segment_views()
+
+    assert len(segment_views) == len(segments)
+    assert len(segment_views) >= 1
+
+    for segment_view, segment_bytes in zip(segment_views, segments):
+        assert not isinstance(segment_view, bytes)
+        view = memoryview(segment_view)
+        try:
+            assert view.readonly is True
+            assert view.tobytes() == segment_bytes
+        finally:
+            view.release()
+
+
+@pytest.mark.skipif(
+    platform.python_implementation() == "PyPy",
+    reason="TODO: Investigate segmented serialization support on PyPy.",
+)
+def test_roundtrip_segment_views(all_types):
+    msg = all_types.TestAllTypes.new_message()
+    test_regression.init_all_types(msg)
+
+    segment_views = msg.to_segment_views()
+    msg = all_types.TestAllTypes.from_segments(segment_views)
+    test_regression.check_all_types(msg)
+
+
+@pytest.mark.skipif(
+    platform.python_implementation() == "PyPy",
+    reason="TODO: Investigate segmented serialization support on PyPy.",
+)
+def test_segment_views_are_not_writable(all_types):
+    msg = all_types.TestAllTypes.new_message()
+    test_regression.init_all_types(msg)
+
+    segment_views = msg.to_segment_views()
+    view = memoryview(segment_views[0])
+    try:
+        assert len(view) > 0
+        with pytest.raises(TypeError):
+            view[0] = 0
+    finally:
+        view.release()
+
+
+@pytest.mark.skipif(
+    platform.python_implementation() == "PyPy",
+    reason="TODO: Investigate segmented serialization support on PyPy.",
+)
+def test_segment_view_keeps_message_alive(all_types):
+    msg = all_types.TestAllTypes.new_message()
+    test_regression.init_all_types(msg)
+
+    segment_views = msg.to_segment_views()
+    segment_view = segment_views[0]
+    view = memoryview(segment_view)
+    expected = view.tobytes()
+
+    del msg
+    del segment_views
+    del segment_view
+    gc.collect()
+
+    try:
+        assert view.tobytes() == expected
+    finally:
+        view.release()
+
+
+def test_segment_views_require_root_struct(all_types):
+    msg = all_types.TestAllTypes.new_message()
+    nested = msg.init("structField")
+
+    with pytest.raises(capnp.KjException):
+        nested.to_segment_views()
+
+
 @pytest.mark.skipif(
     sys.version_info[0] < 3,
     reason="mmap doesn't implement the buffer interface under python 2.",