fix: close tenant session after migration (#1835)

Signed-off-by: Timo Glastra <[email protected]>
openwallet-foundation · Apr 15, 2024 · eb2c513 · eb2c513
1 parent 6ec43eb
commit eb2c513
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 2 deletions.
diff --git a/packages/tenants/src/context/TenantAgentContextProvider.ts b/packages/tenants/src/context/TenantAgentContextProvider.ts
@@ -179,11 +179,14 @@ export class TenantAgentContextProvider implements AgentContextProvider {
    * to update the storage for a tenant manually
    */
   public async updateTenantStorage(tenantRecord: TenantRecord, updateOptions?: UpdateAssistantUpdateOptions) {
-    await this.tenantSessionCoordinator.getContextForSession(tenantRecord, {
+    const agentContext = await this.tenantSessionCoordinator.getContextForSession(tenantRecord, {
       // runInMutex allows us to run the updateTenantStorage method in a mutex lock
       // prevent other sessions from being started while the update is in progress
       runInMutex: (agentContext) => this._updateTenantStorage(tenantRecord, agentContext, updateOptions),
     })
+
+    // End sesion afterwards
+    await agentContext.endSession()
   }
 
   /**

diff --git a/packages/tenants/src/context/TenantSessionCoordinator.ts b/packages/tenants/src/context/TenantSessionCoordinator.ts
@@ -54,6 +54,10 @@ export class TenantSessionCoordinator {
     )
   }
 
+  public getSessionCountForTenant(tenantId: string) {
+    return this.tenantAgentContextMapping[tenantId]?.sessionCount ?? 0
+  }
+
   /**
    * Get agent context to use for a session. If an agent context for this tenant does not exist yet
    * it will create it and store it for later use. If the agent context does already exist it will
@@ -94,7 +98,23 @@ export class TenantSessionCoordinator {
         )
 
         if (runInMutex) {
-          await runInMutex(tenantSessions.agentContext)
+          try {
+            await runInMutex(tenantSessions.agentContext)
+          } catch (error) {
+            // If the runInMutex failed we should release the session again
+            tenantSessions.sessionCount--
+            this.logger.debug(
+              `Decreased agent context session count for tenant '${tenantSessions.agentContext.contextCorrelationId}' to ${tenantSessions.sessionCount} due to failure in mutex script`,
+              error
+            )
+
+            if (tenantSessions.sessionCount <= 0 && tenantSessions.agentContext) {
+              await this.closeAgentContext(tenantSessions.agentContext)
+              delete this.tenantAgentContextMapping[tenantSessions.agentContext.contextCorrelationId]
+            }
+
+            throw error
+          }
         }
 
         return tenantSessions.agentContext

diff --git a/packages/tenants/tests/tenants-storage-update.test.ts b/packages/tenants/tests/tenants-storage-update.test.ts
@@ -14,6 +14,7 @@ import path from 'path'
 import { AskarModule, AskarMultiWalletDatabaseScheme } from '../../askar/src'
 import { ariesAskar } from '../../askar/tests/helpers'
 import { testLogger } from '../../core/tests'
+import { TenantSessionCoordinator } from '../src/context/TenantSessionCoordinator'
 
 import { TenantsModule } from '@credo-ts/tenants'
 
@@ -193,6 +194,9 @@ describe('Tenants Storage Update', () => {
       agent.modules.tenants.getTenantAgent({ tenantId: '1d45d3c2-3480-4375-ac6f-47c322f091b0' })
     ).rejects.toThrow(/Current agent storage for tenant 1d45d3c2-3480-4375-ac6f-47c322f091b0 is not up to date/)
 
+    const tenantSessionCoordinator = agent.dependencyManager.resolve(TenantSessionCoordinator)
+    expect(tenantSessionCoordinator.getSessionCountForTenant(tenant.id)).toBe(0)
+
     // Update tenant
     await agent.modules.tenants.updateTenantStorage({
       tenantId: tenant.id,
@@ -201,6 +205,9 @@ describe('Tenants Storage Update', () => {
       },
     })
 
+    // Should have closed session after upgrade
+    expect(tenantSessionCoordinator.getSessionCountForTenant(tenant.id)).toBe(0)
+
     // Expect tenant storage version to be 0.5
     const updatedTenant = await agent.modules.tenants.getTenantById('1d45d3c2-3480-4375-ac6f-47c322f091b0')
     expect(updatedTenant.storageVersion).toBe('0.5')