fix io_scheduler/task_container circular reference

cpp/mrc/include/mrc/coroutines/io_scheduler.hpp

@@ -35,8 +35,8 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
- 
- #pragma once
+
+#pragma once
 
 #include "mrc/coroutines/detail/poll_info.hpp"
 #include "mrc/coroutines/fd.hpp"
@@ -60,11 +60,15 @@
 #include <vector>
 
 namespace mrc::coroutines {
+
 class IoScheduler : public Scheduler
 {
+  private:
     using timed_events_t = detail::PollInfo::timed_events_t;
 
   public:
+    static std::shared_ptr<IoScheduler> get_instance();
+
     class schedule_operation;
     friend schedule_operation;
 
@@ -124,7 +128,7 @@ class IoScheduler : public Scheduler
     auto operator=(const IoScheduler&) -> IoScheduler& = delete;
     auto operator=(IoScheduler&&) -> IoScheduler&      = delete;
 
-    ~IoScheduler();
+    ~IoScheduler() override;
 
     /**
      * Given a ThreadStrategy::manual this function should be called at regular intervals to

cpp/mrc/include/mrc/coroutines/scheduler.hpp

@@ -18,6 +18,7 @@
 #pragma once
 
 #include "mrc/coroutines/task.hpp"
+#include "mrc/coroutines/time.hpp"
 
 #include <coroutine>
 #include <cstddef>
@@ -44,6 +45,10 @@ class Scheduler : public std::enable_shared_from_this<Scheduler>
      * @brief Suspends the current function and resumes it according to the scheduler's implementation.
      */
     [[nodiscard]] virtual Task<> yield() = 0;
+
+    [[nodiscard]] virtual Task<> yield_for(std::chrono::milliseconds amount) = 0;
+
+    [[nodiscard]] virtual Task<> yield_until(time_point_t time) = 0;
 };
 
 }  // namespace mrc::coroutines

cpp/mrc/src/public/coroutines/io_scheduler.cpp

@@ -58,13 +58,32 @@
 using namespace std::chrono_literals;
 
 namespace mrc::coroutines {
+
+std::shared_ptr<IoScheduler> IoScheduler::get_instance()
+{
+    static std::shared_ptr<IoScheduler> instance;
+    static std::mutex instance_mutex{};
+
+    if (instance == nullptr)
+    {
+        auto lock = std::lock_guard(instance_mutex);
+
+        if (instance == nullptr)
+        {
+            instance = std::make_shared<IoScheduler>();
+        }
+    }
+
+    return instance;
+}
+
 IoScheduler::IoScheduler(Options opts) :
   m_opts(std::move(opts)),
   m_epoll_fd(epoll_create1(EPOLL_CLOEXEC)),
   m_shutdown_fd(eventfd(0, EFD_CLOEXEC | EFD_NONBLOCK)),
   m_timer_fd(timerfd_create(CLOCK_MONOTONIC, TFD_NONBLOCK | TFD_CLOEXEC)),
   m_schedule_fd(eventfd(0, EFD_CLOEXEC | EFD_NONBLOCK)),
-  m_owned_tasks(new mrc::coroutines::TaskContainer(this->shared_from_this()))
+  m_owned_tasks(new mrc::coroutines::TaskContainer(std::shared_ptr<IoScheduler>(this, [](auto _) {})))
 {
     if (opts.execution_strategy == ExecutionStrategy::process_tasks_on_thread_pool)
     {

cpp/mrc/tests/CMakeLists.txt

@@ -17,6 +17,7 @@
 add_executable(test_mrc
   coroutines/test_async_generator.cpp
   coroutines/test_event.cpp
+  coroutines/test_io_scheduler.cpp
   coroutines/test_latch.cpp
   coroutines/test_ring_buffer.cpp
   coroutines/test_task_container.cpp

cpp/mrc/tests/coroutines/test_io_scheduler.cpp

@@ -0,0 +1,43 @@
+/**
+ * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "mrc/coroutines/async_generator.hpp"
+#include "mrc/coroutines/io_scheduler.hpp"
+#include "mrc/coroutines/sync_wait.hpp"
+#include "mrc/coroutines/task.hpp"
+
+#include <gtest/gtest.h>
+
+#include <coroutine>
+
+using namespace mrc;
+using namespace std::chrono_literals;
+
+class TestCoroIoScheduler : public ::testing::Test
+{};
+
+TEST_F(TestCoroIoScheduler, YieldFor)
+{
+    auto scheduler = coroutines::IoScheduler::get_instance();
+
+    auto task = [scheduler]() -> coroutines::Task<> {
+        // co_await scheduler->yield_for(1000ms);
+        co_return;
+    };
+
+    coroutines::sync_wait(task());
+}

python/mrc/_pymrc/include/pymrc/asyncio_scheduler.hpp

@@ -22,9 +22,11 @@
 #include "pymrc/utilities/object_wrappers.hpp"
 
 #include <boost/fiber/future/async.hpp>
+#include <mrc/coroutines/io_scheduler.hpp>
 #include <mrc/coroutines/scheduler.hpp>
 #include <mrc/coroutines/task.hpp>
 #include <mrc/coroutines/task_container.hpp>
+#include <mrc/coroutines/time.hpp>
 #include <pybind11/pybind11.h>
 #include <pybind11/pytypes.h>
 
@@ -90,6 +92,18 @@ class AsyncioScheduler : public mrc::coroutines::Scheduler
         co_await ContinueOnLoopOperation(m_loop);
     }
 
+    [[nodiscard]] coroutines::Task<> yield_for(std::chrono::milliseconds amount) override
+    {
+        co_await coroutines::IoScheduler::get_instance()->yield_for(amount);
+        co_await ContinueOnLoopOperation(m_loop);
+    };
+
+    [[nodiscard]] coroutines::Task<> yield_until(mrc::coroutines::time_point_t time) override
+    {
+        co_await coroutines::IoScheduler::get_instance()->yield_until(time);
+        co_await ContinueOnLoopOperation(m_loop);
+    };
+
   private:
     mrc::pymrc::PyHolder m_loop;
 };