fix memoryview avoid double release buffer

qingshi163 · qingshi163 · commit 0041a81d16fa · 2021-10-25T17:26:54.000+02:00
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,12 @@ impl AsBuffer for PyMemoryView {
     }
 }
 
+impl Drop for PyMemoryView {
+    fn drop(&mut self) {
+        self.release();
+    }
+}
+
 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();
@@ -193,9 +191,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)
         }
