Support python custom message builder and make Data field's type return MemoryView (#380)

This PR is for resolving the following issue:
[issue](#379)

1. Created `_PyCustomMessageBuilder` extends `MessageBuilder`, enabling the ability to customise the `SegmentAllocate` method in Python. This allows allocation and data population within shared memory, and supports zero-copy inter-process data transfer by passing segment offsets.

2. Fields of type `Data` now support being set with a `memoryview`. When retrieving a `Data` field from a `DynamicStructBuilder`, it will return a writable `memoryview`, allowing users to modify the data directly. This enables memory to be pre-allocated and content to be modified in later, eliminating an extra copy. When retrieving a `Data` field from a `DynamicStructReader`, it will return a read-only `memoryview`, allowing user to read data without memory copy.

* add memoryview and custom builder

* support set dynamic field

* add curSize

* add initialSize and lastSize

* change StringPtr name

* add test case

* refine test case

* convert func to py callable object

* add initial value

* refine example

* add copy as_reader and new_message, make structReader's data field return RO memoryView

* rebase master and bugfix

* reformat flake8

* refine test case

* refine test cases for blob

* remove unused import for flake8

* run black .

---------

Co-authored-by: Brian Xu <brian.xu1@bytedance.com>

Author: BrianXu0623

capnp/__init__.py

@@ -48,6 +48,7 @@
     _InterfaceModule,
     _ListSchema,
     _MallocMessageBuilder,
+    _PyCustomMessageBuilder,
     _PackedFdMessageReader,
     _StreamFdMessageReader,
     _StructModule,

capnp/includes/PyCustomMessageBuilder.cpp

@@ -0,0 +1,50 @@
+#include "PyCustomMessageBuilder.h"
+#include <stdexcept>
+
+namespace capnp {
+
+PyCustomMessageBuilder::PyCustomMessageBuilder(
+    PyObject* allocateSegmentCallable, uint firstSegmentWords)
+    : allocateSegmentCallable(allocateSegmentCallable), firstSize(firstSegmentWords)
+{
+  KJ_REQUIRE(PyCallable_Check(allocateSegmentCallable),
+               "allocateSegmentCallable must be callable");
+  Py_INCREF(allocateSegmentCallable);
+}
+
+PyCustomMessageBuilder::~PyCustomMessageBuilder() noexcept(false) {
+  PyGILState_STATE gstate = PyGILState_Ensure();
+
+  for (auto* obj : allocatedBuffers) {
+    Py_DECREF(obj);
+  }
+  allocatedBuffers.clear();
+
+  Py_DECREF(allocateSegmentCallable);
+  PyGILState_Release(gstate);
+}
+
+kj::ArrayPtr<capnp::word> PyCustomMessageBuilder::allocateSegment(capnp::uint minimumSize) {
+  PyGILState_STATE gstate = PyGILState_Ensure();
+  KJ_DEFER({ PyGILState_Release(gstate); });
+  if (curSize == 0) {
+    minimumSize = kj::max(minimumSize, firstSize);
+  }
+  PyObject* pyBufObj = PyObject_CallFunction(allocateSegmentCallable, "I", minimumSize);
+  KJ_REQUIRE(pyBufObj, "PyCustomMessageBuilder: allocateSegment failed");
+  allocatedBuffers.push_back(pyBufObj);
+
+
+  Py_buffer view;
+  int bufRes = PyObject_GetBuffer(pyBufObj, &view, PyBUF_SIMPLE);
+  KJ_REQUIRE(bufRes == 0, "PyCustomMessageBuilder: object does not support buffer protocol");
+  KJ_DEFER({ PyBuffer_Release(&view); });
+
+  size_t byteCount = view.len;
+  size_t wordCount = byteCount / sizeof(capnp::word);
+  KJ_REQUIRE(wordCount >= minimumSize, "PyCustomMessageBuilder: buffer too small for minimumSize");
+  curSize += wordCount;
+  return kj::arrayPtr(reinterpret_cast<capnp::word*>(view.buf), wordCount);
+}
+
+}

capnp/includes/PyCustomMessageBuilder.h

@@ -0,0 +1,28 @@
+#pragma once
+
+#include "Python.h"
+#include <capnp/message.h>
+#include <capnp/serialize.h>
+#include <vector>
+
+namespace capnp {
+
+class PyCustomMessageBuilder : public capnp::MessageBuilder {
+public:
+  explicit PyCustomMessageBuilder(PyObject* allocateSegmentCallable,
+  uint firstSegmentWords = capnp::SUGGESTED_FIRST_SEGMENT_WORDS);
+
+  ~PyCustomMessageBuilder() noexcept(false) override;
+
+  kj::ArrayPtr<capnp::word> allocateSegment(capnp::uint minimumSize) override;
+
+private:
+  PyObject* allocateSegmentCallable;
+
+  uint firstSize;
+  uint curSize = 0;
+
+  std::vector<PyObject*> allocatedBuffers;
+};
+
+}

capnp/includes/schema_cpp.pxd

@@ -714,6 +714,11 @@ cdef extern from "capnp/message.h" namespace " ::capnp":
     enum Void:
         VOID
 
+cdef extern from "PyCustomMessageBuilder.h" namespace " ::capnp":
+    cdef cppclass PyCustomMessageBuilder(MessageBuilder):
+        PyCustomMessageBuilder(PyObject* allocateSegmentCallable)
+        PyCustomMessageBuilder(PyObject* allocateSegmentCallable, int firstSegmentSize)
+
 cdef extern from "capnp/common.h" namespace " ::capnp":
     cdef cppclass word nogil:
         pass

capnp/lib/capnp.pxd

@@ -66,7 +66,7 @@ cdef class _DynamicStructReader:
     cpdef _get_by_field(self, _StructSchemaField field)
     cpdef _has_by_field(self, _StructSchemaField field)
 
-    cpdef as_builder(self, num_first_segment_words=?)
+    cpdef as_builder(self, num_first_segment_words=?, allocate_seg_callable=?)
 
 
 cdef class _DynamicStructBuilder:
@@ -99,7 +99,7 @@ cdef class _DynamicStructBuilder:
     cpdef disown(self, field)
 
     cpdef as_reader(self)
-    cpdef copy(self, num_first_segment_words=?)
+    cpdef copy(self, num_first_segment_words=?, allocate_seg_callable=?)
 
 cdef class _DynamicEnumField:
     cdef object thisptr

capnp/lib/capnp.pyx

@@ -13,8 +13,9 @@ from capnp.helpers.helpers cimport init_capnp_api
 from capnp.includes.capnp_cpp cimport AsyncIoStream, WaitScope, PyPromise, VoidPromise, EventPort, EventLoop, PyAsyncIoStream, PromiseFulfiller, VoidPromiseFulfiller, tryReadMessage, writeMessage, makeException, PythonInterfaceDynamicImpl
 from capnp.includes.schema_cpp cimport (MessageReader,)
 
+from builtins import memoryview as BuiltinsMemoryview
 from cpython cimport array, Py_buffer, PyObject_CheckBuffer
-from cpython.buffer cimport PyBUF_SIMPLE, PyBUF_WRITABLE, PyBUF_WRITE, PyBUF_READ
+from cpython.buffer cimport PyBUF_SIMPLE, PyBUF_WRITABLE, PyBUF_WRITE, PyBUF_READ, PyBUF_CONTIG_RO
 from cpython.memoryview cimport PyMemoryView_FromMemory
 from cpython.exc cimport PyErr_Clear
 from cython.operator cimport dereference as deref
@@ -667,7 +668,7 @@ cdef to_python_reader(C_DynamicValue.Reader self, object parent):
         return (<char*>temp_text.begin())[:temp_text.size()]
     elif type == capnp.TYPE_DATA:
         temp_data = self.asData()
-        return <bytes>((<char*>temp_data.begin())[:temp_data.size()])
+        return PyMemoryView_FromMemory(<char *> temp_data.begin(), temp_data.size(), PyBUF_READ)
     elif type == capnp.TYPE_LIST:
         return _DynamicListReader()._init(self.asList(), parent)
     elif type == capnp.TYPE_STRUCT:
@@ -701,7 +702,7 @@ cdef to_python_builder(C_DynamicValue.Builder self, object parent):
         return (<char*>temp_text.begin())[:temp_text.size()]
     elif type == capnp.TYPE_DATA:
         temp_data = self.asData()
-        return <bytes>((<char*>temp_data.begin())[:temp_data.size()])
+        return PyMemoryView_FromMemory(<char *> temp_data.begin(), temp_data.size(), PyBUF_WRITE)
     elif type == capnp.TYPE_LIST:
         return _DynamicListBuilder()._init(self.asList(), parent)
     elif type == capnp.TYPE_STRUCT:
@@ -766,6 +767,20 @@ cdef _setBytes(_DynamicSetterClasses thisptr, field, value):
     cdef C_DynamicValue.Reader temp = C_DynamicValue.Reader(temp_string)
     thisptr.set(field, temp)
 
+cdef _setMemoryview(_DynamicSetterClasses thisptr, field, value):
+    cdef Py_buffer buf
+    cdef capnp.StringPtr temp_string
+    cdef C_DynamicValue.Reader temp
+    if PyObject_GetBuffer(value, &buf, PyBUF_CONTIG_RO) != 0:
+        raise KjException(
+            "cannot get buffer from memory view, for field '{}'".format(field)
+        )
+    try:
+        temp_string = capnp.StringPtr(<char *> buf.buf, buf.len)
+        temp = C_DynamicValue.Reader(temp_string)
+        thisptr.set(field, temp)
+    finally:
+        PyBuffer_Release(&buf)
 
 cdef _setBaseString(_DynamicSetterClasses thisptr, field, value):
     encoded_value = value.encode('utf-8')
@@ -779,6 +794,20 @@ cdef _setBytesField(DynamicStruct_Builder thisptr, _StructSchemaField field, val
     cdef C_DynamicValue.Reader temp = C_DynamicValue.Reader(temp_string)
     thisptr.setByField(field.thisptr, temp)
 
+cdef _setMemoryviewField(DynamicStruct_Builder thisptr, _StructSchemaField field, value):
+    cdef Py_buffer buf
+    cdef capnp.StringPtr temp_string
+    cdef C_DynamicValue.Reader temp
+    if PyObject_GetBuffer(value, &buf, PyBUF_CONTIG_RO) != 0:
+        raise KjException(
+            "cannot get buffer from memory view, for field '{}'".format(field)
+        )
+    try:
+        temp_string = capnp.StringPtr(<char *>buf.buf, buf.len)
+        temp = C_DynamicValue.Reader(temp_string)
+        thisptr.setByField(field.thisptr, temp)
+    finally:
+        PyBuffer_Release(&buf)
 
 cdef _setBaseStringField(DynamicStruct_Builder thisptr, _StructSchemaField field, value):
     encoded_value = value.encode('utf-8')
@@ -805,6 +834,8 @@ cdef _setDynamicField(_DynamicSetterClasses thisptr, field, value, parent):
         thisptr.set(field, temp)
     elif value_type is bytes:
         _setBytes(thisptr, field, value)
+    elif isinstance(value, BuiltinsMemoryview):
+        _setMemoryview(thisptr, field, value)
     elif isinstance(value, basestring):
         _setBaseString(thisptr, field, value)
     elif value_type is list:
@@ -870,6 +901,8 @@ cdef _setDynamicFieldWithField(DynamicStruct_Builder thisptr, _StructSchemaField
         thisptr.setByField(field.thisptr, temp)
     elif value_type is bytes:
         _setBytesField(thisptr, field, value)
+    elif isinstance(value, BuiltinsMemoryview):
+        _setMemoryviewField(thisptr, field, value)
     elif isinstance(value, basestring):
         _setBaseStringField(thisptr, field, value)
     elif value_type is list:
@@ -1242,19 +1275,28 @@ cdef class _DynamicStructReader:
     def to_dict(self, verbose=False, ordered=False, encode_bytes_as_base64=False):
         return _to_dict(self, verbose, ordered, encode_bytes_as_base64)
 
-    cpdef as_builder(self, num_first_segment_words=None):
+    cpdef as_builder(self, num_first_segment_words=None, allocate_seg_callable=None):
         """A method for casting this Reader to a Builder
 
         This is a copying operation with respect to the message's buffer.
         Changes in the new builder will not reflect in the original reader.
 
         :type num_first_segment_words: int
         :param num_first_segment_words: Size of the first segment to allocate (in words ie. 8 byte increments)
+        
+        :type allocate_seg_callable: Callable[[int], bytearray]
+        :param allocate_seg_callable: A python callable object that takes the minimum number of 8-byte 
+        words to allocate (as an `int`) and returns a `bytearray`. This is used to customize the memory 
+        allocation strategy.
 
         :rtype: :class:`_DynamicStructBuilder`
         """
-        builder = _MallocMessageBuilder(num_first_segment_words)
-        return builder.set_root(self)
+        if allocate_seg_callable is None:
+            builder = _MallocMessageBuilder(num_first_segment_words)
+            return builder.set_root(self)
+        else:
+            builder = _PyCustomMessageBuilder(allocate_seg_callable, num_first_segment_words)
+            return builder.set_root(self)
 
     property total_size:
         def __get__(self):
@@ -1593,19 +1635,28 @@ cdef class _DynamicStructBuilder:
         reader._obj_to_pin = self
         return reader
 
-    cpdef copy(self, num_first_segment_words=None):
+    cpdef copy(self, num_first_segment_words=None, allocate_seg_callable=None):
         """A method for copying this Builder
 
         This is a copying operation with respect to the message's buffer.
         Changes in the new builder will not reflect in the original reader.
 
         :type num_first_segment_words: int
         :param num_first_segment_words: Size of the first segment to allocate (in words ie. 8 byte increments)
+        
+        :type allocate_seg_callable: Callable[[int], bytearray]
+        :param allocate_seg_callable: A python callable object that takes the minimum number of 8-byte 
+        words to allocate (as an `int`) and returns a `bytearray`. This is used to customize the memory 
+        allocation strategy.
 
         :rtype: :class:`_DynamicStructBuilder`
         """
-        builder = _MallocMessageBuilder(num_first_segment_words)
-        return builder.set_root(self)
+        if allocate_seg_callable is None:
+            builder = _MallocMessageBuilder(num_first_segment_words)
+            return builder.set_root(self)
+        else:
+            builder = _PyCustomMessageBuilder(allocate_seg_callable, num_first_segment_words)
+            return builder.set_root(self)
 
     property schema:
         """A property that returns the _StructSchema object matching this writer"""
@@ -3145,8 +3196,12 @@ class _StructABCMeta(type):
         return isinstance(obj, cls.__base__) and obj.schema == cls._schema
 
 
-cdef _new_message(self, kwargs, num_first_segment_words):
-    builder = _MallocMessageBuilder(num_first_segment_words)
+cdef _new_message(self, kwargs, num_first_segment_words, allocate_seg_callable):
+    cdef _MessageBuilder builder
+    if allocate_seg_callable is None:
+        builder = _MallocMessageBuilder(num_first_segment_words)
+    else:
+        builder = _PyCustomMessageBuilder(allocate_seg_callable, num_first_segment_words)
     msg = builder.init_root(self.schema)
     if kwargs is not None:
         msg.from_dict(kwargs)
@@ -3387,12 +3442,17 @@ class _StructModule(object):
     def __call__(self, num_first_segment_words=None, **kwargs):
         return self.new_message(num_first_segment_words=num_first_segment_words, **kwargs)
 
-    def new_message(self, num_first_segment_words=None, **kwargs):
+    def new_message(self, num_first_segment_words=None, allocate_seg_callable=None, **kwargs):
         """Returns a newly allocated builder message.
 
         :type num_first_segment_words: int
         :param num_first_segment_words: Size of the first segment to allocate (in words ie. 8 byte increments)
 
+        :type allocate_seg_callable: Callable[[int], bytearray]
+        :param allocate_seg_callable: A python callable object that takes the minimum number of 8-byte
+        words to allocate (as an `int`) and returns a `bytearray`. This is used to customize the memory
+        allocation strategy.
+
         :type kwargs: dict
         :param kwargs: A list of fields and their values to initialize in the struct.
 
@@ -3401,7 +3461,7 @@ class _StructModule(object):
 
         :rtype: :class:`_DynamicStructBuilder`
         """
-        return _new_message(self, kwargs, num_first_segment_words)
+        return _new_message(self, kwargs, num_first_segment_words, allocate_seg_callable)
 
 
 class _InterfaceModule(object):
@@ -3758,6 +3818,50 @@ cdef class _MallocMessageBuilder(_MessageBuilder):
             self.thisptr = new schema_cpp.MallocMessageBuilder(size)
 
 
+cdef class _PyCustomMessageBuilder(_MessageBuilder):
+    """The class for building Cap'n Proto messages,
+    with customised memory allocation strategy
+
+    You will use this class if you want to customise the allocateSegment method,
+    and define your own memory allocation strategy.
+    """
+    def __init__(self, allocate_seg_callable, size=None):
+        """ The constructor requires you to provide a Python callable object as a parameter.
+        This callable object will be invoked in the allocateSegment method of the MessageBuilder
+        to allocate memory. The allocated memory will be managed within the MessageBuilder.
+
+        :type allocate_seg_callable: Callable[[int], bytearray]
+        :param allocate_seg_callable: A python callable object that takes the minimum number of 8-byte
+        words to allocate (as an `int`) and returns a `bytearray`. This is used to customize the memory
+        allocation strategy.
+
+        Required function signature is like this:
+        def __call__(self, minimum_size: int) -> bytearray:
+        Note that the unit of minimum_size is words, ie. 8 byte increments.
+
+            class Allocator:
+                def __init__(self):
+                    self.cur_size = 0
+                def __call__(self, minimum_size: int) -> bytearray:
+                    size = max(minimum_size, self.cur_size)
+                    self.cur_size += size
+                    WORD_SIZE = 8
+                    byte_count = size * WORD_SIZE
+                    return bytearray(byte_count)
+
+            addressbook = capnp.load('addressbook.capnp')
+            message = capnp._PyCustomMessageBuilder(allocator)
+            person = message.init_root(addressbook.Person)
+
+        :type size: int
+        :param size: Size of the first segment to allocate (in words ie. 8 byte increments)
+        """
+        if size is None:
+            self.thisptr = new schema_cpp.PyCustomMessageBuilder(<PyObject*>allocate_seg_callable)
+        else:
+            self.thisptr = new schema_cpp.PyCustomMessageBuilder(<PyObject*>allocate_seg_callable, size)
+
+
 cdef class _MessageReader:
     """An abstract base class for reading Cap'n Proto messages