diff --git a/src/main/java/io/supertokens/storageLayer/StorageLayer.java b/src/main/java/io/supertokens/storageLayer/StorageLayer.java index 967c538fc..9fe857653 100644 --- a/src/main/java/io/supertokens/storageLayer/StorageLayer.java +++ b/src/main/java/io/supertokens/storageLayer/StorageLayer.java @@ -242,6 +242,8 @@ public static void loadAllTenantStorage(Main main, TenantConfig[] tenants) } main.getResourceDistributor().clearAllResourcesWithResourceKey(RESOURCE_KEY); + Set userPoolsInUse = new HashSet<>(); + for (ResourceDistributor.KeyClass key : resourceKeyToStorageMap.keySet()) { Storage currStorage = resourceKeyToStorageMap.get(key); String userPoolId = currStorage.getUserPoolId(); @@ -256,12 +258,15 @@ public static void loadAllTenantStorage(Main main, TenantConfig[] tenants) main.getResourceDistributor().setResource(key.getTenantIdentifier(), RESOURCE_KEY, new StorageLayer(resourceKeyToStorageMap.get(key))); + + userPoolsInUse.add(userPoolId); } // TODO: should the below code be outside of this locked code cause it takes time // and any other thread that will want access to the resource distributor will have // to wait for this? // we remove storage layers that are no longer being used + for (ResourceDistributor.KeyClass key : existingStorageMap.keySet()) { try { if (((StorageLayer) main.getResourceDistributor() @@ -272,7 +277,11 @@ public static void loadAllTenantStorage(Main main, TenantConfig[] tenants) ((StorageLayer) existingStorageMap.get(key)).storage.stopLogging(); } } catch (TenantOrAppNotFoundException e) { - throw new IllegalStateException(e); + // this means a tenant has been removed but the storage may need closing + if (!userPoolsInUse.contains(((StorageLayer) existingStorageMap.get(key)).storage.getUserPoolId())) { + ((StorageLayer) existingStorageMap.get(key)).storage.close(); + ((StorageLayer) existingStorageMap.get(key)).storage.stopLogging(); + } } } diff --git a/src/test/java/io/supertokens/test/multitenant/StorageLayerTest.java b/src/test/java/io/supertokens/test/multitenant/StorageLayerTest.java index 59785a020..99548cc5a 100644 --- a/src/test/java/io/supertokens/test/multitenant/StorageLayerTest.java +++ b/src/test/java/io/supertokens/test/multitenant/StorageLayerTest.java @@ -21,6 +21,7 @@ import io.supertokens.ProcessState; import io.supertokens.featureflag.EE_FEATURES; import io.supertokens.featureflag.FeatureFlagTestContent; +import io.supertokens.multitenancy.Multitenancy; import io.supertokens.pluginInterface.STORAGE_TYPE; import io.supertokens.pluginInterface.Storage; import io.supertokens.pluginInterface.exceptions.InvalidConfigException; @@ -1819,4 +1820,107 @@ public void testThatStoragePointingToSameDbSharesThInstance() throws Exception { process.kill(); assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STOPPED)); } + + @Test + public void testThatStorageIsClosedAfterTenantDeletion() throws Exception { + String[] args = {"../"}; + + TestingProcessManager.TestingProcess process = TestingProcessManager.start(args, false); + FeatureFlagTestContent.getInstance(process.getProcess()) + .setKeyValue(FeatureFlagTestContent.ENABLED_FEATURES, new EE_FEATURES[]{EE_FEATURES.MULTI_TENANCY}); + process.startProcess(); + assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STARTED)); + + if (StorageLayer.getStorage(process.getProcess()).getType() != STORAGE_TYPE.SQL) { + return; + } + + if (StorageLayer.isInMemDb(process.getProcess())) { + return; + } + + JsonObject config = new JsonObject(); + StorageLayer.getStorage(new TenantIdentifier(null, null, null), process.getProcess()) + .modifyConfigToAddANewUserPoolForTesting(config, 1); + + Multitenancy.addNewOrUpdateAppOrTenant(process.getProcess(), new TenantConfig( + new TenantIdentifier(null, null, "t1"), + new EmailPasswordConfig(true), + new ThirdPartyConfig(true, null), + new PasswordlessConfig(true), + config + ), false); + + Storage storage = StorageLayer.getStorage(new TenantIdentifier(null, null, "t1"), process.getProcess()); + + Multitenancy.deleteTenant(new TenantIdentifier(null, null, "t1"), process.getProcess()); + + // Should not be able to query from the storage + try { + storage.getKeyValue(new TenantIdentifier(null, null, "t1"), "somekey"); + fail(); + } catch (IllegalStateException e) { + assertTrue(e.getMessage().contains("call initPool before getConnection")); + } + process.kill(); + assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STOPPED)); + } + + @Test + public void testThatStorageIsClosedOnlyWhenNoMoreTenantsArePointingToIt() throws Exception { + String[] args = {"../"}; + + TestingProcessManager.TestingProcess process = TestingProcessManager.start(args, false); + FeatureFlagTestContent.getInstance(process.getProcess()) + .setKeyValue(FeatureFlagTestContent.ENABLED_FEATURES, new EE_FEATURES[]{EE_FEATURES.MULTI_TENANCY}); + process.startProcess(); + assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STARTED)); + + if (StorageLayer.getStorage(process.getProcess()).getType() != STORAGE_TYPE.SQL) { + return; + } + + if (StorageLayer.isInMemDb(process.getProcess())) { + return; + } + + JsonObject config = new JsonObject(); + StorageLayer.getStorage(new TenantIdentifier(null, null, null), process.getProcess()) + .modifyConfigToAddANewUserPoolForTesting(config, 1); + + // 2 tenants using the same storage + Multitenancy.addNewOrUpdateAppOrTenant(process.getProcess(), new TenantConfig( + new TenantIdentifier(null, null, "t1"), + new EmailPasswordConfig(true), + new ThirdPartyConfig(true, null), + new PasswordlessConfig(true), + config + ), false); + Multitenancy.addNewOrUpdateAppOrTenant(process.getProcess(), new TenantConfig( + new TenantIdentifier(null, null, "t2"), + new EmailPasswordConfig(true), + new ThirdPartyConfig(true, null), + new PasswordlessConfig(true), + config + ), false); + + Storage storage = StorageLayer.getStorage(new TenantIdentifier(null, null, "t1"), process.getProcess()); + + Multitenancy.deleteTenant(new TenantIdentifier(null, null, "t1"), process.getProcess()); + + // Storage should still be active + storage.getKeyValue(new TenantIdentifier(null, null, "t1"), "somekey"); + + Multitenancy.deleteTenant(new TenantIdentifier(null, null, "t2"), process.getProcess()); + + // Storage should be closed now + try { + storage.getKeyValue(new TenantIdentifier(null, null, "t1"), "somekey"); + fail(); + } catch (IllegalStateException e) { + assertTrue(e.getMessage().contains("call initPool before getConnection")); + } + process.kill(); + assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STOPPED)); + } }