diff --git a/CHANGELOG.md b/CHANGELOG.md index a725396..53fcc48 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ### New +- Implement entity support in the in-memory testing backend ([#341](https://github.com/microsoft/durabletask-js/pull/341)) + ### Fixes diff --git a/packages/durabletask-js/src/testing/in-memory-backend.ts b/packages/durabletask-js/src/testing/in-memory-backend.ts index c1748da..06a9f9b 100644 --- a/packages/durabletask-js/src/testing/in-memory-backend.ts +++ b/packages/durabletask-js/src/testing/in-memory-backend.ts @@ -5,6 +5,7 @@ import * as pb from "../proto/orchestrator_service_pb"; import * as pbh from "../utils/pb-helper.util"; import { OrchestrationStatus as ClientOrchestrationStatus } from "../orchestration/enum/orchestration-status.enum"; import { ParentOrchestrationInstance } from "../types/parent-orchestration-instance.type"; +import { StringValue } from "google-protobuf/google/protobuf/wrappers_pb"; import { randomUUID } from "crypto"; /** Mints a fresh per-execution ID (DTFx `Guid.ToString("N")` idiom: 32 hex chars, no dashes). */ @@ -42,6 +43,41 @@ export interface ActivityWorkItem { completionToken: number; } +/** + * Internal entity state stored by the in-memory backend. + */ +export interface EntityState { + instanceId: string; + serializedState?: string; + lastModifiedAt: Date; + /** Critical section ID currently holding the lock, if any. */ + lockedBy?: string; + /** Operations queued but not yet dispatched to a worker. */ + pendingOperations: pb.HistoryEvent[]; + /** Operations handed to a worker and awaiting a batch result. */ + dispatchedOperations: pb.HistoryEvent[]; + completionToken: number; +} + +/** + * Pending lock request from an orchestration that could not be granted immediately. + */ +interface PendingLockRequest { + criticalSectionId: string; + parentInstanceId: string; + lockSet: string[]; +} + +/** + * Entity work item that needs to be executed. + */ +export interface EntityWorkItem { + instanceId: string; + entityState?: string; + operations: pb.HistoryEvent[]; + completionToken: number; +} + /** * Promise resolver for waiting on orchestration state changes. */ @@ -68,6 +104,11 @@ export class InMemoryOrchestrationBackend { private readonly orchestrationQueue: string[] = []; private readonly orchestrationQueueSet: Set = new Set(); private readonly activityQueue: ActivityWorkItem[] = []; + private readonly entities: Map = new Map(); + private readonly entityQueue: string[] = []; + private readonly entityQueueSet: Set = new Set(); + private readonly entityInFlight: Set = new Set(); + private pendingLockRequests: PendingLockRequest[] = []; private readonly stateWaiters: Map = new Map(); private readonly pendingTimers: Set> = new Set(); private readonly instanceTimers: Map>> = new Map(); @@ -316,6 +357,176 @@ export class InMemoryOrchestrationBackend { return this.activityQueue.shift(); } + /** + * Gets the next entity work item to process, if any. + * + * Drains all operations queued for the entity into a single batch and marks the + * entity in-flight so it is not dispatched again until the batch completes. + */ + getNextEntityWorkItem(): EntityWorkItem | undefined { + const skipped: string[] = []; + let workItem: EntityWorkItem | undefined; + + while (this.entityQueue.length > 0) { + const entityId = this.entityQueue.shift()!; + this.entityQueueSet.delete(entityId); + + const entity = this.entities.get(entityId); + if (!entity || entity.pendingOperations.length === 0) { + continue; + } + + // Already executing a batch; re-queue so it is picked up after completion. + if (this.entityInFlight.has(entityId)) { + skipped.push(entityId); + continue; + } + + this.entityInFlight.add(entityId); + entity.dispatchedOperations = entity.pendingOperations.slice(); + entity.pendingOperations.length = 0; + + workItem = { + instanceId: entity.instanceId, + entityState: entity.serializedState, + operations: entity.dispatchedOperations, + completionToken: entity.completionToken, + }; + break; + } + + for (const entityId of skipped) { + this.enqueueEntity(entityId); + } + + return workItem; + } + + /** + * Completes an entity batch execution: persists the new state, delivers operation + * results back to calling orchestrations, and applies entity side-effect actions. + */ + completeEntityTask(instanceId: string, completionToken: number, result: pb.EntityBatchResult): void { + const entity = this.entities.get(instanceId); + if (!entity || entity.completionToken !== completionToken) { + return; // Entity was reset/purged, or this is a stale completion + } + + const dispatched = entity.dispatchedOperations; + + entity.serializedState = result.getEntitystate()?.getValue() ?? undefined; + entity.lastModifiedAt = new Date(); + entity.dispatchedOperations = []; + entity.completionToken = this.nextCompletionToken++; + this.entityInFlight.delete(instanceId); + + // Deliver results to callers. Results are index-aligned with the dispatched + // operations because the executor pushes exactly one result per operation. + const results = result.getResultsList(); + for (let i = 0; i < dispatched.length; i++) { + this.deliverEntityOperationResult(dispatched[i], results[i]); + } + + for (const action of result.getActionsList()) { + this.processEntityAction(action); + } + + // A batch may have produced signals for this same entity, or the entity may have + // been signalled while its batch was running. + if (entity.pendingOperations.length > 0) { + this.enqueueEntity(instanceId); + } + } + + /** + * Delivers a single entity operation result back to the orchestration that called it. + * Signals are fire-and-forget and produce no response. + */ + private deliverEntityOperationResult(operation: pb.HistoryEvent, result?: pb.OperationResult): void { + const called = operation.getEntityoperationcalled(); + if (!called || !result) { + return; + } + + const parentInstanceId = called.getParentinstanceid()?.getValue(); + if (!parentInstanceId) { + return; + } + + const parent = this.instances.get(parentInstanceId); + if (!parent) { + return; + } + + const event = new pb.HistoryEvent(); + event.setEventid(-1); + event.setTimestamp(pbh.newTimestamp(new Date())); + + const success = result.getSuccess(); + const failure = result.getFailure(); + + if (success) { + const completed = new pb.EntityOperationCompletedEvent(); + completed.setRequestid(called.getRequestid()); + const output = success.getResult(); + if (output) { + completed.setOutput(output); + } + event.setEntityoperationcompleted(completed); + } else if (failure) { + const failed = new pb.EntityOperationFailedEvent(); + failed.setRequestid(called.getRequestid()); + const details = failure.getFailuredetails(); + if (details) { + failed.setFailuredetails(details); + } + event.setEntityoperationfailed(failed); + } else { + return; + } + + parent.pendingEvents.push(event); + parent.lastUpdatedAt = new Date(); + this.enqueueOrchestration(parentInstanceId); + } + + /** + * Applies a side-effect action produced by an entity operation. + */ + private processEntityAction(action: pb.OperationAction): void { + const sendSignal = action.getSendsignal(); + if (sendSignal) { + this.signalEntityInternal(sendSignal.getInstanceid(), sendSignal.getName(), sendSignal.getInput()?.getValue()); + return; + } + + const startNew = action.getStartneworchestration(); + if (startNew) { + const instanceId = startNew.getInstanceid() || randomUUID().replace(/-/g, ""); + if (!this.instances.has(instanceId)) { + this.createInstance(instanceId, startNew.getName(), startNew.getInput()?.getValue()); + } + } + } + + /** + * Signals an entity from outside an orchestration (client-initiated, fire-and-forget). + * + * @param entityId The entity instance ID, in `@name@key` form. + * @param operation The operation name to invoke. + * @param input Optional serialized operation input. + */ + signalEntity(entityId: string, operation: string, input?: string): void { + this.signalEntityInternal(entityId, operation, input); + } + + /** + * Gets the current state of an entity, if it exists. Intended for test assertions. + */ + getEntity(instanceId: string): EntityState | undefined { + return this.entities.get(instanceId); + } + /** * Completes an orchestration execution with the given actions. */ @@ -468,7 +679,7 @@ export class InMemoryOrchestrationBackend { * Checks if there are any pending work items. */ hasPendingWork(): boolean { - return this.orchestrationQueue.length > 0 || this.activityQueue.length > 0; + return this.orchestrationQueue.length > 0 || this.activityQueue.length > 0 || this.entityQueue.length > 0; } /** @@ -479,6 +690,11 @@ export class InMemoryOrchestrationBackend { this.orchestrationQueue.length = 0; this.orchestrationQueueSet.clear(); this.activityQueue.length = 0; + this.entities.clear(); + this.entityQueue.length = 0; + this.entityQueueSet.clear(); + this.entityInFlight.clear(); + this.pendingLockRequests = []; for (const waiters of this.stateWaiters.values()) { for (const waiter of waiters) { waiter.reject(new Error("Backend was reset")); @@ -551,6 +767,9 @@ export class InMemoryOrchestrationBackend { case pb.OrchestratorAction.OrchestratoractiontypeCase.SENDEVENT: this.processSendEventAction(action.getSendevent()!); break; + case pb.OrchestratorAction.OrchestratoractiontypeCase.SENDENTITYMESSAGE: + this.processSendEntityMessageAction(instance, action); + break; case pb.OrchestratorAction.OrchestratoractiontypeCase.REWINDORCHESTRATION: this.processRewindOrchestrationAction(instance, action.getRewindorchestration()!); break; @@ -830,6 +1049,215 @@ export class InMemoryOrchestrationBackend { this.enqueueOrchestration(instance.instanceId); } + /** + * Handles an entity message emitted by an orchestration. + * + * Mirrors the Python in-memory backend: each outbound message is echoed into the + * orchestration history (so replay validation sees the confirmation) and then routed + * to the target entity or the lock manager. + */ + private processSendEntityMessageAction(instance: OrchestrationInstance, action: pb.OrchestratorAction): void { + const entityMessage = action.getSendentitymessage()!; + const actionId = action.getId(); + const messageTypeCase = pb.SendEntityMessageAction.EntitymessagetypeCase; + + switch (entityMessage.getEntitymessagetypeCase()) { + case messageTypeCase.ENTITYOPERATIONSIGNALED: { + const signaled = entityMessage.getEntityoperationsignaled()!; + this.appendEntityHistoryEvent(instance, actionId, (e) => e.setEntityoperationsignaled(signaled)); + + const targetId = signaled.getTargetinstanceid()?.getValue(); + if (targetId) { + const queued = new pb.HistoryEvent(); + queued.setEventid(-1); + queued.setTimestamp(pbh.newTimestamp(new Date())); + queued.setEntityoperationsignaled(signaled); + this.queueEntityOperation(targetId, queued); + } + break; + } + case messageTypeCase.ENTITYOPERATIONCALLED: { + const called = entityMessage.getEntityoperationcalled()!; + this.appendEntityHistoryEvent(instance, actionId, (e) => e.setEntityoperationcalled(called)); + + if (instance.status === pb.OrchestrationStatus.ORCHESTRATION_STATUS_PENDING) { + instance.status = pb.OrchestrationStatus.ORCHESTRATION_STATUS_RUNNING; + } + + const targetId = called.getTargetinstanceid()?.getValue(); + if (targetId) { + const queued = new pb.HistoryEvent(); + queued.setEventid(-1); + queued.setTimestamp(pbh.newTimestamp(new Date())); + queued.setEntityoperationcalled(called); + this.queueEntityOperation(targetId, queued); + } + break; + } + case messageTypeCase.ENTITYLOCKREQUESTED: { + const lockRequested = entityMessage.getEntitylockrequested()!; + this.appendEntityHistoryEvent(instance, actionId, (e) => e.setEntitylockrequested(lockRequested)); + + if (instance.status === pb.OrchestrationStatus.ORCHESTRATION_STATUS_PENDING) { + instance.status = pb.OrchestrationStatus.ORCHESTRATION_STATUS_RUNNING; + } + + const parentId = lockRequested.getParentinstanceid()?.getValue(); + if (parentId) { + this.tryGrantLock({ + criticalSectionId: lockRequested.getCriticalsectionid(), + parentInstanceId: parentId, + lockSet: lockRequested.getLocksetList().slice(), + }); + } + break; + } + case messageTypeCase.ENTITYUNLOCKSENT: { + const unlock = entityMessage.getEntityunlocksent()!; + this.appendEntityHistoryEvent(instance, actionId, (e) => e.setEntityunlocksent(unlock)); + + const targetId = unlock.getTargetinstanceid()?.getValue(); + if (targetId) { + const entity = this.entities.get(targetId); + if (entity && entity.lockedBy === unlock.getCriticalsectionid()) { + entity.lockedBy = undefined; + } + } + + this.tryGrantPendingLocks(); + break; + } + default: + throw new Error( + `Unknown entity message type '${entityMessage.getEntitymessagetypeCase()}' for orchestration ` + + `'${instance.instanceId}'. This likely means the in-memory backend needs to be updated to handle ` + + `a newly introduced entity message type.`, + ); + } + } + + /** + * Appends the confirmation event for an outbound entity message to the orchestration history. + * The event ID must be the originating action ID so replay validation can match it. + */ + private appendEntityHistoryEvent( + instance: OrchestrationInstance, + actionId: number, + populate: (event: pb.HistoryEvent) => void, + ): void { + const event = new pb.HistoryEvent(); + event.setEventid(actionId); + event.setTimestamp(pbh.newTimestamp(new Date())); + populate(event); + instance.history.push(event); + } + + /** + * Gets (or lazily creates) the state record for an entity. + */ + private getOrCreateEntity(entityId: string): EntityState { + let entity = this.entities.get(entityId); + if (!entity) { + entity = { + instanceId: entityId, + lastModifiedAt: new Date(), + pendingOperations: [], + dispatchedOperations: [], + completionToken: this.nextCompletionToken++, + }; + this.entities.set(entityId, entity); + } + return entity; + } + + private queueEntityOperation(entityId: string, event: pb.HistoryEvent): void { + const entity = this.getOrCreateEntity(entityId); + entity.pendingOperations.push(event); + this.enqueueEntity(entityId); + } + + private enqueueEntity(entityId: string): void { + if (!this.entityQueueSet.has(entityId)) { + this.entityQueue.push(entityId); + this.entityQueueSet.add(entityId); + } + } + + /** + * Signals an entity as a side effect of another entity's operation. + */ + private signalEntityInternal(entityId: string, operation: string, input?: string): void { + const signaled = new pb.EntityOperationSignaledEvent(); + signaled.setRequestid(randomUUID().replace(/-/g, "")); + signaled.setOperation(operation); + if (input !== undefined) { + const value = new StringValue(); + value.setValue(input); + signaled.setInput(value); + } + const target = new StringValue(); + target.setValue(entityId); + signaled.setTargetinstanceid(target); + + const event = new pb.HistoryEvent(); + event.setEventid(-1); + event.setTimestamp(pbh.newTimestamp(new Date())); + event.setEntityoperationsignaled(signaled); + + this.queueEntityOperation(entityId, event); + } + + private canGrantLock(pending: PendingLockRequest): boolean { + return !pending.lockSet.some((entityId) => this.entities.get(entityId)?.lockedBy !== undefined); + } + + /** + * Grants a lock to every entity in the lock set and notifies the parent orchestration. + * Assumes availability was already verified via {@link canGrantLock}. + */ + private grantLock(pending: PendingLockRequest): void { + for (const entityId of pending.lockSet) { + this.getOrCreateEntity(entityId).lockedBy = pending.criticalSectionId; + } + + const parent = this.instances.get(pending.parentInstanceId); + if (!parent) { + return; + } + + const granted = new pb.EntityLockGrantedEvent(); + granted.setCriticalsectionid(pending.criticalSectionId); + + const event = new pb.HistoryEvent(); + event.setEventid(-1); + event.setTimestamp(pbh.newTimestamp(new Date())); + event.setEntitylockgranted(granted); + + parent.pendingEvents.push(event); + parent.lastUpdatedAt = new Date(); + this.enqueueOrchestration(pending.parentInstanceId); + } + + private tryGrantLock(pending: PendingLockRequest): void { + if (!this.canGrantLock(pending)) { + this.pendingLockRequests.push(pending); + return; + } + this.grantLock(pending); + } + + private tryGrantPendingLocks(): void { + const stillPending: PendingLockRequest[] = []; + for (const pending of this.pendingLockRequests) { + if (this.canGrantLock(pending)) { + this.grantLock(pending); + } else { + stillPending.push(pending); + } + } + this.pendingLockRequests = stillPending; + } + private processSendEventAction(sendEvent: pb.SendEventAction): void { const targetInstanceId = sendEvent.getInstance()?.getInstanceid(); const eventName = sendEvent.getName(); diff --git a/packages/durabletask-js/src/testing/test-client.ts b/packages/durabletask-js/src/testing/test-client.ts index 09d1538..517fbca 100644 --- a/packages/durabletask-js/src/testing/test-client.ts +++ b/packages/durabletask-js/src/testing/test-client.ts @@ -7,6 +7,8 @@ import { TOrchestrator } from "../types/orchestrator.type"; import { TInput } from "../types/input.type"; import { OrchestrationState } from "../orchestration/orchestration-state"; import { FailureDetails } from "../task/failure-details"; +import { EntityInstanceId } from "../entities/entity-instance-id"; +import { EntityMetadata } from "../entities/entity-metadata"; import { InMemoryOrchestrationBackend, OrchestrationInstance } from "./in-memory-backend"; import * as pb from "../proto/orchestrator_service_pb"; @@ -140,6 +142,49 @@ export class TestOrchestrationClient { this.backend.rewindInstance(instanceId, reason); } + /** + * Signals an entity with a one-way (fire-and-forget) operation. + * + * @param id The entity to signal. + * @param operationName The name of the operation to invoke. + * @param input Optional operation input. Serialized as JSON. + */ + async signalEntity(id: EntityInstanceId, operationName: string, input?: unknown): Promise { + this.backend.signalEntity( + id.toString(), + operationName, + input === undefined ? undefined : JSON.stringify(input), + ); + } + + /** + * Gets the metadata of an entity, or undefined if the entity does not exist. + * + * @param id The entity to look up. + * @param includeState Whether to deserialize and include the entity state. + */ + async getEntity( + id: EntityInstanceId, + includeState: boolean = true, + ): Promise | undefined> { + const entity = this.backend.getEntity(id.toString()); + if (!entity) { + return undefined; + } + + return { + id, + lastModifiedTime: entity.lastModifiedAt, + backlogQueueSize: entity.pendingOperations.length, + lockedBy: entity.lockedBy, + includesState: includeState, + state: + includeState && entity.serializedState !== undefined + ? (JSON.parse(entity.serializedState) as T) + : undefined, + }; + } + /** * Stops the client. No-op for in-memory backend. */ diff --git a/packages/durabletask-js/src/testing/test-worker.ts b/packages/durabletask-js/src/testing/test-worker.ts index ed389b8..b0bacf3 100644 --- a/packages/durabletask-js/src/testing/test-worker.ts +++ b/packages/durabletask-js/src/testing/test-worker.ts @@ -4,11 +4,20 @@ import { Registry } from "../worker/registry"; import { OrchestrationExecutor } from "../worker/orchestration-executor"; import { ActivityExecutor } from "../worker/activity-executor"; +import { TaskEntityShim } from "../worker/entity-executor"; +import { EntityInstanceId } from "../entities/entity-instance-id"; +import { EntityFactory } from "../entities/task-entity"; import { TOrchestrator } from "../types/orchestrator.type"; import { TActivity } from "../types/activity.type"; import { TInput } from "../types/input.type"; import { TOutput } from "../types/output.type"; -import { InMemoryOrchestrationBackend, OrchestrationInstance, ActivityWorkItem } from "./in-memory-backend"; +import { + InMemoryOrchestrationBackend, + OrchestrationInstance, + ActivityWorkItem, + EntityWorkItem, +} from "./in-memory-backend"; +import { StringValue } from "google-protobuf/google/protobuf/wrappers_pb"; import * as pb from "../proto/orchestrator_service_pb"; import * as pbh from "../utils/pb-helper.util"; @@ -73,6 +82,27 @@ export class TestOrchestrationWorker { return name; } + /** + * Registers an entity with the worker. + */ + addEntity(factory: EntityFactory): string { + if (this.isRunning) { + throw new Error("Cannot add entity while worker is running."); + } + return this.registry.addEntity(factory); + } + + /** + * Registers a named entity function with the worker. + */ + addNamedEntity(name: string, factory: EntityFactory): string { + if (this.isRunning) { + throw new Error("Cannot add entity while worker is running."); + } + this.registry.addNamedEntity(name, factory); + return name; + } + /** * Starts the worker processing loop. */ @@ -126,6 +156,13 @@ export class TestOrchestrationWorker { processedAny = true; } + // Then process entities + const entity = this.backend.getNextEntityWorkItem(); + if (entity) { + await this.processEntity(entity); + processedAny = true; + } + // If nothing was processed, yield to allow other async operations if (!processedAny) { await this.yieldToEventLoop(); @@ -175,6 +212,95 @@ export class TestOrchestrationWorker { } } + /** + * Processes a single entity work item by executing the batch of queued operations. + */ + private async processEntity(workItem: EntityWorkItem): Promise { + const { instanceId, completionToken } = workItem; + + const batchResult = await this.executeEntityBatch(workItem); + this.backend.completeEntityTask(instanceId, completionToken, batchResult); + } + + /** + * Builds an EntityBatchRequest from the work item and runs it through the entity shim. + * Any failure is converted into a per-operation failure so callers are never left hanging. + */ + private async executeEntityBatch(workItem: EntityWorkItem): Promise { + const { instanceId, entityState, operations } = workItem; + + try { + const entityId = EntityInstanceId.fromString(instanceId); + const factory = this.registry.getEntity(entityId.name); + + if (!factory) { + return this.createEntityFailureResult(operations, new Error(`No entity task named '${entityId.name}' was found.`)); + } + + const request = new pb.EntityBatchRequest(); + request.setInstanceid(instanceId); + if (entityState !== undefined) { + const state = new StringValue(); + state.setValue(entityState); + request.setEntitystate(state); + } + request.setOperationsList(operations.map((op) => this.toOperationRequest(op))); + + const shim = new TaskEntityShim(factory(), entityId); + return await shim.executeAsync(request); + } catch (error: unknown) { + const err = error instanceof Error ? error : new Error(String(error)); + return this.createEntityFailureResult(operations, err); + } + } + + /** + * Converts a queued entity operation history event into the OperationRequest the shim expects. + */ + private toOperationRequest(event: pb.HistoryEvent): pb.OperationRequest { + const request = new pb.OperationRequest(); + const called = event.getEntityoperationcalled(); + const signaled = event.getEntityoperationsignaled(); + + if (called) { + request.setOperation(called.getOperation()); + request.setRequestid(called.getRequestid()); + const input = called.getInput(); + if (input) { + request.setInput(input); + } + } else if (signaled) { + request.setOperation(signaled.getOperation()); + request.setRequestid(signaled.getRequestid()); + const input = signaled.getInput(); + if (input) { + request.setInput(input); + } + } + + return request; + } + + /** + * Builds a batch result that fails every operation in the batch with the same error. + */ + private createEntityFailureResult(operations: pb.HistoryEvent[], error: Error): pb.EntityBatchResult { + const batchResult = new pb.EntityBatchResult(); + const failureDetails = pbh.newFailureDetails(error); + + batchResult.setResultsList( + operations.map(() => { + const failure = new pb.OperationResultFailure(); + failure.setFailuredetails(failureDetails); + const result = new pb.OperationResult(); + result.setFailure(failure); + return result; + }), + ); + + return batchResult; + } + /** * Yields control to the event loop to allow timers and I/O to process. */ diff --git a/packages/durabletask-js/test/in-memory-backend-entities.spec.ts b/packages/durabletask-js/test/in-memory-backend-entities.spec.ts new file mode 100644 index 0000000..830a026 --- /dev/null +++ b/packages/durabletask-js/test/in-memory-backend-entities.spec.ts @@ -0,0 +1,278 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import { + InMemoryOrchestrationBackend, + TestOrchestrationClient, + TestOrchestrationWorker, + OrchestrationStatus, + OrchestrationContext, + TOrchestrator, + TaskEntity, + EntityInstanceId, +} from "../src"; + +class CounterEntity extends TaskEntity<{ count: number }> { + add(amount: number): number { + this.state.count += amount; + return this.state.count; + } + + get(): number { + return this.state.count; + } + + boom(): void { + throw new Error("Intentional entity failure"); + } + + relay(amount: number): void { + this.context?.signalEntity(new EntityInstanceId("relaytarget", "k"), "add", amount); + } + + protected initializeState(): { count: number } { + return { count: 0 }; + } +} + +class RelayTargetEntity extends TaskEntity<{ count: number }> { + add(amount: number): number { + this.state.count += amount; + return this.state.count; + } + + protected initializeState(): { count: number } { + return { count: 0 }; + } +} + +/** Waits until `predicate` holds or the timeout elapses. Returns whether it held. */ +async function waitFor(predicate: () => boolean, timeoutMs = 5000): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (predicate()) { + return true; + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + return predicate(); +} + +describe("In-Memory Backend - Entities", () => { + let backend: InMemoryOrchestrationBackend; + let client: TestOrchestrationClient; + let worker: TestOrchestrationWorker; + + const counterId = new EntityInstanceId("counter", "mykey"); + + beforeEach(() => { + backend = new InMemoryOrchestrationBackend(); + client = new TestOrchestrationClient(backend); + worker = new TestOrchestrationWorker(backend); + }); + + afterEach(async () => { + try { + await worker.stop(); + } catch { + // Ignore if not running + } + backend.reset(); + }); + + it("should deliver a signal from an orchestration to the entity", async () => { + const orchestrator: TOrchestrator = async (ctx: OrchestrationContext): Promise => { + ctx.entities.signalEntity(counterId, "add", 5); + return "signaled"; + }; + + worker.addOrchestrator(orchestrator); + worker.addNamedEntity("counter", () => new CounterEntity()); + await worker.start(); + + const id = await client.scheduleNewOrchestration(orchestrator); + const state = await client.waitForOrchestrationCompletion(id, true, 10); + expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED); + + // The signal is fire-and-forget, so it lands after the orchestration completes. + const delivered = await waitFor(() => backend.getEntity(counterId.toString())?.serializedState !== undefined); + expect(delivered).toBe(true); + + const entity = await client.getEntity<{ count: number }>(counterId); + expect(entity?.state?.count).toEqual(5); + }); + + it("should batch multiple signals to the same entity", async () => { + const orchestrator: TOrchestrator = async (ctx: OrchestrationContext): Promise => { + ctx.entities.signalEntity(counterId, "add", 1); + ctx.entities.signalEntity(counterId, "add", 2); + ctx.entities.signalEntity(counterId, "add", 3); + return "done"; + }; + + worker.addOrchestrator(orchestrator); + worker.addNamedEntity("counter", () => new CounterEntity()); + await worker.start(); + + const id = await client.scheduleNewOrchestration(orchestrator); + await client.waitForOrchestrationCompletion(id, true, 10); + + const settled = await waitFor(() => { + const entity = backend.getEntity(counterId.toString()); + return entity?.serializedState !== undefined && JSON.parse(entity.serializedState).count === 6; + }); + expect(settled).toBe(true); + }); + + it("should return the result of callEntity to the calling orchestration", async () => { + const orchestrator: TOrchestrator = async function* (ctx: OrchestrationContext): any { + yield ctx.entities.callEntity(counterId, "add", 7); + const total = yield ctx.entities.callEntity(counterId, "add", 3); + return total; + }; + + worker.addOrchestrator(orchestrator); + worker.addNamedEntity("counter", () => new CounterEntity()); + await worker.start(); + + const id = await client.scheduleNewOrchestration(orchestrator); + const state = await client.waitForOrchestrationCompletion(id, true, 30); + + expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED); + expect(state?.serializedOutput).toEqual(JSON.stringify(10)); + }); + + it("should persist entity state across separate callEntity batches", async () => { + const orchestrator: TOrchestrator = async function* (ctx: OrchestrationContext): any { + yield ctx.entities.callEntity(counterId, "add", 4); + return yield ctx.entities.callEntity(counterId, "get"); + }; + + worker.addOrchestrator(orchestrator); + worker.addNamedEntity("counter", () => new CounterEntity()); + await worker.start(); + + const id = await client.scheduleNewOrchestration(orchestrator); + const state = await client.waitForOrchestrationCompletion(id, true, 30); + + expect(state?.serializedOutput).toEqual(JSON.stringify(4)); + }); + + it("should surface an entity operation failure to the calling orchestration", async () => { + const orchestrator: TOrchestrator = async function* (ctx: OrchestrationContext): any { + try { + yield ctx.entities.callEntity(counterId, "boom"); + return "no-error"; + } catch (e) { + return `caught:${(e as Error).message}`; + } + }; + + worker.addOrchestrator(orchestrator); + worker.addNamedEntity("counter", () => new CounterEntity()); + await worker.start(); + + const id = await client.scheduleNewOrchestration(orchestrator); + const state = await client.waitForOrchestrationCompletion(id, true, 30); + + expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED); + expect(state?.serializedOutput).toContain("caught:"); + expect(state?.serializedOutput).not.toContain("no-error"); + }); + + it("should fail a call to an unregistered entity instead of hanging", async () => { + const unknownId = new EntityInstanceId("notregistered", "k"); + + const orchestrator: TOrchestrator = async function* (ctx: OrchestrationContext): any { + try { + yield ctx.entities.callEntity(unknownId, "add", 1); + return "no-error"; + } catch (e) { + return `caught:${(e as Error).message}`; + } + }; + + worker.addOrchestrator(orchestrator); + await worker.start(); + + const id = await client.scheduleNewOrchestration(orchestrator); + const state = await client.waitForOrchestrationCompletion(id, true, 30); + + expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED); + expect(state?.serializedOutput).toContain("caught:"); + }); + + it("should grant and release entity locks", async () => { + const orchestrator: TOrchestrator = async function* (ctx: OrchestrationContext): any { + const lock = yield ctx.entities.lockEntities(counterId); + const value = yield ctx.entities.callEntity(counterId, "add", 2); + lock.release(); + return value; + }; + + worker.addOrchestrator(orchestrator); + worker.addNamedEntity("counter", () => new CounterEntity()); + await worker.start(); + + const id = await client.scheduleNewOrchestration(orchestrator); + const state = await client.waitForOrchestrationCompletion(id, true, 30); + + expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED); + expect(state?.serializedOutput).toEqual(JSON.stringify(2)); + + const released = await waitFor(() => backend.getEntity(counterId.toString())?.lockedBy === undefined); + expect(released).toBe(true); + }); + + it("should deliver a client-initiated signal to the entity", async () => { + worker.addNamedEntity("counter", () => new CounterEntity()); + await worker.start(); + + await client.signalEntity(counterId, "add", 11); + + const delivered = await waitFor(() => { + const entity = backend.getEntity(counterId.toString()); + return entity?.serializedState !== undefined && JSON.parse(entity.serializedState).count === 11; + }); + expect(delivered).toBe(true); + + const metadata = await client.getEntity<{ count: number }>(counterId); + expect(metadata?.id.toString()).toEqual(counterId.toString()); + expect(metadata?.includesState).toBe(true); + expect(metadata?.state?.count).toEqual(11); + }); + + it("should return undefined metadata for an entity that does not exist", async () => { + const metadata = await client.getEntity(new EntityInstanceId("counter", "absent")); + expect(metadata).toBeUndefined(); + }); + + it("should apply a signal sent by one entity to another entity", async () => { + worker.addNamedEntity("counter", () => new CounterEntity()); + worker.addNamedEntity("relaytarget", () => new RelayTargetEntity()); + await worker.start(); + + await client.signalEntity(counterId, "relay", 9); + + const relayed = await waitFor(() => { + const target = backend.getEntity(new EntityInstanceId("relaytarget", "k").toString()); + return target?.serializedState !== undefined && JSON.parse(target.serializedState).count === 9; + }); + expect(relayed).toBe(true); + }); + + it("should clear entity state on reset", async () => { + worker.addNamedEntity("counter", () => new CounterEntity()); + await worker.start(); + + await client.signalEntity(counterId, "add", 3); + const delivered = await waitFor(() => backend.getEntity(counterId.toString())?.serializedState !== undefined); + expect(delivered).toBe(true); + + await worker.stop(); + backend.reset(); + + expect(backend.getEntity(counterId.toString())).toBeUndefined(); + expect(backend.hasPendingWork()).toBe(false); + }); +});