Heap safe destructor

network/windows/src/eventloop/allocator.rs

@@ -2,9 +2,10 @@ use crate::errors::throw_from_windows_err_code;
 use std::ffi::c_void;
 use std::pin::Pin;
 use std::ptr::NonNull;
+use std::thread::panicking;
 use windows::Win32::Foundation::{GetLastError, HANDLE};
 use windows::Win32::System::Memory::{
-    HeapAlloc, HeapCreate, HeapFree, HEAP_FLAGS, HEAP_NO_SERIALIZE,
+    HEAP_FLAGS, HEAP_NO_SERIALIZE, HeapAlloc, HeapCreate, HeapDestroy, HeapFree,
 };
 
 pub(crate) struct Allocator {
@@ -16,6 +17,7 @@ struct Header {
     prev: Option<NonNull<Header>>,
     next: Option<NonNull<Header>>,
     destructor: unsafe fn(NonNull<c_void>),
+    value: NonNull<c_void>,
 }
 
 impl Allocator {
@@ -57,6 +59,14 @@ impl Allocator {
         )
         .unwrap_or_else(|| throw_from_windows_err_code(GetLastError()));
 
+        h.as_mut().value = NonNull::new_unchecked(
+            h.as_ptr()
+                .map_addr(|a| a + Self::_calc_offset::<T>())
+                .cast::<c_void>(),
+        );
+        let mut p = Pin::new_unchecked(h.as_mut().value.cast::<T>().as_mut());
+        heap_free_on_exit::<false, _>(self.win_heap, h.cast(), || init(&mut p));
+
         if let Some(mut last) = self.last {
             last.as_mut().next = Some(h);
         }
@@ -65,15 +75,6 @@ impl Allocator {
         h.as_mut().destructor = Self::_destructor::<T>;
         self.last = Some(h);
 
-        let mut p = Pin::new_unchecked(
-            &mut *(h
-                .as_ptr()
-                .map_addr(|a| a + Self::_calc_offset::<T>())
-                .cast::<T>()),
-        );
-
-        init(&mut p);
-
         return p;
     }
 
@@ -95,16 +96,82 @@ impl Allocator {
         if self.last == Some(NonNull::from_mut(h)) {
             self.last = h.prev;
         }
-
-        (h.destructor)(p.cast::<c_void>());
-
-        match HeapFree(
-            self.win_heap,
-            HEAP_FLAGS::default(),
-            Some((h as *mut Header).cast()),
-        ) {
-            Ok(_) => return,
-            Err(err) => throw_from_windows_err_code(err.code())
-        };
+        heap_free_on_exit::<true, _>(self.win_heap, NonNull::from_mut(h).cast(), || {
+            (h.destructor)(p.cast::<c_void>())
+        });
+    }
+}
+
+fn heap_free_on_exit<const FreeOnSuccess: bool, R>(
+    win_heap: HANDLE,
+    ptr: NonNull<c_void>,
+    scope: impl FnOnce() -> R,
+) -> R {
+    let defer = HeapFreeOnDrop::<FreeOnSuccess> { win_heap, ptr };
+    let ret = scope();
+    drop(defer);
+    return ret;
+}
+
+struct HeapFreeOnDrop<const FreeOnSuccess: bool> {
+    win_heap: HANDLE,
+    ptr: NonNull<c_void>,
+}
+
+impl<const FreeOnSuccess: bool> Drop for HeapFreeOnDrop<FreeOnSuccess> {
+    fn drop(&mut self) {
+        if !panicking() && !FreeOnSuccess {
+            return;
+        }
+        unsafe {
+            match HeapFree(
+                self.win_heap,
+                HEAP_FLAGS::default(),
+                Some(self.ptr.cast().as_mut()),
+            ) {
+                Ok(_) => return,
+                Err(err) => throw_from_windows_err_code(err.code()),
+            };
+        }
+    }
+}
+
+impl Drop for Allocator {
+    fn drop(&mut self) {
+        drop(RecursiveHeapDestructor {
+            win_heap: self.win_heap,
+            next: self.last,
+        });
+    }
+}
+
+struct RecursiveHeapDestructor {
+    win_heap: HANDLE,
+    next: Option<NonNull<Header>>,
+}
+
+impl Drop for RecursiveHeapDestructor {
+    fn drop(&mut self) {
+        unsafe {
+            match self.next {
+                Some(mut h) => {
+                    let h = h.as_mut();
+                    let next_destructor = RecursiveHeapDestructor {
+                        win_heap: self.win_heap,
+                        next: h.prev,
+                    };
+                    heap_free_on_exit::<true, _>(
+                        self.win_heap,
+                        NonNull::from_mut(h).cast(),
+                        || (h.destructor)(h.value),
+                    );
+                    drop(next_destructor);
+                }
+                None => match HeapDestroy(self.win_heap) {
+                    Ok(()) => return,
+                    Err(err) => throw_from_windows_err_code(err.code()),
+                },
+            }
+        }
     }
 }