Lazily initialise TLS on slow paths.

Copying an idea from mimalloc, initialise the TLS variable to a global
allocator that doesn't own any memory and then lazily check when we hit
a slow path (which we always do when using the global allocator, because
it doesn't own any memory) if we are the global allocator and replace
it.

There is a slight complication compared to mimalloc's version of this
idea.  Snmalloc collects outgoing messages and it's possible for the
first operation in a thread to be a free of memory allocated by a
different thread.  We address this by initialising the queues with a
size value indicating that they are full and then do the lazy check when
about to insert a message that would make a queue full.  This will then
trigger lazy creation of an allocator.

Global initialisation doesn't work for the fake allocator, so skip most
of its constructor.

src/mem/alloc.h

@@ -234,6 +234,11 @@ namespace snmalloc
     FastFreeLists() : small_fast_free_lists() {}
   };
 
+  ALWAYSINLINE inline void* no_replacement(void*)
+  {
+    return nullptr;
+  }
+
   /**
    * Allocator.  This class is parameterised on three template parameters.  The
    * `MemoryProvider` defines the source of memory for this allocator.
@@ -245,18 +250,27 @@ namespace snmalloc
    * to associate metadata with large (16MiB, by default) regions, allowing an
    * allocator to find the allocator responsible for that region.
    *
-   * The final template parameter, `IsQueueInline`, defines whether the
+   * The next template parameter, `IsQueueInline`, defines whether the
    * message queue for this allocator should be stored as a field of the
    * allocator (`true`) or provided externally, allowing it to be anywhere else
    * in the address space (`false`).
+   *
+   * The final template parameter provides a hook to allow the allocator in use
+   * to be dynamically modified.  This is used to implement a trick from
+   * mimalloc that avoids a conditional branch on the fast path.  We initialise
+   * the thread-local allocator pointer with the address of a global allocator,
+   * which never owns any memory.  When we try to allocate memory, we call the
+   * replacement function.
    */
   template<
     class MemoryProvider = GlobalVirtual,
     class PageMap = SNMALLOC_DEFAULT_PAGEMAP,
-    bool IsQueueInline = true>
+    bool IsQueueInline = true,
+    void* (*Replacement)(void*) = no_replacement>
   class Allocator
   : public FastFreeLists,
-    public Pooled<Allocator<MemoryProvider, PageMap, IsQueueInline>>
+    public Pooled<
+      Allocator<MemoryProvider, PageMap, IsQueueInline, Replacement>>
   {
     LargeAlloc<MemoryProvider> large_allocator;
     PageMap page_map;
@@ -637,7 +651,14 @@ namespace snmalloc
 
     struct RemoteCache
     {
-      size_t size = 0;
+      /**
+       * The total amount of memory stored awaiting dispatch to other
+       * allocators.  This is initialised to the maximum size that we use
+       * before caching so that, when we hit the slow path and need to dispatch
+       * everything, we can check if we are a real allocator and lazily provide
+       * a real allocator.
+       */
+      size_t size = REMOTE_CACHE;
       RemoteList list[REMOTE_SLOTS];
 
       /// Used to find the index into the array of queues for remote
@@ -645,17 +666,18 @@ namespace snmalloc
       /// r is used for which round of sending this is.
       inline size_t get_slot(size_t id, size_t r)
       {
-        constexpr size_t allocator_size =
+        constexpr size_t allocator_size = sizeof(
-          sizeof(Allocator<MemoryProvider, PageMap, IsQueueInline>);
+          Allocator<MemoryProvider, PageMap, IsQueueInline, Replacement>);
         constexpr size_t initial_shift =
           bits::next_pow2_bits_const(allocator_size);
+        assert((initial_shift - (r * REMOTE_SLOT_BITS)) < 64);
         return (id >> (initial_shift + (r * REMOTE_SLOT_BITS))) & REMOTE_MASK;
       }
 
       SNMALLOC_FAST_PATH void
-      dealloc(alloc_id_t target_id, void* p, sizeclass_t sizeclass)
+      dealloc_sized(alloc_id_t target_id, void* p, size_t objectsize)
       {
-        this->size += sizeclass_to_size(sizeclass);
+        this->size += objectsize;
 
         Remote* r = static_cast<Remote*>(p);
         r->set_target_id(target_id);
@@ -666,6 +688,12 @@ namespace snmalloc
         l->last = r;
       }
 
+      SNMALLOC_FAST_PATH void
+      dealloc(alloc_id_t target_id, void* p, sizeclass_t sizeclass)
+      {
+        dealloc_sized(target_id, p, sizeclass_to_size(sizeclass));
+      }
+
       void post(alloc_id_t id)
       {
         // When the cache gets big, post lists to their target allocators.
@@ -780,7 +808,10 @@ namespace snmalloc
 
   public:
     Allocator(
-      MemoryProvider& m, PageMap&& p = PageMap(), RemoteAllocator* r = nullptr)
+      MemoryProvider& m,
+      PageMap&& p = PageMap(),
+      RemoteAllocator* r = nullptr,
+      bool isFake = false)
     : large_allocator(m), page_map(p)
     {
       if constexpr (IsQueueInline)
@@ -796,6 +827,11 @@ namespace snmalloc
       if (id() >= static_cast<alloc_id_t>(-1))
         error("Id should not be -1");
 
+      // If this is fake, don't do any of the bits of initialisation that may
+      // allocate memory.
+      if (isFake)
+        return;
+
       init_message_queue();
       message_queue().invariant();
 
@@ -1055,6 +1091,11 @@ namespace snmalloc
     template<ZeroMem zero_mem, AllowReserve allow_reserve>
     SNMALLOC_SLOW_PATH void* small_alloc_slow(sizeclass_t sizeclass)
     {
+      if (void* replacement = Replacement(this))
+      {
+        return reinterpret_cast<Allocator*>(replacement)
+          ->template small_alloc_slow<zero_mem, allow_reserve>(sizeclass);
+      }
       handle_message_queue();
       size_t rsize = sizeclass_to_size(sizeclass);
       auto& sl = small_classes[sizeclass];
@@ -1205,6 +1246,12 @@ namespace snmalloc
       }
       else
       {
+        if (void* replacement = Replacement(this))
+        {
+          return reinterpret_cast<Allocator*>(replacement)
+            ->template medium_alloc<zero_mem, allow_reserve>(
+              sizeclass, rsize, size);
+        }
         slab = reinterpret_cast<Mediumslab*>(
           large_allocator.template alloc<NoZero, allow_reserve>(
             0, SUPERSLAB_SIZE));
@@ -1277,6 +1324,12 @@ namespace snmalloc
           zero_mem == YesZero ? "zeromem" : "nozeromem",
           allow_reserve == NoReserve ? "noreserve" : "reserve"));
 
+      if (void* replacement = Replacement(this))
+      {
+        return reinterpret_cast<Allocator*>(replacement)
+          ->template large_alloc<zero_mem, allow_reserve>(size);
+      }
+
       size_t size_bits = bits::next_pow2_bits(size);
       size_t large_class = size_bits - SUPERSLAB_BITS;
       assert(large_class < NUM_LARGE_CLASSES);
@@ -1317,17 +1370,36 @@ namespace snmalloc
     remote_dealloc(RemoteAllocator* target, void* p, sizeclass_t sizeclass)
     {
       MEASURE_TIME(remote_dealloc, 4, 16);
+      assert(target->id() != id());
 
       handle_message_queue();
 
       void* offseted = apply_cache_friendly_offset(p, sizeclass);
 
+      // Check whether this will overflow the cache first.  If we are a fake
+      // allocator, then our cache will always be full and so we will never hit
+      // this path.
+      size_t sz = sizeclass_to_size(sizeclass);
+      if ((remote.size + sz) < REMOTE_CACHE)
+      {
+        stats().remote_free(sizeclass);
+        remote.dealloc_sized(target->id(), offseted, sz);
+        return;
+      }
+      // Now that we've established that we're in the slow path (if we're a
+      // real allocator, we will have to empty our cache now), check if we are
+      // a real allocator and construct one if we aren't.
+      if (void* replacement = Replacement(this))
+      {
+        // We have to do a dealloc, not a remote_dealloc here because this may
+        // have been allocated with the allocator that we've just had returned.
+        reinterpret_cast<Allocator*>(replacement)->dealloc(p);
+        return;
+      }
+
       stats().remote_free(sizeclass);
       remote.dealloc(target->id(), offseted, sizeclass);
 
-      if (remote.size < REMOTE_CACHE)
-        return;
-
       stats().remote_post();
       remote.post(id());
     }

src/mem/globalalloc.h

@@ -6,11 +6,25 @@
 
 namespace snmalloc
 {
+  inline void* lazy_replacement(void*);
+  using Alloc =
+    Allocator<GlobalVirtual, SNMALLOC_DEFAULT_PAGEMAP, true, lazy_replacement>;
+
   template<class MemoryProvider>
-  class AllocPool : Pool<Allocator<MemoryProvider>, MemoryProvider>
+  class AllocPool : Pool<
+                      Allocator<
+                        MemoryProvider,
+                        SNMALLOC_DEFAULT_PAGEMAP,
+                        true,
+                        lazy_replacement>,
+                      MemoryProvider>
   {
-    using Alloc = Allocator<MemoryProvider>;
+    using Alloc = Allocator<
-    using Parent = Pool<Allocator<MemoryProvider>, MemoryProvider>;
+      MemoryProvider,
+      SNMALLOC_DEFAULT_PAGEMAP,
+      true,
+      lazy_replacement>;
+    using Parent = Pool<Alloc, MemoryProvider>;
 
   public:
     static AllocPool* make(MemoryProvider& mp)
@@ -175,5 +189,4 @@ namespace snmalloc
     return AllocPool<MemoryProvider>::make(mp);
   }
 
-  using Alloc = Allocator<GlobalVirtual>;
 } // namespace snmalloc

src/mem/threadalloc.h

@@ -15,6 +15,16 @@ 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
+   * slabs to allocate from, it will discover that it is the placeholder and
+   * replace itself with the thread-local allocator, allocating one if
+   * required.  This avoids a branch on the fast path.
+   */
+  HEADER_GLOBAL Alloc GlobalPlaceHolder(
+    default_memory_provider, SNMALLOC_DEFAULT_PAGEMAP(), nullptr, true);
+
 #ifdef SNMALLOC_EXTERNAL_THREAD_ALLOC
   /**
    * Version of the `ThreadAlloc` interface that does no management of thread
@@ -32,6 +42,7 @@ namespace snmalloc
     {
       return (Alloc*&)ThreadAllocUntyped::get();
     }
+    static void register_cleanup() {}
   };
 #endif
 
@@ -59,7 +70,8 @@ namespace snmalloc
      */
     static inline void exit()
     {
-      if (auto* per_thread = get(false))
+      auto* per_thread = get();
+      if ((per_thread != &GlobalPlaceHolder) && (per_thread != nullptr))
       {
         current_alloc_pool()->release(per_thread);
         per_thread = nullptr;
@@ -76,15 +88,12 @@ namespace snmalloc
      * 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(bool create = true)
+    static SNMALLOC_FAST_PATH Alloc*& get()
     {
-      static thread_local Alloc* per_thread;
+      static thread_local Alloc* per_thread = &GlobalPlaceHolder;
-      if (!per_thread && create)
-      {
-        per_thread = current_alloc_pool()->acquire();
-      }
       return per_thread;
     }
+    static void register_cleanup() {}
   };
   /**
    * Version of the `ThreadAlloc` interface that uses C++ `thread_local`
@@ -109,7 +118,7 @@ namespace snmalloc
      * Constructor.  Acquires a new allocator and associates it with this
      * object.  There should be only one instance of this class per thread.
      */
-    ThreadAllocThreadDestructor() : alloc(current_alloc_pool()->acquire()) {}
+    ThreadAllocThreadDestructor() : alloc(&GlobalPlaceHolder) {}
 
     /**
      * Destructor.  Releases the allocator owned by this thread.
@@ -129,6 +138,7 @@ namespace snmalloc
       static thread_local ThreadAllocThreadDestructor per_thread;
       return per_thread.alloc;
     }
+    static void register_cleanup() {}
   };
   // When targeting the FreeBSD kernel, the pthread header exists, but the
   // pthread symbols do not, so don't compile this because it will fail to
@@ -157,9 +167,13 @@ namespace snmalloc
 #  endif
       thread_alloc_release(void* p)
     {
-      Alloc** pp = static_cast<Alloc**>(p);
+      // Keep pthreads happy
-      current_alloc_pool()->release(*pp);
+      void** pp = reinterpret_cast<void**>(p);
       *pp = nullptr;
+      // Actually destroy the allocator and reset TLS
+      Alloc*& alloc = ThreadAllocExplicitTLSCleanup::get();
+      current_alloc_pool()->release(alloc);
+      alloc = &GlobalPlaceHolder;
     }
 
 #  ifdef _WIN32
@@ -185,9 +199,9 @@ namespace snmalloc
      *
      * This must not be called until after `tls_key_create` has returned.
      */
-    static inline void tls_set_value(tls_key_t key, Alloc** value)
+    static inline void tls_set_value(tls_key_t key, Alloc* value)
     {
-      FlsSetValue(key, static_cast<void*>(value));
+      FlsSetValue(key, value);
     }
 #  else
     /**
@@ -214,9 +228,9 @@ namespace snmalloc
      *
      * This must not be called until after `tls_key_create` has returned.
      */
-    static inline void tls_set_value(tls_key_t key, Alloc** value)
+    static inline void tls_set_value(tls_key_t key, Alloc* value)
     {
-      pthread_setspecific(key, static_cast<void*>(value));
+      pthread_setspecific(key, value);
     }
 #  endif
 
@@ -226,7 +240,7 @@ namespace snmalloc
      */
     static SNMALLOC_FAST_PATH Alloc*& inner_get()
     {
-      static thread_local Alloc* per_thread;
+      static thread_local Alloc* per_thread = &GlobalPlaceHolder;
       return per_thread;
     }
 
@@ -239,42 +253,6 @@ namespace snmalloc
     }
 #  endif
 
-    /**
-     * Private initialiser for the per thread allocator
-     */
-    static SNMALLOC_SLOW_PATH Alloc*& inner_init()
-    {
-      Alloc*& per_thread = inner_get();
-
-      // If we don't have an allocator, construct one.
-      if (!per_thread)
-      {
-        // Construct the allocator and assign it to `per_thread` *before* doing
-        // anything else.  This is important because `tls_key_create` may
-        // allocate memory and if we are providing the `malloc` implementation
-        // then this function must be re-entrant within a single thread.  In
-        // this case, the second call to this function will simply return the
-        // allocator.
-        per_thread = current_alloc_pool()->acquire();
-
-        bool first = false;
-        tls_key_t key = Singleton<tls_key_t, tls_key_create>::get(&first);
-        // Associate the new allocator with the destructor.
-        tls_set_value(key, &per_thread);
-
-#  ifdef USE_SNMALLOC_STATS
-        // Allocator is up and running now, safe to call atexit.
-        if (first)
-        {
-          atexit(print_stats);
-        }
-#  else
-        UNUSED(first);
-#  endif
-      }
-      return per_thread;
-    }
-
   public:
     /**
      * Public interface, returns the allocator for the current thread,
@@ -282,13 +260,25 @@ namespace snmalloc
      */
     static SNMALLOC_FAST_PATH Alloc*& get()
     {
-      Alloc*& per_thread = inner_get();
+      return inner_get();
+    }
+    static void register_cleanup()
+    {
+      // Register the allocator destructor.
+      bool first = false;
+      tls_key_t key = Singleton<tls_key_t, tls_key_create>::get(&first);
+      // Associate the new allocator with the destructor.
+      tls_set_value(key, &GlobalPlaceHolder);
 
-      if (likely(per_thread != nullptr))
+#  ifdef USE_SNMALLOC_STATS
-        return per_thread;
+      // Allocator is up and running now, safe to call atexit.
-
+      if (first)
-      // Slow path that performs initialization
+      {
-      return inner_init();
+        atexit(print_stats);
+      }
+#  else
+      UNUSED(first);
+#  endif
     }
   };
 #endif
@@ -310,4 +300,39 @@ namespace snmalloc
 #else
   using ThreadAlloc = ThreadAllocExplicitTLSCleanup;
 #endif
+
+  /**
+   * Slow path for the placeholder replacement.  The simple check that this is
+   * the global placeholder is inlined, the rest of it is only hit in a very
+   * unusual case and so should go off the fast path.
+   */
+  SNMALLOC_SLOW_PATH inline void* lazy_replacement_slow()
+  {
+    auto*& local_alloc = ThreadAlloc::get();
+    if ((local_alloc != nullptr) && (local_alloc != &GlobalPlaceHolder))
+    {
+      return local_alloc;
+    }
+    local_alloc = current_alloc_pool()->acquire();
+    ThreadAlloc::register_cleanup();
+    return local_alloc;
+  }
+
+  /**
+   * Function passed as a template parameter to `Allocator` to allow lazy
+   * replacement.  This is called on all of the slow paths in `Allocator`.  If
+   * the caller is the global placeholder allocator then this function will
+   * check if we've already allocated a per-thread allocator, returning it if
+   * so.  If we have not allocated a per-thread allocator yet, then this
+   * function will allocate one.
+   */
+  ALWAYSINLINE inline void* lazy_replacement(void* existing)
+  {
+    if (existing != &GlobalPlaceHolder)
+    {
+      return nullptr;
+    }
+    return lazy_replacement_slow();
+  }
+
 } // namespace snmalloc

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++)