增加自定义分配器功能

c++ 就是这样的，要极大的满足用户的自定义需求。而用户非常有可能，是需要定制 promise 对象的内存分配的。因此，学习 STL 容器的做法，让用户可以设定 Allocator 模板参数，从而定制化内存需求。
avplayer · Nov 4, 2024 · b8dcdc8 · b8dcdc8
1 parent 4ab020b
commit b8dcdc8
Show file tree

Hide file tree

Showing 4 changed files with 123 additions and 51 deletions.
diff --git a/include/ucoro/awaitable.hpp b/include/ucoro/awaitable.hpp
@@ -55,12 +55,12 @@ namespace ucoro
 	template<typename T>
 	struct await_transformer;
 
-	template<typename T>
-	struct awaitable;
-
-	template<typename T>
+	template<typename T, typename Allocator>
 	struct awaitable_promise;
 
+	template<typename T, typename Allocator = std::allocator<char>>
+	struct awaitable;
+
 	template<typename T, typename CallbackFunction>
 	struct CallbackAwaiter;
 
@@ -143,27 +143,29 @@ namespace ucoro
 		// template_parameter_of<decltype(local_storage), local_storage_t>;  // void
 		//
 		// 首先定义一个接受 template_parameter_of<Testee, FromTemplate> 这样的一个默认模板萃取
-		template<typename Testee, template<typename> typename FromTemplate>
+		template<typename Testee, template<typename...> typename FromTemplate>
 		struct template_parameter_traits;
 
 		// 接着定义一个偏特化，匹配 template_parameter_traits<模板名<参数>, 模板名>
 		// 这样，这个偏特化的 template_parameter_traits 就有了一个
 		// 名为 template_parameter 的成员类型，其定义的类型就是 _template_parameter
 		// 于是就把 TemplateParameter 这个类型给萃取出来了
-		template<template<typename> typename ClassTemplate, typename TemplateParameter>
-		struct template_parameter_traits<ClassTemplate<TemplateParameter>, ClassTemplate>
+		template<template<typename...> typename ClassTemplate, typename... TemplateParameter>
+		struct template_parameter_traits<ClassTemplate<TemplateParameter...>, ClassTemplate>
 		{
-			using template_parameter = TemplateParameter ;
+			using template_parameter_tuple = std::tuple<TemplateParameter...>;
+			using template_parameter = typename std::tuple_element<0, template_parameter_tuple>::type;
 		};
 
 		// 最后，定义一个简化用法的 using 让用户的地方代码变短点
-		template<typename TesteeType, template<typename> typename FromTemplate>
+		template<typename TesteeType, template<typename...> typename FromTemplate>
 		using template_parameter_of = typename template_parameter_traits<
 									std::decay_t<TesteeType>, FromTemplate>::template_parameter;
 
 		// 利用 通用工具 template_parameter_of 萃取 local_storage_t<T> 里的 T
 		template<concepts::local_storage_type LocalStorage>
 		using local_storage_value_type = template_parameter_of<LocalStorage, local_storage_t>;
+
 
 
 		// 利用 通用工具 template_parameter_of 萃取 awaitable<T> 里的 T
@@ -188,28 +190,30 @@ namespace ucoro
 
 	} // namespace traits
 
-	struct debug_coro_promise
+	template<typename Allocator>
+	struct coro_promise_allocation
 	{
-#if defined(DEBUG_CORO_PROMISE_LEAK)
-
 		void* operator new(std::size_t size)
 		{
-			void* ptr = std::malloc(size);
+			void* ptr = Allocator{}.allocate(size);// std::malloc(size);
 			if (!ptr)
 			{
 				throw std::bad_alloc{};
 			}
+#if defined(DEBUG_CORO_PROMISE_LEAK)
 			debug_coro_leak.insert(ptr);
+#endif // DEBUG_CORO_PROMISE_LEAK
 			return ptr;
 		}
 
 		void operator delete(void* ptr, [[maybe_unused]] std::size_t size)
 		{
+#if defined(DEBUG_CORO_PROMISE_LEAK)
 			debug_coro_leak.erase(ptr);
-			std::free(ptr);
+#endif // DEBUG_CORO_PROMISE_LEAK
+			Allocator{}.deallocate((typename Allocator::value_type*)(ptr), size);
 		}
 
-#endif // DEBUG_CORO_PROMISE_LEAK
 	};
 
 	//////////////////////////////////////////////////////////////////////////
@@ -266,45 +270,42 @@ namespace ucoro
 		}
 	};
 
-	//////////////////////////////////////////////////////////////////////////
-
-	template<typename T>
-	struct final_awaitable : std::suspend_always
-	{
-		std::coroutine_handle<> await_suspend(std::coroutine_handle<awaitable_promise<T>> h) noexcept
-		{
-			if (h.promise().continuation_)
-			{
-				// continuation_ 不为空，则 说明 .detach() 被 co_await
-				// 因此，awaitable_detached 析构的时候会顺便撤销自己，所以这里不用 destory
-				// 返回 continuation_，以便让协程框架调用 continuation_.resume()
-				// 这样就把等它的协程唤醒了.
-				return h.promise().continuation_;
-			}
-			// 并且，如果协程处于 .detach() 而没有被 co_await
-			// 则异常一直存储在 promise 里，并没有代码会去调用他的 await_resume() 重抛异常
-			// 所以这里重新抛出来，避免有被静默吞并的异常
-			h.promise().get_value();
-			// 如果 continuation_ 为空，则说明 .detach() 没有被 co_await
-			// 因此，awaitable_detached 对象其实已经析构
-			// 所以必须主动调用 destroy() 以免内存泄漏.
-			h.destroy();
-			return std::noop_coroutine();
-		}
-	};
-
 	//////////////////////////////////////////////////////////////////////////
 	// 返回 T 的协程 awaitable_promise 实现.
 
 	// Promise 类型实现...
-	template<typename T>
-	struct awaitable_promise : public awaitable_promise_value<T>, public debug_coro_promise
+	template<typename T, typename Allocator>
+	struct awaitable_promise : public awaitable_promise_value<T>, public coro_promise_allocation<Allocator>
 	{
-		awaitable<T> get_return_object();
+		awaitable<T, Allocator> get_return_object();
 
 		auto final_suspend() noexcept
 		{
-			return final_awaitable<T>{};
+			struct final_awaitable : std::suspend_always
+			{
+				std::coroutine_handle<> await_suspend(std::coroutine_handle<awaitable_promise> h) noexcept
+				{
+					if (h.promise().continuation_)
+					{
+						// continuation_ 不为空，则 说明 .detach() 被 co_await
+						// 因此，awaitable_detached 析构的时候会顺便撤销自己，所以这里不用 destory
+						// 返回 continuation_，以便让协程框架调用 continuation_.resume()
+						// 这样就把等它的协程唤醒了.
+						return h.promise().continuation_;
+					}
+					// 并且，如果协程处于 .detach() 而没有被 co_await
+					// 则异常一直存储在 promise 里，并没有代码会去调用他的 await_resume() 重抛异常
+					// 所以这里重新抛出来，避免有被静默吞并的异常
+					h.promise().get_value();
+					// 如果 continuation_ 为空，则说明 .detach() 没有被 co_await
+					// 因此，awaitable_detached 对象其实已经析构
+					// 所以必须主动调用 destroy() 以免内存泄漏.
+					h.destroy();
+					return std::noop_coroutine();
+				}
+			};
+
+			return final_awaitable{};
 		}
 
 		auto initial_suspend()
@@ -371,10 +372,10 @@ namespace ucoro
 	//////////////////////////////////////////////////////////////////////////
 
 	// awaitable 协程包装...
-	template<typename T>
+	template<typename T, typename Allocator>
 	struct awaitable
 	{
-		using promise_type = awaitable_promise<T>;
+		using promise_type = awaitable_promise<T, Allocator>;
 
 		explicit awaitable(std::coroutine_handle<promise_type> h)
 			: current_coro_handle_(h)
@@ -507,10 +508,10 @@ namespace ucoro
 
 	//////////////////////////////////////////////////////////////////////////
 
-	template<typename T>
-	awaitable<T> awaitable_promise<T>::get_return_object()
+	template<typename T, typename Allocator>
+	awaitable<T, Allocator> awaitable_promise<T, Allocator>::get_return_object()
 	{
-		auto result = awaitable<T>{std::coroutine_handle<awaitable_promise<T>>::from_promise(*this)};
+		auto result = awaitable<T, Allocator>{std::coroutine_handle<awaitable_promise<T, Allocator>>::from_promise(*this)};
 		return result;
 	}
 

diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
@@ -40,3 +40,5 @@ if (LIBEVENT_FOUND)
 endif(LIBEVENT_FOUND)
 
 add_subdirectory(test_executor)
+
+add_subdirectory(test_custom_allocator)
diff --git a/tests/test_custom_allocator/CMakeLists.txt b/tests/test_custom_allocator/CMakeLists.txt
@@ -0,0 +1,6 @@
+
+add_executable(test_custom_allocator test.cpp)
+target_link_libraries(test_custom_allocator ucoro)
+
+add_test(NAME test_custom_allocator COMMAND test_custom_allocator)
+set_target_properties(test_custom_allocator PROPERTIES FOLDER "ucoro_tests")
diff --git a/tests/test_custom_allocator/test.cpp b/tests/test_custom_allocator/test.cpp
@@ -0,0 +1,63 @@
+
+#include "ucoro/awaitable.hpp"
+#include <iostream>
+
+struct my_allocator
+{
+	using value_type = void;
+
+	value_type* allocate(std::size_t size)
+	{
+		return std::malloc(size);
+	}
+
+	void deallocate(value_type* ptr, std::size_t size)
+	{
+		free(ptr);
+	}
+
+};
+
+// 只有 test2 用自定义分配器分配
+// 验证混合使用 分配器的awaitable也是没问题的
+ucoro::awaitable<int, my_allocator> test()
+{
+	throw std::runtime_error("test throw");
+    co_return 1;
+}
+
+
+ucoro::awaitable<void> test2()
+{
+	throw std::runtime_error("test throw");
+    co_return;
+}
+
+ucoro::awaitable<int> coro_compute()
+{
+	try
+	{
+		sync_await(test2());
+	}
+	catch(const std::exception& e)
+	{
+		std::cerr << "exception in test2: " << e.what() << '\n';
+	}
+
+	co_return co_await test();
+}
+
+int main(int argc, char** argv)
+{
+	try
+	{
+		std::string str = "hello";
+		sync_await(coro_compute(), str);
+	}
+	catch (std::exception& e)
+	{
+		std::cerr << "exception: " << e.what() << std::endl;
+	}
+
+	return 0;
+}
Original file line number	Diff line number	Diff line change
Expand Up		@@ -40,3 +40,5 @@ if (LIBEVENT_FOUND)
		endif(LIBEVENT_FOUND)

		add_subdirectory(test_executor)

		add_subdirectory(test_custom_allocator)