diff --git a/capnp/includes/schema_cpp.pxd b/capnp/includes/schema_cpp.pxd index 4d869cbf1..cf84f81b8 100644 --- a/capnp/includes/schema_cpp.pxd +++ b/capnp/includes/schema_cpp.pxd @@ -2,7 +2,9 @@ # distutils: language = c++ from libc.stdint cimport * +from libcpp.set cimport set from capnp.helpers.non_circular cimport c_reraise_kj_exception as reraise_kj_exception +from capnp.includes.capnp_cpp cimport TypeWhich from capnp.includes.types cimport * @@ -649,6 +651,13 @@ cdef extern from "capnp/message.h" namespace " ::capnp": uint64_t traversalLimitInWords uint nestingLimit + cdef cppclass LazyZeroSegmentAlloc "BuilderOptions::LazyZeroSegmentAlloc": + bint enableLazyZero + set[TypeWhich] skipLazyZeroTypes + + cdef cppclass BuilderOptions: + LazyZeroSegmentAlloc lazyZeroSegmentAlloc + cdef cppclass MessageBuilder nogil: CodeGeneratorRequest.Builder getRootCodeGeneratorRequest'getRoot< ::capnp::schema::CodeGeneratorRequest>'() CodeGeneratorRequest.Builder initRootCodeGeneratorRequest'initRoot< ::capnp::schema::CodeGeneratorRequest>'() @@ -683,6 +692,9 @@ cdef extern from "capnp/message.h" namespace " ::capnp": DynamicOrphan newOrphan'getOrphanage().newOrphan'(StructSchema) + void setOptions'setOptions'(BuilderOptions options) + BuilderOptions getOptions'getOptions'() const + cdef cppclass MessageReader nogil: CodeGeneratorRequest.Reader getRootCodeGeneratorRequest'getRoot< ::capnp::schema::CodeGeneratorRequest>'() InterfaceNode.Reader getRootInterfaceNode'getRoot< ::capnp::schema::InterfaceNode>'() diff --git a/capnp/lib/capnp.pxd b/capnp/lib/capnp.pxd index d84576750..29096be11 100644 --- a/capnp/lib/capnp.pxd +++ b/capnp/lib/capnp.pxd @@ -151,6 +151,15 @@ cdef class _MessageBuilder: cpdef set_root(self, value) cpdef get_segments_for_output(self) cpdef new_orphan(self, schema) + cpdef get_options(self) + cpdef set_options(self, BuilderOptions py_opts) + +cdef class LazyZeroSegmentAlloc: + cdef public bint enableLazyZero + cdef public object skipLazyZeroTypes + +cdef class BuilderOptions: + cdef public LazyZeroSegmentAlloc lazyZeroSegmentAlloc cdef to_python_reader(C_DynamicValue.Reader self, object parent) cdef to_python_builder(C_DynamicValue.Builder self, object parent) diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx index b69e0926d..ce424df43 100644 --- a/capnp/lib/capnp.pyx +++ b/capnp/lib/capnp.pyx @@ -3196,12 +3196,14 @@ class _StructABCMeta(type): return isinstance(obj, cls.__base__) and obj.schema == cls._schema -cdef _new_message(self, kwargs, num_first_segment_words, allocate_seg_callable): +cdef _new_message(self, kwargs, num_first_segment_words, allocate_seg_callable, builder_options): cdef _MessageBuilder builder if allocate_seg_callable is None: builder = _MallocMessageBuilder(num_first_segment_words) + builder.set_options(builder_options) else: builder = _PyCustomMessageBuilder(allocate_seg_callable, num_first_segment_words) + builder.set_options(builder_options) msg = builder.init_root(self.schema) if kwargs is not None: msg.from_dict(kwargs) @@ -3442,7 +3444,7 @@ 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, allocate_seg_callable=None, **kwargs): + def new_message(self, num_first_segment_words=None, allocate_seg_callable=None, builder_options=BuilderOptions(), **kwargs): """Returns a newly allocated builder message. :type num_first_segment_words: int @@ -3453,6 +3455,9 @@ class _StructModule(object): words to allocate (as an `int`) and returns a `bytearray`. This is used to customize the memory allocation strategy. + :type builder_options: BuilderOptions + :param builder_options: This is for configuring message builder options, such as lazy zero segment alloc. + :type kwargs: dict :param kwargs: A list of fields and their values to initialize in the struct. @@ -3461,7 +3466,7 @@ class _StructModule(object): :rtype: :class:`_DynamicStructBuilder` """ - return _new_message(self, kwargs, num_first_segment_words, allocate_seg_callable) + return _new_message(self, kwargs, num_first_segment_words, allocate_seg_callable, builder_options) class _InterfaceModule(object): @@ -3795,6 +3800,51 @@ cdef class _MessageBuilder: ptr = s._thisptr() return _DynamicOrphan()._init(self.thisptr.newOrphan(ptr), self) + cpdef set_options(self, BuilderOptions py_opts): + if py_opts is None: + raise ValueError("options must be an BuilderOptions instance, got None") + + cdef schema_cpp.BuilderOptions opts + opts.lazyZeroSegmentAlloc.enableLazyZero = py_opts.lazyZeroSegmentAlloc.enableLazyZero + + if py_opts.lazyZeroSegmentAlloc is not None: + py_set = py_opts.lazyZeroSegmentAlloc.skipLazyZeroTypes + if py_set is not None: + try: + if types.Data in py_set: + opts.lazyZeroSegmentAlloc.skipLazyZeroTypes.insert(capnp.TypeWhichDATA) + except TypeError: + raise TypeError("skipLazyZeroTypes must be an iterable (e.g. set)") + + self.thisptr.setOptions(opts) + + + cpdef get_options(self): + cdef schema_cpp.BuilderOptions opts = self.thisptr.getOptions() + + cdef BuilderOptions py_opts = BuilderOptions() + py_opts.lazyZeroSegmentAlloc = LazyZeroSegmentAlloc() + py_opts.lazyZeroSegmentAlloc.enableLazyZero = opts.lazyZeroSegmentAlloc.enableLazyZero + py_opts.lazyZeroSegmentAlloc.skipLazyZeroTypes = set() + + if opts.lazyZeroSegmentAlloc.skipLazyZeroTypes.count(capnp.TypeWhichDATA): + py_opts.lazyZeroSegmentAlloc.skipLazyZeroTypes.add(types.Data) + + return py_opts + +cdef class LazyZeroSegmentAlloc: + def __init__(self, enableLazyZero=False, skipLazyZeroTypes=None): + self.enableLazyZero = enableLazyZero + self.skipLazyZeroTypes = skipLazyZeroTypes + if skipLazyZeroTypes == None: + self.skipLazyZeroTypes = set() + +cdef class BuilderOptions: + def __init__(self, lazyZeroSegmentAlloc=None): + self.lazyZeroSegmentAlloc = lazyZeroSegmentAlloc + if lazyZeroSegmentAlloc == None: + self.lazyZeroSegmentAlloc = LazyZeroSegmentAlloc() + cdef class _MallocMessageBuilder(_MessageBuilder): """The main class for building Cap'n Proto messages diff --git a/examples/lazy_zero_allocation.py b/examples/lazy_zero_allocation.py new file mode 100644 index 000000000..00463da27 --- /dev/null +++ b/examples/lazy_zero_allocation.py @@ -0,0 +1,41 @@ +import capnp +import random + +import addressbook_capnp + + +WORD_SIZE = 8 + + +class Allocator: + def __init__(self): + return + + def __call__(self, minimum_size: int) -> bytearray: + # dirty the memory on purpose, to simulate reusing shared memory + return bytearray(random.getrandbits(8) for _ in range(minimum_size * WORD_SIZE)) + + +lazy_zero = capnp.LazyZeroSegmentAlloc(enableLazyZero=True, skipLazyZeroTypes={capnp.types.Data}) +builder_options=capnp.BuilderOptions(lazyZeroSegmentAlloc=lazy_zero) + +# message creation method 1 +person = addressbook_capnp.Person.new_message(allocate_seg_callable=Allocator(), builder_options=builder_options) +print(person.name) # guaranteed empty string +person.name = "test name 1" +print(person.name) +person.init("extraData", 100) # random dirty bytes +print(bytes(person.extraData)) +print() + +# message creation method 2 +builder = capnp._PyCustomMessageBuilder(allocate_seg_callable=Allocator()) +builder.set_options(builder_options) +person = builder.init_root(addressbook_capnp.Person) +print(person.name) # guaranteed empty string +person.name = "test name 2" +print(person.name) +person.init("extraData", 100) # random dirty bytes +print(bytes(person.extraData)) +builder.get_options() +print(len(builder.get_options().lazyZeroSegmentAlloc.skipLazyZeroTypes)) # return a set of length 1