From ce489cfffe84d64f562fe43cfbf6b1bcafb304a8 Mon Sep 17 00:00:00 2001
From: Matthew Parkinson <mjp41@users.noreply.github.com>
Date: Tue, 20 Jun 2023 12:00:34 -0400
Subject: [PATCH] Conditional range (#617)

* Make a conditional range

This range allows for a contained range to be disabled at runtime.
This allows for thread local caching to be disabled if the initial fixed
size heap is below a threshold.
---
 src/snmalloc/backend/fixedglobalconfig.h      |  6 ++
 src/snmalloc/backend/standard_range.h         | 10 ++-
 .../backend_helpers/backend_helpers.h         |  1 +
 .../backend_helpers/staticconditionalrange.h  | 78 +++++++++++++++++++
 src/snmalloc/ds/allocconfig.h                 |  6 ++
 5 files changed, 99 insertions(+), 2 deletions(-)
 create mode 100644 src/snmalloc/backend_helpers/staticconditionalrange.h

diff --git a/src/snmalloc/backend/fixedglobalconfig.h b/src/snmalloc/backend/fixedglobalconfig.h
index b45db72e..c6784e70 100644
--- a/src/snmalloc/backend/fixedglobalconfig.h
+++ b/src/snmalloc/backend/fixedglobalconfig.h
@@ -87,6 +87,12 @@ namespace snmalloc
       auto [heap_base, heap_length] =
         Pagemap::concretePagemap.init(base, length);
 
+      // Make this a alloc_config constant.
+      if (length < MIN_HEAP_SIZE_FOR_THREAD_LOCAL_BUDDY)
+      {
+        LocalState::set_small_heap();
+      }
+
       Authmap::arena = capptr::Arena<void>::unsafe_from(heap_base);
 
       Pagemap::register_range(Authmap::arena, heap_length);
diff --git a/src/snmalloc/backend/standard_range.h b/src/snmalloc/backend/standard_range.h
index 40c7cd21..78609ed2 100644
--- a/src/snmalloc/backend/standard_range.h
+++ b/src/snmalloc/backend/standard_range.h
@@ -49,11 +49,11 @@ namespace snmalloc
     // Use buddy allocators to cache locally.
     using LargeObjectRange = Pipe<
       Stats,
-      LargeBuddyRange<
+      StaticConditionalRange<LargeBuddyRange<
         LocalCacheSizeBits,
         LocalCacheSizeBits,
         Pagemap,
-        page_size_bits>>;
+        page_size_bits>>>;
 
   private:
     using ObjectRange = Pipe<LargeObjectRange, SmallBuddyRange>;
@@ -85,5 +85,11 @@ namespace snmalloc
       // Use the object range to service meta-data requests.
       return object_range;
     }
+
+    static void set_small_heap()
+    {
+      // This disables the thread local caching of large objects.
+      LargeObjectRange::disable_range();
+    }
   };
 } // namespace snmalloc
diff --git a/src/snmalloc/backend_helpers/backend_helpers.h b/src/snmalloc/backend_helpers/backend_helpers.h
index 32b10a4a..2104e681 100644
--- a/src/snmalloc/backend_helpers/backend_helpers.h
+++ b/src/snmalloc/backend_helpers/backend_helpers.h
@@ -14,5 +14,6 @@
 #include "palrange.h"
 #include "range_helpers.h"
 #include "smallbuddyrange.h"
+#include "staticconditionalrange.h"
 #include "statsrange.h"
 #include "subrange.h"
diff --git a/src/snmalloc/backend_helpers/staticconditionalrange.h b/src/snmalloc/backend_helpers/staticconditionalrange.h
new file mode 100644
index 00000000..46637135
--- /dev/null
+++ b/src/snmalloc/backend_helpers/staticconditionalrange.h
@@ -0,0 +1,78 @@
+#pragma once
+#include "../pal/pal.h"
+#include "empty_range.h"
+#include "range_helpers.h"
+
+namespace snmalloc
+{
+  template<typename OptionalRange>
+  struct StaticConditionalRange
+  {
+    // This is a range that can bypass the OptionalRange if it is disabled.
+    // Disabling is global, and not local.
+    // This is used to allow disabling thread local buddy allocators when the
+    // initial fixed size heap is small.
+    //
+    // The range builds a more complex parent
+    //    Pipe<ParentRange, OptionalRange>
+    // and uses the ancestor functions to bypass the OptionalRange if the flag
+    // has been set.
+    template<typename ParentRange>
+    class Type : public ContainsParent<Pipe<ParentRange, OptionalRange>>
+    {
+      // This contains connects the optional range to the parent range.
+      using ActualParentRange = Pipe<ParentRange, OptionalRange>;
+
+      using ContainsParent<ActualParentRange>::parent;
+
+      // Global flag specifying if the optional range should be disabled.
+      static inline bool disable_range_{false};
+
+    public:
+      // Both parent and grandparent must be aligned for this range to be
+      // aligned.
+      static constexpr bool Aligned =
+        ActualParentRange::Aligned && ParentRange::Aligned;
+
+      // Both parent and grandparent must be aligned for this range to be
+      // concurrency safe.
+      static constexpr bool ConcurrencySafe =
+        ActualParentRange::ConcurrencySafe && ParentRange::ConcurrencySafe;
+
+      using ChunkBounds = typename ActualParentRange::ChunkBounds;
+
+      static_assert(
+        std::is_same_v<ChunkBounds, typename ParentRange::ChunkBounds>,
+        "Grandparent and optional parent range chunk bounds must be equal");
+
+      constexpr Type() = default;
+
+      CapPtr<void, ChunkBounds> alloc_range(size_t size)
+      {
+        if (disable_range_)
+        {
+          // Use ancestor to bypass the optional range.
+          return this->template ancestor<ParentRange>()->alloc_range(size);
+        }
+
+        return parent.alloc_range(size);
+      }
+
+      void dealloc_range(CapPtr<void, ChunkBounds> base, size_t size)
+      {
+        if (disable_range_)
+        {
+          // Use ancestor to bypass the optional range.
+          this->template ancestor<ParentRange>()->dealloc_range(base, size);
+          return;
+        }
+        parent.dealloc_range(base, size);
+      }
+
+      static void disable_range()
+      {
+        disable_range_ = true;
+      }
+    };
+  };
+} // namespace snmalloc
diff --git a/src/snmalloc/ds/allocconfig.h b/src/snmalloc/ds/allocconfig.h
index 7f1e55e9..858940f0 100644
--- a/src/snmalloc/ds/allocconfig.h
+++ b/src/snmalloc/ds/allocconfig.h
@@ -92,4 +92,10 @@ namespace snmalloc
     1 << MIN_CHUNK_BITS
 #endif
     ;
+
+  // Used to configure when the backend should use thread local buddies.
+  // This only basically is used to disable some buddy allocators on small
+  // fixed heap scenarios like OpenEnclave.
+  static constexpr size_t MIN_HEAP_SIZE_FOR_THREAD_LOCAL_BUDDY =
+    bits::one_at_bit(27);
 } // namespace snmalloc