From 3f65343170898a44d1870370b7287da7f22a8318 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Wed, 28 Oct 2026 10:52:42 -0400 Subject: [PATCH] fix: hll_union kept memory of input sketches' allocators union_impl copied input sketches with copyAs() and copy_or_downsample(), which allocate with the source's allocator, and update(hll_sketch&&) took over the input's memory outright. With a stateful allocator the union's gadget then held memory owned by another allocator instance, which could be released or reset before the union. Copy only when the input uses the union's allocator; otherwise use the existing paths that build into the union's own memory (inserting coupons, or merging into a new HLL_8 array allocated with the union's allocator). Allocators that always compare equal, such as std::allocator, keep the current behavior. Co-Authored-By: Claude Opus 5.5 --- hll/include/HllUnion-internal.hpp | 27 ++++++++---- hll/include/hll.hpp | 2 +- hll/test/HllUnionTest.cpp | 73 +++++++++++++++++++++++++++++++ 3 files changed, 92 insertions(+), 10 deletions(-) diff --git a/hll/include/HllUnion-internal.hpp b/hll/include/HllUnion-internal.hpp index d83340de..d102b792 100644 --- a/hll/include/HllUnion-internal.hpp +++ b/hll/include/HllUnion-internal.hpp @@ -54,7 +54,9 @@ void hll_union_alloc::update(const hll_sketch_alloc& sketch) { template void hll_union_alloc::update(hll_sketch_alloc&& sketch) { if (sketch.is_empty()) { return; } - if (gadget_.is_empty() && sketch.get_target_type() == HLL_8 && sketch.get_lg_config_k() <= lg_max_k_) { + // taking over the sketch is only possible if its memory can be released with the union's allocator + if (gadget_.is_empty() && sketch.get_target_type() == HLL_8 && sketch.get_lg_config_k() <= lg_max_k_ + && sketch.sketch_impl->getAllocator() == gadget_.sketch_impl->getAllocator()) { if (sketch.get_current_mode() == HLL || sketch.get_lg_config_k() == lg_max_k_) { gadget_ = std::move(sketch); } @@ -208,17 +210,20 @@ double hll_union_alloc::get_rel_err(bool upper_bound, bool unioned, } template -HllSketchImpl* hll_union_alloc::copy_or_downsample(const HllSketchImpl* src_impl, uint8_t tgt_lg_k) { +HllSketchImpl* hll_union_alloc::copy_or_downsample(const HllSketchImpl* src_impl, uint8_t tgt_lg_k, + const A& allocator) { if (src_impl->getCurMode() != HLL) { throw std::logic_error("Attempt to downsample non-HLL sketch"); } const HllArray* src = static_cast*>(src_impl); - const int src_lg_k = src->getLgConfigK(); - if (src_lg_k <= tgt_lg_k) { + const uint8_t src_lg_k = src->getLgConfigK(); + // a copy keeps the source's allocator, so it is only made if that is the union's allocator + if (src_lg_k <= tgt_lg_k && src->getAllocator() == allocator) { return src->copyAs(HLL_8); } typedef typename std::allocator_traits::template rebind_alloc> hll8Alloc; - Hll8Array* tgtHllArr = new (hll8Alloc(src->getAllocator()).allocate(1)) Hll8Array(tgt_lg_k, src->getAllocator()); + const uint8_t lg_k = std::min(src_lg_k, tgt_lg_k); + Hll8Array* tgtHllArr = new (hll8Alloc(allocator).allocate(1)) Hll8Array(lg_k, allocator); tgtHllArr->mergeHll(*src); //both of these are required for isomorphism tgtHllArr->putHipAccum(src->getHipAccum()); @@ -239,8 +244,12 @@ template void hll_union_alloc::union_impl(const hll_sketch_alloc& sketch, uint8_t lg_max_k) { const HllSketchImpl* src_impl = sketch.sketch_impl; //default HllSketchImpl* dst_impl = gadget_.sketch_impl; //default + // the gadget allocates with the union's allocator: the input sketch may use another allocator + // instance (e.g. a stateful one), and copies of it would keep that allocator + const A allocator = dst_impl->getAllocator(); if (src_impl->getCurMode() == LIST || src_impl->getCurMode() == SET) { - if (dst_impl->isEmpty() && src_impl->getLgConfigK() == dst_impl->getLgConfigK()) { + if (dst_impl->isEmpty() && src_impl->getLgConfigK() == dst_impl->getLgConfigK() + && src_impl->getAllocator() == allocator) { dst_impl = src_impl->copyAs(HLL_8); gadget_.sketch_impl->get_deleter()(gadget_.sketch_impl); // gadget to be replaced } else { @@ -254,12 +263,12 @@ void hll_union_alloc::union_impl(const hll_sketch_alloc& sketch, uint8_t l // swap so that src is LIST or SET, tgt is HLL // use lg_max_k because LIST has effective K of 2^26 const CouponList* src = static_cast*>(dst_impl); - dst_impl = copy_or_downsample(src_impl, lg_max_k); + dst_impl = copy_or_downsample(src_impl, lg_max_k, allocator); static_cast*>(dst_impl)->mergeList(*src); gadget_.sketch_impl->get_deleter()(gadget_.sketch_impl); // gadget to be replaced } else { // gadget is HLL if (src_impl->getLgConfigK() < dst_impl->getLgConfigK()) { - dst_impl = copy_or_downsample(dst_impl, sketch.get_lg_config_k()); + dst_impl = copy_or_downsample(dst_impl, sketch.get_lg_config_k(), allocator); gadget_.sketch_impl->get_deleter()(gadget_.sketch_impl); // gadget to be replaced } const HllArray* src = static_cast*>(src_impl); @@ -268,7 +277,7 @@ void hll_union_alloc::union_impl(const hll_sketch_alloc& sketch, uint8_t l static_cast*>(dst_impl)->putHipAccum(0); } } else { // src is HLL, gadget is empty - dst_impl = copy_or_downsample(src_impl, lg_max_k); + dst_impl = copy_or_downsample(src_impl, lg_max_k, allocator); gadget_.sketch_impl->get_deleter()(gadget_.sketch_impl); // gadget to be replaced } gadget_.sketch_impl = dst_impl; // gadget replaced diff --git a/hll/include/hll.hpp b/hll/include/hll.hpp index 9950eec0..f160fb14 100644 --- a/hll/include/hll.hpp +++ b/hll/include/hll.hpp @@ -640,7 +640,7 @@ class hll_union_alloc { */ inline void union_impl(const hll_sketch_alloc& sketch, uint8_t lg_max_k); - static HllSketchImpl* copy_or_downsample(const HllSketchImpl* src_impl, uint8_t tgt_lg_k); + static HllSketchImpl* copy_or_downsample(const HllSketchImpl* src_impl, uint8_t tgt_lg_k, const A& allocator); void coupon_update(uint32_t coupon); diff --git a/hll/test/HllUnionTest.cpp b/hll/test/HllUnionTest.cpp index 1d065d19..1dba6f62 100644 --- a/hll/test/HllUnionTest.cpp +++ b/hll/test/HllUnionTest.cpp @@ -384,4 +384,77 @@ TEST_CASE("hll union: serialization is grouping independent", "[hll_union]") { REQUIRE(serialize_nested_union(a, b, c) == serialize_flat_union(a, b, c)); } +// A stateful allocator: instances with different ids are different allocators, and the bytes each one +// has outstanding are counted. +static long long pool_bytes[2] = {0, 0}; + +template +struct pool_allocator { + using value_type = T; + int id; + explicit pool_allocator(int id): id(id) {} + template + pool_allocator(const pool_allocator& other): id(other.id) {} + T* allocate(std::size_t n) { + pool_bytes[id] += n * sizeof(T); + return static_cast(::operator new(n * sizeof(T))); + } + void deallocate(T* p, std::size_t n) { + pool_bytes[id] -= n * sizeof(T); + ::operator delete(p); + } +}; +template +bool operator==(const pool_allocator& a, const pool_allocator& b) { return a.id == b.id; } +template +bool operator!=(const pool_allocator& a, const pool_allocator& b) { return a.id != b.id; } + +using pool_hll_sketch = hll_sketch_alloc>; +using pool_hll_union = hll_union_alloc>; + +static pool_hll_sketch make_pool_sketch(int id, uint8_t lg_k, target_hll_type type, uint64_t n) { + pool_hll_sketch sketch(lg_k, type, false, pool_allocator(id)); + for (uint64_t i = 0; i < n; ++i) sketch.update(i); + return sketch; +} + +TEST_CASE("hll union: keeps no memory of an input sketch's allocator", "[hll_union]") { + const uint8_t lg_max_k = 10; + // LIST, SET and HLL inputs, with a smaller, the same and a larger lg_k than the union + const uint64_t sizes[] = {5, 100, 10000}; + const uint8_t lg_ks[] = {8, 10, 12}; + for (uint64_t n: sizes) { + for (uint8_t lg_k: lg_ks) { + for (target_hll_type type: {HLL_4, HLL_6, HLL_8}) { + for (bool gadget_empty: {true, false}) { + for (bool move: {false, true}) { + hll_union u(lg_max_k); + pool_hll_union pool_union(lg_max_k, pool_allocator(0)); + if (!gadget_empty) { + for (uint64_t i = 0; i < 3 * n; ++i) { + u.update(i + n); + pool_union.update(i + n); + } + } + { + hll_sketch control(lg_k, type); + for (uint64_t i = 0; i < n; ++i) control.update(i); + u.update(control); + pool_hll_sketch input = make_pool_sketch(1, lg_k, type, n); + if (move) { + pool_union.update(std::move(input)); + } else { + pool_union.update(input); + } + } + REQUIRE(pool_bytes[1] == 0); + REQUIRE(pool_union.get_estimate() == u.get_estimate()); + } + } + } + } + } + REQUIRE(pool_bytes[0] == 0); +} + } /* namespace datasketches */