Add link time interface "get_stack_holder" and "stack_deleter::operator()(char* ptr) const noexcept" for Uthread

async_simple/uthread/internal/thread.cc

@@ -131,13 +131,18 @@ thread_context::thread_context(std::function<void()> func, size_t stack_size)
 
 thread_context::~thread_context() {}
 
-thread_context::stack_holder thread_context::make_stack() {
-    auto stack = stack_holder(new char[stack_size_]);
-    return stack;
+__attribute__((weak))
+void stack_deleter::operator()(char* ptr) const noexcept {
+    delete[] ptr;
 }
 
-void thread_context::stack_deleter::operator()(char* ptr) const noexcept {
-    delete[] ptr;
+__attribute__((weak))
+stack_holder get_stack_holder(unsigned stack_size) {
+    return stack_holder(new char[stack_size]);
+}
+
+stack_holder thread_context::make_stack() {
+    return get_stack_holder(stack_size_);
 }
 
 void thread_context::setup() {

async_simple/uthread/internal/thread.h

@@ -39,12 +39,12 @@ namespace internal {
 inline constexpr size_t default_base_stack_size = 512 * 1024;
 size_t get_base_stack_size();
 
-class thread_context {
-    struct stack_deleter {
-        void operator()(char* ptr) const noexcept;
-    };
-    using stack_holder = std::unique_ptr<char[], stack_deleter>;
+struct stack_deleter {
+    void operator()(char* ptr) const noexcept;
+};
+using stack_holder = std::unique_ptr<char[], stack_deleter>;
 
+class thread_context {
     const size_t stack_size_;
     stack_holder stack_{make_stack()};
     std::function<void()> func_;

async_simple/uthread/test/CMakeLists.txt

@@ -1,4 +1,4 @@
-file(GLOB uthread_test_src "*.cpp")
+file(GLOB uthread_test_src "UthreadTest.cpp")
 
 if (CMAKE_BUILD_TYPE STREQUAL "Debug" AND CMAKE_CXX_COMPILER_ID STREQUAL "Clang" AND
     CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL "13")
@@ -10,3 +10,10 @@ target_link_libraries(async_simple_uthread_test async_simple ${deplibs} ${testde
 
 add_test(NAME run_async_simple_uthread_test COMMAND async_simple_uthread_test)
 
+# We have to create new tests as linking will change the behavior
+add_executable(async_simple_uthread_alloc_test UthreadAllocTest.cpp ${PROJECT_SOURCE_DIR}/async_simple/test/dotest.cpp)
+
+target_link_libraries(async_simple_uthread_alloc_test async_simple ${deplibs} ${testdeplibs})
+
+add_test(NAME run_async_simple_uthread_alloc_test COMMAND async_simple_uthread_alloc_test)
+

async_simple/uthread/test/UthreadAllocTest.cpp

@@ -0,0 +1,120 @@
+#include <exception>
+#include <functional>
+#include <iostream>
+#include <type_traits>
+#include "async_simple/Common.h"
+#include "async_simple/executors/SimpleExecutor.h"
+#include "async_simple/test/unittest.h"
+#include "async_simple/uthread/Await.h"
+#include "async_simple/uthread/Uthread.h"
+
+using namespace std;
+
+static std::atomic<unsigned> get_stack_holder_count = 0;
+static std::atomic<unsigned> delete_stack_holder_count = 0;
+
+namespace async_simple {
+namespace uthread {
+namespace internal {
+
+stack_holder get_stack_holder(unsigned stack_size) {
+    get_stack_holder_count++;
+    return stack_holder(new char[stack_size]);
+}
+
+void stack_deleter::operator()(char* ptr) const noexcept {
+    delete_stack_holder_count++;
+    delete[] ptr;
+}
+}
+}
+}
+
+namespace async_simple {
+namespace uthread {
+class UthreadAllocSwapTest : public FUTURE_TESTBASE {
+public:
+    UthreadAllocSwapTest() : _executor(4) {}
+    void caseSetUp() override {}
+    void caseTearDown() override {}
+
+    template <class Func>
+    void delayedTask(Func&& func, std::size_t ms) {
+        std::thread(
+            [f = std::move(func), ms](Executor* ex) {
+                std::this_thread::sleep_for(std::chrono::milliseconds(ms));
+                ex->schedule(std::move(f));
+            },
+            &_executor)
+            .detach();
+    }
+
+    template <class T>
+    struct Awaiter {
+        Executor* ex;
+        T value;
+
+        Awaiter(Executor* e, T v) : ex(e), value(v) {}
+
+        bool await_ready() { return false; }
+        void await_suspend(std::coroutine_handle<> handle) noexcept {
+            auto ctx = ex->checkout();
+            std::thread([handle, e = ex, ctx]() mutable {
+                std::this_thread::sleep_for(std::chrono::milliseconds(100));
+                Executor::Func f = [handle]() mutable { handle.resume(); };
+                e->checkin(std::move(f), ctx);
+            }).detach();
+        }
+        T await_resume() noexcept { return value; }
+    };
+    template <class T>
+    coro::Lazy<T> lazySum(T x, T y) {
+        co_return co_await Awaiter{&_executor, x + y};
+    }
+
+protected:
+    executors::SimpleExecutor _executor;
+};
+
+TEST_F(UthreadAllocSwapTest, testSwitch) {
+    Executor* ex = &_executor;
+    auto show = [&](const std::string& message) mutable {
+        std::cout << message << "\n";
+    };
+
+    auto ioJob = [&]() -> Future<int> {
+        Promise<int> p;
+        auto f = p.getFuture().via(&_executor);
+        delayedTask(
+            [p = std::move(p)]() mutable {
+                auto value = 1024;
+                p.setValue(value);
+            },
+            100);
+        return f;
+    };
+
+    std::atomic<int> running = 1;
+    _executor.schedule([ex, &running, &show, &ioJob]() mutable {
+        Uthread task1(Attribute{ex}, [&show, &ioJob]() {
+            show("task1 start");
+            auto value = await(ioJob());
+            EXPECT_EQ(1024, value);
+            show("task1 done");
+        });
+        task1.join([&]() { 
+            running--;
+        });
+    });
+
+    while (running) {
+    }
+
+    EXPECT_NE(get_stack_holder_count, 0);
+    while (!delete_stack_holder_count) {
+    }
+
+    EXPECT_EQ(delete_stack_holder_count, get_stack_holder_count);
+}
+}
+}

docs/docs.cn/Uthread.md

@@ -185,3 +185,35 @@ uthread::await(p.getFuture());
 ## Sanitizer
 
 高版本的 Compiler-rt 默认会检测 `Use-After-Return` 的情况。而因为 Uthread 不能很好的处理这个情况，所以当我们在开启高版本 Compiler-rt 后使用 Uthread 时，需要使用 `-fsanitize-address-use-after-return=never` 这个选项来禁止检测 `Use-After-Return`.
+
+## 内存分配
+
+用户可以通过自定义 `stack_holder get_stack_holder(unsigned stack_size)` 和 `void stack_deleter::operator()(char* ptr) const noexcept` 强符号接口
+在链接时覆盖 Uthread 默认的内存分配器。例如：
+
+```C++
+namespace async_simple {
+namespace uthread {
+namespace internal {
+
+stack_holder get_stack_holder(unsigned stack_size) {
+    // user logics
+}
+
+void stack_deleter::operator()(char* ptr) const noexcept {
+    // user logics
+}
+}
+}
+}
+```
+
+其中 `stack_holder` 的定义为:
+
+```
+struct stack_deleter {
+    void operator()(char* ptr) const noexcept;
+};
+using stack_holder = std::unique_ptr<char[], stack_deleter>;
+```
+

docs/docs.en/Uthread.md

@@ -184,3 +184,36 @@ uthread::await(p.getFuture());
 ## Sanitizer
 
 Newer Compiler-rt will enable `Use-After-Return` by default. But it can't take care of Uthread well, so when we use Uthread with newer compiler-rt, we need to disable `Use-After-Return` explicitly by `-fsanitize-address-use-after-return=never`.
+
+## Memory Allocation
+
+
+Users can use self defined `stack_holder get_stack_holder(unsigned stack_size)` and `void stack_deleter::operator()(char* ptr) const noexcept` interfaces
+with strong symbols to override the default allocators and deallocators at link time:
+
+```C++
+namespace async_simple {
+namespace uthread {
+namespace internal {
+
+stack_holder get_stack_holder(unsigned stack_size) {
+    // user logics
+}
+
+void stack_deleter::operator()(char* ptr) const noexcept {
+    // user logics
+}
+}
+}
+}
+```
+
+the definition of  `stack_holder` is:
+
+```
+struct stack_deleter {
+    void operator()(char* ptr) const noexcept;
+};
+using stack_holder = std::unique_ptr<char[], stack_deleter>;
+```
+