fix memoryview avoid double release buffer

qingshi163 · qingshi163 · commit aec309c76b61 · 2021-10-25T21:05:13.000+02:00
diff --git a/vm/src/builtins/bytearray.rs b/vm/src/builtins/bytearray.rs
@@ -115,6 +115,12 @@ impl PyByteArray {
         Ok(())
     }
 
+    #[cfg(debug_assertions)]
+    #[pyproperty]
+    fn exports(&self) -> usize {
+        self.exports.load()
+    }
+
     #[inline]
     fn inner(&self) -> PyRwLockReadGuard<'_, PyBytesInner> {
         self.inner.read()
diff --git a/vm/src/builtins/memory.rs b/vm/src/builtins/memory.rs
@@ -19,7 +19,7 @@ use crate::{
 use crossbeam_utils::atomic::AtomicCell;
 use itertools::Itertools;
 use num_traits::ToPrimitive;
-use std::fmt::Debug;
+use std::{fmt::Debug, mem::ManuallyDrop};
 
 #[derive(FromArgs)]
 pub struct PyMemoryViewNewArgs {
@@ -29,7 +29,8 @@ pub struct PyMemoryViewNewArgs {
 #[pyclass(module = false, name = "memoryview")]
 #[derive(Debug)]
 pub struct PyMemoryView {
-    buffer: PyBuffer,
+    // avoid double release when memoryview had released the buffer before drop
+    buffer: ManuallyDrop<PyBuffer>,
     // the released memoryview does not mean the buffer is destoryed
     // because the possible another memeoryview is viewing from it
     released: AtomicCell<bool>,
@@ -93,7 +94,7 @@ impl PyMemoryView {
         let stop = options.len * options.itemsize;
         let format_spec = Self::parse_format(&options.format, vm)?;
         Ok(PyMemoryView {
-            buffer,
+            buffer: ManuallyDrop::new(buffer),
             released: AtomicCell::new(false),
             start: 0,
             stop,
@@ -114,7 +115,7 @@ impl PyMemoryView {
         let itemsize = options.itemsize;
         let format_spec = Self::parse_format(&options.format, vm)?;
         Ok(PyMemoryView {
-            buffer,
+            buffer: ManuallyDrop::new(buffer),
             released: AtomicCell::new(false),
             start: range.start * itemsize,
             stop: range.end * itemsize,
@@ -132,7 +133,6 @@ impl PyMemoryView {
     #[pymethod]
     pub fn release(&self) {
         if self.released.compare_exchange(false, true).is_ok() {
-            self.buffer.manually_release = true;
             self.buffer.release();
         }
     }
@@ -689,6 +689,16 @@ impl AsBuffer for PyMemoryView {
     }
 }
 
+impl Drop for PyMemoryView {
+    fn drop(&mut self) {
+        if self.released.load() {
+            unsafe { self.buffer.drop_without_release() };
+        } else {
+            unsafe { ManuallyDrop::drop(&mut self.buffer) };
+        }
+    }
+}
+
 impl AsMapping for PyMemoryView {
     fn as_mapping(_zelf: &PyObjectView<Self>, _vm: &VirtualMachine) -> PyMappingMethods {
         PyMappingMethods {
diff --git a/vm/src/protocol/buffer.rs b/vm/src/protocol/buffer.rs
@@ -12,6 +12,7 @@ use crate::{
 };
 use std::{borrow::Cow, fmt::Debug};
 
+#[allow(clippy::type_complexity)]
 pub struct BufferMethods {
     // always reflecting the whole bytes of the most top object
     pub obj_bytes: fn(&PyBuffer) -> BorrowedValue<[u8]>,
@@ -45,8 +46,6 @@ impl Debug for BufferMethods {
 pub struct PyBuffer {
     pub obj: PyObjectRef,
     pub options: BufferOptions,
-    // if true, don't call release when drop
-    pub manually_release: bool,
     methods: &'static BufferMethods,
 }
 
@@ -55,7 +54,6 @@ impl PyBuffer {
         let zelf = Self {
             obj,
             options,
-            manually_release: false,
             methods,
         };
         zelf.retain();
@@ -132,12 +130,23 @@ impl PyBuffer {
             .unwrap_or_else(|| self.obj_bytes_mut())
     }
 
+    // WARNING: should always try to clone from the contiguous first
     pub(crate) fn _collect_bytes(&self, buf: &mut Vec<u8>) {
         self.methods
             .collect_bytes
             .map(|f| f(self, buf))
             .unwrap_or_else(|| buf.extend_from_slice(&self.obj_bytes()))
     }
+
+    // drop PyBuffer without calling release
+    // after this function, the owner should use forget()
+    // or wrap PyBuffer in the ManaullyDrop to prevent drop()
+    pub(crate) unsafe fn drop_without_release(&mut self) {
+        // self.obj = PyObjectRef::from_raw(0 as *const PyObject);
+        // self.options = BufferOptions::default();
+        std::ptr::drop_in_place(&mut self.obj);
+        std::ptr::drop_in_place(&mut self.options);
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -193,9 +202,7 @@ impl TryFromBorrowedObject for PyBuffer {
 // but it is not supported by Rust
 impl Drop for PyBuffer {
     fn drop(&mut self) {
-        if !self.manually_release {
-            self.release();
-        }
+        self.release();
     }
 }
 
diff --git a/vm/src/stdlib/io.rs b/vm/src/stdlib/io.rs
@@ -876,9 +876,12 @@ mod _io {
             } else {
                 let v = std::mem::take(&mut self.buffer);
                 let writebuf = VecBuffer::new(v).into_ref(vm);
-                let memobj =
-                    PyMemoryView::from_buffer_range(writebuf.clone().into_readonly_pybuffer(), buf_range, vm)?
-                        .into_ref(vm);
+                let memobj = PyMemoryView::from_buffer_range(
+                    writebuf.clone().into_readonly_pybuffer(),
+                    buf_range,
+                    vm,
+                )?
+                .into_ref(vm);
 
                 // TODO: loop if write() raises an interrupt
                 let res = vm.call_method(self.raw.as_ref().unwrap(), "write", (memobj.clone(),));
@@ -3286,11 +3289,7 @@ mod _io {
                 len: self.buffer.read().cursor.get_ref().len(),
                 ..Default::default()
             };
-            let buffer = PyBuffer::new(
-                self.clone().into_object(),
-                options,
-                &BYTES_IO_BUFFER_METHODS,
-            );
+            let buffer = PyBuffer::new(self.into_object(), options, &BYTES_IO_BUFFER_METHODS);
             let view = PyMemoryView::from_buffer(buffer, vm)?;
             Ok(view)
         }
