feat: support individual async handler for each op (#3690)

3 files changed, 19 insertions, 16 deletions

diff --git a/core/examples/http_bench.js b/core/examples/http_bench.js
index db57401cc..dc204dfca 100644
--- a/core/examples/http_bench.js
+++ b/core/examples/http_bench.js
@@ -123,8 +123,10 @@ async function serve(rid) {
 let ops;
 
 async function main() {
-  Deno.core.setAsyncHandler(handleAsyncMsgFromRust);
   ops = Deno.core.ops();
+  for (const opName in ops) {
+    Deno.core.setAsyncHandler(ops[opName], handleAsyncMsgFromRust);
+  }
 
   Deno.core.print("http_bench.js start\n");
 
diff --git a/core/isolate.rs b/core/isolate.rs
index 50fcadb48..9c13f0e4d 100644
--- a/core/isolate.rs
+++ b/core/isolate.rs
@@ -900,7 +900,7 @@ pub mod tests {
         "setup2.js",
         r#"
          let nrecv = 0;
-         Deno.core.setAsyncHandler((opId, buf) => {
+         Deno.core.setAsyncHandler(1, (buf) => {
            nrecv++;
          });
          "#,
@@ -1021,7 +1021,7 @@ pub mod tests {
       "overflow_req_sync.js",
       r#"
         let asyncRecv = 0;
-        Deno.core.setAsyncHandler((opId, buf) => { asyncRecv++ });
+        Deno.core.setAsyncHandler(1, (buf) => { asyncRecv++ });
         // Large message that will overflow the shared space.
         let control = new Uint8Array(100 * 1024 * 1024);
         let response = Deno.core.dispatch(1, control);
@@ -1043,7 +1043,7 @@ pub mod tests {
       "overflow_res_sync.js",
       r#"
         let asyncRecv = 0;
-        Deno.core.setAsyncHandler((opId, buf) => { asyncRecv++ });
+        Deno.core.setAsyncHandler(1, (buf) => { asyncRecv++ });
         // Large message that will overflow the shared space.
         let control = new Uint8Array([42]);
         let response = Deno.core.dispatch(1, control);
@@ -1064,8 +1064,7 @@ pub mod tests {
         "overflow_req_async.js",
         r#"
          let asyncRecv = 0;
-         Deno.core.setAsyncHandler((opId, buf) => {
-           assert(opId == 1);
+         Deno.core.setAsyncHandler(1, (buf) => {
            assert(buf.byteLength === 4);
            assert(buf[0] === 43);
            asyncRecv++;
@@ -1097,8 +1096,7 @@ pub mod tests {
         "overflow_res_async.js",
         r#"
          let asyncRecv = 0;
-         Deno.core.setAsyncHandler((opId, buf) => {
-           assert(opId == 1);
+         Deno.core.setAsyncHandler(1, (buf) => {
            assert(buf.byteLength === 100 * 1024 * 1024);
            assert(buf[0] === 4);
            asyncRecv++;
@@ -1126,8 +1124,7 @@ pub mod tests {
         "overflow_res_multiple_dispatch_async.js",
         r#"
          let asyncRecv = 0;
-         Deno.core.setAsyncHandler((opId, buf) => {
-           assert(opId === 1);
+         Deno.core.setAsyncHandler(1, (buf) => {
            assert(buf.byteLength === 100 * 1024 * 1024);
            assert(buf[0] === 4);
            asyncRecv++;
diff --git a/core/shared_queue.js b/core/shared_queue.js
index 430261511..093cc223f 100644
--- a/core/shared_queue.js
+++ b/core/shared_queue.js
@@ -38,6 +38,9 @@ SharedQueue Binary Layout
 
   let sharedBytes;
   let shared32;
+
+  let asyncHandlers;
+
   let initialized = false;
 
   function maybeInit() {
@@ -54,6 +57,7 @@ SharedQueue Binary Layout
     assert(shared32 == null);
     sharedBytes = new Uint8Array(shared);
     shared32 = new Int32Array(shared);
+    asyncHandlers = [];
     // Callers should not call Deno.core.recv, use setAsyncHandler.
     Deno.core.recv(handleAsyncMsgFromRust);
   }
@@ -157,24 +161,24 @@ SharedQueue Binary Layout
     return [opId, buf];
   }
 
-  let asyncHandler;
-  function setAsyncHandler(cb) {
+  function setAsyncHandler(opId, cb) {
     maybeInit();
-    assert(asyncHandler == null);
-    asyncHandler = cb;
+    assert(opId != null);
+    asyncHandlers[opId] = cb;
   }
 
   function handleAsyncMsgFromRust(opId, buf) {
     if (buf) {
       // This is the overflow_response case of deno::Isolate::poll().
-      asyncHandler(opId, buf);
+      asyncHandlers[opId](buf);
     } else {
       while (true) {
         const opIdBuf = shift();
         if (opIdBuf == null) {
           break;
         }
-        asyncHandler(...opIdBuf);
+        assert(asyncHandlers[opIdBuf[0]] != null);
+        asyncHandlers[opIdBuf[0]](opIdBuf[1]);
       }
     }
   }