Make HLL sketch and union exception-safe on allocation failure

hll/include/AuxHashMap-internal.hpp

@@ -35,14 +35,23 @@ auxCount(0),
 entries(1ULL << lgAuxArrInts, 0, allocator)
 {}
 
+template<typename A>
+AuxHashMap<A>::AuxHashMap(const AuxHashMap& that, const A& allocator):
+lgConfigK(that.lgConfigK),
+lgAuxArrInts(that.lgAuxArrInts),
+auxCount(that.auxCount),
+entries(that.entries, allocator)
+{}
+
 template<typename A>
 AuxHashMap<A>* AuxHashMap<A>::newAuxHashMap(uint8_t lgAuxArrInts, uint8_t lgConfigK, const A& allocator) {
-  return new (ahmAlloc(allocator).allocate(1)) AuxHashMap<A>(lgAuxArrInts, lgConfigK, allocator);
+  return HllUtil<A>::template newObject<AuxHashMap<A>>(allocator, lgAuxArrInts, lgConfigK, allocator);
 }
 
 template<typename A>
 AuxHashMap<A>* AuxHashMap<A>::newAuxHashMap(const AuxHashMap& that) {
-  return new (ahmAlloc(that.entries.get_allocator()).allocate(1)) AuxHashMap<A>(that);
+  const A allocator(that.entries.get_allocator());
+  return HllUtil<A>::template newObject<AuxHashMap<A>>(allocator, that, allocator);
 }
 
 template<typename A>
@@ -60,12 +69,14 @@ AuxHashMap<A>* AuxHashMap<A>::deserialize(const void* bytes, size_t len,
   const uint32_t configKmask = (1 << lgConfigK) - 1;
 
   AuxHashMap<A>* auxHashMap;
+  std::unique_ptr<AuxHashMap<A>, std::function<void(AuxHashMap<A>*)>> aux_ptr(nullptr, make_deleter());
   const uint32_t* auxPtr = static_cast<const uint32_t*>(bytes);
   if (srcCompact) {
     if (len < auxCount * sizeof(int)) {
       throw std::out_of_range("Input array too small to hold AuxHashMap image");
     }
-    auxHashMap = new (ahmAlloc(allocator).allocate(1)) AuxHashMap<A>(lgArrInts, lgConfigK, allocator);
+    auxHashMap = newAuxHashMap(lgArrInts, lgConfigK, allocator);
+    aux_ptr.reset(auxHashMap);
     for (uint32_t i = 0; i < auxCount; ++i) {
       const uint32_t pair = auxPtr[i];
       const uint32_t slotNo = HllUtil<A>::getLow26(pair) & configKmask;
@@ -77,7 +88,8 @@ AuxHashMap<A>* AuxHashMap<A>::deserialize(const void* bytes, size_t len,
     if (len < itemsToRead * sizeof(uint32_t)) {
       throw std::out_of_range("Input array too small to hold AuxHashMap image");
     }
-    auxHashMap = new (ahmAlloc(allocator).allocate(1)) AuxHashMap<A>(lgArrInts, lgConfigK, allocator);
+    auxHashMap = newAuxHashMap(lgArrInts, lgConfigK, allocator);
+    aux_ptr.reset(auxHashMap);
     for (uint32_t i = 0; i < itemsToRead; ++i) {
       const uint32_t pair = auxPtr[i];
       if (pair == hll_constants::EMPTY) { continue; }
@@ -88,11 +100,10 @@ AuxHashMap<A>* AuxHashMap<A>::deserialize(const void* bytes, size_t len,
   }
 
   if (auxHashMap->getAuxCount() != auxCount) {
-    make_deleter()(auxHashMap);
     throw std::invalid_argument("Deserialized AuxHashMap has wrong number of entries");
   }
 
-  return auxHashMap;                                    
+  return aux_ptr.release();
 }
 
 template<typename A>
@@ -106,7 +117,7 @@ AuxHashMap<A>* AuxHashMap<A>::deserialize(std::istream& is, uint8_t lgConfigK,
     lgArrInts = lgAuxArrInts;
   }
 
-  AuxHashMap<A>* auxHashMap = new (ahmAlloc(allocator).allocate(1)) AuxHashMap<A>(lgArrInts, lgConfigK, allocator);
+  AuxHashMap<A>* auxHashMap = newAuxHashMap(lgArrInts, lgConfigK, allocator);
   typedef std::unique_ptr<AuxHashMap<A>, std::function<void(AuxHashMap<A>*)>> aux_hash_map_ptr;
   aux_hash_map_ptr aux_ptr(auxHashMap, auxHashMap->make_deleter());
 
@@ -131,7 +142,6 @@ AuxHashMap<A>* AuxHashMap<A>::deserialize(std::istream& is, uint8_t lgConfigK,
   }
 
   if (auxHashMap->getAuxCount() != auxCount) {
-    make_deleter()(auxHashMap);
     throw std::invalid_argument("Deserialized AuxHashMap has wrong number of entries");
   }
 
@@ -149,7 +159,7 @@ std::function<void(AuxHashMap<A>*)> AuxHashMap<A>::make_deleter() {
 
 template<typename A>
 AuxHashMap<A>* AuxHashMap<A>::copy() const {
-  return new (ahmAlloc(entries.get_allocator()).allocate(1)) AuxHashMap<A>(*this);
+  return newAuxHashMap(*this);
 }
 
 template<typename A>
@@ -189,7 +199,13 @@ void AuxHashMap<A>::mustAdd(uint32_t slotNo, uint8_t value) {
   // found empty entry
   entries[~index] = entry_pair;
   ++auxCount;
-  checkGrow();
+  try {
+    checkGrow();
+  } catch (...) {
+    entries[~index] = hll_constants::EMPTY;
+    --auxCount;
+    throw;
+  }
 }
 
 template<typename A>
@@ -224,17 +240,19 @@ void AuxHashMap<A>::checkGrow() {
 template<typename A>
 void AuxHashMap<A>::growAuxSpace() {
   const int configKmask = (1 << lgConfigK) - 1;
-  const int newArrLen = 1 << ++lgAuxArrInts;
+  const uint8_t newLgAuxArrInts = static_cast<uint8_t>(lgAuxArrInts + 1);
+  const int newArrLen = 1 << newLgAuxArrInts;
   vector_int entries_new(newArrLen, 0, entries.get_allocator());
   for (size_t i = 0; i < entries.size(); ++i) {
     const uint32_t fetched = entries[i];
     if (fetched != hll_constants::EMPTY) {
       // find empty in new array
-      const int32_t idx = find(entries_new.data(), lgAuxArrInts, lgConfigK, fetched & configKmask);
+      const int32_t idx = find(entries_new.data(), newLgAuxArrInts, lgConfigK, fetched & configKmask);
       entries_new[~idx] = fetched;
     }
   }
   entries = std::move(entries_new);
+  lgAuxArrInts = newLgAuxArrInts;
 }
 
 //Searches the Aux arr hash table for an empty or a matching slotNo depending on the context.

hll/include/AuxHashMap.hpp

@@ -32,6 +32,7 @@ template<typename A>
 class AuxHashMap final {
   public:
     AuxHashMap(uint8_t lgAuxArrInts, uint8_t lgConfigK, const A& allocator);
+    AuxHashMap(const AuxHashMap& that, const A& allocator);
     static AuxHashMap* newAuxHashMap(uint8_t lgAuxArrInts, uint8_t lgConfigK, const A& allocator);
     static AuxHashMap* newAuxHashMap(const AuxHashMap<A>& that);
 

hll/include/CouponHashSet-internal.hpp

@@ -83,12 +83,20 @@ CouponHashSet<A>* CouponHashSet<A>::newSet(const void* bytes, size_t len, const
   if (lgK <= 7) {
     throw std::invalid_argument("Attempt to deserialize invalid CouponHashSet with lgConfigK <= 7. Found: "
                                 + std::to_string(lgK));
-  }   
+  }
+  HllUtil<A>::checkLgK(lgK);
   uint8_t lgArrInts = data[hll_constants::LG_ARR_BYTE];
   const bool compactFlag = ((data[hll_constants::FLAGS_BYTE] & hll_constants::COMPACT_FLAG_MASK) ? true : false);
 
   uint32_t couponCount;
   std::memcpy(&couponCount, data + hll_constants::HASH_SET_COUNT_INT, sizeof(couponCount));
+  const uint32_t maxCouponCount = hll_constants::RESIZE_NUMER * (uint32_t(1) << (lgK - 3))
+      / hll_constants::RESIZE_DENOM;
+  if (compactFlag && couponCount > maxCouponCount) {
+    throw std::invalid_argument("Attempt to deserialize invalid compact CouponHashSet with couponCount > capacity. Found couponCount: "
+                                + std::to_string(couponCount)
+                                + ", capacity: " + std::to_string(maxCouponCount));
+  }
   if (lgArrInts < hll_constants::LG_INIT_SET_SIZE) {
     lgArrInts = HllUtil<>::computeLgArrInts(SET, couponCount, lgK);
   }
@@ -101,8 +109,8 @@ CouponHashSet<A>* CouponHashSet<A>::newSet(const void* bytes, size_t len, const
                                 + ", found: " + std::to_string(len));
   }
 
-  ChsAlloc chsa(allocator);
-  CouponHashSet<A>* sketch = new (chsa.allocate(1)) CouponHashSet<A>(lgK, tgtHllType, allocator);
+  auto sketch = HllUtil<A>::template newObject<CouponHashSet<A>>(allocator, lgK, tgtHllType, allocator);
+  std::unique_ptr<CouponHashSet<A>, std::function<void(HllSketchImpl<A>*)>> ptr(sketch, sketch->get_deleter());
 
   if (compactFlag) {
     const uint8_t* curPos = data + hll_constants::HASH_SET_INT_ARR_START;
@@ -119,7 +127,7 @@ CouponHashSet<A>* CouponHashSet<A>::newSet(const void* bytes, size_t len, const
                 couponsInArray * sizeof(uint32_t));
   }
 
-  return sketch;
+  return ptr.release();
 }
 
 template<typename A>
@@ -149,16 +157,23 @@ CouponHashSet<A>* CouponHashSet<A>::newSet(std::istream& is, const A& allocator)
     throw std::invalid_argument("Attempt to deserialize invalid CouponHashSet with lgConfigK <= 7. Found: "
                                 + std::to_string(lgK));
   }
+  HllUtil<A>::checkLgK(lgK);
   uint8_t lgArrInts = listHeader[hll_constants::LG_ARR_BYTE];
   const bool compactFlag = ((listHeader[hll_constants::FLAGS_BYTE] & hll_constants::COMPACT_FLAG_MASK) ? true : false);
 
   const auto couponCount = read<uint32_t>(is);
+  const uint32_t maxCouponCount = hll_constants::RESIZE_NUMER * (uint32_t(1) << (lgK - 3))
+      / hll_constants::RESIZE_DENOM;
+  if (compactFlag && couponCount > maxCouponCount) {
+    throw std::invalid_argument("Attempt to deserialize invalid compact CouponHashSet with couponCount > capacity. Found couponCount: "
+                                + std::to_string(couponCount)
+                                + ", capacity: " + std::to_string(maxCouponCount));
+  }
   if (lgArrInts < hll_constants::LG_INIT_SET_SIZE) {
     lgArrInts = HllUtil<>::computeLgArrInts(SET, couponCount, lgK);
   }
 
-  ChsAlloc chsa(allocator);
-  CouponHashSet<A>* sketch = new (chsa.allocate(1)) CouponHashSet<A>(lgK, tgtHllType, allocator);
+  auto sketch = HllUtil<A>::template newObject<CouponHashSet<A>>(allocator, lgK, tgtHllType, allocator);
   typedef std::unique_ptr<CouponHashSet<A>, std::function<void(HllSketchImpl<A>*)>> coupon_hash_set_ptr;
   coupon_hash_set_ptr ptr(sketch, sketch->get_deleter());
 
@@ -183,14 +198,12 @@ CouponHashSet<A>* CouponHashSet<A>::newSet(std::istream& is, const A& allocator)
 
 template<typename A>
 CouponHashSet<A>* CouponHashSet<A>::copy() const {
-  ChsAlloc chsa(this->coupons_.get_allocator());
-  return new (chsa.allocate(1)) CouponHashSet<A>(*this);
+  return HllUtil<A>::template newObject<CouponHashSet<A>>(this->getAllocator(), *this);
 }
 
 template<typename A>
 CouponHashSet<A>* CouponHashSet<A>::copyAs(target_hll_type tgtHllType) const {
-  ChsAlloc chsa(this->coupons_.get_allocator());
-  return new (chsa.allocate(1)) CouponHashSet<A>(*this, tgtHllType);
+  return HllUtil<A>::template newObject<CouponHashSet<A>>(this->getAllocator(), *this, tgtHllType);
 }
 
 template<typename A>
@@ -202,8 +215,8 @@ HllSketchImpl<A>* CouponHashSet<A>::couponUpdate(uint32_t coupon) {
   }
   this->coupons_[~index] = coupon; // found empty
   ++this->couponCount_;
-  if (checkGrowOrPromote()) {
-    return this->promoteHeapListOrSetToHll(*this);
+  if (static_cast<size_t>(hll_constants::RESIZE_DENOM * this->couponCount_) > (hll_constants::RESIZE_NUMER * this->coupons_.size())) {
+    return growOrPromoteWithRollback(static_cast<uint32_t>(~index));
   }
   return this;
 }
@@ -219,15 +232,19 @@ uint8_t CouponHashSet<A>::getPreInts() const {
 }
 
 template<typename A>
-bool CouponHashSet<A>::checkGrowOrPromote() {
-  if (static_cast<size_t>(hll_constants::RESIZE_DENOM * this->couponCount_) > (hll_constants::RESIZE_NUMER * this->coupons_.size())) {
+HllSketchImpl<A>* CouponHashSet<A>::growOrPromoteWithRollback(uint32_t index) {
+  try {
     const uint8_t lgCouponArrInts = count_trailing_zeros_in_u32(static_cast<uint32_t>(this->coupons_.size()));
     if (lgCouponArrInts == (this->lgConfigK_ - 3)) { // at max size
-      return true; // promote to HLL
+      return this->promoteHeapListOrSetToHll(*this);
     }
-    growHashSet(lgCouponArrInts + 1);
+    growHashSet(static_cast<uint8_t>(lgCouponArrInts + 1));
+  } catch (...) {
+    this->coupons_[index] = hll_constants::EMPTY;
+    --this->couponCount_;
+    throw;
   }
-  return false;
+  return this;
 }
 
 template<typename A>

hll/include/CouponHashSet.hpp

@@ -50,7 +50,12 @@ class CouponHashSet : public CouponList<A> {
 
   private:
     using ChsAlloc = typename std::allocator_traits<A>::template rebind_alloc<CouponHashSet<A>>;
-    bool checkGrowOrPromote();
+#if defined(_MSC_VER)
+    __declspec(noinline)
+#elif defined(__GNUC__) || defined(__clang__)
+    __attribute__((noinline))
+#endif
+    HllSketchImpl<A>* growOrPromoteWithRollback(uint32_t index);
     void growHashSet(uint8_t tgtLgCoupArrSize);
 };
 

hll/include/CouponList-internal.hpp

@@ -59,14 +59,12 @@ std::function<void(HllSketchImpl<A>*)> CouponList<A>::get_deleter() const {
 
 template<typename A>
 CouponList<A>* CouponList<A>::copy() const {
-  ClAlloc cla(coupons_.get_allocator());
-  return new (cla.allocate(1)) CouponList<A>(*this);
+  return HllUtil<A>::template newObject<CouponList<A>>(getAllocator(), *this);
 }
 
 template<typename A>
 CouponList<A>* CouponList<A>::copyAs(target_hll_type tgtHllType) const {
-  ClAlloc cla(coupons_.get_allocator());
-  return new (cla.allocate(1)) CouponList<A>(*this, tgtHllType);
+  return HllUtil<A>::template newObject<CouponList<A>>(getAllocator(), *this, tgtHllType);
 }
 
 template<typename A>
@@ -113,8 +111,7 @@ CouponList<A>* CouponList<A>::newList(const void* bytes, size_t len, const A& al
                                 + ", found: " + std::to_string(len));
   }
 
-  ClAlloc cla(allocator);
-  CouponList<A>* sketch = new (cla.allocate(1)) CouponList<A>(lgK, tgtHllType, mode, allocator);
+  auto sketch = HllUtil<A>::template newObject<CouponList<A>>(allocator, lgK, tgtHllType, mode, allocator);
   sketch->couponCount_ = couponCount;
   sketch->putOutOfOrderFlag(oooFlag); // should always be false for LIST
 
@@ -162,8 +159,7 @@ CouponList<A>* CouponList<A>::newList(std::istream& is, const A& allocator) {
                                 + ", capacity: " + std::to_string(listCapacity));
   }
 
-  ClAlloc cla(allocator);
-  CouponList<A>* sketch = new (cla.allocate(1)) CouponList<A>(lgK, tgtHllType, mode, allocator);
+  auto sketch = HllUtil<A>::template newObject<CouponList<A>>(allocator, lgK, tgtHllType, mode, allocator);
   using coupon_list_ptr = std::unique_ptr<CouponList<A>, std::function<void(HllSketchImpl<A>*)>>;
   coupon_list_ptr ptr(sketch, sketch->get_deleter());
   sketch->couponCount_ = couponCount;
@@ -287,10 +283,7 @@ HllSketchImpl<A>* CouponList<A>::couponUpdate(uint32_t coupon) {
       coupons_[i] = coupon; // the actual update
       ++couponCount_;
       if (couponCount_ == static_cast<uint32_t>(coupons_.size())) { // array full
-        if (this->lgConfigK_ < 8) {
-          return promoteHeapListOrSetToHll(*this);
-        }
-        return promoteHeapListToSet(*this);
+        return promoteWithRollback(static_cast<uint32_t>(i));
       }
       return this;
     }
@@ -303,6 +296,20 @@ HllSketchImpl<A>* CouponList<A>::couponUpdate(uint32_t coupon) {
   throw std::runtime_error("Array invalid: no empties and no duplicates");
 }
 
+template<typename A>
+HllSketchImpl<A>* CouponList<A>::promoteWithRollback(uint32_t index) {
+  try {
+    if (this->lgConfigK_ < 8) {
+      return promoteHeapListOrSetToHll(*this);
+    }
+    return promoteHeapListToSet(*this);
+  } catch (...) {
+    coupons_[index] = hll_constants::EMPTY;
+    --couponCount_;
+    throw;
+  }
+}
+
 template<typename A>
 double CouponList<A>::getCompositeEstimate() const { return getEstimate(); }
 

hll/include/CouponList.hpp

@@ -85,6 +85,14 @@ class CouponList : public HllSketchImpl<A> {
     vector_int coupons_;
 
     friend class HllSketchImplFactory<A>;
+
+  private:
+#if defined(_MSC_VER)
+    __declspec(noinline)
+#elif defined(__GNUC__) || defined(__clang__)
+    __attribute__((noinline))
+#endif
+    HllSketchImpl<A>* promoteWithRollback(uint32_t index);
 };
 
 }