From 3af9d3509942ae33d83dd97067a6f44469f1cadd Mon Sep 17 00:00:00 2001
From: Nathaniel Wesley Filardo <nfilardo@microsoft.com>
Date: Tue, 24 Aug 2021 16:34:30 +0100
Subject: [PATCH] Plumb LocalState ptrs through to Pagemap accessors

David points out that we might not have a static way to get at the pagemap, so
it is potentially useful to pass pointers to state objects down from the
Allocators.
---
 src/backend/address_space.h                | 27 +++++++-----
 src/backend/address_space_core.h           | 49 ++++++++++++++--------
 src/backend/backend.h                      | 41 ++++++++++++------
 src/backend/backend_concept.h              | 21 ++++++----
 src/backend/fixedglobalconfig.h            |  5 ++-
 src/mem/corealloc.h                        | 34 +++++++++++----
 src/mem/localalloc.h                       | 14 +++----
 src/mem/localcache.h                       |  5 ++-
 src/mem/remotecache.h                      | 13 +++---
 src/mem/slaballocator.h                    |  2 +-
 src/test/func/fixed_region/fixed_region.cc |  2 +-
 src/test/func/two_alloc_types/alloc1.cc    |  3 +-
 12 files changed, 141 insertions(+), 75 deletions(-)

diff --git a/src/backend/address_space.h b/src/backend/address_space.h
index cbdc0fd..ada5295 100644
--- a/src/backend/address_space.h
+++ b/src/backend/address_space.h
@@ -43,7 +43,8 @@ namespace snmalloc
      * arena_map for use in subsequent amplification.
      */
     template<bool committed, SNMALLOC_CONCEPT(ConceptBackendMetaRange) Pagemap>
-    CapPtr<void, CBChunk> reserve(size_t size)
+    CapPtr<void, CBChunk>
+    reserve(typename Pagemap::LocalState* local_state, size_t size)
     {
 #ifdef SNMALLOC_TRACING
       std::cout << "ASM reserve request:" << size << std::endl;
@@ -63,7 +64,7 @@ namespace snmalloc
         {
           auto base = CapPtr<void, CBChunk>(
             PAL::template reserve_aligned<committed>(size));
-          Pagemap::register_range(address_cast(base), size);
+          Pagemap::register_range(local_state, address_cast(base), size);
           return base;
         }
       }
@@ -71,7 +72,7 @@ namespace snmalloc
       CapPtr<void, CBChunk> res;
       {
         FlagLock lock(spin_lock);
-        res = core.template reserve<PAL, Pagemap>(size);
+        res = core.template reserve<PAL, Pagemap>(local_state, size);
         if (res == nullptr)
         {
           // Allocation failed ask OS for more memory
@@ -133,12 +134,12 @@ namespace snmalloc
             return nullptr;
           }
 
-          Pagemap::register_range(address_cast(block), block_size);
+          Pagemap::register_range(local_state, address_cast(block), block_size);
 
-          core.template add_range<PAL, Pagemap>(block, block_size);
+          core.template add_range<PAL, Pagemap>(local_state, block, block_size);
 
           // still holding lock so guaranteed to succeed.
-          res = core.template reserve<PAL, Pagemap>(size);
+          res = core.template reserve<PAL, Pagemap>(local_state, size);
         }
       }
 
@@ -157,7 +158,8 @@ namespace snmalloc
      * used, by smaller objects.
      */
     template<bool committed, SNMALLOC_CONCEPT(ConceptBackendMetaRange) Pagemap>
-    CapPtr<void, CBChunk> reserve_with_left_over(size_t size)
+    CapPtr<void, CBChunk> reserve_with_left_over(
+      typename Pagemap::LocalState* local_state, size_t size)
     {
       SNMALLOC_ASSERT(size >= sizeof(void*));
 
@@ -165,7 +167,7 @@ namespace snmalloc
 
       size_t rsize = bits::next_pow2(size);
 
-      auto res = reserve<false, Pagemap>(rsize);
+      auto res = reserve<false, Pagemap>(local_state, rsize);
 
       if (res != nullptr)
       {
@@ -173,7 +175,7 @@ namespace snmalloc
         {
           FlagLock lock(spin_lock);
           core.template add_range<PAL, Pagemap>(
-            pointer_offset(res, size), rsize - size);
+            local_state, pointer_offset(res, size), rsize - size);
         }
 
         if constexpr (committed)
@@ -194,10 +196,13 @@ namespace snmalloc
      * Divides blocks into power of two sizes with natural alignment
      */
     template<SNMALLOC_CONCEPT(ConceptBackendMeta) Pagemap>
-    void add_range(CapPtr<void, CBChunk> base, size_t length)
+    void add_range(
+      typename Pagemap::LocalState* local_state,
+      CapPtr<void, CBChunk> base,
+      size_t length)
     {
       FlagLock lock(spin_lock);
-      core.add_range<PAL, Pagemap>(base, length);
+      core.add_range<PAL, Pagemap>(local_state, base, length);
     }
   };
 } // namespace snmalloc
diff --git a/src/backend/address_space_core.h b/src/backend/address_space_core.h
index 189771c..0aa5872 100644
--- a/src/backend/address_space_core.h
+++ b/src/backend/address_space_core.h
@@ -70,6 +70,7 @@ namespace snmalloc
      */
     template<SNMALLOC_CONCEPT(ConceptBackendMeta) Pagemap>
     void set_next(
+      typename Pagemap::LocalState* local_state,
       size_t align_bits,
       CapPtr<FreeChunk, CBChunk> base,
       CapPtr<FreeChunk, CBChunk> next)
@@ -84,7 +85,7 @@ namespace snmalloc
         // external_pointer, for example) will not attempt to follow this
         // "Metaslab" pointer.
         MetaEntry t(reinterpret_cast<Metaslab*>(next.unsafe_ptr()), nullptr, 0);
-        Pagemap::set_metaentry(address_cast(base), 1, t);
+        Pagemap::set_metaentry(local_state, address_cast(base), 1, t);
         return;
       }
 
@@ -101,13 +102,15 @@ namespace snmalloc
      * particular size.
      */
     template<SNMALLOC_CONCEPT(ConceptBackendMeta) Pagemap>
-    CapPtr<FreeChunk, CBChunk>
-    get_next(size_t align_bits, CapPtr<FreeChunk, CBChunk> base)
+    CapPtr<FreeChunk, CBChunk> get_next(
+      typename Pagemap::LocalState* local_state,
+      size_t align_bits,
+      CapPtr<FreeChunk, CBChunk> base)
     {
       if (align_bits >= MIN_CHUNK_BITS)
       {
-        const MetaEntry& t =
-          Pagemap::template get_metaentry<false>(address_cast(base));
+        const MetaEntry& t = Pagemap::template get_metaentry<false>(
+          local_state, address_cast(base));
         return CapPtr<FreeChunk, CBChunk>(
           reinterpret_cast<FreeChunk*>(t.get_metaslab()));
       }
@@ -121,12 +124,15 @@ namespace snmalloc
     template<
       SNMALLOC_CONCEPT(ConceptPAL) PAL,
       SNMALLOC_CONCEPT(ConceptBackendMeta) Pagemap>
-    void add_block(size_t align_bits, CapPtr<FreeChunk, CBChunk> base)
+    void add_block(
+      typename Pagemap::LocalState* local_state,
+      size_t align_bits,
+      CapPtr<FreeChunk, CBChunk> base)
     {
       check_block(base, align_bits);
       SNMALLOC_ASSERT(align_bits < 64);
 
-      set_next<Pagemap>(align_bits, base, ranges[align_bits]);
+      set_next<Pagemap>(local_state, align_bits, base, ranges[align_bits]);
       ranges[align_bits] = base.as_static<FreeChunk>();
     }
 
@@ -137,7 +143,8 @@ namespace snmalloc
     template<
       SNMALLOC_CONCEPT(ConceptPAL) PAL,
       SNMALLOC_CONCEPT(ConceptBackendMeta) Pagemap>
-    CapPtr<void, CBChunk> remove_block(size_t align_bits)
+    CapPtr<void, CBChunk>
+    remove_block(typename Pagemap::LocalState* local_state, size_t align_bits)
     {
       CapPtr<FreeChunk, CBChunk> first = ranges[align_bits];
       if (first == nullptr)
@@ -150,7 +157,7 @@ namespace snmalloc
 
         // Look for larger block and split up recursively
         CapPtr<void, CBChunk> bigger =
-          remove_block<PAL, Pagemap>(align_bits + 1);
+          remove_block<PAL, Pagemap>(local_state, align_bits + 1);
         if (bigger != nullptr)
         {
           // This block is going to be broken up into sub CHUNK_SIZE blocks
@@ -165,6 +172,7 @@ namespace snmalloc
           auto left_over = pointer_offset(bigger, left_over_size);
 
           add_block<PAL, Pagemap>(
+            local_state,
             align_bits,
             Aal::capptr_bound<FreeChunk, CBChunk>(left_over, left_over_size));
           check_block(left_over.as_static<FreeChunk>(), align_bits);
@@ -174,7 +182,7 @@ namespace snmalloc
       }
 
       check_block(first, align_bits);
-      ranges[align_bits] = get_next<Pagemap>(align_bits, first);
+      ranges[align_bits] = get_next<Pagemap>(local_state, align_bits, first);
       return first.as_void();
     }
 
@@ -186,7 +194,10 @@ namespace snmalloc
     template<
       SNMALLOC_CONCEPT(ConceptPAL) PAL,
       SNMALLOC_CONCEPT(ConceptBackendMeta) Pagemap>
-    void add_range(CapPtr<void, CBChunk> base, size_t length)
+    void add_range(
+      typename Pagemap::LocalState* local_state,
+      CapPtr<void, CBChunk> base,
+      size_t length)
     {
       // For start and end that are not chunk sized, we need to
       // commit the pages to track the allocations.
@@ -211,7 +222,7 @@ namespace snmalloc
         auto b = base.as_static<FreeChunk>();
 
         check_block(b, align_bits);
-        add_block<PAL, Pagemap>(align_bits, b);
+        add_block<PAL, Pagemap>(local_state, align_bits, b);
 
         base = pointer_offset(base, align);
         length -= align;
@@ -246,7 +257,8 @@ namespace snmalloc
     template<
       SNMALLOC_CONCEPT(ConceptPAL) PAL,
       SNMALLOC_CONCEPT(ConceptBackendMeta) Pagemap>
-    CapPtr<void, CBChunk> reserve(size_t size)
+    CapPtr<void, CBChunk>
+    reserve(typename Pagemap::LocalState* local_state, size_t size)
     {
 #ifdef SNMALLOC_TRACING
       std::cout << "ASM Core reserve request:" << size << std::endl;
@@ -255,7 +267,8 @@ namespace snmalloc
       SNMALLOC_ASSERT(bits::is_pow2(size));
       SNMALLOC_ASSERT(size >= sizeof(void*));
 
-      return remove_block<PAL, Pagemap>(bits::next_pow2_bits(size));
+      return remove_block<PAL, Pagemap>(
+        local_state, bits::next_pow2_bits(size));
     }
 
     /**
@@ -268,7 +281,8 @@ namespace snmalloc
     template<
       SNMALLOC_CONCEPT(ConceptPAL) PAL,
       SNMALLOC_CONCEPT(ConceptBackendMeta) Pagemap>
-    CapPtr<void, CBChunk> reserve_with_left_over(size_t size)
+    CapPtr<void, CBChunk> reserve_with_left_over(
+      typename Pagemap::LocalState* local_state, size_t size)
     {
       SNMALLOC_ASSERT(size >= sizeof(void*));
 
@@ -276,13 +290,14 @@ namespace snmalloc
 
       size_t rsize = bits::next_pow2(size);
 
-      auto res = reserve<PAL, Pagemap>(rsize);
+      auto res = reserve<PAL, Pagemap>(local_state, rsize);
 
       if (res != nullptr)
       {
         if (rsize > size)
         {
-          add_range<PAL, Pagemap>(pointer_offset(res, size), rsize - size);
+          add_range<PAL, Pagemap>(
+            local_state, pointer_offset(res, size), rsize - size);
         }
       }
       return res;
diff --git a/src/backend/backend.h b/src/backend/backend.h
index 49bbcc4..9a39645 100644
--- a/src/backend/backend.h
+++ b/src/backend/backend.h
@@ -67,6 +67,12 @@ namespace snmalloc
         concretePagemap;
 
     public:
+      /**
+       * Provide a type alias for LocalState so that we can refer to it without
+       * needing the whole BackendAllocator type at hand.
+       */
+      using LocalState = BackendAllocator::LocalState;
+
       /**
        * Get the metadata associated with a chunk.
        *
@@ -74,8 +80,10 @@ namespace snmalloc
        * to access a location that is not backed by a chunk.
        */
       template<bool potentially_out_of_range = false>
-      SNMALLOC_FAST_PATH static const MetaEntry& get_metaentry(address_t p)
+      SNMALLOC_FAST_PATH static const MetaEntry&
+      get_metaentry(LocalState* ls, address_t p)
       {
+        UNUSED(ls);
         return concretePagemap.template get<potentially_out_of_range>(p);
       }
 
@@ -83,16 +91,19 @@ namespace snmalloc
        * Set the metadata associated with a chunk.
        */
       SNMALLOC_FAST_PATH
-      static void set_metaentry(address_t p, size_t size, MetaEntry t)
+      static void
+      set_metaentry(LocalState* ls, address_t p, size_t size, MetaEntry t)
       {
+        UNUSED(ls);
         for (address_t a = p; a < p + size; a += MIN_CHUNK_SIZE)
         {
           concretePagemap.set(a, t);
         }
       }
 
-      static void register_range(address_t p, size_t sz)
+      static void register_range(LocalState* ls, address_t p, size_t sz)
       {
+        UNUSED(ls);
         concretePagemap.register_range(p, sz);
       }
     };
@@ -107,14 +118,15 @@ namespace snmalloc
     }
 
     template<bool fixed_range_ = fixed_range>
-    static std::enable_if_t<fixed_range_> init(void* base, size_t length)
+    static std::enable_if_t<fixed_range_>
+    init(LocalState* local_state, void* base, size_t length)
     {
       static_assert(fixed_range_ == fixed_range, "Don't set SFINAE parameter!");
 
       auto [heap_base, heap_length] =
         Pagemap::concretePagemap.init(base, length);
       address_space.template add_range<Pagemap>(
-        CapPtr<void, CBChunk>(heap_base), heap_length);
+        local_state, CapPtr<void, CBChunk>(heap_base), heap_length);
     }
 
   private:
@@ -172,14 +184,16 @@ namespace snmalloc
         auto& local = local_state->local_address_space;
 #endif
 
-        p = local.template reserve_with_left_over<PAL, Pagemap>(size);
+        p = local.template reserve_with_left_over<PAL, Pagemap>(
+          local_state, size);
         if (p != nullptr)
         {
           return p;
         }
 
         auto refill_size = LOCAL_CACHE_BLOCK;
-        auto refill = global.template reserve<false, Pagemap>(refill_size);
+        auto refill =
+          global.template reserve<false, Pagemap>(local_state, refill_size);
         if (refill == nullptr)
           return nullptr;
 
@@ -191,10 +205,12 @@ namespace snmalloc
         }
 #endif
         PAL::template notify_using<NoZero>(refill.unsafe_ptr(), refill_size);
-        local.template add_range<PAL, Pagemap>(refill, refill_size);
+        local.template add_range<PAL, Pagemap>(
+          local_state, refill, refill_size);
 
         // This should succeed
-        return local.template reserve_with_left_over<PAL, Pagemap>(size);
+        return local.template reserve_with_left_over<PAL, Pagemap>(
+          local_state, size);
       }
 
 #ifdef SNMALLOC_CHECK_CLIENT
@@ -205,7 +221,7 @@ namespace snmalloc
         size_t rsize = bits::max(OS_PAGE_SIZE, bits::next_pow2(size));
         size_t size_request = rsize * 64;
 
-        p = global.template reserve<false, Pagemap>(size_request);
+        p = global.template reserve<false, Pagemap>(local_state, size_request);
         if (p == nullptr)
           return nullptr;
 
@@ -221,7 +237,8 @@ namespace snmalloc
       SNMALLOC_ASSERT(!is_meta);
 #endif
 
-      p = global.template reserve_with_left_over<true, Pagemap>(size);
+      p = global.template reserve_with_left_over<true, Pagemap>(
+        local_state, size);
       return p;
     }
 
@@ -285,7 +302,7 @@ namespace snmalloc
       }
 
       MetaEntry t(meta, remote, sizeclass);
-      Pagemap::set_metaentry(address_cast(p), size, t);
+      Pagemap::set_metaentry(local_state, address_cast(p), size, t);
       return {p, meta};
     }
   };
diff --git a/src/backend/backend_concept.h b/src/backend/backend_concept.h
index 0717c49..96c7319 100644
--- a/src/backend/backend_concept.h
+++ b/src/backend/backend_concept.h
@@ -17,14 +17,18 @@ namespace snmalloc
    */
   template<typename Meta>
   concept ConceptBackendMeta =
-    requires(address_t addr, size_t sz, MetaEntry t)
+    requires(
+      typename Meta::LocalState* ls,
+      address_t addr,
+      size_t sz,
+      MetaEntry t)
   {
-    { Meta::set_metaentry(addr, sz, t) } -> ConceptSame<void>;
+    { Meta::set_metaentry(ls, addr, sz, t) } -> ConceptSame<void>;
 
-    { Meta::template get_metaentry<true>(addr) }
+    { Meta::template get_metaentry<true>(ls, addr) }
       -> ConceptSame<const MetaEntry&>;
 
-    { Meta::template get_metaentry<false>(addr) }
+    { Meta::template get_metaentry<false>(ls, addr) }
       -> ConceptSame<const MetaEntry&>;
   };
 
@@ -37,9 +41,9 @@ namespace snmalloc
    */
   template<typename Meta>
   concept ConceptBackendMeta_Range =
-    requires(address_t addr, size_t sz)
+    requires(typename Meta::LocalState* ls, address_t addr, size_t sz)
   {
-    { Meta::register_range(addr, sz) } -> ConceptSame<void>;
+    { Meta::register_range(ls, addr, sz) } -> ConceptSame<void>;
   };
 
   /**
@@ -64,7 +68,7 @@ namespace snmalloc
    *  * inherit from CommonConfig (see commonconfig.h)
    *  * specify which PAL is in use via T::Pal
    *  * have static pagemap accessors via T::Pagemap
-   *  * define a T::LocalState type
+   *  * define a T::LocalState type (and alias it as T::Pagemap::LocalState)
    *  * define T::Options of type snmalloc::Flags
    *  * expose the global allocator pool via T::pool()
    *
@@ -74,6 +78,9 @@ namespace snmalloc
     std::is_base_of<CommonConfig, Globals>::value &&
     ConceptPAL<typename Globals::Pal> &&
     ConceptBackendMetaRange<typename Globals::Pagemap> &&
+    ConceptSame<
+      typename Globals::LocalState,
+      typename Globals::Pagemap::LocalState> &&
     requires()
     {
       typename Globals::LocalState;
diff --git a/src/backend/fixedglobalconfig.h b/src/backend/fixedglobalconfig.h
index dcb6e9f..7864723 100644
--- a/src/backend/fixedglobalconfig.h
+++ b/src/backend/fixedglobalconfig.h
@@ -46,9 +46,10 @@ namespace snmalloc
       snmalloc::register_clean_up();
     }
 
-    static void init(void* base, size_t length)
+    static void
+    init(typename Backend::LocalState* local_state, void* base, size_t length)
     {
-      Backend::init(base, length);
+      Backend::init(local_state, base, length);
     }
   };
 }
diff --git a/src/mem/corealloc.h b/src/mem/corealloc.h
index af3b19f..c708127 100644
--- a/src/mem/corealloc.h
+++ b/src/mem/corealloc.h
@@ -114,6 +114,21 @@ namespace snmalloc
       }
     }
 
+    /**
+     * Return a pointer to the backend state.
+     */
+    LocalState* backend_state_ptr()
+    {
+      if constexpr (SharedStateHandle::Options.CoreAllocOwnsLocalState)
+      {
+        return &backend_state;
+      }
+      else
+      {
+        return backend_state;
+      }
+    }
+
     /**
      * Return this allocator's "truncated" ID, an integer useful as a hash
      * value of this allocator.
@@ -376,8 +391,8 @@ namespace snmalloc
       for (size_t i = 0; i < REMOTE_BATCH; i++)
       {
         auto p = message_queue().peek();
-        auto& entry =
-          SharedStateHandle::Pagemap::get_metaentry(snmalloc::address_cast(p));
+        auto& entry = SharedStateHandle::Pagemap::get_metaentry(
+          backend_state_ptr(), snmalloc::address_cast(p));
 
         auto r = message_queue().dequeue(key_global);
 
@@ -516,9 +531,10 @@ namespace snmalloc
     SNMALLOC_FAST_PATH bool post()
     {
       // stats().remote_post();  // TODO queue not in line!
-      bool sent_something = attached_cache->remote_dealloc_cache
-                              .post<sizeof(CoreAllocator), SharedStateHandle>(
-                                public_state()->trunc_id(), key_global);
+      bool sent_something =
+        attached_cache->remote_dealloc_cache
+          .post<sizeof(CoreAllocator), SharedStateHandle>(
+            backend_state_ptr(), public_state()->trunc_id(), key_global);
 
       return sent_something;
     }
@@ -538,8 +554,8 @@ namespace snmalloc
 
     SNMALLOC_FAST_PATH void dealloc_local_object(void* p)
     {
-      auto entry =
-        SharedStateHandle::Pagemap::get_metaentry(snmalloc::address_cast(p));
+      auto entry = SharedStateHandle::Pagemap::get_metaentry(
+        backend_state_ptr(), snmalloc::address_cast(p));
       if (likely(dealloc_local_object_fast(entry, p, entropy)))
         return;
 
@@ -666,7 +682,7 @@ namespace snmalloc
           bool need_post = true; // Always going to post, so ignore.
           auto n = p->atomic_read_next(key_global);
           auto& entry = SharedStateHandle::Pagemap::get_metaentry(
-            snmalloc::address_cast(p));
+            backend_state_ptr(), snmalloc::address_cast(p));
           handle_dealloc_remote(entry, p, need_post);
           p = n;
         }
@@ -681,7 +697,7 @@ namespace snmalloc
 
       auto posted =
         attached_cache->flush<sizeof(CoreAllocator), SharedStateHandle>(
-          [&](auto p) { dealloc_local_object(p); });
+          backend_state_ptr(), [&](auto p) { dealloc_local_object(p); });
 
       // We may now have unused slabs, return to the global allocator.
       for (sizeclass_t sizeclass = 0; sizeclass < NUM_SIZECLASSES; sizeclass++)
diff --git a/src/mem/localalloc.h b/src/mem/localalloc.h
index 02c210c..36f4544 100644
--- a/src/mem/localalloc.h
+++ b/src/mem/localalloc.h
@@ -260,8 +260,8 @@ namespace snmalloc
         std::cout << "Remote dealloc post" << p << " size " << alloc_size(p)
                   << std::endl;
 #endif
-        MetaEntry entry =
-          SharedStateHandle::Pagemap::get_metaentry(address_cast(p));
+        MetaEntry entry = SharedStateHandle::Pagemap::get_metaentry(
+          core_alloc->backend_state_ptr(), address_cast(p));
         local_cache.remote_dealloc_cache.template dealloc<sizeof(CoreAlloc)>(
           entry.get_remote()->trunc_id(), CapPtr<void, CBAlloc>(p), key_global);
         post_remote_cache();
@@ -472,8 +472,8 @@ namespace snmalloc
       //  before init, that maps null to a remote_deallocator that will never be
       //  in thread local state.
 
-      const MetaEntry& entry =
-        SharedStateHandle::Pagemap::get_metaentry(address_cast(p));
+      const MetaEntry& entry = SharedStateHandle::Pagemap::get_metaentry(
+        core_alloc->backend_state_ptr(), address_cast(p));
       if (likely(local_cache.remote_allocator == entry.get_remote()))
       {
         if (likely(CoreAlloc::dealloc_local_object_fast(
@@ -583,8 +583,8 @@ namespace snmalloc
       // To handle this case we require the uninitialised pagemap contain an
       // entry for the first chunk of memory, that states it represents a large
       // object, so we can pull the check for null off the fast path.
-      MetaEntry entry =
-        SharedStateHandle::Pagemap::get_metaentry(address_cast(p_raw));
+      MetaEntry entry = SharedStateHandle::Pagemap::get_metaentry(
+        core_alloc->backend_state_ptr(), address_cast(p_raw));
 
       if (likely(entry.get_remote() != SharedStateHandle::fake_large_remote))
         return sizeclass_to_size(entry.get_sizeclass());
@@ -611,7 +611,7 @@ namespace snmalloc
       // TODO bring back the CHERI bits. Wes to review if required.
       MetaEntry entry =
         SharedStateHandle::Pagemap::template get_metaentry<true>(
-          address_cast(p_raw));
+          core_alloc->backend_state_ptr(), address_cast(p_raw));
       auto sizeclass = entry.get_sizeclass();
       if (likely(entry.get_remote() != SharedStateHandle::fake_large_remote))
       {
diff --git a/src/mem/localcache.h b/src/mem/localcache.h
index 7e7687f..74a35fb 100644
--- a/src/mem/localcache.h
+++ b/src/mem/localcache.h
@@ -74,7 +74,8 @@ namespace snmalloc
       size_t allocator_size,
       typename SharedStateHandle,
       typename DeallocFun>
-    bool flush(DeallocFun dealloc)
+    bool flush(
+      typename SharedStateHandle::LocalState* local_state, DeallocFun dealloc)
     {
       auto& key = entropy.get_free_list_key();
 
@@ -92,7 +93,7 @@ namespace snmalloc
       }
 
       return remote_dealloc_cache.post<allocator_size, SharedStateHandle>(
-        remote_allocator->trunc_id(), key_global);
+        local_state, remote_allocator->trunc_id(), key_global);
     }
 
     template<ZeroMem zero_mem, typename SharedStateHandle, typename Slowpath>
diff --git a/src/mem/remotecache.h b/src/mem/remotecache.h
index cee3274..4b88d23 100644
--- a/src/mem/remotecache.h
+++ b/src/mem/remotecache.h
@@ -76,7 +76,10 @@ namespace snmalloc
     }
 
     template<size_t allocator_size, typename SharedStateHandle>
-    bool post(RemoteAllocator::alloc_id_t id, const FreeListKey& key)
+    bool post(
+      typename SharedStateHandle::LocalState* local_state,
+      RemoteAllocator::alloc_id_t id,
+      const FreeListKey& key)
     {
       SNMALLOC_ASSERT(initialised);
       size_t post_round = 0;
@@ -94,8 +97,8 @@ namespace snmalloc
           if (!list[i].empty())
           {
             auto [first, last] = list[i].extract_segment(key);
-            MetaEntry entry =
-              SharedStateHandle::Pagemap::get_metaentry(address_cast(first));
+            MetaEntry entry = SharedStateHandle::Pagemap::get_metaentry(
+              local_state, address_cast(first));
             entry.get_remote()->enqueue(first, last, key);
             sent_something = true;
           }
@@ -117,8 +120,8 @@ namespace snmalloc
           // Use the next N bits to spread out remote deallocs in our own
           // slot.
           auto r = resend.take(key);
-          MetaEntry entry =
-            SharedStateHandle::Pagemap::get_metaentry(address_cast(r));
+          MetaEntry entry = SharedStateHandle::Pagemap::get_metaentry(
+            local_state, address_cast(r));
           auto i = entry.get_remote()->trunc_id();
           size_t slot = get_slot<allocator_size>(i, post_round);
           list[slot].add(r, key);
diff --git a/src/mem/slaballocator.h b/src/mem/slaballocator.h
index 31f6e94..f714c07 100644
--- a/src/mem/slaballocator.h
+++ b/src/mem/slaballocator.h
@@ -108,7 +108,7 @@ namespace snmalloc
 #endif
         MetaEntry entry{meta, remote, sizeclass};
         SharedStateHandle::Pagemap::set_metaentry(
-          address_cast(slab), slab_size, entry);
+          &local_state, address_cast(slab), slab_size, entry);
         return {slab, meta};
       }
 
diff --git a/src/test/func/fixed_region/fixed_region.cc b/src/test/func/fixed_region/fixed_region.cc
index db75ccf..5273772 100644
--- a/src/test/func/fixed_region/fixed_region.cc
+++ b/src/test/func/fixed_region/fixed_region.cc
@@ -29,7 +29,7 @@ int main()
   std::cout << "Allocated region " << oe_base << " - "
             << pointer_offset(oe_base, size) << std::endl;
 
-  CustomGlobals::init(oe_base, size);
+  CustomGlobals::init(nullptr, oe_base, size);
   FixedAlloc a;
 
   size_t object_size = 128;
diff --git a/src/test/func/two_alloc_types/alloc1.cc b/src/test/func/two_alloc_types/alloc1.cc
index feb47bd..8a21c80 100644
--- a/src/test/func/two_alloc_types/alloc1.cc
+++ b/src/test/func/two_alloc_types/alloc1.cc
@@ -19,5 +19,6 @@ namespace snmalloc
 
 extern "C" void oe_allocator_init(void* base, void* end)
 {
-  snmalloc::CustomGlobals::init(base, address_cast(end) - address_cast(base));
+  snmalloc::CustomGlobals::init(
+    nullptr, base, address_cast(end) - address_cast(base));
 }