io_,task_scheduler: Add scheduleOnceNoAlloc.

Author: rtrbt

software/src/modules/io_scheduler/io_scheduler.cpp

@@ -96,10 +96,10 @@ void IoScheduler::drive(RoundStateBase *state)
 
     state->in_flight = true;
 
-    this->scheduleOnce([state]() {
+    this->scheduleOnceNoAlloc(&state->during_io_task_buf, [state]() {
         state->run_during_io();
 
-        task_scheduler.scheduleOnce([state]() {
+        task_scheduler.scheduleOnceNoAlloc(&state->after_io_task_buf, [state]() {
             state->run_after_io();
             state->in_flight = false;
         }, 0_ms, state->src_location);

software/src/modules/io_scheduler/io_scheduler.h

@@ -52,6 +52,11 @@ class IoScheduler final : public IModule
         return scheduler.scheduleOnce(std::move(fn), delay_ms, src_location);
     }
 
+    uint64_t scheduleOnceNoAlloc(aligned_storage<Task> *task_buf, std::function<void(void)> &&fn, millis_t delay_ms = 0_ms, const std::source_location &src_location = std::source_location::current())
+    {
+        return scheduler.scheduleOnceNoAlloc(task_buf, std::move(fn), delay_ms, src_location);
+    }
+
     // Registers an uncancelable driven round.
     // before_io and after_io run on the main task, during_io runs on the IO task.
     // Pass nullptr for before_io/after_io if unused.
@@ -124,6 +129,8 @@ class IoScheduler final : public IModule
     [[noreturn]] static void task_fn(void *arg);
 
     struct RoundStateBase {
+        aligned_storage<Task> during_io_task_buf;
+        aligned_storage<Task> after_io_task_buf;
         const std::source_location src_location;
         bool in_flight = false; // Only accessed on the main task.
 

software/src/modules/task_scheduler/task_scheduler.cpp

@@ -248,14 +248,21 @@ bool TaskScheduler::nextTaskReady()
     return next_deadline <= now;
 }
 
-uint64_t TaskScheduler::scheduleOnce(std::function<void(void)> &&fn, millis_t delay_ms, const std::source_location &src_location)
+uint64_t TaskScheduler::scheduleOnceNoAlloc(aligned_storage<Task> *task_buffer, std::function<void(void)> &&fn, millis_t delay_ms, const std::source_location &src_location, bool transfer_task_ownership)
 {
     std::lock_guard<std::mutex> lock{this->task_mutex};
     uint64_t task_id = ++last_task_id;
-    tasks.emplace(new Task(std::move(fn), task_id, delay_ms, 0_us, src_location.file_name(), src_location.line(), true));
+    Task *ptr = new(task_buffer) Task(std::move(fn), task_id, delay_ms, 0_us, src_location.file_name(), src_location.line(), true);
+    ptr->owned = transfer_task_ownership;
+    tasks.emplace(ptr);
     return task_id;
 }
 
+uint64_t TaskScheduler::scheduleOnce(std::function<void(void)> &&fn, millis_t delay_ms, const std::source_location &src_location)
+{
+    return this->scheduleOnceNoAlloc(new aligned_storage<Task>, std::move(fn), delay_ms, src_location, true);
+}
+
 uint64_t TaskScheduler::scheduleWithFixedDelay(std::function<void(void)> &&fn, millis_t first_delay_ms, millis_t delay_ms, const std::source_location &src_location)
 {
     std::lock_guard<std::mutex> lock{this->task_mutex};
@@ -268,7 +275,9 @@ uint64_t TaskScheduler::scheduleUncancelable(std::function<void(void)> &&fn, mil
 {
     std::lock_guard<std::mutex> lock{this->task_mutex};
     uint64_t task_id = ++last_task_id | (1ull << 62ull);
-    tasks.emplace(perm_new<Task>(DRAM, std::move(fn), task_id, first_delay_ms, delay_ms, src_location.file_name(), src_location.line(), false));
+    auto *task = perm_new<Task>(DRAM, std::move(fn), task_id, first_delay_ms, delay_ms, src_location.file_name(), src_location.line(), false);
+    task->owned = false;
+    tasks.emplace(task);
     return task_id;
 }
 

software/src/modules/task_scheduler/task_scheduler.h

@@ -46,10 +46,22 @@ struct Task {
     // Task is currently running while another thread (or fn itself) called rescheduleNow with this task's ID.
     // Will be executed as next task again.
     bool immediate_reschedule;
+    // Task was allocated by the task scheduler.
+    bool owned = true;
 
     Task(std::function<void(void)> &&fn, uint64_t task_id, micros_t first_run_delay, micros_t delay, const char *file, uint_least32_t line, bool once);
 };
 
+// https://stackoverflow.com/a/50957671
+// This specializes the default deleter used by std::unique_ptr to only delete a Task if it is owned.
+template <>
+struct std::default_delete<Task> {
+    default_delete() = default;
+    template <class U>
+    constexpr default_delete(default_delete<U>) noexcept {}
+    void operator()(Task* p) const noexcept { if (p->owned) delete p; }
+};
+
 #define IS_WALL_CLOCK_TASK_ID(task_id) (task_id & (1ull << 63))
 #define IS_UNCANCELABLE_TASK_ID(task_id) (task_id & (1ull << 62))
 
@@ -123,6 +135,7 @@ class TaskScheduler final : public IModule
     bool nextTaskReady();
 
     uint64_t scheduleOnce(std::function<void(void)> &&fn, millis_t delay_ms = 0_ms, const std::source_location &src_location = std::source_location::current());
+    uint64_t scheduleOnceNoAlloc(aligned_storage<Task> *task_buffer, std::function<void(void)> &&fn, millis_t delay_ms = 0_ms, const std::source_location &src_location = std::source_location::current(), bool transfer_task_ownership = false);
 
     [[nodiscard("Use scheduleUncancelable if you don't need the returned task ID to cancel this task later. Cast to void if you intend to write a self-canceling task that will use task_scheduler.currentTaskId()")]]
     inline uint64_t scheduleWithFixedDelay(std::function<void(void)> &&fn, millis_t delay_ms, const std::source_location &src_location = std::source_location::current()) {return this->scheduleWithFixedDelay(std::move(fn), 0_ms, delay_ms, src_location);}

software/src/tools.h

@@ -297,3 +297,11 @@ struct PointerWithBits {
     [[gnu::always_inline]]
     T *get_ptr() const { return reinterpret_cast<T *>(reinterpret_cast<uintptr_t>(this->ptr) & (~BITMASK)); }
 };
+
+// std::aligned_storage is deprecated in C++23.
+// https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2021/p1413r3.pdf
+// Use this as an replacement.
+template <typename T>
+struct aligned_storage {
+    alignas(T) std::byte buf[sizeof(T)];
+};