Merge pull request #71 from microsoft/tls-2

Improved ThreadAlloc::get API

src/ds/helpers.h

@@ -87,4 +87,17 @@ namespace snmalloc
       return array[i & (rlength - 1)];
     }
   };
+
+  /**
+   * Helper class to execute a specified function on destruction.
+   */
+  template<void f()>
+  class OnDestruct
+  {
+  public:
+    ~OnDestruct()
+    {
+      f();
+    }
+  };
 } // namespace snmalloc

src/mem/alloc.h

@@ -1056,7 +1056,7 @@ namespace snmalloc
     }
 
     template<ZeroMem zero_mem, AllowReserve allow_reserve>
-    inline void* small_alloc(size_t size)
+    SNMALLOC_FAST_PATH void* small_alloc(size_t size)
     {
       MEASURE_TIME_MARKERS(
         small_alloc,

src/mem/threadalloc.h

@@ -7,10 +7,10 @@
 #error At most one out of SNMALLOC_USE_THREAD_CLEANUP and SNMALLOC_USE_THREAD_DESTRUCTOR may be defined.
 #endif
 
+extern "C" void _malloc_thread_cleanup();
+
 namespace snmalloc
 {
-  extern "C" void _malloc_thread_cleanup(void);
-
   /**
    * A global fake allocator object.  This never allocates memory and, as a
    * result, never owns any slabs.  On the slow paths, where it would fetch
@@ -33,15 +33,106 @@ namespace snmalloc
    */
   class ThreadAllocUntypedWrapper
   {
+  protected:
+    static void register_cleanup() {}
+
   public:
-    static SNMALLOC_FAST_PATH Alloc*& get()
+    static SNMALLOC_FAST_PATH Alloc* get_noncachable()
+    {
+      return (Alloc*&)ThreadAllocUntyped::get();
+    }
+
+    static SNMALLOC_FAST_PATH Alloc* get()
     {
       return (Alloc*&)ThreadAllocUntyped::get();
     }
-    static void register_cleanup() {}
   };
 #endif
 
+  /**
+   * Common aspects of thread local allocator. Subclasses handle how releasing
+   * the allocator is triggered.
+   */
+  class ThreadAllocCommon
+  {
+    friend void* lazy_replacement_slow();
+
+  protected:
+    static inline void inner_release()
+    {
+      auto& per_thread = get_reference();
+      if (per_thread != &GlobalPlaceHolder)
+      {
+        current_alloc_pool()->release(per_thread);
+        per_thread = &GlobalPlaceHolder;
+      }
+    }
+
+    /**
+     * Default clean up does nothing except print statistics if enabled.
+     **/
+    static void register_cleanup()
+    {
+#ifdef USE_SNMALLOC_STATS
+      Singleton<int, atexit_print_stats>::get();
+#endif
+    }
+
+#ifdef USE_SNMALLOC_STATS
+    static void print_stats()
+    {
+      Stats s;
+      current_alloc_pool()->aggregate_stats(s);
+      s.print<Alloc>(std::cout);
+    }
+
+    static int atexit_print_stats() noexcept
+    {
+      return atexit(print_stats);
+    }
+#endif
+
+    static inline Alloc*& get_reference()
+    {
+      static thread_local Alloc* alloc = &GlobalPlaceHolder;
+      return alloc;
+    }
+
+  public:
+    /**
+     * Public interface, returns the allocator for this thread, constructing
+     * one if necessary.
+     *
+     * If no operations have been performed on an allocator returned by either
+     * `get()` nor `get_noncachable()`, then the value contained in the return
+     * will be an Alloc* that will always use the slow path.
+     *
+     * Only use this API if you intend to use the returned allocator just once
+     * per call, or if you know other calls have already been made to the
+     * allocator.
+     */
+    static inline Alloc* get_noncachable()
+    {
+      return get_reference();
+    }
+
+    /**
+     * Public interface, returns the allocator for this thread, constructing
+     * one if necessary.
+     *
+     * The returned Alloc* is guaranteed to be initialised.  This incurs a cost,
+     * so use `get_noncachable` if you can meet its criteria.
+     */
+    static SNMALLOC_FAST_PATH Alloc* get()
+    {
+      auto alloc = get_reference();
+      auto new_alloc = lazy_replacement(alloc);
+      return (likely(new_alloc == nullptr)) ?
+        alloc :
+        reinterpret_cast<Alloc*>(new_alloc);
+    }
+  };
+
   /**
    * Version of the `ThreadAlloc` interface that uses a hook provided by libc
    * to destroy thread-local state.  This is the ideal option, because it
@@ -51,58 +142,14 @@ namespace snmalloc
    * This class is used only when snmalloc is compiled as part of a compatible
    * libc (for example, FreeBSD libc).
    */
-  class ThreadAllocLibcCleanup
+  class ThreadAllocLibcCleanup : public ThreadAllocCommon
   {
     /**
      * Libc will call `_malloc_thread_cleanup` just before a thread terminates.
      * This function must be allowed to call back into this class to destroy
      * the state.
      */
-    friend void _malloc_thread_cleanup(void);
-
-    /**
-     * Function called when the thread exits.  This is guaranteed to be called
-     * precisely once per thread and releases the current allocator.
-     */
-    static inline void exit()
-    {
-      auto* per_thread = get();
-      if ((per_thread != &GlobalPlaceHolder) && (per_thread != nullptr))
-      {
-        current_alloc_pool()->release(per_thread);
-        per_thread = nullptr;
-      }
-    }
-
-  public:
-    /**
-     * Returns a pointer to the allocator associated with this thread.  If
-     * `create` is true, it will create an allocator if one does not exist,
-     * otherwise it will return `nullptr` in this case.  This should be called
-     * with `create == false` only during thread teardown.
-     *
-     * The non-create case exists so that the `per_thread` variable can be a
-     * local static and not a global, allowing ODR to deduplicate it.
-     */
-    static SNMALLOC_FAST_PATH Alloc*& get()
-    {
-      static thread_local Alloc* per_thread = &GlobalPlaceHolder;
-      return per_thread;
-    }
-    static void register_cleanup() {}
-  };
-
-  /**
-   * Helper class to execute a specified function on destruction.
-   */
-  template<void f()>
-  class OnDestruct
-  {
-  public:
-    ~OnDestruct()
-    {
-      f();
-    }
+    friend void ::_malloc_thread_cleanup();
   };
 
   /**
@@ -117,56 +164,17 @@ namespace snmalloc
    * environment and so should be the simplest for initial bringup on an
    * unsupported platform.  It is currently used in the FreeBSD kernel version.
    */
-  class ThreadAllocThreadDestructor
+  class ThreadAllocThreadDestructor : public ThreadAllocCommon
   {
     template<void f()>
     friend class OnDestruct;
 
-    /**
-     * Releases the allocator owned by this thread.
-     */
-    static void inner_release()
-    {
-      if (get() != &GlobalPlaceHolder)
-      {
-        current_alloc_pool()->release(get());
-        get() = &GlobalPlaceHolder;
-      }
-    }
-
-#ifdef USE_SNMALLOC_STATS
-    static void print_stats()
-    {
-      Stats s;
-      current_alloc_pool()->aggregate_stats(s);
-      s.print<Alloc>(std::cout);
-    }
-
-    static int atexit_print_stats()
-    {
-      return atexit(print_stats);
-    }
-#endif
-
   public:
-    /**
-     * Public interface, returns the allocator for this thread, constructing
-     * one if necessary.
-     */
-    static inline Alloc*& get()
-    {
-      static thread_local Alloc* alloc = &GlobalPlaceHolder;
-      return alloc;
-    }
-
     static void register_cleanup()
     {
-      static thread_local OnDestruct<ThreadAllocThreadDestructor::inner_release>
-        tidier;
+      static thread_local OnDestruct<ThreadAllocCommon::inner_release> tidier;
 
-#ifdef USE_SNMALLOC_STATS
-      Singleton<int, atexit_print_stats>::get();
-#endif
+      ThreadAllocCommon::register_cleanup();
     }
   };
 
@@ -175,13 +183,11 @@ namespace snmalloc
    * Entry point the allows libc to call into the allocator for per-thread
    * cleanup.
    */
-  extern "C" void _malloc_thread_cleanup(void)
+  extern "C" void _malloc_thread_cleanup()
   {
-    ThreadAllocLibcCleanup::exit();
+    ThreadAllocLibcCleanup::inner_release();
   }
   using ThreadAlloc = ThreadAllocLibcCleanup;
-#elif defined(SNMALLOC_USE_THREAD_DESTRUCTOR)
-  using ThreadAlloc = ThreadAllocThreadDestructor;
 #elif defined(SNMALLOC_EXTERNAL_THREAD_ALLOC)
   using ThreadAlloc = ThreadAllocUntypedWrapper;
 #else
@@ -195,11 +201,8 @@ namespace snmalloc
    */
   SNMALLOC_SLOW_PATH inline void* lazy_replacement_slow()
   {
-    auto*& local_alloc = ThreadAlloc::get();
-    if ((local_alloc != nullptr) && (local_alloc != &GlobalPlaceHolder))
-    {
-      return local_alloc;
-    }
+    auto*& local_alloc = ThreadAlloc::get_reference();
+    assert(local_alloc == &GlobalPlaceHolder);
     local_alloc = current_alloc_pool()->acquire();
     assert(local_alloc != &GlobalPlaceHolder);
     ThreadAlloc::register_cleanup();

src/test/func/fixed_region/fixed_region.cc

@@ -37,7 +37,7 @@ int main()
   oe_end = (uint8_t*)oe_base + size;
   std::cout << "Allocated region " << oe_base << " - " << oe_end << std::endl;
 
-  auto& a = ThreadAlloc::get();
+  auto a = ThreadAlloc::get();
 
   for (size_t i = 0; i < 1000; i++)
   {

src/test/func/memory/memory.cc

@@ -7,7 +7,7 @@ using namespace snmalloc;
 
 void test_alloc_dealloc_64k()
 {
-  auto& alloc = ThreadAlloc::get();
+  auto alloc = ThreadAlloc::get();
 
   constexpr size_t count = 1 << 12;
   constexpr size_t outer_count = 12;
@@ -39,7 +39,7 @@ void test_alloc_dealloc_64k()
 
 void test_random_allocation()
 {
-  auto& alloc = ThreadAlloc::get();
+  auto alloc = ThreadAlloc::get();
   std::unordered_set<void*> allocated;
 
   constexpr size_t count = 10000;
@@ -91,7 +91,7 @@ void test_random_allocation()
 
 void test_calloc()
 {
-  auto& alloc = ThreadAlloc::get();
+  auto alloc = ThreadAlloc::get();
 
   for (size_t size = 16; size <= (1 << 24); size <<= 1)
   {
@@ -162,7 +162,7 @@ void test_double_alloc()
 void test_external_pointer()
 {
   // Malloc does not have an external pointer querying mechanism.
-  auto& alloc = ThreadAlloc::get();
+  auto alloc = ThreadAlloc::get();
 
   for (uint8_t sc = 0; sc < NUM_SIZECLASSES; sc++)
   {
@@ -208,7 +208,7 @@ void test_external_pointer_large()
 {
   xoroshiro::p128r64 r;
 
-  auto& alloc = ThreadAlloc::get();
+  auto alloc = ThreadAlloc::get();
 
   constexpr size_t count_log = snmalloc::bits::is64() ? 5 : 3;
   constexpr size_t count = 1 << count_log;
@@ -244,7 +244,7 @@ void test_external_pointer_large()
 
 void test_external_pointer_dealloc_bug()
 {
-  auto& alloc = ThreadAlloc::get();
+  auto alloc = ThreadAlloc::get();
   constexpr size_t count = (SUPERSLAB_SIZE / SLAB_SIZE) * 2;
   void* allocs[count];
 
@@ -268,7 +268,7 @@ void test_external_pointer_dealloc_bug()
 
 void test_alloc_16M()
 {
-  auto& alloc = ThreadAlloc::get();
+  auto alloc = ThreadAlloc::get();
   // sizes >= 16M use large_alloc
   const size_t size = 16'000'000;
 
@@ -279,7 +279,7 @@ void test_alloc_16M()
 
 void test_calloc_16M()
 {
-  auto& alloc = ThreadAlloc::get();
+  auto alloc = ThreadAlloc::get();
   // sizes >= 16M use large_alloc
   const size_t size = 16'000'000;
 

src/test/perf/contention/contention.cc

@@ -75,7 +75,7 @@ size_t swapcount;
 
 void test_tasks_f(size_t id)
 {
-  Alloc*& a = ThreadAlloc::get();
+  Alloc* a = ThreadAlloc::get();
   xoroshiro::p128r32 r(id + 5000);
 
   for (size_t n = 0; n < swapcount; n++)
@@ -100,7 +100,7 @@ void test_tasks_f(size_t id)
 
 void test_tasks(size_t num_tasks, size_t count, size_t size)
 {
-  Alloc*& a = ThreadAlloc::get();
+  Alloc* a = ThreadAlloc::get();
 
   contention = new std::atomic<size_t*>[size];
   xoroshiro::p128r32 r;

src/test/perf/external_pointer/externalpointer.cc

@@ -49,7 +49,7 @@ namespace test
 
   void test_external_pointer(xoroshiro::p128r64& r)
   {
-    auto& alloc = ThreadAlloc::get();
+    auto alloc = ThreadAlloc::get();
 #ifdef NDEBUG
     static constexpr size_t iterations = 10000000;
 #else
@@ -79,9 +79,6 @@ namespace test
 int main(int, char**)
 {
   xoroshiro::p128r64 r;
-  // Force a per-thread allocator to actually exist.
-  void* p = ThreadAlloc::get()->alloc(16);
-  ThreadAlloc::get()->dealloc(p);
 #if NDEBUG
   size_t nn = 30;
 #else

src/test/perf/singlethread/singlethread.cc

@@ -7,7 +7,7 @@ using namespace snmalloc;
 template<ZeroMem zero_mem>
 void test_alloc_dealloc(size_t count, size_t size, bool write)
 {
-  auto*& alloc = ThreadAlloc::get();
+  auto* alloc = ThreadAlloc::get();
 
   DO_TIME(
     "Count: " << std::setw(6) << count << ", Size: " << std::setw(6) << size