feat(core): Ops can take several zero copy buffers (#4788)

1 files changed, 60 insertions, 5 deletions

diff --git a/core/core_isolate.rs b/core/core_isolate.rs
index dff887ab3..5ddc6571b 100644
--- a/core/core_isolate.rs
+++ b/core/core_isolate.rs
@@ -366,7 +366,7 @@ impl CoreIsolate {
   /// Requires runtime to explicitly ask for op ids before using any of the ops.
   pub fn register_op<F>(&mut self, name: &str, op: F) -> OpId
   where
-    F: Fn(&mut CoreIsolateState, &[u8], Option<ZeroCopyBuf>) -> Op + 'static,
+    F: Fn(&mut CoreIsolateState, &[u8], &mut [ZeroCopyBuf]) -> Op + 'static,
   {
     let state_rc = Self::state(self);
     let mut state = state_rc.borrow_mut();
@@ -484,7 +484,7 @@ impl CoreIsolateState {
   /// Requires runtime to explicitly ask for op ids before using any of the ops.
   pub fn register_op<F>(&mut self, name: &str, op: F) -> OpId
   where
-    F: Fn(&mut CoreIsolateState, &[u8], Option<ZeroCopyBuf>) -> Op + 'static,
+    F: Fn(&mut CoreIsolateState, &[u8], &mut [ZeroCopyBuf]) -> Op + 'static,
   {
     self.op_registry.register(name, op)
   }
@@ -504,10 +504,10 @@ impl CoreIsolateState {
     scope: &mut impl v8::ToLocal<'s>,
     op_id: OpId,
     control_buf: &[u8],
-    zero_copy_buf: Option<ZeroCopyBuf>,
+    zero_copy_bufs: &mut [ZeroCopyBuf],
   ) -> Option<(OpId, Box<[u8]>)> {
     let op = if let Some(dispatcher) = self.op_registry.get(op_id) {
-      dispatcher(self, control_buf, zero_copy_buf)
+      dispatcher(self, control_buf, zero_copy_bufs)
     } else {
       let message =
         v8::String::new(scope, &format!("Unknown op id: {}", op_id)).unwrap();
@@ -718,6 +718,7 @@ pub mod tests {
   pub enum Mode {
     Async,
     AsyncUnref,
+    AsyncZeroCopy(u8),
     OverflowReqSync,
     OverflowResSync,
     OverflowReqAsync,
@@ -732,7 +733,7 @@ pub mod tests {
 
     let dispatcher = move |_state: &mut CoreIsolateState,
                            control: &[u8],
-                           _zero_copy: Option<ZeroCopyBuf>|
+                           zero_copy: &mut [ZeroCopyBuf]|
           -> Op {
       dispatch_count_.fetch_add(1, Ordering::Relaxed);
       match mode {
@@ -752,6 +753,18 @@ pub mod tests {
           };
           Op::AsyncUnref(fut.boxed())
         }
+        Mode::AsyncZeroCopy(count) => {
+          assert_eq!(control.len(), 1);
+          assert_eq!(control[0], 24);
+          assert_eq!(zero_copy.len(), count as usize);
+          zero_copy.iter().enumerate().for_each(|(idx, buf)| {
+            assert_eq!(buf.len(), 1);
+            assert_eq!(idx, buf[0] as usize);
+          });
+
+          let buf = vec![43u8].into_boxed_slice();
+          Op::Async(futures::future::ready(buf).boxed())
+        }
         Mode::OverflowReqSync => {
           assert_eq!(control.len(), 100 * 1024 * 1024);
           let buf = vec![43u8].into_boxed_slice();
@@ -817,6 +830,48 @@ pub mod tests {
   }
 
   #[test]
+  fn test_dispatch_no_zero_copy_buf() {
+    let (mut isolate, dispatch_count) = setup(Mode::AsyncZeroCopy(0));
+    js_check(isolate.execute(
+      "filename.js",
+      r#"
+        let control = new Uint8Array([24]);
+        Deno.core.send(1, control);
+        "#,
+    ));
+    assert_eq!(dispatch_count.load(Ordering::Relaxed), 1);
+  }
+
+  #[test]
+  fn test_dispatch_one_zero_copy_buf() {
+    let (mut isolate, dispatch_count) = setup(Mode::AsyncZeroCopy(1));
+    js_check(isolate.execute(
+      "filename.js",
+      r#"
+        let control = new Uint8Array([24]);
+        let zero_copy = new Uint8Array([0]);
+        Deno.core.send(1, control, zero_copy);
+        "#,
+    ));
+    assert_eq!(dispatch_count.load(Ordering::Relaxed), 1);
+  }
+
+  #[test]
+  fn test_dispatch_two_zero_copy_bufs() {
+    let (mut isolate, dispatch_count) = setup(Mode::AsyncZeroCopy(2));
+    js_check(isolate.execute(
+      "filename.js",
+      r#"
+        let control = new Uint8Array([24]);
+        let zero_copy_a = new Uint8Array([0]);
+        let zero_copy_b = new Uint8Array([1]);
+        Deno.core.send(1, control, zero_copy_a, zero_copy_b);
+        "#,
+    ));
+    assert_eq!(dispatch_count.load(Ordering::Relaxed), 1);
+  }
+
+  #[test]
   fn test_poll_async_delayed_ops() {
     run_in_task(|cx| {
       let (mut isolate, dispatch_count) = setup(Mode::Async);