free the right way

Signed-off-by: Andrew Duffy <andrew@a10y.dev>

Author: a10y

vortex-cuda/cudf-test/src/lib.rs

@@ -5,10 +5,11 @@
 
 #![allow(clippy::unwrap_used)]
 
+use std::sync::LazyLock;
+
 use arrow_schema::DataType;
 use arrow_schema::ffi::FFI_ArrowSchema;
 use futures::executor::block_on;
-use std::sync::LazyLock;
 use vortex::array::IntoArray;
 use vortex::array::arrays::PrimitiveArray;
 use vortex::array::session::ArraySession;

vortex-cuda/src/arrow/canonical.rs

@@ -60,10 +60,18 @@ async fn export_primitive(
     ctx: &mut CudaExecutionCtx,
 ) -> VortexResult<ArrowArray> {
     unsafe extern "C" fn release(array: *mut ArrowArray) {
-        // SAFETY: this is only safe if the caller provides a valid pointer to an `ArrowArray`.
-        drop(unsafe { Box::from_raw(array) });
-    }
+        // SAFETY: this is only safe if we're dropping an ArrowArray that was created from Rust
+        //  code. This is necessary to ensure that the fields inside the CudaPrivateData
+        //  get dropped to free native/GPU memory.
+        unsafe {
+            let private_data_ptr =
+                std::ptr::replace(&raw mut (*array).private_data, std::ptr::null_mut());
+            drop(Box::from_raw(private_data_ptr.cast::<CudaPrivateData>()));
 
+            // update the release function to NULL to avoid any possibility of double-frees.
+            (*array).release = None;
+        }
+    }
     let len = array.len();
     let PrimitiveArrayParts {
         buffer, validity, ..

vortex-cuda/src/arrow/mod.rs

@@ -103,15 +103,6 @@ impl ArrowArray {
     }
 }
 
-impl Drop for ArrowArray {
-    fn drop(&mut self) {
-        // SAFETY: this is only safe if we're dropping an ArrowArray that was created from Rust
-        //  code. This is necessary to ensure that the fields inside of the CudaPrivateData
-        //  get dropped to free native/GPU memory.
-        drop(unsafe { Box::from_raw(self.private_data.cast::<CudaPrivateData>()) })
-    }
-}
-
 #[expect(
     unused,
     reason = "cuda_stream and cuda_buffers need to have deferred drop"