From 32cbe2f9e7e27544499d3b50625492f5a62a00b0 Mon Sep 17 00:00:00 2001
From: Yuta Saito <kateinoigakukun@gmail.com>
Date: Wed, 17 May 2023 13:44:20 +0000
Subject: [PATCH] Prohibit await in non-async context and remove unnecessary
 loop fiber

Some part of this change are derived from
https://github.com/ruby/ruby.wasm/pull/205 by @tmtm

Co-authored-by: TOMITA Masahiro <tommy@tmtm.org>
---
 ext/js/lib/js.rb                              | 36 +++++--------------
 .../ruby-wasm-wasi/test/eval_async.test.ts    |  7 ++--
 .../ruby-wasm-wasi/test/unit/test_async.rb    | 10 ++++++
 3 files changed, 21 insertions(+), 32 deletions(-)

diff --git a/ext/js/lib/js.rb b/ext/js/lib/js.rb
index 9249860e1..06a31619c 100644
--- a/ext/js/lib/js.rb
+++ b/ext/js/lib/js.rb
@@ -39,42 +39,22 @@ module JS
   Null = JS.eval("return null")
 
   class PromiseScheduler
-    Task = Struct.new(:fiber, :status, :value)
-
-    def initialize(main_fiber)
-      @tasks = []
-      @is_spinning = false
-      @loop_fiber =
-        Fiber.new do
-          loop do
-            while task = @tasks.shift
-              task.fiber.transfer(task.value, task.status)
-            end
-            @is_spinning = false
-            main_fiber.transfer
-          end
-        end
+    def initialize(loop)
+      @loop = loop
     end
 
     def await(promise)
       current = Fiber.current
       promise.call(
         :then,
-        ->(value) { enqueue Task.new(current, :success, value) },
-        ->(value) { enqueue Task.new(current, :failure, value) }
+        ->(value) { current.transfer(value, :success) },
+        ->(value) { current.transfer(value, :failure) }
       )
-      value, status = @loop_fiber.transfer
+      raise "JS::Object#await can be called only from evalAsync" if @loop == current
+      value, status = @loop.transfer
       raise JS::Error.new(value) if status == :failure
       value
     end
-
-    def enqueue(task)
-      @tasks << task
-      unless @is_spinning
-        @is_spinning = true
-        JS.global.queueMicrotask -> { @loop_fiber.transfer }
-      end
-    end
   end
 
   @promise_scheduler = PromiseScheduler.new Fiber.current
@@ -120,8 +100,8 @@ def respond_to_missing?(sym, include_private)
   # This method looks like a synchronous method, but it actually runs asynchronously using fibers.
   # In other words, the next line to the `await` call at Ruby source will be executed after the
   # promise will be resolved. However, it does not block JavaScript event loop, so the next line
-  # to the `RubyVM.eval` or `RubyVM.evalAsync` (in the case when no `await` operator before the
-  # call expression) at JavaScript source will be executed without waiting for the promise.
+  # to the RubyVM.evalAsync` (in the case when no `await` operator before the call expression)
+  # at JavaScript source will be executed without waiting for the promise.
   #
   # The below example shows how the execution order goes. It goes in the order of "step N"
   #
diff --git a/packages/npm-packages/ruby-wasm-wasi/test/eval_async.test.ts b/packages/npm-packages/ruby-wasm-wasi/test/eval_async.test.ts
index 14f4ef750..71e3dbc01 100644
--- a/packages/npm-packages/ruby-wasm-wasi/test/eval_async.test.ts
+++ b/packages/npm-packages/ruby-wasm-wasi/test/eval_async.test.ts
@@ -32,9 +32,8 @@ describe("Async Ruby code evaluation", () => {
 
   test("await outside of evalAsync", async () => {
     const vm = await initRubyVM();
-    const result = vm.eval(
-      `require "js"; JS.global[:Promise].resolve(42).await`
-    );
-    expect(result.call("nil?").toString()).toBe("true");
+    expect(() => {
+      vm.eval(`require "js"; JS.global[:Promise].resolve(42).await`);
+    }).toThrowError("JS::Object#await can be called only from evalAsync");
   });
 });
diff --git a/packages/npm-packages/ruby-wasm-wasi/test/unit/test_async.rb b/packages/npm-packages/ruby-wasm-wasi/test/unit/test_async.rb
index 73a65a5d3..4ae15d1cf 100644
--- a/packages/npm-packages/ruby-wasm-wasi/test/unit/test_async.rb
+++ b/packages/npm-packages/ruby-wasm-wasi/test/unit/test_async.rb
@@ -42,4 +42,14 @@ def test_concurrent_promises
     assert_equal 43, pac1[:promise].await.to_i
     assert_equal 42, pac0[:promise].await.to_i
   end
+
+  def test_await_in_fiber
+    fiber_ended = false
+    Fiber.new do
+      promise = JS.eval("return Promise.resolve(42)")
+      assert_equal 42, promise.await.to_i
+      fiber_ended = true
+    end.resume
+    assert_equal true, fiber_ended
+  end
 end