Allow reading current task from Task::current()

CMakeLists.txt

@@ -11,8 +11,8 @@ add_library(cosched STATIC scheduler.cpp include/scheduler.hpp)
 target_link_libraries(cosched PUBLIC async)
 target_include_directories(cosched PUBLIC include/)
 
-#add_executable(test test.cpp)
-#target_link_libraries(test PRIVATE cosched)
+add_executable(test test.cpp)
+target_link_libraries(test PRIVATE cosched)
 
 install(TARGETS cosched
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR})

include/scheduler.hpp

@@ -31,6 +31,8 @@ enum class TaskState {
 class Task {
     friend class Scheduler;
 
+    static thread_local class Task *current;
+
     class Scheduler *scheduler;
     async::recurring_event resume;
 
@@ -41,12 +43,19 @@ class Task {
     TaskState state = TaskState::running;
     bool suspended = false;
 
+    void kill();
+
 public:
     Task(Scheduler *scheduler, const std::string& name)
         : scheduler(scheduler), name(name) {}
     Task(const Task&) = delete;
     Task(Task&&) = delete;
 
+    static inline
+    Task& get_current() {
+        return *current;
+    }
+
     const std::string& get_name() const {
         return name;
     }
@@ -88,36 +97,6 @@ public:
 };
 
 
-class TaskPtr {
-    unsigned *ref_count;
-    Task *task;
-
-    void finalize();
-
-public:
-    TaskPtr(Task *task) : task(task) {
-        ref_count = new unsigned(1);
-    }
-    ~TaskPtr() {
-        if (!task) return;
-        if (--*ref_count == 0) {
-            delete ref_count;
-            finalize();
-        }
-    }
-    TaskPtr(const TaskPtr& o) : ref_count(o.ref_count), task(o.task) {
-        ++*ref_count;
-    }
-    TaskPtr(TaskPtr&& o) : ref_count(o.ref_count), task(o.task) {
-        o.task = nullptr;
-    }
-
-    auto operator ->() {
-        return task;
-    }
-};
-
-
 class Scheduler {
     friend class Task;
     friend class TaskPtr;
@@ -138,17 +117,13 @@ public:
         return tasks;
     }
 
-    TaskPtr create_task(const std::string& name) {
+    // Creates new task, returns it and switches to it
+    void create_task(const std::string& name) {
         std::scoped_lock L(tasks_mutex);
-        return TaskPtr(tasks.emplace_back(std::make_unique<Task>(this, name)).get());
+        Task::current = tasks.emplace_back(std::make_unique<Task>(this, name)).get();
     }
 
     void run();
 };
-
-
-inline void TaskPtr::finalize() {
-    task->get_scheduler().delete_task(task);
-}
 }
 #endif // _SCHEDULER_HPP

scheduler.cpp

@@ -5,6 +5,10 @@
 
 
 
+void CoSched::Task::kill() {
+    get_scheduler().delete_task(this);
+}
+
 async::result<bool> CoSched::Task::yield() {
     if (state == TaskState::terminating) {
         // If it was terminating, it can finally be declared dead now
@@ -63,10 +67,17 @@ CoSched::Task *CoSched::Scheduler::get_next_task() {
 
 void CoSched::Scheduler::run() {
     while (!tasks.empty()) {
-        // Get next task
-        auto next_task = get_next_task();
+        // If current task isn't sleeping, it is considered a zombie so removed from list
+        if (Task::current && Task::current->state != TaskState::sleeping) {
+            delete_task(std::exchange(Task::current, nullptr));
+        }
+
+        // Get new task
+        Task::current = get_next_task();
 
         // Resume task if any
-        if (next_task) next_task->resume.raise();
+        if (Task::current) Task::current->resume.raise();
     }
 }
+
+thread_local CoSched::Task *CoSched::Task::current;

test.cpp

@@ -5,20 +5,22 @@
 
 
 
-async::result<void> test_task(CoSched::TaskPtr task) {
-    for (unsigned x = 100; co_await task->yield(); x--) {
-        std::cout << task->get_name() << ": " << x << '\n';
-        if (x == 10) task->terminate();
+async::result<void> test_task() {
+    auto& task = CoSched::Task::get_current();
+    for (unsigned x = 100; co_await task.yield(); x--) {
+        std::cout << task.get_name() << ": " << x << '\n';
+        if (x == 10) task.terminate();
     }
 }
 
 int main () {
     CoSched::Scheduler scheduler;
     for (const auto& name : {"A", "B", "C", "D", "E", "F"}) {
-        auto task = scheduler.create_task(name);
-        async::detach(test_task(task));
-        if (task->get_name() == "B" || task->get_name() == "D") {
-            task->set_priority(CoSched::PRIO_HIGH);
+        scheduler.create_task(name);
+        async::detach(test_task());
+        auto& task = CoSched::Task::get_current();
+        if (task.get_name() == "B" || task.get_name() == "D") {
+            task.set_priority(CoSched::PRIO_HIGH);
         }
     }
     scheduler.run();