New PAL for open enclave

CMakeLists.txt

@@ -51,6 +51,10 @@ macro(warnings_high)
   endif()
 endmacro()
 
+macro(oe_simulate target)
+  target_compile_definitions(${target} PRIVATE IS_REALLY_ADDRESS_SPACE_CONSTRAINED)
+endmacro()
+
 macro(clangformat_targets)
   # The clang-format tool is installed under a variety of different names.  Try
   # to find a sensible one.  Only look for versions 9 explicitly - we don't
@@ -244,6 +248,9 @@ if(NOT DEFINED SNMALLOC_ONLY_HEADER_LIBRARY)
     add_shim(snmallocshim SHARED ${SHARED_FILES})
     add_shim(snmallocshim-1mib SHARED ${SHARED_FILES})
     target_compile_definitions(snmallocshim-1mib PRIVATE IS_ADDRESS_SPACE_CONSTRAINED)
+    # Build a shim with some settings from oe.
+    add_shim(snmallocshim-oe SHARED ${SHARED_FILES})
+    oe_simulate(snmallocshim-oe)
   endif()
 
   if(SNMALLOC_RUST_SUPPORT)
@@ -260,7 +267,7 @@ if(NOT DEFINED SNMALLOC_ONLY_HEADER_LIBRARY)
   foreach(TEST_CATEGORY ${TEST_CATEGORIES})
     subdirlist(TESTS ${TESTDIR}/${TEST_CATEGORY})
     foreach(TEST ${TESTS})
-      foreach(SUPER_SLAB_SIZE 1;16)
+      foreach(SUPER_SLAB_SIZE 1;16;oe)
         unset(SRC)
         aux_source_directory(${TESTDIR}/${TEST_CATEGORY}/${TEST} SRC)
         set(TESTNAME "${TEST_CATEGORY}-${TEST}-${SUPER_SLAB_SIZE}")
@@ -269,6 +276,9 @@ if(NOT DEFINED SNMALLOC_ONLY_HEADER_LIBRARY)
         if (${SUPER_SLAB_SIZE} EQUAL 1)
           target_compile_definitions(${TESTNAME} PRIVATE IS_ADDRESS_SPACE_CONSTRAINED)
         endif()
+        if (${SUPER_SLAB_SIZE} EQUAL oe)
+          oe_simulate(${TESTNAME})
+        endif()
         target_link_libraries(${TESTNAME} snmalloc_lib)
         if (${TEST} MATCHES "release-.*")
           message(STATUS "Adding test: ${TESTNAME} only for release configs")

src/mem/allocconfig.h

@@ -50,6 +50,15 @@ namespace snmalloc
 #endif
     ;
 
+  // Specifies even smaller slab and super slab sizes for open enclave.
+  static constexpr size_t REALLY_ADDRESS_SPACE_CONSTRAINED =
+#ifdef IS_REALLY_ADDRESS_SPACE_CONSTRAINED
+    true
+#else
+    false
+#endif
+    ;
+
   static constexpr size_t RESERVE_MULTIPLE =
 #ifdef USE_RESERVE_MULTIPLE
     USE_RESERVE_MULTIPLE
@@ -100,14 +109,17 @@ namespace snmalloc
   static constexpr size_t MIN_ALLOC_BITS = bits::ctz_const(MIN_ALLOC_SIZE);
 
   // Slabs are 64 KiB unless constrained to 16 KiB.
-  static constexpr size_t SLAB_BITS = ADDRESS_SPACE_CONSTRAINED ? 14 : 16;
+  static constexpr size_t SLAB_BITS = REALLY_ADDRESS_SPACE_CONSTRAINED ?
+    13 :
+    (ADDRESS_SPACE_CONSTRAINED ? 14 : 16);
   static constexpr size_t SLAB_SIZE = 1 << SLAB_BITS;
   static constexpr size_t SLAB_MASK = ~(SLAB_SIZE - 1);
 
   // Superslabs are composed of this many slabs. Slab offsets are encoded as
   // a byte, so the maximum count is 256. This must be a power of two to
   // allow fast masking to find a superslab start address.
-  static constexpr size_t SLAB_COUNT_BITS = ADDRESS_SPACE_CONSTRAINED ? 6 : 8;
+  static constexpr size_t SLAB_COUNT_BITS =
+    REALLY_ADDRESS_SPACE_CONSTRAINED ? 5 : (ADDRESS_SPACE_CONSTRAINED ? 6 : 8);
   static constexpr size_t SLAB_COUNT = 1 << SLAB_COUNT_BITS;
   static constexpr size_t SUPERSLAB_SIZE = SLAB_SIZE * SLAB_COUNT;
   static constexpr size_t SUPERSLAB_MASK = ~(SUPERSLAB_SIZE - 1);

src/pal/pal_open_enclave.h

@@ -1,7 +1,10 @@
 #pragma once
 
 #include "ds/address.h"
+#include "ds/flaglock.h"
 #include "pal_plain.h"
+
+#include <array>
 #ifdef OPEN_ENCLAVE
 extern "C" void* oe_memset_s(void* p, size_t p_size, int c, size_t size);
 extern "C" [[noreturn]] void oe_abort();
@@ -10,24 +13,98 @@ namespace snmalloc
 {
   class PALOpenEnclave
   {
-    static inline std::atomic<void*> oe_base;
-    static inline void* oe_end = nullptr;
+    /**
+     * Implements a power of two allocator, where all blocks are aligned to the
+     * same power of two as their size. This is what snmalloc uses to get
+     * alignment of very large sizeclasses.
+     *
+     * Pals are not required to unreserve memory, so this does not require the
+     * usual complexity of a buddy allocator.
+     */
+
+    // There are a maximum of two blocks for any size/align in a range.
+    // One before the point of maximum alignment, and one after.
+    static inline std::array<std::array<void*, 2>, bits::BITS> ranges;
+    // This is infrequently used code, a spin lock simplifies the code
+    // considerably, and should never be on the fast path.
+    static inline std::atomic_flag spin_lock;
+
+    static void add_block(size_t align_bits, void* base)
+    {
+      if (ranges[align_bits][0] == nullptr)
+      {
+        ranges[align_bits][0] = base;
+        return;
+      }
+
+      if (ranges[align_bits][1] != nullptr)
+        error("Critical assumption violated!");
+
+      ranges[align_bits][1] = base;
+    }
+
+    static void* remove_block(size_t align_bits)
+    {
+      auto first = ranges[align_bits][0];
+      if (first == nullptr)
+      {
+        if (align_bits < (bits::BITS - 1))
+        {
+          // Look for larger block and split up recursively
+          void* bigger = remove_block(align_bits + 1);
+          if (bigger == nullptr)
+          {
+            // Out of memory.
+            return bigger;
+          }
+          void* left_over =
+            pointer_offset(bigger, bits::one_at_bit(align_bits));
+          ranges[align_bits][0] = left_over;
+          return bigger;
+        }
+        // Out of memory
+        return nullptr;
+      }
+
+      auto second = ranges[align_bits][1];
+      if (second != nullptr)
+      {
+        ranges[align_bits][1] = nullptr;
+        return second;
+      }
+
+      ranges[align_bits][0] = nullptr;
+      return first;
+    }
 
   public:
     /**
      * This will be called by oe_allocator_init to set up enclave heap bounds.
      */
     static void setup_initial_range(void* base, void* end)
     {
-      oe_base = base;
-      oe_end = end;
+      // Find the minimum set of maximally aligned blocks in this range.
+      // Each block's alignment and size are equal.
+      size_t length = pointer_diff(base, end);
+      while (length != 0)
+      {
+        size_t base_align_bits = bits::ctz(address_cast(base));
+        size_t length_align_bits = (bits::BITS - 1) - bits::clz(length);
+        size_t align_bits = bits::min(base_align_bits, length_align_bits);
+        size_t align = bits::one_at_bit(align_bits);
+
+        add_block(align_bits, base);
+
+        base = pointer_offset(base, align);
+        length -= align;
+      }
     }
 
     /**
      * Bitmap of PalFeatures flags indicating the optional features that this
      * PAL supports.
      */
-    static constexpr uint64_t pal_features = 0;
+    static constexpr uint64_t pal_features = AlignedAllocation;
 
     static constexpr size_t page_size = 0x1000;
 
@@ -38,22 +115,18 @@ namespace snmalloc
     }
 
     template<bool committed>
-    void* reserve(size_t size) noexcept
+    static void* reserve(size_t size, size_t align) noexcept
     {
-      void* old_base = oe_base;
-      void* next_base;
-      auto end = oe_end;
-      do
-      {
-        auto new_base = old_base;
-        next_base = pointer_offset(new_base, size);
-
-        if (next_base > end)
-          return nullptr;
-
-      } while (!oe_base.compare_exchange_strong(old_base, next_base));
+      // The following are all true from the current way snmalloc uses the PAL.
+      // The implementation here is depending on them.
+      SNMALLOC_ASSERT(size == bits::next_pow2(size));
+      SNMALLOC_ASSERT(align == bits::next_pow2(align));
+      if (size != align)
+        error("Critical assumption violated!");
 
-      return old_base;
+      FlagLock lock(spin_lock);
+      size_t align_bits = bits::next_pow2_bits(align);
+      return remove_block(align_bits);
     }
 
     template<bool page_aligned = false>

src/test/func/two_alloc_types/alloc2.cc

@@ -1,4 +1,7 @@
-#undef IS_ADDRESS_SPACE_CONSTRAINED
+// Remove parameters feed from test harness
+#undef ADDRESS_SPACE_CONSTRAINED
+#undef REALLY_ADDRESS_SPACE_CONSTRAINED
+
 #define SNMALLOC_NAME_MANGLE(a) host_##a
 #define NO_BOOTSTRAP_ALLOCATOR
 #define SNMALLOC_EXPOSE_PAGEMAP