From 484fc215a95232b861b81d6def6e42260ac8a1f9 Mon Sep 17 00:00:00 2001 From: Niels Klomp Date: Sat, 9 Mar 2024 02:04:25 +0100 Subject: [PATCH] feat: Add support to automatically cleanup on final states, as well as to cleanup all other instances when starting a machine --- .../MachineStatePersistenceAgentLogic.ts | 7 +++- .../src/agent/MachineStatePersistence.ts | 39 +++++++++++++------ .../src/functions/machineRegistration.ts | 27 +++++++++++-- .../xstate-persistence/src/types/types.ts | 6 +-- 4 files changed, 60 insertions(+), 19 deletions(-) diff --git a/packages/xstate-persistence/src/__tests__/shared/MachineStatePersistenceAgentLogic.ts b/packages/xstate-persistence/src/__tests__/shared/MachineStatePersistenceAgentLogic.ts index 87701ee5e..35fe77808 100644 --- a/packages/xstate-persistence/src/__tests__/shared/MachineStatePersistenceAgentLogic.ts +++ b/packages/xstate-persistence/src/__tests__/shared/MachineStatePersistenceAgentLogic.ts @@ -135,7 +135,12 @@ export default (testContext: { getAgent: () => ConfiguredAgent; setup: () => Pro it('should automatically store xstate state changes', async (): Promise => { instance.start() - const init = await machineStatePersistRegistration({ context, interpreter: instance, machineName: instance.machine.id }) + const init = await machineStatePersistRegistration({ + context, + interpreter: instance, + machineName: instance.machine.id, + cleanupOnFinalState: false, + }) console.log(JSON.stringify(init, null, 2)) if (!init) { return Promise.reject(new Error('No init')) diff --git a/packages/xstate-persistence/src/agent/MachineStatePersistence.ts b/packages/xstate-persistence/src/agent/MachineStatePersistence.ts index d328ee2f3..1539fe040 100644 --- a/packages/xstate-persistence/src/agent/MachineStatePersistence.ts +++ b/packages/xstate-persistence/src/agent/MachineStatePersistence.ts @@ -79,31 +79,43 @@ export class MachineStatePersistence implements IAgentPlugin { await context.agent.machineStateInit({ ...event.data }) break case MachineStatePersistEventType.EVERY: - if (event.data.state.done) { - // TODO: Cleanup on done - } // We are keeping track of the update counter in the events, ensuring we do not process out of order - await context.agent.machineStatePersist({ ...event.data, updatedCount: event.data._eventCounter ?? event.data.updatedCount }) + await context.agent.machineStatePersist({ + ...event.data, + cleanupOnFinalState: event.data.cleanupOnFinalState ?? event.data._cleanupOnFinalState, + updatedCount: event.data._eventCounter ?? event.data.updatedCount, + }) break default: - return Promise.reject(Error('Event type not supported')) + return Promise.reject(Error(`Event type ${event.type} not supported`)) } } - private async machineStateInit(args: InitMachineStateArgs): Promise { - const { tenantId, machineName, expiresAt, customInstanceId, existingInstanceId } = args + private async machineStateInit(args: InitMachineStateArgs, context: RequiredContext): Promise { + const { tenantId, machineName, expiresAt, customInstanceId, existingInstanceId, cleanupAllOtherInstances } = args debug( `machineStateInit for machine name ${machineName}, tenant ${tenantId}, custom instance ${customInstanceId}, existing id ${existingInstanceId}` ) - let machineInit: MachineStateInit | undefined = undefined if (customInstanceId && existingInstanceId) { return Promise.reject(new Error(`Cannot have both a custom and existing instance id at the same time`)) - } else if (existingInstanceId) { + } + if (cleanupAllOtherInstances !== false) { + await context.agent.machineStatesDeleteExpired({ machineName, tenantId, deleteDoneStates: true }) + await context.agent.machineStatesDeleteExpired({ machineName, tenantId, deleteDoneStates: false }) + const activeMachineStates = (await context.agent.machineStatesFindActive({ machineName, tenantId })).filter( + (state) => !existingInstanceId || state.instanceId !== existingInstanceId + ) + await Promise.all(activeMachineStates.map((state) => context.agent.machineStateDelete({ instanceId: state.instanceId, tenantId }))) + } + let machineInit: MachineStateInit | undefined = undefined + + if (existingInstanceId) { // A existing instanceId is provided. First lookup whether this id is persisted, if not an error is thrown debug(`machineStateInit is using a previously persisted instance id (${existingInstanceId})`) const state = await this.store.getMachineState({ tenantId, instanceId: existingInstanceId }) machineInit = storeInfoToMachineInit({ ...state, stateType: 'existing' }) - } else if (customInstanceId) { + } + if (customInstanceId) { // A custom instanceId is provided. debug(`machineStateInit is using a custom instance id (${customInstanceId})`) } @@ -120,8 +132,9 @@ export class MachineStatePersistence implements IAgentPlugin { debug(`machineStateInit result: ${JSON.stringify(machineInit)}`) return machineInit } - private async machineStatePersist(args: MachineStatePersistArgs): Promise { + private async machineStatePersist(args: MachineStatePersistArgs, context: RequiredContext): Promise { const { instanceId, tenantId, machineName, updatedCount } = args + const cleanupOnFinalState = args.cleanupOnFinalState !== false debug(`machineStatePersist for machine name ${machineName}, updateCount: ${updatedCount}, instance ${instanceId} and tenant ${tenantId}...`) const queriedStates = await this.store.findMachineStates({ filter: [{ instanceId, tenantId }] }) const existingState = queriedStates.length === 1 ? queriedStates[0] : undefined @@ -131,6 +144,10 @@ export class MachineStatePersistence implements IAgentPlugin { debug( `machineStatePersist success for machine name ${machineName}, instance ${instanceId}, update count ${machineStateInfo.updatedCount}, tenant ${tenantId}, last event: ${machineStateInfo.latestEventType}, last state: ${machineStateInfo.latestStateName}` ) + if (cleanupOnFinalState && machineStateInfo.state.done) { + debug(`reached final state for machine ${machineName} instance ${instanceId} and auto cleanup was enabled. Deleting machine state`) + await context.agent.machineStateDelete(machineStateInfo) + } return machineStateInfo } diff --git a/packages/xstate-persistence/src/functions/machineRegistration.ts b/packages/xstate-persistence/src/functions/machineRegistration.ts index a5a712b8e..8086ed750 100644 --- a/packages/xstate-persistence/src/functions/machineRegistration.ts +++ b/packages/xstate-persistence/src/functions/machineRegistration.ts @@ -59,8 +59,9 @@ export const machineStatePersistOnTransition = async < interpreter: Interpreter context: IAgentContext // We use any as this method could be called from an agent with access to, but not exposing this plugin init: MachineStateInit + cleanupOnFinalState?: boolean }): Promise => { - const { context, init, interpreter } = opts + const { cleanupOnFinalState, context, init, interpreter } = opts if (!(context.agent.availableMethods().includes('machineStatePersist') && 'machineStatePersist' in context.agent)) { console.log(`IMachineStatePersistence was not exposed in the current agent. Disabling machine state persistence events`) return @@ -78,6 +79,7 @@ export const machineStatePersistOnTransition = async < state, _eventCounter: _eventCounter++, _eventDate: new Date(), + _cleanupOnFinalState: cleanupOnFinalState !== false, }, }, context @@ -105,6 +107,8 @@ export const machineStatePersistRegistration = async < args: Omit & Partial> & MachineStatePersistenceOpts & { + cleanupOnFinalState?: boolean + cleanupAllOtherInstances?: boolean interpreter: Interpreter context: IAgentContext // We use any as this method could be called from an agent with access to, but not exposing this plugin } @@ -208,11 +212,12 @@ export const interpreterStartOrResumeFromInit = async < TResolvedTypesMeta = TypegenDisabled >(args: { init: MachineStateInit & { stateType?: MachineStateInitType } + cleanupAllOtherInstances?: boolean noRegistration?: boolean interpreter: Interpreter context: IAgentContext }): Promise> => { - const { init, noRegistration, interpreter, context } = args + const { init, noRegistration, interpreter, cleanupAllOtherInstances, context } = args const { stateType, instanceId, machineName, tenantId, expiresAt } = init if (init.machineName !== interpreter.id) { throw new Error(`Machine state init machine name ${init.machineName} does not match name from state machine interpreter ${interpreter.id}`) @@ -225,6 +230,7 @@ export const interpreterStartOrResumeFromInit = async < tenantId, ...(stateType === 'existing' && { existingInstanceId: instanceId }), ...(stateType === 'new' && { customInstanceId: instanceId }), + cleanupAllOtherInstances, context, interpreter, }) @@ -278,16 +284,27 @@ export const interpreterStartOrResume = async < tenantId?: string singletonCheck: boolean noRegistration?: boolean + cleanupAllOtherInstances?: boolean interpreter: Interpreter context: IAgentContext }): Promise> => { - const { stateType, singletonCheck, instanceId, tenantId, noRegistration, context, interpreter } = args + const { stateType, singletonCheck, instanceId, tenantId, noRegistration, context, interpreter, cleanupAllOtherInstances } = args const machineName = args.machineName ?? interpreter.id - const activeStates = await context.agent.machineStatesFindActive({ + let activeStates = await context.agent.machineStatesFindActive({ machineName, tenantId, instanceId, }) + if (activeStates.length > 0 && cleanupAllOtherInstances) { + // We cleanup here to not influence the logic below. Normally the agent machineStateInit method does the cleanup + await Promise.all(activeStates.map((state) => context.agent.machineStateDelete({ tenantId: args.tenantId, instanceId: state.instanceId }))) + // We search again, given the delete is using the passed in tenantId, instead of relying on the persisted tenantId. Should not matter, but just making sure + activeStates = await context.agent.machineStatesFindActive({ + machineName, + tenantId, + instanceId, + }) + } if (singletonCheck && activeStates.length > 0) { if (stateType === 'new' || activeStates.every((state) => state.instanceId !== instanceId)) { return Promise.reject(new Error(`Found ${activeStates.length} active '${machineName}' instances, but only one is allows at the same time`)) @@ -305,6 +322,7 @@ export const interpreterStartOrResume = async < customInstanceId: instanceId, machineName: machineName ?? interpreter.id, tenantId, + cleanupAllOtherInstances, }) return await interpreterStartOrResumeFromInit({ init, noRegistration, interpreter, context }) } @@ -317,6 +335,7 @@ export const interpreterStartOrResume = async < customInstanceId: instanceId, machineName: machineName ?? interpreter.id, tenantId, + cleanupAllOtherInstances, }) return await interpreterStartOrResumeFromInit({ init, noRegistration, interpreter, context }) } diff --git a/packages/xstate-persistence/src/types/types.ts b/packages/xstate-persistence/src/types/types.ts index 0983d627b..69eefcd63 100644 --- a/packages/xstate-persistence/src/types/types.ts +++ b/packages/xstate-persistence/src/types/types.ts @@ -66,7 +66,7 @@ export type DeleteStateResult = number */ export type MachineStatePersistEvent = { type: MachineStatePersistEventType - data: MachineStatePersistArgs & { _eventCounter: number; _eventDate: Date } + data: MachineStatePersistArgs & { _eventCounter: number; _eventDate: Date; _cleanupOnFinalState: boolean } } /** @@ -120,14 +120,14 @@ export type MachineStateInit = Pick, 'instanceId'> & Pick & - Pick + Pick & { cleanupAllOtherInstances?: boolean } /** * Represents the arguments required to persist the machine state. */ export type MachineStatePersistArgs = Omit & Pick & - Partial> + Partial> & { cleanupOnFinalState?: boolean } /** * Represents the arguments required to get machine state.