feat: support task deadlines

bytecodealliance · Jun 5, 2024 · 3fb8933 · 3fb8933
1 parent aa9ff14
commit 3fb8933
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 9 deletions.
diff --git a/builtins/web/fetch/fetch-api.cpp b/builtins/web/fetch/fetch-api.cpp
@@ -20,7 +20,7 @@ class ResponseFutureTask final : public api::AsyncTask {
     handle_ = res.unwrap();
   }
 
-  [[nodiscard]] bool run(api::Engine *engine) override {
+  [[nodiscard]] bool run(api::Engine *engine, bool deadline) override {
     // MOZ_ASSERT(ready());
     JSContext *cx = engine->cx();
 

diff --git a/builtins/web/fetch/request-response.cpp b/builtins/web/fetch/request-response.cpp
@@ -70,7 +70,7 @@ class BodyFutureTask final : public api::AsyncTask {
     handle_ = res.unwrap();
   }
 
-  [[nodiscard]] bool run(api::Engine *engine) override {
+  [[nodiscard]] bool run(api::Engine *engine, bool deadline) override {
     // MOZ_ASSERT(ready());
     JSContext *cx = engine->cx();
     RootedObject owner(cx, streams::NativeStreamSource::owner(body_source_));
@@ -142,7 +142,7 @@ class ResponseFutureTask final : public api::AsyncTask {
 
   ResponseFutureTask(const RootedObject &rooted);
 
-  [[nodiscard]] bool run(api::Engine *engine) override {
+  [[nodiscard]] bool run(api::Engine *engine, bool deadline) override {
     // MOZ_ASSERT(ready());
     JSContext *cx = engine->cx();
 

diff --git a/builtins/web/timers.cpp b/builtins/web/timers.cpp
@@ -34,9 +34,12 @@ class TimerTask final : public api::AsyncTask {
     }
 
     handle_ = host_api::MonotonicClock::subscribe(deadline_, true);
+    if (handle_ == INVALID_POLLABLE_HANDLE) {
+      deadline = deadline_;
+    }
   }
 
-  [[nodiscard]] bool run(api::Engine *engine) override {
+  [[nodiscard]] bool run(api::Engine *engine, bool deadline) override {
     JSContext *cx = engine->cx();
 
     const RootedObject callback(cx, callback_);

diff --git a/host-apis/wasi-0.2.0/host_api.cpp b/host-apis/wasi-0.2.0/host_api.cpp
@@ -477,7 +477,7 @@ class BodyAppendTask final : public api::AsyncTask {
     state_ = State::BlockedOnBoth;
   }
 
-  [[nodiscard]] bool run(api::Engine *engine) override {
+  [[nodiscard]] bool run(api::Engine *engine, bool deadline) override {
     // If run is called while we're blocked on the incoming stream, that means that stream's
     // pollable has resolved, so the stream must be ready.
     if (state_ == State::BlockedOnBoth || state_ == State::BlockedOnIncoming) {

diff --git a/include/extension-api.h b/include/extension-api.h
@@ -117,9 +117,10 @@ class AsyncTask {
   PollableHandle handle_ = -1;
 
 public:
+  uint64_t deadline = 0;
   virtual ~AsyncTask() = default;
 
-  virtual bool run(Engine *engine) = 0;
+  virtual bool run(Engine *engine, bool deadline) = 0;
   virtual bool cancel(Engine *engine) = 0;
 
   [[nodiscard]] virtual PollableHandle id() {

diff --git a/runtime/event_loop.cpp b/runtime/event_loop.cpp
@@ -67,14 +67,37 @@ bool EventLoop::run_event_loop(api::Engine *engine, double total_compute) {
       exit_event_loop();
       return false;
     }
+
+    const auto tasks = &queue.get().tasks;
+    size_t tasks_size = tasks->size();
+
+    // If there is any task with a deadline, let the deadline run it
+    uint64_t now = 0;
+    for (size_t task_idx = 0; task_idx < tasks_size; ++task_idx) {
+      const auto task = tasks->at(task_idx);
+      if (task->deadline > 0) {
+        if (now == 0) {
+          now = host_api::MonotonicClock::now();
+        }
+        if (task->deadline <= now) {
+          fprintf(stderr, "DEADLINE RUN");
+          bool success = task->run(engine, true);
+          tasks->erase(tasks->begin() + task_idx);
+          if (!success) {
+            exit_event_loop();
+            return false;
+          }
+          continue;
+        }
+      }
+    }
+
     // if there is no interest in the event loop at all, just run one tick
     if (interest_complete()) {
       exit_event_loop();
       return true;
     }
 
-    const auto tasks = &queue.get().tasks;
-    size_t tasks_size = tasks->size();
     if (tasks_size == 0) {
       exit_event_loop();
       MOZ_ASSERT(!interest_complete());
@@ -87,7 +110,7 @@ bool EventLoop::run_event_loop(api::Engine *engine, double total_compute) {
     size_t task_idx = api::AsyncTask::select(tasks);
 
     auto task = tasks->at(task_idx);
-    bool success = task->run(engine);
+    bool success = task->run(engine, false);
     tasks->erase(tasks->begin() + task_idx);
     if (!success) {
       exit_event_loop();