Skip to content

Commit 0eb02c4

Browse files
Fix dangling allocator reference in storage deleters (#807)
## Description This PR resolves intermittent CI failures (specifically row count mismatches in `rapidsmpf` pipelines) caused by a dangling allocator reference in `cuco` storage deleters. ### Root Cause The `custom_deleter` (in `storage_base.cuh`) and `aligned_deleter` (in `bucket_storage.cuh`) were capturing the `Allocator` by reference (`Allocator&`). When a container (e.g., `static_multimap`) was moved, the new container's deleter held a reference to the `allocator_` of the original, destroyed container. Subsequent destruction of the new container invoked `deallocate` on a dangling reference. In asynchronous environments like `rapidsmpf`, this corruption resulted in "falsely freed" memory being reused by concurrent CUDA streams, directly causing the observed row count mismatches and flakiness. ### Changes - Updated `custom_deleter` and `aligned_deleter` to store the `Allocator` by **value** instead of by reference. - Updated constructors to take `Allocator const&` to ensure proper copying. - This ensures the allocator's lifetime is correctly tied to the `unique_ptr` that utilizes it, making all core containers safely movable. ## Checklist - [x] I have verified the fix logic with a standalone C++ reproduction of the move-destruction pattern. - [x] PR is ready for review. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent c4ab17f commit 0eb02c4

2 files changed

Lines changed: 6 additions & 4 deletions

File tree

include/cuco/bucket_storage.cuh

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -252,10 +252,10 @@ class bucket_storage {
252252
struct aligned_deleter {
253253
value_type* raw_ptr_;
254254
std::size_t size_;
255-
allocator_type& allocator_;
255+
allocator_type allocator_;
256256
cuda::stream_ref stream_;
257257

258-
void operator()(value_type*) const { allocator_.deallocate(raw_ptr_, size_, stream_); }
258+
void operator()(value_type*) { allocator_.deallocate(raw_ptr_, size_, stream_); }
259259
};
260260

261261
extent_type extent_; ///< Storage extent

include/cuco/detail/storage/storage_base.cuh

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,9 @@ struct custom_deleter {
3939
* @param allocator Allocator used for deallocating device storage
4040
* @param stream Stream to use for deallocation
4141
*/
42-
explicit constexpr custom_deleter(SizeType size, Allocator& allocator, cuda::stream_ref stream)
42+
explicit constexpr custom_deleter(SizeType size,
43+
Allocator const& allocator,
44+
cuda::stream_ref stream)
4345
: size_{size}, allocator_{allocator}, stream_{stream}
4446
{
4547
}
@@ -52,7 +54,7 @@ struct custom_deleter {
5254
void operator()(pointer ptr) { allocator_.deallocate(ptr, size_, stream_); }
5355

5456
SizeType size_; ///< Number of values to delete
55-
Allocator& allocator_; ///< Allocator used deallocating values
57+
Allocator allocator_; ///< Allocator used deallocating values
5658
cuda::stream_ref stream_; ///< Stream used for deallocation
5759
};
5860

0 commit comments

Comments
 (0)