set_photon_thread_stack_allocator able to hook photon threads alloc…

…ation, helps with GC adaption Signed-off-by: Coldwings <coldwings@me.com>
alibaba · Jul 4, 2023 · f01c76a · f01c76a
1 parent 5058a7a
commit f01c76a
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 14 deletions.
diff --git a/thread/thread.cpp b/thread/thread.cpp
@@ -111,6 +111,30 @@ namespace photon
         }
     };
 
+    void* default_photon_thread_stack_alloc(void*, size_t stack_size) {
+        char* ptr = nullptr;
+        int err = posix_memalign((void**)&ptr, PAGE_SIZE, stack_size);
+        if (unlikely(err))
+            LOG_ERROR_RETURN(err, nullptr, "Failed to allocate photon stack! ",
+                             ERRNO(err));
+#if defined(__linux__)
+        madvise(ptr, stack_size, MADV_NOHUGEPAGE);
+#endif
+        return ptr;
+    }
+
+    void default_photon_thread_stack_dealloc(void*, void* ptr, size_t size) {
+#ifndef __aarch64__
+        madvise(ptr, size, MADV_DONTNEED);
+#endif
+        free(ptr);
+    }
+
+    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);
+
     struct vcpu_t;
     struct thread;
     class Stack
@@ -257,11 +281,9 @@ namespace photon
         }
         void dispose() {
             assert(state == states::DONE);
-            register auto b = buf; //store in register to prevent from being deleted by madvise
-#ifndef __aarch64__
-            madvise(b, stack_size, MADV_DONTNEED);
-#endif
-            free(b);
+            // `buf` and `stack_size` will always store on register
+            // when calling deallocating.
+            photon_thread_dealloc(buf, stack_size);
         }
     };
 
@@ -816,13 +838,7 @@ R"(
             LOG_ERROR_RETURN(ENOSYS, nullptr, "Photon not initialized in this vCPU (OS thread)");
         size_t randomizer = (rand() % 32) * (1024 + 8);
         stack_size = align_up(randomizer + stack_size + sizeof(thread), PAGE_SIZE);
-        char* ptr = nullptr;
-        int err = posix_memalign((void**)&ptr, PAGE_SIZE, stack_size);
-        if (unlikely(err))
-            LOG_ERROR_RETURN(err, nullptr, "Failed to allocate photon stack! ", ERRNO(err));
-#if defined(__linux__)
-        madvise(ptr, stack_size, MADV_NOHUGEPAGE);
-#endif
+        char* ptr = (char*)photon_thread_alloc(stack_size);
         auto p = ptr + stack_size - sizeof(thread) - randomizer;
         (uint64_t&)p &= ~63;
         auto th = new (p) thread;

diff --git a/thread/thread.h b/thread/thread.h
@@ -32,8 +32,6 @@ namespace photon
     int wait_all();
     int timestamp_updater_init();
     int timestamp_updater_fini();
-    void* stackful_malloc(size_t size);
-    void stackful_free(void* ptr);
 
 
     struct thread;
@@ -433,6 +431,23 @@ namespace photon
     bool is_master_event_engine_default();
     void reset_master_event_engine_default();
 
+    // alloc space on rear end of current thread stack,
+    // helps allocating when using hybrid C++20 style coroutine
+    void* stackful_malloc(size_t size);
+    void stackful_free(void* ptr);
+
+    // 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);
+    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});
+
     // Saturating addition, primarily for timeout caculation
     __attribute__((always_inline)) inline uint64_t sat_add(uint64_t x,
                                                            uint64_t y) {