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 */