Repository: Azure/azure-functions-durable-js
Pull request #168 (2020-02-25 23:57:34)

Handling async entities
This addresses issue #144 but bubbling the async handling of entities all the way up to the entity shim.

I've added a test and some samples to it as well.

aaronpowell: @anthonychu There was a few more places than you noted in #144 that needed tweaking, but I think I've got that bug sorted now.
anthonychu: Thanks for this @aaronpowell. I told you it would be easy to trick you into helping us. It worked!
aaronpowell: @anthonychu you know how to spike my interest 🤣
aaronpowell: @cgillum The tests have both sync and async entities in them, the only differences is you _always_ have to `await` the created entity from `shim.ts`.

I didn't test the new sample I added though as I couldn't work out how to trigger it.
ConnorMcMahon: @aaronpowell 

To clarify, there are no breaking changes for customers who have written their entities as is today, right? The new `await` is handled purely in `shim.ts`, and does not affect existing customer code.

Patch:
diff --git a/samples/AsyncCounterEntity/functions.json b/samples/AsyncCounterEntity/functions.json
new file mode 100644
index 0000000..b79bed5
--- /dev/null
+++ b/samples/AsyncCounterEntity/functions.json
@@ -0,0 +1,10 @@
+{
+    "bindings": [
+      {
+        "name": "context",
+        "type": "entityTrigger",
+        "direction": "in"
+      }
+    ],
+    "disabled": false
+  }
\ No newline at end of file
diff --git a/samples/AsyncCounterEntity/index.js b/samples/AsyncCounterEntity/index.js
new file mode 100644
index 0000000..3066d71
--- /dev/null
+++ b/samples/AsyncCounterEntity/index.js
@@ -0,0 +1,21 @@
+const df = require("durable-functions");
+
+module.exports = df.entity(async function(context) {
+    await Promise.resolve();
+    let currentValue = context.df.getState(() => 0);
+
+    switch (context.df.operationName) {
+        case "add":
+            const amount = context.df.getInput();
+            currentValue += amount;
+            break;
+        case "reset":
+            currentValue = 0;
+            break;
+        case "get":
+            context.df.return(currentValue);
+            break;
+    }
+
+    context.df.setState(currentValue);
+});
\ No newline at end of file
diff --git a/samples/AsyncCounterOrchestration/function.json b/samples/AsyncCounterOrchestration/function.json
new file mode 100644
index 0000000..acf1ac7
--- /dev/null
+++ b/samples/AsyncCounterOrchestration/function.json
@@ -0,0 +1,10 @@
+{
+  "bindings": [
+    {
+      "name": "context",
+      "type": "orchestrationTrigger",
+      "direction": "in"
+    }
+  ],
+  "disabled": false
+}
\ No newline at end of file
diff --git a/samples/AsyncCounterOrchestration/index.js b/samples/AsyncCounterOrchestration/index.js
new file mode 100644
index 0000000..40b6164
--- /dev/null
+++ b/samples/AsyncCounterOrchestration/index.js
@@ -0,0 +1,10 @@
+const df = require("durable-functions");
+
+module.exports = df.orchestrator(function*(context){
+    const entityId = new df.EntityId("AsyncCounterEntity", "myAsyncCounter");
+
+    currentValue = yield context.df.callEntity(entityId, "get");
+    if (currentValue < 10) {
+        yield context.df.callEntity(entityId, "add", 1);
+    }
+});
diff --git a/src/entity.ts b/src/entity.ts
index 4579090..62bbac0 100644
--- a/src/entity.ts
+++ b/src/entity.ts
@@ -30,7 +30,7 @@ export class Entity {
             context.df = this.getCurrentDurableEntityContext(entityBinding, returnState, i, startTime);
 
             try {
-                this.fn(context);
+                await Promise.resolve(this.fn(context));
                 if (!returnState.results[i]) {
                     const elapsedMs = this.computeElapsedMilliseconds(startTime);
                     returnState.results[i] = new OperationResult(false, elapsedMs);
diff --git a/src/orchestrator.ts b/src/orchestrator.ts
index 2ecdc9d..2a4afca 100644
--- a/src/orchestrator.ts
+++ b/src/orchestrator.ts
@@ -146,6 +146,13 @@ export class Orchestrator {
                     return;
                 }
 
+                // The first time a task is marked as complete, the history event that finally marked the task as completed
+                // should not yet have been played by the Durable Task framework, resulting in isReplaying being false.
+                // On replays, the event will have already been processed by the framework, and IsPlayed will be marked as true.
+                if (state[partialResult.completionIndex] !== undefined) {
+                  context.df.isReplaying = state[partialResult.completionIndex].IsPlayed;
+                }
+
                 if (TaskFilter.isFailedTask(partialResult)) {
                     if (!gen.throw) {
                         throw new Error("Cannot properly throw the execption returned by customer code");
@@ -177,13 +184,6 @@ export class Orchestrator {
                 decisionStartedEvent = newDecisionStartedEvent || decisionStartedEvent;
                 context.df.currentUtcDateTime = this.currentUtcDateTime = new Date(decisionStartedEvent.Timestamp);
 
-                // The first time a task is marked as complete, the history event that finally marked the task as completed
-                // should not yet have been played by the Durable Task framework, resulting in isReplaying being false.
-                // On replays, the event will have already been processed by the framework, and IsPlayed will be marked as true.
-                if (state[partialResult.completionIndex] !== undefined) {
-                  context.df.isReplaying = state[partialResult.completionIndex].IsPlayed;
-                }
-
                 g = gen.next(partialResult.result);
             }
         } catch (error) {
diff --git a/src/shim.ts b/src/shim.ts
index ffdcdda..8dcf374 100644
--- a/src/shim.ts
+++ b/src/shim.ts
@@ -25,7 +25,7 @@ export function orchestrator(fn: (context: IOrchestrationFunctionContext) => Ite
 export function entity(fn: (context: IEntityFunctionContext) => unknown)
     : (context: IEntityFunctionContext) => void {
     const listener = new Entity(fn).listen();
-    return (context: IEntityFunctionContext) => {
-        listener(context);
+    return async (context: IEntityFunctionContext) => {
+        await listener(context);
     };
 }


Test patch:
diff --git a/test/integration/entity-spec.ts b/test/integration/entity-spec.ts
index 1bc2c54..c3aeda1 100644
--- a/test/integration/entity-spec.ts
+++ b/test/integration/entity-spec.ts
@@ -18,7 +18,7 @@ describe("Entity", () => {
         const mockContext = new MockContext({
             context: testData.input,
         });
-        entity(mockContext);
+        await entity(mockContext);
 
         expect(mockContext.doneValue).to.not.equal(undefined);
 
@@ -36,7 +36,28 @@ describe("Entity", () => {
         const mockContext = new MockContext({
             context: testData.input,
         });
-        entity(mockContext);
+        await entity(mockContext);
+
+        expect(mockContext.doneValue).to.not.equal(undefined);
+
+        if (mockContext.doneValue) {
+            entityStateMatchesExpected(mockContext.doneValue, testData.output);
+        }
+    });
+
+    it("AsyncStringStore entity with no initial state.", async () => {
+        const entity = TestEntities.AsyncStringStore;
+        const operations: StringStoreOperation[] = [];
+        operations.push({ kind: "set", value: "set 1"});
+        operations.push({ kind: "get"});
+        operations.push({ kind: "set", value: "set 2"});
+        operations.push({ kind: "get"});
+
+        const testData = TestEntityBatches.GetAsyncStringStoreBatch(operations, undefined);
+        const mockContext = new MockContext({
+            context: testData.input,
+        });
+        await entity(mockContext);
 
         expect(mockContext.doneValue).to.not.equal(undefined);
 
diff --git a/test/testobjects/testentities.ts b/test/testobjects/testentities.ts
index 3444074..2a589a9 100644
--- a/test/testobjects/testentities.ts
+++ b/test/testobjects/testentities.ts
@@ -38,4 +38,20 @@ export class TestEntities {
                 throw Error("Invalid operation");
         }
     });
+
+    public static AsyncStringStore: any = df.entity(async (context: IEntityFunctionContext) => {
+        await new Promise((resolve) => setTimeout(() => resolve(), 0)); // force onto the event loop and result in a no-op delay
+        switch (context.df.operationName) {
+            case "set":
+                context.df.setState(context.df.getInput());
+                break;
+            case "get":
+                context.df.return(context.df.getState());
+                break;
+            default:
+                throw new Error("No such operation exists");
+        }
+
+
+    });
 }
diff --git a/test/testobjects/testentitybatches.ts b/test/testobjects/testentitybatches.ts
index e592e65..5dc11af 100644
--- a/test/testobjects/testentitybatches.ts
+++ b/test/testobjects/testentitybatches.ts
@@ -132,4 +132,50 @@ export class TestEntityBatches {
             output,
         };
     }
+
+    public static GetAsyncStringStoreBatch(operations: StringStoreOperation[], existingState: string | undefined): EntityInputsAndOutputs {
+        const id = new df.EntityId("asyncstringstore", "asyncstringstorekey");
+
+        const entityExists = existingState !== undefined;
+        const output = new EntityState([], []);
+        if (entityExists)  {
+            output.entityState = JSON.stringify(existingState);
+            output.entityExists = entityExists;
+        }
+
+        const batch: RequestMessage[] = [];
+        let operationCount = 0;
+        for (const operation of operations) {
+            batch[operationCount] = new RequestMessage();
+            switch (operation.kind) {
+                case "get":
+                    // Handle inputs
+                    batch[operationCount].id = JSON.stringify(operationCount);
+                    batch[operationCount].name = "get";
+                    batch[operationCount].signal = false;
+
+                    // Handle outputs
+                    output.results[operationCount] = new OperationResult(false, -1, output.entityState);
+                    break;
+                case "set":
+                    // Handle inputs
+                    const value = JSON.stringify(operation.value);
+                    batch[operationCount].id = JSON.stringify(operationCount);
+                    batch[operationCount].name = "set";
+                    batch[operationCount].signal = false;
+                    batch[operationCount].input = value;
+
+                    // Handle outputs
+                    output.results[operationCount] = new OperationResult(false, -1);
+                    output.entityExists = true;
+                    output.entityState = value;
+                    break;
+            }
+            operationCount++;
+        }
+        return  {
+            input: new DurableEntityBindingInfo(id, entityExists, JSON.stringify(existingState), batch),
+            output,
+        };
+    }
 }