[K/N] custom-alloc: Fix off-by-one in threshold

The existing threshold for determining between medium and large page
does not take the dummy block into account, meaning that an allocation
on the threshold will lead to a loop that only terminates when the
process goes out of memory.


Co-authored-by: Troels Lund <troels@google.com>

Merge-request: KOTLIN-MR-618
Merged-by: Alexander Shabalin <alexander.shabalin@jetbrains.com>

kotlin-native/runtime/src/custom_alloc/cpp/CustomAllocConstants.hpp

@@ -24,7 +24,8 @@ inline constexpr const size_t MEDIUM_PAGE_SIZE = (256 * KiB);
 inline constexpr const size_t MEDIUM_PAGE_CELL_COUNT =
         ((MEDIUM_PAGE_SIZE - sizeof(kotlin::alloc::MediumPage)) / sizeof(kotlin::alloc::Cell));
 
-inline constexpr const size_t LARGE_PAGE_SIZE_THRESHOLD = (MEDIUM_PAGE_CELL_COUNT - 1);
+// MEDIUM_PAGE_CELL_COUNT minus one cell for header minus another for the 0-sized dummy block at cells_[0]
+inline constexpr const size_t MEDIUM_PAGE_MAX_BLOCK_SIZE = (MEDIUM_PAGE_CELL_COUNT - 2);
 
 inline constexpr const size_t EXTRA_OBJECT_PAGE_SIZE = 64 * KiB;
 inline constexpr const int EXTRA_OBJECT_COUNT =

kotlin-native/runtime/src/custom_alloc/cpp/CustomAllocator.cpp

@@ -12,6 +12,7 @@
 #include <new>
 
 #include "ConcurrentMarkAndSweep.hpp"
+#include "CustomAllocConstants.hpp"
 #include "CustomLogging.hpp"
 #include "ExtraObjectData.hpp"
 #include "ExtraObjectPage.hpp"
@@ -126,7 +127,7 @@ uint8_t* CustomAllocator::Allocate(uint64_t size) noexcept {
     uint8_t* ptr;
     if (cellCount <= SMALL_PAGE_MAX_BLOCK_SIZE) {
         ptr = AllocateInSmallPage(cellCount);
-    } else if (cellCount > LARGE_PAGE_SIZE_THRESHOLD) {
+    } else if (cellCount > MEDIUM_PAGE_MAX_BLOCK_SIZE) {
         ptr = AllocateInLargePage(cellCount);
     } else {
         ptr = AllocateInMediumPage(cellCount);

kotlin-native/runtime/src/custom_alloc/cpp/CustomAllocatorTest.cpp

@@ -6,6 +6,7 @@
 #include <cstdint>
 #include <random>
 
+#include "CustomAllocConstants.hpp"
 #include "CustomAllocator.hpp"
 #include "GCScheduler.hpp"
 #include "Memory.h"
@@ -55,6 +56,32 @@ TEST(CustomAllocTest, SmallAllocSameSmallPage) {
     }
 }
 
+TEST(CustomAllocTest, SmallPageThreshold) {
+    Heap heap;
+    kotlin::gc::GCSchedulerConfig config;
+    kotlin::gc::GCSchedulerThreadData schedulerData(config, [](auto&) {});
+    CustomAllocator ca(heap, schedulerData);
+    const int FROM = SMALL_PAGE_MAX_BLOCK_SIZE - 10;
+    const int TO = SMALL_PAGE_MAX_BLOCK_SIZE + 10;
+    for (int blocks = FROM; blocks <= TO; ++blocks) {
+        TypeInfo fakeType = {.instanceSize_ = 8 * blocks, .flags_ = 0};
+        ca.CreateObject(&fakeType);
+    }
+}
+
+TEST(CustomAllocTest, MediumPageThreshold) {
+    Heap heap;
+    kotlin::gc::GCSchedulerConfig config;
+    kotlin::gc::GCSchedulerThreadData schedulerData(config, [](auto&) {});
+    CustomAllocator ca(heap, schedulerData);
+    const int FROM = MEDIUM_PAGE_MAX_BLOCK_SIZE - 10;
+    const int TO = MEDIUM_PAGE_MAX_BLOCK_SIZE + 10;
+    for (int blocks = FROM; blocks <= TO; ++blocks) {
+        TypeInfo fakeType = {.instanceSize_ = 8 * blocks, .flags_ = 0};
+        ca.CreateObject(&fakeType);
+    }
+}
+
 TEST(CustomAllocTest, TwoAllocatorsDifferentPages) {
     for (int blocks = MIN_BLOCK_SIZE; blocks < 2000; ++blocks) {
         Heap heap;

kotlin-native/runtime/src/custom_alloc/cpp/LargePage.cpp

@@ -16,7 +16,7 @@ namespace kotlin::alloc {
 
 LargePage* LargePage::Create(uint64_t cellCount) noexcept {
     CustomAllocInfo("LargePage::Create(%" PRIu64 ")", cellCount);
-    RuntimeAssert(cellCount > LARGE_PAGE_SIZE_THRESHOLD, "blockSize too small for large page");
+    RuntimeAssert(cellCount > MEDIUM_PAGE_MAX_BLOCK_SIZE, "blockSize too small for large page");
     uint64_t size = sizeof(LargePage) + cellCount * sizeof(uint64_t);
     return new (SafeAlloc(size)) LargePage();
 }