reland JsRuntime/Worker is not a Future (#7924)

1 files changed, 55 insertions, 46 deletions

diff --git a/core/runtime.rs b/core/runtime.rs
index 193e33420..26b62cbe3 100644
--- a/core/runtime.rs
+++ b/core/runtime.rs
@@ -23,6 +23,7 @@ use crate::shared_queue::SharedQueue;
 use crate::shared_queue::RECOMMENDED_SIZE;
 use crate::BufVec;
 use crate::OpState;
+use futures::future::poll_fn;
 use futures::stream::FuturesUnordered;
 use futures::stream::StreamExt;
 use futures::stream::StreamFuture;
@@ -442,28 +443,24 @@ impl JsRuntime {
         .remove_near_heap_limit_callback(cb, heap_limit);
     }
   }
-}
-
-extern "C" fn near_heap_limit_callback<F>(
-  data: *mut c_void,
-  current_heap_limit: usize,
-  initial_heap_limit: usize,
-) -> usize
-where
-  F: FnMut(usize, usize) -> usize,
-{
-  let callback = unsafe { &mut *(data as *mut F) };
-  callback(current_heap_limit, initial_heap_limit)
-}
 
-impl Future for JsRuntime {
-  type Output = Result<(), AnyError>;
+  /// Runs event loop to completion
+  ///
+  /// This future resolves when:
+  ///  - there are no more pending dynamic imports
+  ///  - there are no more pending ops
+  pub async fn run_event_loop(&mut self) -> Result<(), AnyError> {
+    poll_fn(|cx| self.poll_event_loop(cx)).await
+  }
 
-  fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
-    let runtime = self.get_mut();
-    runtime.shared_init();
+  /// Runs a single tick of event loop
+  pub fn poll_event_loop(
+    &mut self,
+    cx: &mut Context,
+  ) -> Poll<Result<(), AnyError>> {
+    self.shared_init();
 
-    let state_rc = Self::state(runtime.v8_isolate());
+    let state_rc = Self::state(self.v8_isolate());
     {
       let state = state_rc.borrow();
       state.waker.register(cx.waker());
@@ -471,21 +468,21 @@ impl Future for JsRuntime {
 
     // Dynamic module loading - ie. modules loaded using "import()"
     {
-      let poll_imports = runtime.prepare_dyn_imports(cx)?;
+      let poll_imports = self.prepare_dyn_imports(cx)?;
       assert!(poll_imports.is_ready());
 
-      let poll_imports = runtime.poll_dyn_imports(cx)?;
+      let poll_imports = self.poll_dyn_imports(cx)?;
       assert!(poll_imports.is_ready());
 
-      runtime.check_promise_exceptions()?;
+      self.check_promise_exceptions()?;
     }
 
     // Ops
     {
-      let overflow_response = runtime.poll_pending_ops(cx);
-      runtime.async_op_response(overflow_response)?;
-      runtime.drain_macrotasks()?;
-      runtime.check_promise_exceptions()?;
+      let overflow_response = self.poll_pending_ops(cx);
+      self.async_op_response(overflow_response)?;
+      self.drain_macrotasks()?;
+      self.check_promise_exceptions()?;
     }
 
     let state = state_rc.borrow();
@@ -509,6 +506,18 @@ impl Future for JsRuntime {
   }
 }
 
+extern "C" fn near_heap_limit_callback<F>(
+  data: *mut c_void,
+  current_heap_limit: usize,
+  initial_heap_limit: usize,
+) -> usize
+where
+  F: FnMut(usize, usize) -> usize,
+{
+  let callback = unsafe { &mut *(data as *mut F) };
+  callback(current_heap_limit, initial_heap_limit)
+}
+
 impl JsRuntimeState {
   // Called by V8 during `Isolate::mod_instantiate`.
   pub fn module_resolve_cb(
@@ -1215,13 +1224,13 @@ pub mod tests {
     futures::executor::block_on(lazy(move |cx| f(cx)));
   }
 
-  fn poll_until_ready<F>(future: &mut F, max_poll_count: usize) -> F::Output
-  where
-    F: Future + Unpin,
-  {
+  fn poll_until_ready(
+    runtime: &mut JsRuntime,
+    max_poll_count: usize,
+  ) -> Result<(), AnyError> {
     let mut cx = Context::from_waker(futures::task::noop_waker_ref());
     for _ in 0..max_poll_count {
-      match future.poll_unpin(&mut cx) {
+      match runtime.poll_event_loop(&mut cx) {
         Poll::Pending => continue,
         Poll::Ready(val) => return val,
       }
@@ -1437,7 +1446,7 @@ pub mod tests {
         )
         .unwrap();
       assert_eq!(dispatch_count.load(Ordering::Relaxed), 1);
-      assert!(matches!(runtime.poll_unpin(cx), Poll::Ready(Ok(_))));
+      assert!(matches!(runtime.poll_event_loop(cx), Poll::Ready(Ok(_))));
       assert_eq!(dispatch_count.load(Ordering::Relaxed), 1);
       runtime
         .execute(
@@ -1450,11 +1459,11 @@ pub mod tests {
         )
         .unwrap();
       assert_eq!(dispatch_count.load(Ordering::Relaxed), 2);
-      assert!(matches!(runtime.poll_unpin(cx), Poll::Ready(Ok(_))));
+      assert!(matches!(runtime.poll_event_loop(cx), Poll::Ready(Ok(_))));
       runtime.execute("check3.js", "assert(nrecv == 2)").unwrap();
       assert_eq!(dispatch_count.load(Ordering::Relaxed), 2);
       // We are idle, so the next poll should be the last.
-      assert!(matches!(runtime.poll_unpin(cx), Poll::Ready(Ok(_))));
+      assert!(matches!(runtime.poll_event_loop(cx), Poll::Ready(Ok(_))));
     });
   }
 
@@ -1478,7 +1487,7 @@ pub mod tests {
       assert_eq!(dispatch_count.load(Ordering::Relaxed), 1);
       // The above op never finish, but runtime can finish
       // because the op is an unreffed async op.
-      assert!(matches!(runtime.poll_unpin(cx), Poll::Ready(Ok(_))));
+      assert!(matches!(runtime.poll_event_loop(cx), Poll::Ready(Ok(_))));
     })
   }
 
@@ -1608,7 +1617,7 @@ pub mod tests {
         )
         .unwrap();
       assert_eq!(dispatch_count.load(Ordering::Relaxed), 1);
-      assert!(matches!(runtime.poll_unpin(cx), Poll::Ready(Ok(_))));
+      assert!(matches!(runtime.poll_event_loop(cx), Poll::Ready(Ok(_))));
       runtime
         .execute("check.js", "assert(asyncRecv == 1);")
         .unwrap();
@@ -1700,7 +1709,7 @@ pub mod tests {
          "#,
         )
         .unwrap();
-      if let Poll::Ready(Err(_)) = runtime.poll_unpin(&mut cx) {
+      if let Poll::Ready(Err(_)) = runtime.poll_event_loop(&mut cx) {
         unreachable!();
       }
     });
@@ -1713,7 +1722,7 @@ pub mod tests {
       runtime
         .execute("core_test.js", include_str!("core_test.js"))
         .unwrap();
-      if let Poll::Ready(Err(_)) = runtime.poll_unpin(&mut cx) {
+      if let Poll::Ready(Err(_)) = runtime.poll_event_loop(&mut cx) {
         unreachable!();
       }
     });
@@ -1739,7 +1748,7 @@ pub mod tests {
           include_str!("encode_decode_test.js"),
         )
         .unwrap();
-      if let Poll::Ready(Err(_)) = runtime.poll_unpin(&mut cx) {
+      if let Poll::Ready(Err(_)) = runtime.poll_event_loop(&mut cx) {
         unreachable!();
       }
     });
@@ -2047,7 +2056,7 @@ pub mod tests {
 
       assert_eq!(count.load(Ordering::Relaxed), 0);
       // We should get an error here.
-      let result = runtime.poll_unpin(cx);
+      let result = runtime.poll_event_loop(cx);
       if let Poll::Ready(Ok(_)) = result {
         unreachable!();
       }
@@ -2140,14 +2149,14 @@ pub mod tests {
         .unwrap();
 
       // First poll runs `prepare_load` hook.
-      assert!(matches!(runtime.poll_unpin(cx), Poll::Pending));
+      assert!(matches!(runtime.poll_event_loop(cx), Poll::Pending));
       assert_eq!(prepare_load_count.load(Ordering::Relaxed), 1);
 
       // Second poll actually loads modules into the isolate.
-      assert!(matches!(runtime.poll_unpin(cx), Poll::Ready(Ok(_))));
+      assert!(matches!(runtime.poll_event_loop(cx), Poll::Ready(Ok(_))));
       assert_eq!(resolve_count.load(Ordering::Relaxed), 4);
       assert_eq!(load_count.load(Ordering::Relaxed), 2);
-      assert!(matches!(runtime.poll_unpin(cx), Poll::Ready(Ok(_))));
+      assert!(matches!(runtime.poll_event_loop(cx), Poll::Ready(Ok(_))));
       assert_eq!(resolve_count.load(Ordering::Relaxed), 4);
       assert_eq!(load_count.load(Ordering::Relaxed), 2);
     })
@@ -2179,10 +2188,10 @@ pub mod tests {
         )
         .unwrap();
       // First poll runs `prepare_load` hook.
-      let _ = runtime.poll_unpin(cx);
+      let _ = runtime.poll_event_loop(cx);
       assert_eq!(prepare_load_count.load(Ordering::Relaxed), 1);
       // Second poll triggers error
-      let _ = runtime.poll_unpin(cx);
+      let _ = runtime.poll_event_loop(cx);
     })
   }
 
@@ -2313,7 +2322,7 @@ main();
     at async error_async_stack.js:10:5
 "#;
 
-      match runtime.poll_unpin(cx) {
+      match runtime.poll_event_loop(cx) {
         Poll::Ready(Err(e)) => {
           assert_eq!(e.to_string(), expected_error);
         }