python - 使用 numpy/ctypes 公开 C 分配的内存缓冲区的更安全方法？

Question

我正在为 C 库编写 Python 绑定(bind)，该库使用共享内存缓冲区来存储其内部状态。这些缓冲区的分配和释放是由库本身在 Python 之外完成的，但我可以通过从 Python 中调用包装的构造函数/析构函数来间接控制何时发生这种情况。我想向 Python 公开一些缓冲区，以便我可以从中读取，并在某些情况下将值推送给它们。性能和内存使用是重要的问题，因此我希望尽可能避免复制数据。

我目前的方法是创建一个 numpy 数组，它可以直接查看 ctypes 指针:

import numpy as np
import ctypes as C

libc = C.CDLL('libc.so.6')

class MyWrapper(object):

    def __init__(self, n=10):
        # buffer allocated by external library
        addr = libc.malloc(C.sizeof(C.c_int) * n)
        self._cbuf = (C.c_int * n).from_address(addr)

    def __del__(self):
        # buffer freed by external library
        libc.free(C.addressof(self._cbuf))
        self._cbuf = None

    @property
    def buffer(self):
        return np.ctypeslib.as_array(self._cbuf)

除了避免复制之外，这还意味着我可以使用 numpy 的索引和赋值语法，并将其直接传递给其他 numpy 函数:

wrap = MyWrapper()
buf = wrap.buffer       # buf is now a writeable view of a C-allocated buffer

buf[:] = np.arange(10)  # this is pretty cool!
buf[::2] += 10

print(wrap.buffer)
# [10  1 12  3 14  5 16  7 18  9]

然而，它本身也很危险:

del wrap                # free the pointer

print(buf)              # this is bad!
# [1852404336 1969367156  538978662  538976288  538976288  538976288
#  1752440867 1763734377 1633820787       8548]

# buf[0] = 99           # uncomment this line if you <3 segfaults

为了更安全，我需要能够在尝试读取/写入数组内容之前检查底层 C 指针是否已被释放。我对如何做到这一点有一些想法:

• 一种方法是生成一个 np.ndarray 的子类，该子类包含对 MyWrapper 的 _cbuf 属性的引用，检查它是否在对其底层内存进行任何读取/写入之前为 None，如果是这种情况，则会引发异常。
• 我可以轻松地在同一个缓冲区上生成多个 View ，例如通过 .view 强制转换或切片，因此每个都需要继承对 _cbuf 的引用和执行检查的方法。我怀疑这可以通过覆盖 __array_finalize__ 来实现，但我不确定具体如何。
• 还需要在读取和/或写入数组内容的任何操作之前调用“指针检查”方法。我对 numpy 的内部结构知之甚少，无法提供详尽的覆盖方法列表。

如何实现执行此检查的 np.ndarray 的子类？任何人都可以提出更好的方法吗？

---

更新:这个类做了我想要的大部分:

class SafeBufferView(np.ndarray):

    def __new__(cls, get_buffer, shape=None, dtype=None):
        obj = np.ctypeslib.as_array(get_buffer(), shape).view(cls)
        if dtype is not None:
            obj.dtype = dtype
        obj._get_buffer = get_buffer
        return obj

    def __array_finalize__(self, obj):
        if obj is None: return
        self._get_buffer = getattr(obj, "_get_buffer", None)

    def __array_prepare__(self, out_arr, context=None):
        if not self._get_buffer(): raise Exception("Dangling pointer!")
        return out_arr

    # this seems very heavy-handed - surely there must be a better way?
    def __getattribute__(self, name):
        if name not in ["__new__", "__array_finalize__", "__array_prepare__",
                        "__getattribute__", "_get_buffer"]:
            if not self._get_buffer(): raise Exception("Dangling pointer!")
        return super(np.ndarray, self).__getattribute__(name)

例如:

wrap = MyWrapper()
sb = SafeBufferView(lambda: wrap._cbuf)
sb[:] = np.arange(10)

print(repr(sb))
# SafeBufferView([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)

print(repr(sb[::2]))
# SafeBufferView([0, 2, 4, 6, 8], dtype=int32)

sbv = sb.view(np.double)
print(repr(sbv))
# SafeBufferView([  2.12199579e-314,   6.36598737e-314,   1.06099790e-313,
#          1.48539705e-313,   1.90979621e-313])

# we have to call the destructor method of `wrap` explicitly - `del wrap` won't
# do anything because `sb` and `sbv` both hold references to `wrap`
wrap.__del__()

print(sb)                # Exception: Dangling pointer!
print(sb + 1)            # Exception: Dangling pointer!
print(sbv)               # Exception: Dangling pointer!
print(np.sum(sb))        # Exception: Dangling pointer!
print(sb.dot(sb))        # Exception: Dangling pointer!

print(np.dot(sb, sb))    # oops...
# -70104698

print(np.extract(np.ones(10), sb))
# array([251019024,     32522, 498870232,     32522,         4,         5,
#               6,         7,        48,         0], dtype=int32)

# np.copyto(sb, np.ones(10, np.int32))    # don't try this at home, kids!

我确定我还漏掉了其他边缘情况。

---

更新 2:按照 @ivan_pozdeev 的建议，我尝试了 weakref.proxy。这是一个好主意，但不幸的是我看不出它如何与 numpy 数组一起工作。我可以尝试为 .buffer:

返回的 numpy 数组创建一个弱引用

wrap = MyWrapper()
wr = weakref.proxy(wrap.buffer)
print(wr)
# ReferenceError: weakly-referenced object no longer exists
# <weakproxy at 0x7f6fe715efc8 to NoneType at 0x91a870>

我认为这里的问题是 wrap.buffer 返回的 np.ndarray 实例立即超出范围。一种解决方法是让类在初始化时实例化数组，持有对它的强引用，并让 .buffer() getter 返回一个 weakref.proxy 到数组:

class MyWrapper2(object):

    def __init__(self, n=10):
        # buffer allocated by external library
        addr = libc.malloc(C.sizeof(C.c_int) * n)
        self._cbuf = (C.c_int * n).from_address(addr)
        self._buffer = np.ctypeslib.as_array(self._cbuf)

    def __del__(self):
        # buffer freed by external library
        libc.free(C.addressof(self._cbuf))
        self._cbuf = None
        self._buffer = None

    @property
    def buffer(self):
        return weakref.proxy(self._buffer)

但是，如果我在同一个数组上创建第二个 View ，而缓冲区仍处于分配状态，则会中断:

wrap2 = MyWrapper2()
buf = wrap2.buffer
buf[:] = np.arange(10)

buf2 = buf[:]   # create a second view onto the contents of buf

print(repr(buf))
# <weakproxy at 0x7fec3e709b50 to numpy.ndarray at 0x210ac80>
print(repr(buf2))
# array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)

wrap2.__del__()

print(buf2[:])  # this is bad
# [1291716568    32748 1291716568    32748        0        0        0
#         0       48        0] 

print(buf[:])   # WTF?!
# [34525664        0        0        0        0        0        0        0
#         0        0]  

这是 严重 坏了 - 在调用 wrap2.__del__() 之后，我不仅可以读取和写入 buf2 这是一个 numpy 数组查看 wrap2._cbuf，但我什至可以读取和写入 buf，考虑到 wrap2.__del__() 集，这应该是不可能的wrap2._buffer 到 None。

Answer 1

当存在任何 numpy 数组时，您必须保留对 Wrapper 的引用。实现这一点的最简单方法是将此引用保存在 ctype-buffer 的属性中:

class MyWrapper(object):
    def __init__(self, n=10):
        # buffer allocated by external library
        self.size = n
        self.addr = libc.malloc(C.sizeof(C.c_int) * n)

    def __del__(self):
        # buffer freed by external library
        libc.free(self.addr)

    @property
    def buffer(self):
        buf = (C.c_int * self.size).from_address(self.addr)
        buf._wrapper = self
        return np.ctypeslib.as_array(buf)

当最后一个引用(例如最后一个 numpy 数组)被垃圾收集时，您的包装器会自动释放。