feat(kernel): process callbacks while EndRequest is awaiting promise

Author: mrgrain

CONTRIBUTING.md

@@ -1,15 +1,12 @@
 # Contributing to jsii
-
 Thanks for your interest in contributing to AWS JSII! :heart:
 
 This document describes how to set up a development environment and submit your
 contributions. Please read it carefully and let us know if it's not up-to date
 (or even better, submit a pull request with your corrections! :wink:).
 
 ## Pre-requisites
-
 ### Setup Docker image
-
 Due to the polyglot nature of `jsii`, the toolchain requirements are somewhat
 more complicated than for most projects. In order to locally develop `jsii`, you
 will need a number of tools.
@@ -22,38 +19,37 @@ our own CI/CD: the ["superchain" image][superchain] from.
 The image can be built for local usage, too:
 
 ```console
-IMAGE=superchain
-docker build -t ${IMAGE} ./superchain
+$ IMAGE=superchain
+$ docker build -t ${IMAGE} ./superchain
 ```
 
 In order to get an interactive shell within a Docker container using the
 *superchain* image you just built:
 
 ```console
-cd jsii # go to the root of the jsii repo
-docker run --rm --net=host -it -v $PWD:$PWD -w $PWD ${IMAGE}
+$ cd jsii # go to the root of the jsii repo
+$ docker run --rm --net=host -it -v $PWD:$PWD -w $PWD ${IMAGE}
 ```
 
 In the shell that pops up, the `npm run` commands in the following sections must
 be executed.
 
 ### Alternative: Manually install the toolchain
-
 The following tools need to be installed to develop on JSII locally. We recommend
 using the docker image from the above section, but if you wish to, you can install
 in your development environment.
 
 - [Node `14.6.0`] or later
 - [Yarn `1.19.1`] or later
 - An OpenJDK-8 distribution (e.g: [Oracle's OpenJDK8], [Amazon Corretto 8])
-  - [`maven >= 3.0.5`](https://maven.apache.org)
+  + [`maven >= 3.0.5`](https://maven.apache.org)
 - [.NET Core `3.1`] or later
-  - *Recommended:* [`mono >= 5`](https://www.mono-project.com)
+  + *Recommended:* [`mono >= 5`](https://www.mono-project.com)
 - [Python `3.7.3`] or later
-  - [`pip`](https://pip.pypa.io/en/stable/installing/)
-  - [`setuptools >= 38.6.0`](https://pypi.org/project/setuptools/)
-  - [`wheel`](https://pypi.org/project/wheel/)
-  - *Recommended:* [`twine`](https://pypi.org/project/twine/)
+  + [`pip`](https://pip.pypa.io/en/stable/installing/)
+  + [`setuptools >= 38.6.0`](https://pypi.org/project/setuptools/)
+  + [`wheel`](https://pypi.org/project/wheel/)
+  + *Recommended:* [`twine`](https://pypi.org/project/twine/)
 - [Go] `1.18` or newer
 
 [Node `14.6.0`]: https://nodejs.org/download/release/v14.6.0/
@@ -65,15 +61,14 @@ in your development environment.
 [Go]: https://go.dev/dl/
 
 ## Getting Started
-
 ### Bootstrapping
 
 The project is managed as a [monorepo] using [lerna].
 
 [monorepo]: https://github.com/babel/babel/blob/main/doc/design/monorepo.md
 [lerna]: https://github.com/lerna/lerna
 
-1. Check out this repository and change directory to its root.
+1. Check out this respository and change directory to its root.
 2. Run `yarn install && yarn build` to install lerna, bootstrap the repository
    and perform an initial build and test cycle.
 
@@ -145,10 +140,10 @@ The runtime client library should be implemented as a module under
 
 The jsii runtime client library usually includes the following components:
 
-- Child process manager: responsible to start/stop the __@jsii/runtime__ child
+- Child process manager: responsible to start/stop the **@jsii/runtime** child
   process.
 - Protocol layer: implements the STDIN/STDOUT protocol that interacts with the
-  __@jsii/runtime__.
+  **@jsii/runtime**.
 - Proxy layer: includes base classes and serialization utilities to implement
   the generated proxy classes.
 
@@ -168,7 +163,6 @@ The [Python](./packages/jsii-pacmak/lib/targets/python.ts) target is a good
 example to work from.
 
 ## Releasing
-
 ### The `jsii/superchain` Docker image
 
 Upon merging new changes to the `main` branch, the `jsii/superchain:nightly`

packages/@jsii/kernel/src/api.ts

@@ -73,6 +73,8 @@ export function isPropertyOverride(value: Override): value is PropertyOverride {
 
 export interface Callback {
   readonly cbid: string;
+  /** Whether this callback is synchronous. */
+  readonly sync: boolean;
   readonly cookie: string | undefined;
   readonly invoke?: InvokeRequest;
   readonly get?: GetRequest;

packages/@jsii/kernel/src/kernel.test.ts

@@ -2212,6 +2212,36 @@ defineTest('invokeBinScript() accepts arguments', (sandbox) => {
   });
 });
 
+// defineTest('ImplementationFromAsyncContext compliance', async (sandbox) => {
+//   const producer = sandbox.create({
+//     fqn: 'Object',
+//     overrides: [{ method: 'produce', cookie: 'produce1234' }],
+//     interfaces: ['jsii-calc.IPromiseProducer'],
+//   });
+
+//   const obj = sandbox.create({
+//     fqn: 'jsii-calc.ImplementationFromAsyncContext',
+//     args: [producer],
+//   });
+
+//   const promise1 = sandbox.begin({
+//     objref: obj,
+//     method: 'doAsyncWork',
+//   });
+
+//   const callbacks1 = sandbox.callbacks();
+//   expect(callbacks1.callbacks.length).toBe(1);
+//   expect(callbacks1.callbacks[0].cookie).toBe('produce1234');
+
+//   sandbox.complete({
+//     cbid: callbacks1.callbacks[0].cbid,
+//     result: 'test-string',
+//   });
+
+//   const result = (await sandbox.end(promise1)).result;
+//   expect(result).toBe('test-string');
+// });
+
 // =================================================================================================
 
 const testNames: { [name: string]: boolean } = {};

packages/@jsii/kernel/src/kernel.ts

@@ -447,7 +447,52 @@ export class Kernel {
 
     let result;
     try {
-      result = await promise;
+      let settled = false;
+      /**
+       * Poll for new callback requests until the promise is resolved. This is
+       * to allow any promises necessary for the promise to be able to settle.
+       * We use setImmediate so the next poll happens on the next run loop tick,
+       * after other microtasks might have been paused on a pending callback.
+       */
+      // eslint-disable-next-line no-inner-declarations
+      function pollForCallbacks(kernel: Kernel) {
+        // Promise has settled already, not going any further...
+        if (settled) {
+          return;
+        }
+
+        for (const [cbid, cb] of kernel.cbs.entries()) {
+          kernel.waiting.set(cbid, cb);
+          kernel.cbs.delete(cbid);
+          try {
+            cb.succeed(
+              kernel.callbackHandler({
+                cbid,
+                sync: false,
+                cookie: cb.override.cookie,
+                invoke: {
+                  objref: cb.objref,
+                  method: cb.override.method,
+                  args: cb.args,
+                },
+              }),
+            );
+          } catch (err) {
+            cb.fail(err);
+          } finally {
+            kernel.waiting.delete(cbid);
+          }
+        }
+        if (!settled) {
+          setImmediate(pollForCallbacks, kernel);
+        }
+      }
+      pollForCallbacks(this);
+
+      result = await promise.finally(() => {
+        settled = true;
+      });
+
       this._debug('promise result:', result);
     } catch (e: any) {
       this._debug('promise error:', e);
@@ -475,14 +520,16 @@ export class Kernel {
     };
   }
 
+  /** @deprecated the flow should be handled directly by "end" */
   public callbacks(_req?: api.CallbacksRequest): api.CallbacksResponse {
     this._debug('callbacks');
     const ret = Array.from(this.cbs.entries()).map(([cbid, cb]) => {
       this.waiting.set(cbid, cb); // move to waiting
       this.cbs.delete(cbid); // remove from created
       const callback: api.Callback = {
-        cbid,
         cookie: cb.override.cookie,
+        cbid,
+        sync: false,
         invoke: {
           objref: cb.objref,
           method: cb.override.method,
@@ -758,6 +805,7 @@ export class Kernel {
         const result = this.callbackHandler({
           cookie: override.cookie,
           cbid: this._makecbid(),
+          sync: true,
           get: { objref, property: propertyName },
         });
         this._debug('callback returned', result);
@@ -774,6 +822,7 @@ export class Kernel {
         this.callbackHandler({
           cookie: override.cookie,
           cbid: this._makecbid(),
+          sync: true,
           set: {
             objref,
             property: propertyName,
@@ -910,6 +959,7 @@ export class Kernel {
           const result = this.callbackHandler({
             cookie: override.cookie,
             cbid: this._makecbid(),
+            sync: true,
             invoke: {
               objref,
               method: methodName,

packages/@jsii/python-runtime/src/jsii/_kernel/__init__.py

@@ -1,6 +1,7 @@
 import datetime
 import inspect
 import itertools
+import time
 from types import FunctionType, MethodType, BuiltinFunctionType, LambdaType
 
 from typing import Callable, cast, Any, List, Optional, Sequence, Type
@@ -28,6 +29,7 @@
     CreateResponse,
     DeleteRequest,
     EndRequest,
+    EndResponse,
     EnumRef,
     GetRequest,
     GetResponse,
@@ -466,26 +468,11 @@ def ainvoke(self, obj: Any, method: str, args: Optional[List[Any]] = None) -> An
         if isinstance(promise, Callback):
             promise = _callback_till_result(self, promise, BeginResponse)
 
-        callbacks = self.provider.callbacks(CallbacksRequest()).callbacks
-        while callbacks:
-            for callback in callbacks:
-                try:
-                    result = _handle_callback(self, callback)
-                except Exception as exc:
-                    # TODO: Maybe we want to print the whole traceback here?
-                    complete = self.provider.complete(
-                        CompleteRequest(cbid=callback.cbid, err=str(exc))
-                    )
-                else:
-                    complete = self.provider.complete(
-                        CompleteRequest(cbid=callback.cbid, result=result)
-                    )
-
-                assert complete.cbid == callback.cbid
-
-            callbacks = self.provider.callbacks(CallbacksRequest()).callbacks
-
-        return self.provider.end(EndRequest(promiseid=promise.promiseid)).result
+        response = self.provider.end(EndRequest(promiseid=promise.promiseid))
+        if isinstance(response, Callback):
+            return _callback_till_result(self, response, EndResponse).result
+        else:
+            return response.result
 
     def stats(self):
         resp = self.provider.stats(StatsRequest())

packages/@jsii/python-runtime/src/jsii/_kernel/providers/process.py

@@ -291,7 +291,7 @@ def stop(self):
         assert self._process.stdin is not None
         if not self._process.stdin.closed:
             self._process.stdin.write(b'{"exit":0}\n')
-            # Close the process' STDIN, singaling we are done with it
+            # Close the process' STDIN, signaling we are done with it
             self._process.stdin.close()
 
         try:

packages/@jsii/python-runtime/src/jsii/_kernel/types.py

@@ -264,6 +264,7 @@ class StatsResponse:
     LoadResponse,
     CreateResponse,
     DeleteResponse,
+    EndResponse,
     GetResponse,
     InvokeResponse,
     InvokeScriptResponse,

packages/@jsii/runtime/lib/host.ts

@@ -82,7 +82,7 @@ export class KernelHost {
       return this.processRequest(
         req,
         completeCallback.bind(this),
-        /* sync */ true,
+        callback.sync,
       );
     }
   }
@@ -126,7 +126,7 @@ export class KernelHost {
       // promises. see the kernel test 'async overrides: two overrides'
       // for an example for this use case.
       if (apiReq.api === 'begin' || apiReq.api === 'complete') {
-        checkIfAsyncIsAllowed();
+        assertAsyncIsAllowed();
 
         this.debug('processing pending promises before responding');
 
@@ -141,7 +141,7 @@ export class KernelHost {
       // if this is an async method, return immediately and
       // call next only when the promise is fulfilled.
       if (this.isPromise(ret)) {
-        checkIfAsyncIsAllowed();
+        assertAsyncIsAllowed();
 
         this.debug('waiting for promise to be fulfilled');
 
@@ -169,7 +169,7 @@ export class KernelHost {
     // indicate this request was processed (synchronously).
     return next();
 
-    function checkIfAsyncIsAllowed() {
+    function assertAsyncIsAllowed() {
       if (sync) {
         throw new JsiiFault(
           'Cannot handle async operations while waiting for a sync callback to return',