From 27d9c797c2f4c3e1543cba33df3567379e577b69 Mon Sep 17 00:00:00 2001
From: Miuna <809711+Mishura4@users.noreply.github.com>
Date: Mon, 16 Dec 2024 15:13:28 -0500
Subject: [PATCH] fix: coro timers ez pz lemon squeezy (#1351)

---
 include/dpp/cluster.h          | 27 +++++++++++++++++++++++++++
 include/dpp/zlibcontext.h      |  2 +-
 src/dpp/cluster/timer.cpp      | 11 ++++++-----
 src/dpp/socketengines/poll.cpp |  1 +
 src/unittest/test.cpp          |  3 +--
 5 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/include/dpp/cluster.h b/include/dpp/cluster.h
index 144b64b0c5..283d622a05 100644
--- a/include/dpp/cluster.h
+++ b/include/dpp/cluster.h
@@ -445,6 +445,33 @@ class DPP_EXPORT cluster {
 	 */
 	timer start_timer(timer_callback_t on_tick, uint64_t frequency, timer_callback_t on_stop = {});
 
+#ifdef DPP_CORO
+	/**
+	 * @brief Start a coroutine timer. Every `frequency` seconds, the callback is called.
+	 * 
+	 * @param on_tick The callback lambda to call for this timer when ticked
+	 * @param on_stop The callback lambda to call for this timer when it is stopped
+	 * @param frequency How often to tick the timer in seconds
+	 * @return timer A handle to the timer, used to remove that timer later
+	 */
+	template <std::invocable<timer> T, std::invocable<timer> U = std::function<void(timer)>>
+	requires (dpp::awaitable_type<typename std::invoke_result<T, timer>::type>)
+	timer start_timer(T&& on_tick, uint64_t frequency, U&& on_stop = {}) {
+		std::function<void(timer)> ticker = [fun = std::forward<T>(on_tick)](timer t) mutable -> dpp::job {
+			co_await std::invoke(fun, t);
+		};
+		std::function<void(timer)> stopper;
+		if constexpr (dpp::awaitable_type<typename std::invoke_result<U, timer>::type>) {
+			stopper = [fun = std::forward<U>(on_stop)](timer t) mutable -> dpp::job {
+				co_await std::invoke(fun, t);
+			};
+		} else {
+			stopper = std::forward<U>(on_stop);
+		}
+		return start_timer(std::move(ticker), frequency, std::move(stopper));
+	}
+#endif
+
 	/**
 	 * @brief Stop a ticking timer
 	 * 
diff --git a/include/dpp/zlibcontext.h b/include/dpp/zlibcontext.h
index a1a7f48350..e3396cd348 100644
--- a/include/dpp/zlibcontext.h
+++ b/include/dpp/zlibcontext.h
@@ -42,7 +42,7 @@ constexpr size_t DECOMP_BUFFER_SIZE = 512 * 1024;
  * This wraps the C pointers needed for zlib with unique_ptr and gives us a nice
  * buffer abstraction so we don't need to wrestle with raw pointers.
  */
-class zlibcontext {
+class DPP_EXPORT zlibcontext {
 public:
 	/**
 	 * @brief Zlib stream struct. The actual type is defined in zlib.h
diff --git a/src/dpp/cluster/timer.cpp b/src/dpp/cluster/timer.cpp
index 2fa610c002..e7bbfea2c2 100644
--- a/src/dpp/cluster/timer.cpp
+++ b/src/dpp/cluster/timer.cpp
@@ -32,8 +32,8 @@ timer cluster::start_timer(timer_callback_t on_tick, uint64_t frequency, timer_c
 
 	new_timer.handle = next_handle++;
 	new_timer.next_tick = time(nullptr) + frequency;
-	new_timer.on_tick = on_tick;
-	new_timer.on_stop = on_stop;
+	new_timer.on_tick = std::move(on_tick);
+	new_timer.on_stop = std::move(on_stop);
 	new_timer.frequency = frequency;
 
 	std::lock_guard<std::mutex> l(timer_guard);
@@ -64,16 +64,17 @@ void cluster::tick_timers() {
 		timer_t cur_timer;
 		{
 			std::lock_guard<std::mutex> l(timer_guard);
-			cur_timer = next_timer.top();
-			if (cur_timer.next_tick > now) {
+			if (next_timer.top().next_tick > now) {
 				/* Nothing to do */
 				break;
 			}
+			cur_timer = std::move(next_timer.top());
 			next_timer.pop();
 		}
 		timers_deleted_t::iterator deleted_iter{};
 		bool deleted{};
 		{
+			std::lock_guard<std::mutex> l(timer_guard);
 			deleted_iter = deleted_timers.find(cur_timer.handle);
 			deleted = deleted_iter != deleted_timers.end();
 		}
@@ -83,7 +84,7 @@ void cluster::tick_timers() {
 			cur_timer.next_tick += cur_timer.frequency;
 			{
 				std::lock_guard<std::mutex> l(timer_guard);
-				next_timer.emplace(cur_timer);
+				next_timer.emplace(std::move(cur_timer));
 			}
 		} else {
 			/* Deleted timers are not reinserted into the priority queue and their on_stop is called */
diff --git a/src/dpp/socketengines/poll.cpp b/src/dpp/socketengines/poll.cpp
index beb168ed9e..aca03bd2b3 100644
--- a/src/dpp/socketengines/poll.cpp
+++ b/src/dpp/socketengines/poll.cpp
@@ -20,6 +20,7 @@
  ************************************************************************************/
 
 #include <dpp/socketengine.h>
+#include <dpp/exception.h>
 #include <vector>
 #include <shared_mutex>
 #ifdef _WIN32
diff --git a/src/unittest/test.cpp b/src/unittest/test.cpp
index 4d5d308ee0..9837d934fa 100644
--- a/src/unittest/test.cpp
+++ b/src/unittest/test.cpp
@@ -2341,8 +2341,7 @@ Markdown lol \\|\\|spoiler\\|\\| \\~\\~strikethrough\\~\\~ \\`small \\*code\\* b
 		}
 
 		set_test(TIMERSTART, false);
-		static uint32_t ticks = 0;
-		dpp::timer th = bot.start_timer([](dpp::timer timer_handle) {
+		dpp::timer th = bot.start_timer([ticks = 0](dpp::timer timer_handle) mutable {
 			if (ticks == 2) {
 				/* The simple test timer ticks every second.
 				 * If we get to 2 seconds, we know the timer is working.