Fix memory leaks when parallel merging hash tables by transferring memory blocks from thread-local tables into the global table.

Author: pmenon

src/codegen/util/hash_table.cpp

@@ -76,6 +76,30 @@ HashTable::Entry *HashTable::EntryBuffer::NextFree() {
   return entry;
 }
 
+void HashTable::EntryBuffer::TransferMemoryBlocks(
+    HashTable::EntryBuffer &target) {
+  // Find end of our memory block chain
+  MemoryBlock *tail = block_;
+
+  // Check if there is anything to transfer
+  if (tail == nullptr) {
+    return;
+  }
+
+  // Move to the end
+  while (tail->next != nullptr) {
+    tail = tail->next;
+  }
+
+  // Transfer everything to the target entry buffer
+  do {
+    tail->next = target.block_;
+  } while (!::peloton::atomic_cas(&target.block_, target.block_, block_));
+
+  // Success
+  block_ = nullptr;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 ///
 /// Hash Table
@@ -224,7 +248,7 @@ void HashTable::ReserveLazy(
   PELOTON_MEMSET(directory_, 0, alloc_size);
 }
 
-void HashTable::MergeLazyUnfinished(const HashTable &other) {
+void HashTable::MergeLazyUnfinished(HashTable &other) {
   // Begin with the head of the linked list of entries, stored in the first
   // directory entry
   PELOTON_ASSERT(other.directory_[0] != nullptr);
@@ -248,6 +272,10 @@ void HashTable::MergeLazyUnfinished(const HashTable &other) {
 
   // Increment number of elements
   ::peloton::atomic_add(&num_elems_, other.NumElements());
+
+  // Transfer all allocated memory blocks in the other table into this one
+  other.num_elems_ = other.capacity_ = 0;
+  other.entry_buffer_.TransferMemoryBlocks(entry_buffer_);
 }
 
 void HashTable::Resize() {

src/include/codegen/util/hash_table.h

@@ -112,7 +112,7 @@ class HashTable {
    *
    * @param The hash table whose contents we will merge into this one..
    */
-  void MergeLazyUnfinished(const HashTable &other);
+  void MergeLazyUnfinished(HashTable &other);
 
   //////////////////////////////////////////////////////////////////////////////
   ///
@@ -216,21 +216,34 @@ class HashTable {
      */
     Entry *NextFree();
 
+    /**
+     * Transfer all allocated memory blocks in this entry buffer into the target
+     * buffer.
+     *
+     * @param target Buffer where all blocks are transferred to
+     */
+    void TransferMemoryBlocks(EntryBuffer &target);
+
    private:
-    // This struct allows us to chain together chunks of memory
+    // This struct represents a chunk of heap memory. We chain together these
+    // chunks to avoid the need for a std::vector.
     struct MemoryBlock {
       MemoryBlock *next;
       char data[0];
     };
 
     // The memory pool where block allocations are sourced
     ::peloton::type::AbstractPool &memory_;
+
     // The sizes of each entry
     uint32_t entry_size_;
+
     // The current active block
     MemoryBlock *block_;
+
     // A pointer into the block where the next position is
     char *next_entry_;
+
     // The number of available bytes left in the block
     uint64_t available_bytes_;
   };

test/codegen/hash_table_test.cpp

@@ -195,22 +195,24 @@ TEST_F(HashTableTest, CanInsertLazilyWithDups) {
 }
 
 TEST_F(HashTableTest, ParallelMerge) {
-  codegen::util::HashTable global_table{GetMemPool(), sizeof(Key),
-                                        sizeof(Value)};
-
   constexpr uint32_t num_threads = 4;
   constexpr uint32_t to_insert = 20000;
 
-  std::mutex keys_mutex;
-  std::vector<Key> keys;
-
+  // Allocate hash tables for each thread
   executor::ExecutorContext exec_ctx{nullptr};
 
-  // Allocate hash tables for each thread
   auto &thread_states = exec_ctx.GetThreadStates();
   thread_states.Reset(sizeof(codegen::util::HashTable));
   thread_states.Allocate(num_threads);
 
+  // The keys we insert
+  std::mutex keys_mutex;
+  std::vector<Key> keys;
+
+  // The global hash table
+  codegen::util::HashTable global_table{*exec_ctx.GetPool(), sizeof(Key),
+                                        sizeof(Value)};
+
   auto add_key = [&keys_mutex, &keys](const Key &k) {
     std::lock_guard<std::mutex> lock{keys_mutex};
     keys.emplace_back(k);
@@ -260,6 +262,13 @@ TEST_F(HashTableTest, ParallelMerge) {
   // Now merge thread-local tables into global table in parallel
   LaunchParallelTest(num_threads, merge_fn);
 
+  // Clean up local tables
+  for (uint32_t tid = 0; tid < num_threads; tid++) {
+    auto *table = reinterpret_cast<codegen::util::HashTable *>(
+        thread_states.AccessThreadState(tid));
+    codegen::util::HashTable::Destroy(*table);
+  }
+
   // Now probe global
   EXPECT_EQ(to_insert * num_threads, global_table.NumElements());
   EXPECT_LE(global_table.NumElements(), global_table.Capacity());