Refactor dispatch handling (#2452)

Promise id is now created in core and passed back to JS.

2 files changed, 59 insertions, 52 deletions

diff --git a/js/dispatch.ts b/js/dispatch.ts
index 0c6e70709..36f97363f 100644
--- a/js/dispatch.ts
+++ b/js/dispatch.ts
@@ -5,36 +5,30 @@ import * as msg from "gen/cli/msg_generated";
 import * as errors from "./errors";
 import * as util from "./util";
 import {
-  nextPromiseId,
   recordFromBufMinimal,
   handleAsyncMsgFromRustMinimal
 } from "./dispatch_minimal";
 
 const promiseTable = new Map<number, util.Resolvable<msg.Base>>();
 
-interface FlatbufferRecord {
-  promiseId: number;
-  base: msg.Base;
-}
-
-function flatbufferRecordFromBuf(buf: Uint8Array): FlatbufferRecord {
+function flatbufferRecordFromBuf(buf: Uint8Array): msg.Base {
   const bb = new flatbuffers.ByteBuffer(buf);
   const base = msg.Base.getRootAsBase(bb);
-  return {
-    promiseId: base.cmdId(),
-    base
-  };
+  return base;
 }
 
-export function handleAsyncMsgFromRust(ui8: Uint8Array): void {
+export function handleAsyncMsgFromRust(
+  promiseId: number,
+  ui8: Uint8Array
+): void {
   const buf32 = new Int32Array(ui8.buffer, ui8.byteOffset, ui8.byteLength / 4);
   const recordMin = recordFromBufMinimal(buf32);
   if (recordMin) {
     // Fast and new
-    handleAsyncMsgFromRustMinimal(ui8, recordMin);
+    handleAsyncMsgFromRustMinimal(promiseId, ui8, recordMin);
   } else {
     // Legacy
-    let { promiseId, base } = flatbufferRecordFromBuf(ui8);
+    let base = flatbufferRecordFromBuf(ui8);
     const promise = promiseTable.get(promiseId);
     util.assert(promise != null, `Expecting promise in table. ${promiseId}`);
     promiseTable.delete(promiseId);
@@ -56,14 +50,26 @@ function sendInternal(
   innerType: msg.Any,
   inner: flatbuffers.Offset,
   zeroCopy: undefined | ArrayBufferView,
-  sync = true
-): [number, null | Uint8Array] {
-  const cmdId = nextPromiseId();
+  isSync: true
+): Uint8Array | null;
+function sendInternal(
+  builder: flatbuffers.Builder,
+  innerType: msg.Any,
+  inner: flatbuffers.Offset,
+  zeroCopy: undefined | ArrayBufferView,
+  isSync: false
+): Promise<msg.Base>;
+function sendInternal(
+  builder: flatbuffers.Builder,
+  innerType: msg.Any,
+  inner: flatbuffers.Offset,
+  zeroCopy: undefined | ArrayBufferView,
+  isSync: boolean
+): Promise<msg.Base> | Uint8Array | null {
   msg.Base.startBase(builder);
+  msg.Base.addSync(builder, isSync);
   msg.Base.addInner(builder, inner);
   msg.Base.addInnerType(builder, innerType);
-  msg.Base.addSync(builder, sync);
-  msg.Base.addCmdId(builder, cmdId);
   builder.finish(msg.Base.endBase(builder));
 
   const control = builder.asUint8Array();
@@ -74,7 +80,25 @@ function sendInternal(
   );
 
   builder.inUse = false;
-  return [cmdId, response];
+
+  if (typeof response === "number") {
+    const promise = util.createResolvable<msg.Base>();
+    promiseTable.set(response, promise);
+    util.assert(!isSync);
+    return promise;
+  } else {
+    if (!isSync) {
+      util.assert(response !== null);
+      const base = flatbufferRecordFromBuf(response as Uint8Array);
+      const err = errors.maybeError(base);
+      if (err != null) {
+        return Promise.reject(err);
+      } else {
+        return Promise.resolve(base);
+      }
+    }
+    return response;
+  }
 }
 
 // @internal
@@ -84,16 +108,7 @@ export function sendAsync(
   inner: flatbuffers.Offset,
   data?: ArrayBufferView
 ): Promise<msg.Base> {
-  const [cmdId, response] = sendInternal(
-    builder,
-    innerType,
-    inner,
-    data,
-    false
-  );
-  util.assert(response == null); // null indicates async.
-  const promise = util.createResolvable<msg.Base>();
-  promiseTable.set(cmdId, promise);
+  const promise = sendInternal(builder, innerType, inner, data, false);
   return promise;
 }
 
@@ -104,10 +119,8 @@ export function sendSync(
   inner: flatbuffers.Offset,
   data?: ArrayBufferView
 ): null | msg.Base {
-  const [cmdId, response] = sendInternal(builder, innerType, inner, data, true);
-  util.assert(cmdId >= 0);
-  util.assert(response != null); // null indicates async.
-  if (response!.length === 0) {
+  const response = sendInternal(builder, innerType, inner, data, true);
+  if (response == null || response.length === 0) {
     return null;
   } else {
     const bb = new flatbuffers.ByteBuffer(response!);
diff --git a/js/dispatch_minimal.ts b/js/dispatch_minimal.ts
index 17d328110..bf9065f56 100644
--- a/js/dispatch_minimal.ts
+++ b/js/dispatch_minimal.ts
@@ -5,14 +5,8 @@ import { core } from "./core";
 
 const DISPATCH_MINIMAL_TOKEN = 0xcafe;
 const promiseTableMin = new Map<number, util.Resolvable<number>>();
-let _nextPromiseId = 0;
-
-export function nextPromiseId(): number {
-  return _nextPromiseId++;
-}
 
 export interface RecordMinimal {
-  promiseId: number;
   opId: number;
   arg: number;
   result: number;
@@ -28,10 +22,9 @@ export function hasMinimalToken(i32: Int32Array): boolean {
 export function recordFromBufMinimal(buf32: Int32Array): null | RecordMinimal {
   if (hasMinimalToken(buf32)) {
     return {
-      promiseId: buf32[1],
-      opId: buf32[2],
-      arg: buf32[3],
-      result: buf32[4]
+      opId: buf32[1],
+      arg: buf32[2],
+      result: buf32[3]
     };
   }
   return null;
@@ -46,12 +39,13 @@ const scratchBytes = new Uint8Array(
 util.assert(scratchBytes.byteLength === scratch32.length * 4);
 
 export function handleAsyncMsgFromRustMinimal(
+  promiseId: number,
   ui8: Uint8Array,
   record: RecordMinimal
 ): void {
   // Fast and new
   util.log("minimal handleAsyncMsgFromRust ", ui8.length);
-  const { promiseId, result } = record;
+  const { result } = record;
   const promise = promiseTableMin.get(promiseId);
   promiseTableMin.delete(promiseId);
   promise!.resolve(result);
@@ -62,16 +56,16 @@ export function sendAsyncMinimal(
   arg: number,
   zeroCopy: Uint8Array
 ): Promise<number> {
-  const promiseId = nextPromiseId(); // AKA cmdId
-
   scratch32[0] = DISPATCH_MINIMAL_TOKEN;
-  scratch32[1] = promiseId;
-  scratch32[2] = opId;
-  scratch32[3] = arg;
+  scratch32[1] = opId;
+  scratch32[2] = arg;
+
+  const promiseId = core.dispatch(scratchBytes, zeroCopy);
+
+  util.assert(typeof promiseId == "number");
 
   const promise = util.createResolvable<number>();
-  promiseTableMin.set(promiseId, promise);
+  promiseTableMin.set(promiseId as number, promise);
 
-  core.dispatch(scratchBytes, zeroCopy);
   return promise;
 }