From b7fe8ea65436c106ae9c0311527c909465d3ecf6 Mon Sep 17 00:00:00 2001
From: Istvan Haller <ihaller@microsoft.com>
Date: Tue, 24 Aug 2021 15:32:51 +0100
Subject: [PATCH] More comments and improved test

---
 src/mem/corealloc.h        |  3 +++
 src/mem/pool.h             | 27 ++++++++++++++++++++++++-
 src/test/func/pool/pool.cc | 41 ++++++++++++++++++++++++++++----------
 3 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/src/mem/corealloc.h b/src/mem/corealloc.h
index 30dba53..b67df18 100644
--- a/src/mem/corealloc.h
+++ b/src/mem/corealloc.h
@@ -786,6 +786,9 @@ namespace snmalloc
     }
   };
 
+  /**
+   * Use this alias to access the pool of allocators throughout snmalloc.
+   */
   template<typename SharedStateHandle>
   using AllocPool = Pool<
     CoreAllocator<SharedStateHandle>,
diff --git a/src/mem/pool.h b/src/mem/pool.h
index 34f58e9..aecc550 100644
--- a/src/mem/pool.h
+++ b/src/mem/pool.h
@@ -37,12 +37,20 @@ namespace snmalloc
   };
 
   /**
-   * Class used to instantiate a global non-allocator PoolState.
+   * Helper class used to instantiate a global PoolState.
+   *
+   * SingletonPoolState::pool is the default provider for the PoolState within
+   * the Pool class.
    */
   template<typename T, typename SharedStateHandle>
   class SingletonPoolState
   {
     static inline PoolState<T> state;
+    /**
+     * Thread-local initialization marker for a contention-free way to skip
+     * initialization. SharedStateHandle::ensure_init allows for multiple calls,
+     * but it uses a reentrency-safe check which is more expensive.
+     */
     static thread_local inline bool initialized;
 
     /**
@@ -71,6 +79,10 @@ namespace snmalloc
     }
 
   public:
+    /**
+     * Returns a reference for the global PoolState for the given type.
+     * Also forces the initialization of the backend state, if needed.
+     */
     SNMALLOC_FAST_PATH static PoolState<T>& pool()
     {
       if (unlikely(!initialized))
@@ -82,6 +94,19 @@ namespace snmalloc
     }
   };
 
+  /**
+   * Wrapper class to access a pool of a particular type of object.
+   *
+   * The third template argument is a method to retrieve the actual PoolState.
+   *
+   * For the pool of allocators, refer to the AllocPool alias defined in
+   * corealloc.h.
+   *
+   * For a pool of another type it is recommended to leave the
+   * third template argument with its default value. The SingletonPoolState
+   * class is used as a helper to provide a default PoolState management for
+   * this use case.
+   */
   template<
     typename T,
     typename SharedStateHandle,
diff --git a/src/test/func/pool/pool.cc b/src/test/func/pool/pool.cc
index 550dd4e..8288c29 100644
--- a/src/test/func/pool/pool.cc
+++ b/src/test/func/pool/pool.cc
@@ -7,7 +7,7 @@ using namespace snmalloc;
 
 struct PoolAEntry : Pooled<PoolAEntry>
 {
-  size_t field;
+  int field;
 
   PoolAEntry() : field(1){};
 };
@@ -16,7 +16,7 @@ using PoolA = Pool<PoolAEntry, Alloc::StateHandle>;
 
 struct PoolBEntry : Pooled<PoolBEntry>
 {
-  size_t field;
+  int field;
 
   PoolBEntry() : field(0){};
   PoolBEntry(size_t f) : field(f){};
@@ -30,6 +30,7 @@ void test_alloc()
   SNMALLOC_CHECK(ptr != nullptr);
   // Pool allocations should not be visible to debug_check_empty.
   snmalloc::debug_check_empty<Alloc::StateHandle>();
+  PoolA::release(ptr);
 }
 
 void test_constructor()
@@ -45,6 +46,10 @@ void test_constructor()
   auto ptr3 = PoolB::acquire(1);
   SNMALLOC_CHECK(ptr3 != nullptr);
   SNMALLOC_CHECK(ptr3->field == 1);
+
+  PoolA::release(ptr1);
+  PoolB::release(ptr2);
+  PoolB::release(ptr3);
 }
 
 void test_alloc_many()
@@ -66,13 +71,6 @@ void test_alloc_many()
   }
 }
 
-void test_alloc_dealloc()
-{
-  auto ptr = PoolA::acquire();
-  SNMALLOC_CHECK(ptr != nullptr);
-  PoolA::release(ptr);
-}
-
 void test_double_alloc()
 {
   auto ptr1 = PoolA::acquire();
@@ -83,6 +81,8 @@ void test_double_alloc()
   PoolA::release(ptr2);
   auto ptr3 = PoolA::acquire();
   SNMALLOC_CHECK(ptr2 == ptr3);
+  PoolA::release(ptr1);
+  PoolA::release(ptr3);
 }
 
 void test_different_alloc()
@@ -93,6 +93,27 @@ void test_different_alloc()
   auto ptr2 = PoolB::acquire();
   SNMALLOC_CHECK(ptr2 != nullptr);
   SNMALLOC_CHECK(static_cast<void*>(ptr1) != static_cast<void*>(ptr2));
+  PoolB::release(ptr2);
+}
+
+void test_iterator()
+{
+  PoolAEntry* before_iteration_ptr = PoolA::acquire();
+
+  PoolAEntry* ptr = nullptr;
+  while ((ptr = PoolA::iterate(ptr)) != nullptr)
+  {
+    ptr->field = 2;
+  }
+
+  SNMALLOC_CHECK(before_iteration_ptr->field == 2);
+
+  PoolAEntry* after_iteration_ptr = PoolA::acquire();
+
+  SNMALLOC_CHECK(after_iteration_ptr->field == 2);
+
+  PoolA::release(before_iteration_ptr);
+  PoolA::release(after_iteration_ptr);
 }
 
 int main(int argc, char** argv)
@@ -110,8 +131,8 @@ int main(int argc, char** argv)
   test_alloc();
   test_constructor();
   test_alloc_many();
-  test_alloc_dealloc();
   test_double_alloc();
   test_different_alloc();
+  test_iterator();
   return 0;
 }