Fix: get_data_as_view retaion the backing reader/builder even after the view is released, leading to memory leakage (#406)
* fix: return empty memoryview for uninitialized DATA fields Use a module-level sentinel when Cap'n Proto reports a NULL pointer with zero size so PyBuffer_FillInfo receives a valid address for unset fields. Co-authored-by: Cursor <cursoragent@cursor.com> * fix: release buffer info if memoryview construction fails PyBuffer_FillInfo pins `self` via buf.obj; call PyBuffer_Release on failure so that reference is not leaked. This is safe for sentinel-backed empty views: PyBuffer_Release only decrements buf.obj and does not free buf.buf. Co-authored-by: Cursor <cursoragent@cursor.com> * docs: clarify lifetime rules for zero-copy buffer views Document borrowing semantics, mutation hazards, and empty DATA field behavior for get_data_as_view and to_segment_views. Co-authored-by: Cursor <cursoragent@cursor.com> * fix: pin DATA field views via shared buffer exporter Replace PyMemoryView_FromBuffer with a _BorrowedBufferView holder and PyMemoryView_FromObject so get_data_as_view() correctly pins the struct reader/builder for the memoryview lifetime. Generalize the same exporter for to_segment_views() and add regression tests for packed payload release. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: bigtailfox <leoherz.liu@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -16,7 +16,7 @@ from capnp.includes.schema_cpp cimport (MessageReader,)
|
|||||||
from builtins import memoryview as BuiltinsMemoryview
|
from builtins import memoryview as BuiltinsMemoryview
|
||||||
from cpython cimport array, Py_buffer, PyObject_CheckBuffer
|
from cpython cimport array, Py_buffer, PyObject_CheckBuffer
|
||||||
from cpython.buffer cimport PyBUF_SIMPLE, PyBUF_WRITABLE, PyBUF_WRITE, PyBUF_READ, PyBUF_CONTIG_RO, PyBuffer_FillInfo
|
from cpython.buffer cimport PyBUF_SIMPLE, PyBUF_WRITABLE, PyBUF_WRITE, PyBUF_READ, PyBUF_CONTIG_RO, PyBuffer_FillInfo
|
||||||
from cpython.memoryview cimport PyMemoryView_FromMemory, PyMemoryView_FromBuffer
|
from cpython.memoryview cimport PyMemoryView_FromMemory, PyMemoryView_FromObject
|
||||||
from cpython.bytes cimport PyBytes_FromStringAndSize
|
from cpython.bytes cimport PyBytes_FromStringAndSize
|
||||||
from cpython.exc cimport PyErr_Clear
|
from cpython.exc cimport PyErr_Clear
|
||||||
from cpython.pyport cimport PY_SSIZE_T_MAX
|
from cpython.pyport cimport PY_SSIZE_T_MAX
|
||||||
@@ -54,6 +54,8 @@ _CAPNP_VERSION_MINOR = capnp.CAPNP_VERSION_MINOR
|
|||||||
_CAPNP_VERSION_MICRO = capnp.CAPNP_VERSION_MICRO
|
_CAPNP_VERSION_MICRO = capnp.CAPNP_VERSION_MICRO
|
||||||
_CAPNP_VERSION = capnp.CAPNP_VERSION
|
_CAPNP_VERSION = capnp.CAPNP_VERSION
|
||||||
|
|
||||||
|
cdef char _EMPTY_DATA_VIEW_SENTINEL = 0
|
||||||
|
|
||||||
cdef dict _type_registry = {}
|
cdef dict _type_registry = {}
|
||||||
|
|
||||||
|
|
||||||
@@ -1176,20 +1178,24 @@ cdef class _MessageSize:
|
|||||||
|
|
||||||
|
|
||||||
@cython.internal
|
@cython.internal
|
||||||
cdef class _SegmentView:
|
cdef class _BorrowedBufferView:
|
||||||
cdef object _builder
|
"""Buffer-protocol exporter that pins an owner while a view borrows its memory."""
|
||||||
|
cdef object _owner
|
||||||
cdef const char* _ptr
|
cdef const char* _ptr
|
||||||
cdef Py_ssize_t _size
|
cdef Py_ssize_t _size
|
||||||
|
cdef bint _readonly
|
||||||
|
|
||||||
cdef _init(self, object builder, const char* ptr, Py_ssize_t size):
|
cdef _init(self, object owner, const void* ptr, Py_ssize_t size, bint readonly):
|
||||||
self._builder = builder
|
self._owner = owner
|
||||||
self._ptr = ptr
|
self._ptr = <const char*>ptr
|
||||||
self._size = size
|
self._size = size
|
||||||
|
self._readonly = readonly
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def __getbuffer__(self, Py_buffer *buffer, int flags):
|
def __getbuffer__(self, Py_buffer *buffer, int flags):
|
||||||
if PyBuffer_FillInfo(buffer, self, <void*>self._ptr, self._size, 1, flags) < 0:
|
if PyBuffer_FillInfo(buffer, self, <void*>self._ptr, self._size,
|
||||||
raise BufferError("Failed to create segment buffer view")
|
self._readonly, flags) < 0:
|
||||||
|
raise BufferError("Failed to create borrowed buffer view")
|
||||||
|
|
||||||
def __releasebuffer__(self, Py_buffer *buffer):
|
def __releasebuffer__(self, Py_buffer *buffer):
|
||||||
pass
|
pass
|
||||||
@@ -1198,7 +1204,56 @@ cdef class _SegmentView:
|
|||||||
return self._size
|
return self._size
|
||||||
|
|
||||||
def __repr__(self):
|
def __repr__(self):
|
||||||
return '<capnp segment view size=%d>' % self._size
|
if self._readonly:
|
||||||
|
return '<capnp borrowed buffer view size=%d read-only>' % self._size
|
||||||
|
return '<capnp borrowed buffer view size=%d writable>' % self._size
|
||||||
|
|
||||||
|
|
||||||
|
cdef inline object _memoryview_borrowing(object owner, void* ptr, Py_ssize_t size,
|
||||||
|
bint readonly):
|
||||||
|
cdef _BorrowedBufferView exporter
|
||||||
|
exporter = _BorrowedBufferView()._init(owner, ptr, size, readonly)
|
||||||
|
return PyMemoryView_FromObject(exporter)
|
||||||
|
|
||||||
|
|
||||||
|
cdef void _data_field_ptr_reader(_DynamicStructReader self, field,
|
||||||
|
void** data_ptr, size_t* data_size) except *:
|
||||||
|
cdef C_DynamicValue.Reader val
|
||||||
|
cdef capnp.Data.Reader temp_data
|
||||||
|
|
||||||
|
try:
|
||||||
|
val = self.thisptr.get(field)
|
||||||
|
except KjException as e:
|
||||||
|
raise e._to_python() from None
|
||||||
|
|
||||||
|
if val.getType() != capnp.TYPE_DATA:
|
||||||
|
raise TypeError("Field '{}' is not a DATA field".format(field))
|
||||||
|
|
||||||
|
temp_data = val.asData()
|
||||||
|
data_ptr[0] = <void*>temp_data.begin()
|
||||||
|
data_size[0] = temp_data.size()
|
||||||
|
if data_size[0] == 0 and data_ptr[0] == NULL:
|
||||||
|
data_ptr[0] = <void*>&_EMPTY_DATA_VIEW_SENTINEL
|
||||||
|
|
||||||
|
|
||||||
|
cdef void _data_field_ptr_builder(_DynamicStructBuilder self, field,
|
||||||
|
void** data_ptr, size_t* data_size) except *:
|
||||||
|
cdef C_DynamicValue.Builder val
|
||||||
|
cdef capnp.Data.Builder temp_data
|
||||||
|
|
||||||
|
try:
|
||||||
|
val = self.thisptr.get(field)
|
||||||
|
except KjException as e:
|
||||||
|
raise e._to_python() from None
|
||||||
|
|
||||||
|
if val.getType() != capnp.TYPE_DATA:
|
||||||
|
raise TypeError("Field '{}' is not a DATA field".format(field))
|
||||||
|
|
||||||
|
temp_data = val.asData()
|
||||||
|
data_ptr[0] = <void*>temp_data.begin()
|
||||||
|
data_size[0] = temp_data.size()
|
||||||
|
if data_size[0] == 0 and data_ptr[0] == NULL:
|
||||||
|
data_ptr[0] = <void*>&_EMPTY_DATA_VIEW_SENTINEL
|
||||||
|
|
||||||
|
|
||||||
@cython.internal
|
@cython.internal
|
||||||
@@ -1219,10 +1274,11 @@ cdef class _SegmentViews:
|
|||||||
if word_count > <size_t>(PY_SSIZE_T_MAX // 8):
|
if word_count > <size_t>(PY_SSIZE_T_MAX // 8):
|
||||||
raise OverflowError("segment is too large to expose as a Python buffer")
|
raise OverflowError("segment is too large to expose as a Python buffer")
|
||||||
byte_count = <Py_ssize_t>(8 * word_count)
|
byte_count = <Py_ssize_t>(8 * word_count)
|
||||||
self._views.append(_SegmentView()._init(
|
self._views.append(_BorrowedBufferView()._init(
|
||||||
builder,
|
builder,
|
||||||
<const char*>segments[i].begin(),
|
segments[i].begin(),
|
||||||
byte_count))
|
byte_count,
|
||||||
|
True))
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def __getitem__(self, index):
|
def __getitem__(self, index):
|
||||||
@@ -1292,27 +1348,21 @@ cdef class _DynamicStructReader:
|
|||||||
return self.thisptr.hasByField(field.thisptr)
|
return self.thisptr.hasByField(field.thisptr)
|
||||||
|
|
||||||
cpdef get_data_as_view(self, field):
|
cpdef get_data_as_view(self, field):
|
||||||
|
"""Efficiently get a read-only memoryview for a DATA field without copying.
|
||||||
|
|
||||||
|
.. warning::
|
||||||
|
The returned memoryview *borrows* memory owned by this message. It stays valid while
|
||||||
|
the memoryview (and any object derived from it) is alive; an internal exporter pins
|
||||||
|
this reader for that duration. Do not let the message be mutated underneath an
|
||||||
|
outstanding view.
|
||||||
|
|
||||||
|
An unset/empty DATA field yields a valid, zero-length view (it does not raise).
|
||||||
"""
|
"""
|
||||||
Efficiently get a read-only memoryview for a DATA field without copying.
|
cdef void* data_ptr
|
||||||
"""
|
cdef size_t data_size
|
||||||
cdef C_DynamicValue.Reader val
|
|
||||||
cdef capnp.Data.Reader temp_data
|
|
||||||
|
|
||||||
try:
|
_data_field_ptr_reader(self, field, &data_ptr, &data_size)
|
||||||
val = self.thisptr.get(field)
|
return _memoryview_borrowing(self, data_ptr, <Py_ssize_t>data_size, True)
|
||||||
except KjException as e:
|
|
||||||
raise e._to_python() from None
|
|
||||||
|
|
||||||
if val.getType() != capnp.TYPE_DATA:
|
|
||||||
raise TypeError("Field '{}' is not a DATA field".format(field))
|
|
||||||
|
|
||||||
temp_data = val.asData()
|
|
||||||
|
|
||||||
# Return read-only memoryview
|
|
||||||
cdef Py_buffer buf
|
|
||||||
if PyBuffer_FillInfo(&buf, self, <void*>temp_data.begin(), temp_data.size(), 1, PyBUF_CONTIG_RO) < 0:
|
|
||||||
raise KjException("Failed to create buffer info")
|
|
||||||
return PyMemoryView_FromBuffer(&buf)
|
|
||||||
|
|
||||||
cpdef _which_str(self):
|
cpdef _which_str(self):
|
||||||
try:
|
try:
|
||||||
@@ -1531,8 +1581,21 @@ cdef class _DynamicStructBuilder:
|
|||||||
cpdef to_segment_views(_DynamicStructBuilder self):
|
cpdef to_segment_views(_DynamicStructBuilder self):
|
||||||
"""Returns the struct's containing message as zero-copy, read-only segment views.
|
"""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 returned object is a sequence of read-only buffer-protocol views, one per output
|
||||||
the builder while the views are still in use.
|
segment. Each view borrows memory owned by the message builder; the segment pointers and
|
||||||
|
sizes are captured eagerly at call time (a snapshot).
|
||||||
|
|
||||||
|
.. warning::
|
||||||
|
The views (and any buffer exported from them, e.g. by ``memoryview()`` or by a
|
||||||
|
consumer that holds them) keep the builder pinned and remain valid only while no
|
||||||
|
mutation happens. Do NOT mutate, re-set, reset, or reuse the builder while any view or
|
||||||
|
exported buffer is still alive -- this includes calls that may allocate (e.g. getting
|
||||||
|
an unset pointer/struct/list field). Mutating after the snapshot can grow/relocate
|
||||||
|
segments, leaving the views pointing at stale or truncated data. Sharing the views
|
||||||
|
across threads or ``await`` points while the builder may change is a data race.
|
||||||
|
|
||||||
|
Lifetime is enforced only by buffer-protocol reference counting (memory is not freed
|
||||||
|
while a view is held); correctness of the *contents* is the caller's responsibility.
|
||||||
|
|
||||||
:rtype: sequence
|
:rtype: sequence
|
||||||
"""
|
"""
|
||||||
@@ -1731,30 +1794,29 @@ cdef class _DynamicStructBuilder:
|
|||||||
return _DynamicOrphan()._init(self.thisptr.disown(field), self._parent)
|
return _DynamicOrphan()._init(self.thisptr.disown(field), self._parent)
|
||||||
|
|
||||||
cpdef get_data_as_view(self, field):
|
cpdef get_data_as_view(self, field):
|
||||||
"""
|
"""Efficiently get a writable memoryview for a DATA field without copying.
|
||||||
Efficiently get a writable memoryview for a DATA field without copying.
|
|
||||||
|
This allows in-place modification of the underlying buffer::
|
||||||
|
|
||||||
This allows in-place modification of the underlying buffer:
|
|
||||||
msg.get_data_as_view('myField')[0] = 0xFF
|
msg.get_data_as_view('myField')[0] = 0xFF
|
||||||
|
|
||||||
|
.. warning::
|
||||||
|
The returned memoryview *borrows* mutable memory owned by this message builder. It is
|
||||||
|
valid while the memoryview (and any object derived from it) is alive; an internal
|
||||||
|
exporter pins this builder for that duration. Do NOT mutate, re-set, reset, or reuse the
|
||||||
|
builder while a view is outstanding -- including calls that may allocate (e.g. getting
|
||||||
|
an unset pointer/struct/list field), since those can relocate or stale the borrowed
|
||||||
|
memory. Sharing a view across threads or ``await`` points while the builder may change
|
||||||
|
is a data race.
|
||||||
|
|
||||||
|
An unset/empty DATA field yields a valid but zero-length view (writes are no-ops); to write
|
||||||
|
into the field, initialize it to the desired size first.
|
||||||
"""
|
"""
|
||||||
cdef C_DynamicValue.Builder val
|
cdef void* data_ptr
|
||||||
cdef capnp.Data.Builder temp_data
|
cdef size_t data_size
|
||||||
|
|
||||||
try:
|
_data_field_ptr_builder(self, field, &data_ptr, &data_size)
|
||||||
val = self.thisptr.get(field)
|
return _memoryview_borrowing(self, data_ptr, <Py_ssize_t>data_size, False)
|
||||||
except KjException as e:
|
|
||||||
raise e._to_python() from None
|
|
||||||
|
|
||||||
if val.getType() != capnp.TYPE_DATA:
|
|
||||||
raise TypeError("Field '{}' is not a DATA field".format(field))
|
|
||||||
|
|
||||||
temp_data = val.asData()
|
|
||||||
|
|
||||||
# Return writable memoryview
|
|
||||||
cdef Py_buffer buf
|
|
||||||
if PyBuffer_FillInfo(&buf, self, <void*>temp_data.begin(), temp_data.size(), 0, PyBUF_WRITABLE) < 0:
|
|
||||||
raise KjException("Failed to create buffer info")
|
|
||||||
return PyMemoryView_FromBuffer(&buf)
|
|
||||||
|
|
||||||
cpdef as_reader(self):
|
cpdef as_reader(self):
|
||||||
"""A method for casting this Builder to a Reader
|
"""A method for casting this Builder to a Reader
|
||||||
|
|||||||
@@ -1,4 +1,8 @@
|
|||||||
import os
|
import os
|
||||||
|
import tempfile
|
||||||
|
import weakref
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
import capnp
|
import capnp
|
||||||
import sys
|
import sys
|
||||||
@@ -152,6 +156,33 @@ def test_corner_cases_values(all_types):
|
|||||||
assert msg.get_data_as_view("dataField").tobytes() == binary_data
|
assert msg.get_data_as_view("dataField").tobytes() == binary_data
|
||||||
|
|
||||||
|
|
||||||
|
def test_uninitialized_data_get_view(all_types):
|
||||||
|
"""
|
||||||
|
Default DATA fields should expose an empty memoryview instead of failing on a NULL buffer pointer.
|
||||||
|
"""
|
||||||
|
builder = all_types.TestAllTypes.new_message()
|
||||||
|
builder_view = builder.get_data_as_view("dataField")
|
||||||
|
|
||||||
|
assert isinstance(builder_view, memoryview)
|
||||||
|
assert builder_view.readonly is False
|
||||||
|
assert len(builder_view) == 0
|
||||||
|
assert builder_view.tobytes() == b""
|
||||||
|
|
||||||
|
reader = all_types.TestAllTypes.new_message().as_reader()
|
||||||
|
reader_view = reader.get_data_as_view("dataField")
|
||||||
|
|
||||||
|
assert isinstance(reader_view, memoryview)
|
||||||
|
assert reader_view.readonly is True
|
||||||
|
assert len(reader_view) == 0
|
||||||
|
assert reader_view.tobytes() == b""
|
||||||
|
|
||||||
|
with pytest.raises(IndexError):
|
||||||
|
builder_view[0] = 0xFF
|
||||||
|
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
builder_view[0:1] = b"\xff"
|
||||||
|
|
||||||
|
|
||||||
def test_error_wrong_type(all_types):
|
def test_error_wrong_type(all_types):
|
||||||
"""
|
"""
|
||||||
Test error handling: Calling get_data_as_view on non-Data fields.
|
Test error handling: Calling get_data_as_view on non-Data fields.
|
||||||
@@ -206,3 +237,56 @@ def test_view_keeps_message_alive(all_types):
|
|||||||
gc.collect()
|
gc.collect()
|
||||||
|
|
||||||
assert view.tobytes() == expected_data
|
assert view.tobytes() == expected_data
|
||||||
|
|
||||||
|
|
||||||
|
def test_data_view_exports_through_buffer_exporter(all_types):
|
||||||
|
"""Returned memoryviews should pin an internal exporter, not bare pointers."""
|
||||||
|
msg = all_types.TestAllTypes.new_message()
|
||||||
|
msg.dataField = b"exporter_check"
|
||||||
|
view = msg.get_data_as_view("dataField")
|
||||||
|
|
||||||
|
assert isinstance(view, memoryview)
|
||||||
|
assert view.obj is not None
|
||||||
|
assert len(view.obj) == len(view)
|
||||||
|
|
||||||
|
|
||||||
|
def test_data_view_survives_del_builder(all_types):
|
||||||
|
msg = all_types.TestAllTypes.new_message()
|
||||||
|
msg.dataField = b"persistence_check"
|
||||||
|
view = msg.get_data_as_view("dataField")
|
||||||
|
|
||||||
|
del msg
|
||||||
|
gc.collect()
|
||||||
|
|
||||||
|
assert view.tobytes() == b"persistence_check"
|
||||||
|
|
||||||
|
|
||||||
|
def test_data_view_releases_packed_payload():
|
||||||
|
schema_text = """
|
||||||
|
@0x9d7d4f087df9b6e1;
|
||||||
|
struct BlobMsg {
|
||||||
|
data @0 :Data;
|
||||||
|
}
|
||||||
|
"""
|
||||||
|
|
||||||
|
class Payload(bytearray):
|
||||||
|
pass
|
||||||
|
|
||||||
|
td = tempfile.TemporaryDirectory()
|
||||||
|
path = Path(td.name) / "blob.capnp"
|
||||||
|
path.write_text(schema_text)
|
||||||
|
schema = capnp.load(str(path))
|
||||||
|
try:
|
||||||
|
payload = Payload(schema.BlobMsg.new_message(data=b"x" * 4096).to_bytes_packed())
|
||||||
|
payload_ref = weakref.ref(payload)
|
||||||
|
|
||||||
|
reader = schema.BlobMsg.from_bytes_packed(payload)
|
||||||
|
view = reader.get_data_as_view("data")
|
||||||
|
view.release()
|
||||||
|
|
||||||
|
del view, reader, payload
|
||||||
|
gc.collect()
|
||||||
|
|
||||||
|
assert payload_ref() is None
|
||||||
|
finally:
|
||||||
|
td.cleanup()
|
||||||
|
|||||||
Reference in New Issue
Block a user