diff --git a/benches/bench_dict.rs b/benches/bench_dict.rs index 620550da2fe..114f38eb0f2 100644 --- a/benches/bench_dict.rs +++ b/benches/bench_dict.rs @@ -9,7 +9,7 @@ use pyo3::{prelude::*, types::IntoPyDict}; fn iter_dict(b: &mut Bencher) { let gil = Python::acquire_gil(); let py = gil.python(); - const LEN: usize = 1_000_000; + const LEN: usize = 1_000_00; let dict = (0..LEN as u64).map(|i| (i, i * 2)).into_py_dict(py); let mut sum = 0; b.iter(|| { diff --git a/src/pythonrun.rs b/src/pythonrun.rs index 3b98827d1e5..357641ce954 100644 --- a/src/pythonrun.rs +++ b/src/pythonrun.rs @@ -108,8 +108,8 @@ impl Drop for GILGuard { /// Release pool struct ReleasePool { - owned: Vec>, - borrowed: Vec>, + owned: ArrayList>, + borrowed: ArrayList>, pointers: *mut Vec>, obj: Vec>, p: spin::Mutex<*mut Vec>>, @@ -118,8 +118,8 @@ struct ReleasePool { impl ReleasePool { fn new() -> ReleasePool { ReleasePool { - owned: Vec::with_capacity(256), - borrowed: Vec::with_capacity(256), + owned: ArrayList::new(), + borrowed: ArrayList::new(), pointers: Box::into_raw(Box::new(Vec::with_capacity(256))), obj: Vec::with_capacity(8), p: spin::Mutex::new(Box::into_raw(Box::new(Vec::with_capacity(256)))), @@ -128,20 +128,16 @@ impl ReleasePool { unsafe fn release_pointers(&mut self) { let mut v = self.p.lock(); - - // vec of pointers - let ptr = *v; - let vec: &'static mut Vec<_> = &mut *ptr; + let vec = &mut **v; if vec.is_empty() { return; } // switch vectors - *v = self.pointers; - self.pointers = ptr; + std::mem::swap(&mut self.pointers, &mut *v); drop(v); - // release py objects + // release PyObjects for ptr in vec.iter_mut() { ffi::Py_DECREF(ptr.as_ptr()); } @@ -149,18 +145,13 @@ impl ReleasePool { } pub unsafe fn drain(&mut self, owned: usize, borrowed: usize, pointers: bool) { - let len = self.owned.len(); - if owned < len { - for ptr in &mut self.owned[owned..len] { - ffi::Py_DECREF(ptr.as_ptr()); - } - self.owned.set_len(owned); - } - - let len = self.borrowed.len(); - if borrowed < len { - self.borrowed.set_len(borrowed); + // Release owned objects(call decref) + while owned < self.owned.len() { + let last = self.owned.pop_back().unwrap(); + ffi::Py_DECREF(last.as_ptr()); } + // Release borrowed objects(don't call decref) + self.borrowed.truncate(borrowed); if pointers { self.release_pointers(); @@ -232,23 +223,18 @@ pub unsafe fn register_any<'p, T: 'static>(obj: T) -> &'p T { } pub unsafe fn register_pointer(obj: NonNull) { - let pool: &'static mut ReleasePool = &mut *POOL; - - let mut v = pool.p.lock(); - let pool: &'static mut Vec<_> = &mut *(*v); - pool.push(obj); + let pool = &mut *POOL; + (**pool.p.lock()).push(obj); } pub unsafe fn register_owned(_py: Python, obj: NonNull) -> &PyObjectRef { - let pool: &'static mut ReleasePool = &mut *POOL; - pool.owned.push(obj); - &*(&pool.owned[pool.owned.len() - 1] as *const _ as *const PyObjectRef) + let pool = &mut *POOL; + &*(pool.owned.push_back(obj) as *const _ as *const PyObjectRef) } pub unsafe fn register_borrowed(_py: Python, obj: NonNull) -> &PyObjectRef { - let pool: &'static mut ReleasePool = &mut *POOL; - pool.borrowed.push(obj); - &*(&pool.borrowed[pool.borrowed.len() - 1] as *const _ as *const PyObjectRef) + let pool = &mut *POOL; + &*(pool.borrowed.push_back(obj) as *const _ as *const PyObjectRef) } impl GILGuard { @@ -278,6 +264,64 @@ impl GILGuard { } } +use self::array_list::ArrayList; + +mod array_list { + use std::collections::LinkedList; + use std::mem; + + const BLOCK_SIZE: usize = 256; + + /// A container type for Release Pool + /// See #271 for why this is crated + pub(super) struct ArrayList { + inner: LinkedList<[T; BLOCK_SIZE]>, + length: usize, + } + + impl ArrayList { + pub fn new() -> Self { + ArrayList { + inner: LinkedList::new(), + length: 0, + } + } + pub fn push_back(&mut self, item: T) -> &T { + let next_idx = self.next_idx(); + if next_idx == 0 { + self.inner.push_back(unsafe { mem::uninitialized() }); + } + self.inner.back_mut().unwrap()[next_idx] = item; + self.length += 1; + &self.inner.back().unwrap()[next_idx] + } + pub fn pop_back(&mut self) -> Option { + self.length -= 1; + let current_idx = self.next_idx(); + if self.length >= BLOCK_SIZE && current_idx == 0 { + let last_list = self.inner.pop_back()?; + return Some(last_list[0].clone()); + } + self.inner.back().map(|arr| arr[current_idx].clone()) + } + pub fn len(&self) -> usize { + self.length + } + pub fn truncate(&mut self, new_len: usize) { + if self.length <= new_len { + return; + } + while self.inner.len() > (new_len + BLOCK_SIZE - 1) / BLOCK_SIZE { + self.inner.pop_back(); + } + self.length = new_len; + } + fn next_idx(&self) -> usize { + self.length % BLOCK_SIZE + } + } +} + #[cfg(test)] mod test { use super::{GILPool, NonNull, ReleasePool, POOL}; diff --git a/tests/test_dict_iter.rs b/tests/test_dict_iter.rs new file mode 100644 index 00000000000..f4e36017b02 --- /dev/null +++ b/tests/test_dict_iter.rs @@ -0,0 +1,17 @@ +extern crate pyo3; + +use pyo3::{prelude::*, types::IntoPyDict}; + +#[test] +fn iter_dict_nosegv() { + let gil = Python::acquire_gil(); + let py = gil.python(); + const LEN: usize = 10_000_000; + let dict = (0..LEN as u64).map(|i| (i, i * 2)).into_py_dict(py); + let mut sum = 0; + for (k, _v) in dict.iter() { + let i: u64 = k.extract().unwrap(); + sum += i; + } + assert_eq!(sum, 49999995000000); +}