From 884ba04d06bd8b70c35f5f31f6af78d9acff9cb2 Mon Sep 17 00:00:00 2001
From: NITHIN <nithinvenkatesh@rockset.com>
Date: Tue, 22 Nov 2022 01:37:08 +0000
Subject: [PATCH] [folly] Use allocateOverAligned to allocate over-aligned
 ConcurrentHashMap nodes

Summary:
ConcurrentHashMap uses alignment to ensure that there is sufficient
paddint to avoid false sharing, but doesn't ensure that the alignment
condition is actually met.  This causes problems under UBSAN.

Test Plan:
* existing unit tests
* no UBSAN failures
---
 folly/concurrency/ConcurrentHashMap.h | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/folly/concurrency/ConcurrentHashMap.h b/folly/concurrency/ConcurrentHashMap.h
index 205c4b2403e..b792668e953 100644
--- a/folly/concurrency/ConcurrentHashMap.h
+++ b/folly/concurrency/ConcurrentHashMap.h
@@ -156,6 +156,7 @@ class ConcurrentHashMap {
   using EnableHeterogeneousFind = std::enable_if_t<
       detail::EligibleForHeterogeneousFind<KeyType, HashFn, KeyEqual, K>::value,
       T>;
+  using SegAlloc = typename std::allocator_traits<Allocator>::template rebind_alloc<SegmentT>;
 
   float load_factor_ = SegmentT::kDefaultLoadFactor;
 
@@ -235,7 +236,7 @@ class ConcurrentHashMap {
       auto seg = segments_[i].load(std::memory_order_relaxed);
       if (seg) {
         seg->~SegmentT();
-        Allocator().deallocate((uint8_t*)seg, sizeof(SegmentT));
+        deallocateOverAligned(SegAlloc(), seg, 1);
       }
       segments_[i].store(
           o.segments_[i].load(std::memory_order_relaxed),
@@ -263,7 +264,7 @@ class ConcurrentHashMap {
       auto seg = segments_[i].load(std::memory_order_relaxed);
       if (seg) {
         seg->~SegmentT();
-        Allocator().deallocate((uint8_t*)seg, sizeof(SegmentT));
+        deallocateOverAligned(SegAlloc(), seg, 1);
       }
     }
     cohort_shutdown_cleanup();
@@ -727,13 +728,13 @@ class ConcurrentHashMap {
     SegmentT* seg = segments_[i].load(std::memory_order_acquire);
     if (!seg) {
       auto b = ensureCohort();
-      SegmentT* newseg = (SegmentT*)Allocator().allocate(sizeof(SegmentT));
+      SegmentT* newseg = allocateOverAligned(SegAlloc(), 1);
       newseg = new (newseg)
           SegmentT(size_ >> ShardBits, load_factor_, max_size_ >> ShardBits, b);
       if (!segments_[i].compare_exchange_strong(seg, newseg)) {
         // seg is updated with new value, delete ours.
         newseg->~SegmentT();
-        Allocator().deallocate((uint8_t*)newseg, sizeof(SegmentT));
+        deallocateOverAligned(SegAlloc(), newseg, 1);
       } else {
         seg = newseg;
         updateBeginAndEndSegments(i);