Fix thread-local PooledStackAllocator init order

CMakeLists.txt

@@ -212,9 +212,9 @@ file(GLOB PHOTON_SRC RELATIVE "${PROJECT_SOURCE_DIR}"
 if (APPLE)
     list(APPEND PHOTON_SRC io/kqueue.cpp)
 else ()
-    file(GLOB CACHE_SRC 
+    file(GLOB CACHE_SRC
         fs/cache/*.cpp
-        fs/cache/full_file_cache/*.cpp 
+        fs/cache/full_file_cache/*.cpp
         fs/cache/persistent_cache/*.cpp
     )
     list(APPEND PHOTON_SRC

photon.cpp

@@ -86,6 +86,8 @@ static int init_event_engine(uint64_t engine, uint64_t flags, const PhotonOption
 int __photon_init(uint64_t event_engine, uint64_t io_engine, const PhotonOptions& options) {
     if (options.use_pooled_stack_allocator) {
         use_pooled_stack_allocator();
+    } else {
+        use_default_stack_allocator();
     }
     if (options.bypass_threadpool) {
         set_bypass_threadpool(true);

thread/stack-allocator.cpp

@@ -18,9 +18,11 @@ limitations under the License.
 #include <linux/mman.h>
 #endif
 #include <errno.h>
+#include <mutex>
 #include <photon/common/alog.h>
 #include <photon/common/utility.h>
 #include <photon/thread/arch.h>
+#include <photon/thread/thread.h>
 #include <stdlib.h>
 #include <sys/mman.h>
 #include <unistd.h>
@@ -29,6 +31,8 @@ limitations under the License.
 
 namespace photon {
 
+static std::once_flag init_once_flag;
+
 template <size_t MIN_ALLOCATION_SIZE = 4UL * 1024,
           size_t MAX_ALLOCATION_SIZE = 64UL * 1024 * 1024>
 class PooledStackAllocator {
@@ -169,8 +173,13 @@ size_t pooled_stack_trim_threshold(size_t x) {
     return get_pooled_stack_allocator().threshold(x);
 }
 
-size_t pooled_stack_trim_current_vcpu(size_t keep_size);
-size_t pooled_stack_trim_threshold(size_t x);
+void use_pooled_stack_allocator() {
+    get_pooled_stack_allocator();
+    std::call_once(init_once_flag, [] {
+        photon::set_thread_stack_allocator({&pooled_stack_alloc, nullptr},
+                                           {&pooled_stack_dealloc, nullptr});
+    });
+}
 
 void* default_photon_thread_stack_alloc(void*, size_t stack_size) {
     char* ptr = nullptr;
@@ -193,4 +202,11 @@ void default_photon_thread_stack_dealloc(void*, void* ptr, size_t size) {
     free(ptr);
 }
 
+void use_default_stack_allocator() {
+    std::call_once(init_once_flag, [] {
+        photon::set_thread_stack_allocator({&default_photon_thread_stack_alloc, nullptr},
+                                           {&default_photon_thread_stack_dealloc, nullptr});
+    });
+}
+
 }  // namespace photon

thread/stack-allocator.h

@@ -18,33 +18,25 @@ limitations under the License.
 #include <photon/thread/thread.h>
 #include <stddef.h>
 
+// Set photon allocator/deallocator for photon thread stack.
+// This is a hook for thread allocation, both alloc and dealloc.
+// Helps user to do more works like mark GC while allocating
+
 namespace photon {
 
-// Set photon allocator/deallocator for photon thread stack
-// this is a hook for thread allocation, both alloc and dealloc
-// helps user to do more works like mark GC while allocating
-void* default_photon_thread_stack_alloc(void*, size_t stack_size);
-void default_photon_thread_stack_dealloc(void*, void* stack_ptr,
-                                            size_t stack_size);
+// A memory aligned allocator with stack-overflow protection.
+// This is the default option. No need to set explicitly.
+void use_default_stack_allocator();
 
-// Threadlocal Pooled stack allocator
-// better performance, and keep thread safe
-void* pooled_stack_alloc(void*, size_t stack_size);
-void pooled_stack_dealloc(void*, void* stack_ptr, size_t stack_size);
+// Thread-local Pooled stack allocator.
+// Better performance, and keep thread safe.
+void use_pooled_stack_allocator();
 
 // Free memory in pooled stack allocator till in-pool memory size less than
 // `keep_size` for current vcpu
 size_t pooled_stack_trim_current_vcpu(size_t keep_size);
+
 // Pooled stack allocator set keep-in-pool size
 size_t pooled_stack_trim_threshold(size_t threshold);
 
-void set_photon_thread_stack_allocator(
-    Delegate<void*, size_t> photon_thread_alloc = {
-        &default_photon_thread_stack_alloc, nullptr},
-    Delegate<void, void*, size_t> photon_thread_dealloc = {
-        &default_photon_thread_stack_dealloc, nullptr});
-inline void use_pooled_stack_allocator() {
-    set_photon_thread_stack_allocator({&pooled_stack_alloc, nullptr},
-                                      {&pooled_stack_dealloc, nullptr});
-}
 }  // namespace photon

thread/thread.cpp

@@ -125,10 +125,8 @@ namespace photon
         }
     };
 
-    static Delegate<void*, size_t> photon_thread_alloc(
-        &default_photon_thread_stack_alloc, nullptr);
-    static Delegate<void, void*, size_t> photon_thread_dealloc(
-        &default_photon_thread_stack_dealloc, nullptr);
+    static Delegate<void*, size_t> photon_thread_alloc;
+    static Delegate<void, void*, size_t> photon_thread_dealloc;
 
     struct vcpu_t;
     struct thread;
@@ -2095,6 +2093,10 @@ R"(
     }
 
     int vcpu_init(uint64_t flags) {
+        if (unlikely(!photon_thread_alloc || !photon_thread_dealloc)) {
+            // For some test cases that don't use photon::init
+            use_default_stack_allocator();
+        }
         uint64_t FLAGS = VCPU_ENABLE_ACTIVE_WORK_STEALING |
                          VCPU_ENABLE_PASSIVE_WORK_STEALING;
         if (unlikely(flags & ~FLAGS))
@@ -2148,11 +2150,9 @@ R"(
         CURRENT->stackful_free(ptr);
     }
 
-    void set_photon_thread_stack_allocator(
-        Delegate<void *, size_t> _photon_thread_alloc,
-        Delegate<void, void *, size_t> _photon_thread_dealloc) {
-        photon_thread_alloc = _photon_thread_alloc;
-        photon_thread_dealloc = _photon_thread_dealloc;
+    void set_thread_stack_allocator(Delegate<void *, size_t> alloc, Delegate<void, void *, size_t> dealloc) {
+        photon_thread_alloc = alloc;
+        photon_thread_dealloc = dealloc;
     }
 
     extern "C" {

thread/thread.h

@@ -538,6 +538,8 @@ namespace photon
     // helps allocating when using hybrid C++20 style coroutine
     void* stackful_malloc(size_t size);
     void stackful_free(void* ptr);
+
+    void set_thread_stack_allocator(Delegate<void*, size_t> alloc, Delegate<void, void*, size_t> dealloc);
 };
 
 /*