Skip to content

Commit

Permalink
Add test for leaking allocator during clonning
Browse files Browse the repository at this point in the history
  • Loading branch information
@JustForFun88
JustForFun88 committed Aug 24, 2023
1 parent 91390f3 commit 50a9e7b
Show file tree
Hide file tree
Showing 2 changed files with 120 additions and 115 deletions.
153 changes: 120 additions & 33 deletions src/map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6658,7 +6658,7 @@ mod test_map {
use allocator_api2::alloc::{AllocError, Allocator, Global};
use core::alloc::Layout;
use core::ptr::NonNull;
use core::sync::atomic::{AtomicBool, Ordering};
use core::sync::atomic::{AtomicI8, Ordering};
use rand::{rngs::SmallRng, Rng, SeedableRng};
use std::borrow::ToOwned;
use std::cell::RefCell;
Expand Down Expand Up @@ -8510,47 +8510,44 @@ mod test_map {
let _map2 = map1.clone();
}

#[test]
fn test_hashmap_into_iter_bug() {
let dropped: Arc<AtomicBool> = Arc::new(AtomicBool::new(false));
struct MyAllocInner {
drop_count: Arc<AtomicI8>,
}

{
struct MyAllocInner {
drop_flag: Arc<AtomicBool>,
}
#[derive(Clone)]
struct MyAlloc {
_inner: Arc<MyAllocInner>,
}

#[derive(Clone)]
struct MyAlloc {
_inner: Arc<MyAllocInner>,
}
impl Drop for MyAllocInner {
fn drop(&mut self) {
println!("MyAlloc freed.");
self.drop_count.fetch_sub(1, Ordering::SeqCst);
}
}

impl Drop for MyAllocInner {
fn drop(&mut self) {
println!("MyAlloc freed.");
self.drop_flag.store(true, Ordering::SeqCst);
}
}
unsafe impl Allocator for MyAlloc {
fn allocate(&self, layout: Layout) -> std::result::Result<NonNull<[u8]>, AllocError> {
let g = Global;
g.allocate(layout)
}

unsafe impl Allocator for MyAlloc {
fn allocate(
&self,
layout: Layout,
) -> std::result::Result<NonNull<[u8]>, AllocError> {
let g = Global;
g.allocate(layout)
}
unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
let g = Global;
g.deallocate(ptr, layout)
}
}

unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
let g = Global;
g.deallocate(ptr, layout)
}
}
#[test]
fn test_hashmap_into_iter_bug() {
let dropped: Arc<AtomicI8> = Arc::new(AtomicI8::new(1));

{
let mut map = crate::HashMap::with_capacity_in(
10,
MyAlloc {
_inner: Arc::new(MyAllocInner {
drop_flag: dropped.clone(),
drop_count: dropped.clone(),
}),
},
);
Expand All @@ -8563,6 +8560,96 @@ mod test_map {
}
}

assert!(dropped.load(Ordering::SeqCst));
// All allocator clones should already be dropped.
assert_eq!(dropped.load(Ordering::SeqCst), 0);
}

#[test]
#[should_panic = "panic in clone"]
fn test_clone_memory_leaks_and_double_drop() {
let dropped: Arc<AtomicI8> = Arc::new(AtomicI8::new(2));

{
let mut map = crate::HashMap::with_capacity_in(
10,
MyAlloc {
_inner: Arc::new(MyAllocInner {
drop_count: dropped.clone(),
}),
},
);

const DISARMED: bool = false;
const ARMED: bool = true;

struct CheckedCloneDrop {
panic_in_clone: bool,
dropped: bool,
need_drop: Vec<i32>,
}

impl Clone for CheckedCloneDrop {
fn clone(&self) -> Self {
if self.panic_in_clone {
panic!("panic in clone")
}
Self {
panic_in_clone: self.panic_in_clone,
dropped: self.dropped,
need_drop: self.need_drop.clone(),
}
}
}

impl Drop for CheckedCloneDrop {
fn drop(&mut self) {
if self.dropped {
panic!("double drop");
}
self.dropped = true;
}
}

let armed_flags = [
DISARMED, DISARMED, ARMED, DISARMED, DISARMED, DISARMED, DISARMED,
];

// Hash and Key must be equal to each other for controlling the elements placement
// so that we can be sure that we first clone elements that don't panic during cloning.
for (idx, &panic_in_clone) in armed_flags.iter().enumerate() {
let idx = idx as u64;
map.table.insert(
idx,
(
idx,
CheckedCloneDrop {
panic_in_clone,
dropped: false,
need_drop: vec![0, 1, 2, 3],
},
),
|(k, _)| *k,
);
}

let mut count = 0;
// Let's check that all elements are located as we wanted
//
// SAFETY: We know for sure that `RawTable` will outlive
// the returned `RawIter` iterator.
for ((key, value), panic_in_clone) in map.iter().zip(armed_flags) {
assert_eq!(*key, count);
assert_eq!(value.panic_in_clone, panic_in_clone);
count += 1;
}
assert_eq!(map.len(), 7);
assert_eq!(count, 7);

// Clone should normally clone a few elements, and then (when the
// clone function panics), deallocate both its own memory, memory
// of `dropped: Arc<AtomicI8>` and the memory of already cloned
// elements (Vec<i32> memory inside CheckedCloneDrop).
let _table2 = map.clone();
}
}
}
82 changes: 0 additions & 82 deletions src/raw/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3393,88 +3393,6 @@ mod test_map {
}
}

#[test]
#[should_panic = "panic in clone"]
fn test_clone_memory_leaks_and_double_drop() {
use ::alloc::vec::Vec;

const DISARMED: bool = false;
const ARMED: bool = true;

struct CheckedCloneDrop {
panic_in_clone: bool,
dropped: bool,
need_drop: Vec<i32>,
}

impl Clone for CheckedCloneDrop {
fn clone(&self) -> Self {
if self.panic_in_clone {
panic!("panic in clone")
}
Self {
panic_in_clone: self.panic_in_clone,
dropped: self.dropped,
need_drop: self.need_drop.clone(),
}
}
}

impl Drop for CheckedCloneDrop {
fn drop(&mut self) {
if self.dropped {
panic!("double drop");
}
self.dropped = true;
}
}

let mut table: RawTable<(u64, CheckedCloneDrop)> = RawTable::with_capacity(7);

let armed_flags = [
DISARMED, DISARMED, ARMED, DISARMED, DISARMED, DISARMED, DISARMED,
];

// Hash and Key must be equal to each other for controlling the elements placement
// so that we can be sure that we first clone elements that don't panic during cloning.
for (idx, &panic_in_clone) in armed_flags.iter().enumerate() {
let idx = idx as u64;
table.insert(
idx,
(
idx,
CheckedCloneDrop {
panic_in_clone,
dropped: false,
need_drop: vec![0, 1, 2, 3],
},
),
|(k, _)| *k,
);
}
let mut count = 0;
unsafe {
// Let's check that all elements are located as we wanted
//
// SAFETY: We know for sure that `RawTable` will outlive
// the returned `RawIter` iterator.
for (bucket, panic_in_clone) in table.iter().zip(armed_flags) {
// SAFETY: It's OK, we get bucket from RawIter, so it is safe to read
let (key, value) = bucket.as_ref();
assert_eq!(*key, count);
assert_eq!(value.panic_in_clone, panic_in_clone);
count += 1;
}
}
assert_eq!(table.len(), 7);
assert_eq!(count, 7);

// Clone should normally clone a few elements, and then (when the clone
// function panics), deallocate both its own memory and the memory of
// already cloned elements (Vec<i32> memory inside CheckedCloneDrop).
let _table2 = table.clone();
}

/// CHECKING THAT WE ARE NOT TRYING TO READ THE MEMORY OF
/// AN UNINITIALIZED TABLE DURING THE DROP
#[test]
Expand Down

0 comments on commit 50a9e7b

Please sign in to comment.