Fixes issue with memory leak during testing (#35)

* fixes issue #34

* updates changelog

Co-authored-by: Rishabh <rishabh@supertokens.io>

Author: rishabhpoddar

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+- Fixes https://github.com/supertokens/supertokens-postgresql-plugin/issues/34
 
 ## [1.14.0] - 2022-02-23
 

src/main/java/io/supertokens/storage/postgresql/ConnectionPool.java

@@ -24,6 +24,7 @@
 import io.supertokens.storage.postgresql.config.PostgreSQLConfig;
 import io.supertokens.storage.postgresql.output.Logging;
 
+import java.io.ObjectInputFilter;
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.text.DecimalFormat;
@@ -33,12 +34,24 @@
 public class ConnectionPool extends ResourceDistributor.SingletonResource {
 
     private static final String RESOURCE_KEY = "io.supertokens.storage.postgresql.ConnectionPool";
-    private final HikariDataSource ds;
+    private static HikariDataSource hikariDataSource = null;
 
     private ConnectionPool(Start start) {
         if (!start.enabled) {
             throw new RuntimeException("Connection to refused"); // emulates exception thrown by Hikari
         }
+
+        if (ConnectionPool.hikariDataSource != null) {
+            // This implies that it was already created before and that
+            // there is no need to create Hikari again.
+
+            // If ConnectionPool.hikariDataSource == null, it implies that
+            // either the config file had changed somehow (which means the plugin JAR was reloaded, resulting in static
+            // variables to be set to null), or it means that this is the first time we are trying to connect to a db
+            // (applicable only for testing).
+            return;
+        }
+
         HikariConfig config = new HikariConfig();
         PostgreSQLConfig userConfig = Config.getConfig(start);
         config.setDriverClassName("org.postgresql.Driver");
@@ -81,7 +94,7 @@ private ConnectionPool(Start start) {
         // - Failed to validate connection org.mariadb.jdbc.MariaDbConnection@79af83ae (Connection.setNetworkTimeout
         // cannot be called on a closed connection). Possibly consider using a shorter maxLifetime value.
         config.setPoolName("SuperTokens");
-        ds = new HikariDataSource(config);
+        hikariDataSource = new HikariDataSource(config);
     }
 
     private static int getTimeToWaitToInit(Start start) {
@@ -166,13 +179,14 @@ public static Connection getConnection(Start start) throws SQLException {
         if (!start.enabled) {
             throw new SQLException("Storage layer disabled");
         }
-        return getInstance(start).ds.getConnection();
+        return ConnectionPool.hikariDataSource.getConnection();
     }
 
     static void close(Start start) {
         if (getInstance(start) == null) {
             return;
         }
-        getInstance(start).ds.close();
+        ConnectionPool.hikariDataSource.close();
+        ConnectionPool.hikariDataSource = null;
     }
 }

src/main/java/io/supertokens/storage/postgresql/output/CustomLayout.java

@@ -28,11 +28,11 @@
 
 class CustomLayout extends LayoutBase<ILoggingEvent> {
 
-    private final Start start;
+    private final String processID;
 
-    CustomLayout(Start start) {
+    CustomLayout(String processID) {
         super();
-        this.start = start;
+        this.processID = processID;
     }
 
     @Override
@@ -47,7 +47,7 @@ public String doLayout(ILoggingEvent event) {
         sbuf.append(" | ");
 
         sbuf.append("pid: ");
-        sbuf.append(start.getProcessId());
+        sbuf.append(this.processID);
         sbuf.append(" | ");
 
         sbuf.append("[");

src/main/java/io/supertokens/storage/postgresql/output/LayoutWrappingEncoder.java

@@ -28,8 +28,8 @@ class LayoutWrappingEncoder extends EncoderBase<ILoggingEvent> {
 
     private Layout<ILoggingEvent> layout;
 
-    LayoutWrappingEncoder(Start start) {
-        layout = new CustomLayout(start);
+    LayoutWrappingEncoder(String processID) {
+        layout = new CustomLayout(processID);
     }
 
     @Override

src/main/java/io/supertokens/storage/postgresql/output/Logging.java

@@ -35,13 +35,11 @@ public class Logging extends ResourceDistributor.SingletonResource {
 
     private Logging(Start start, String infoLogPath, String errorLogPath) {
         this.infoLogger = infoLogPath.equals("null")
-                ? createLoggerForConsole(start, "io.supertokens.storage.postgresql.Info." + start.getProcessId())
-                : createLoggerForFile(start, infoLogPath,
-                        "io.supertokens.storage.postgresql.Info." + start.getProcessId());
+                ? createLoggerForConsole(start, "io.supertokens.storage.postgresql.Info")
+                : createLoggerForFile(start, infoLogPath, "io.supertokens.storage.postgresql.Info");
         this.errorLogger = errorLogPath.equals("null")
-                ? createLoggerForConsole(start, "io.supertokens.storage.postgresql.Error." + start.getProcessId())
-                : createLoggerForFile(start, errorLogPath,
-                        "io.supertokens.storage.postgresql.Error." + start.getProcessId());
+                ? createLoggerForConsole(start, "io.supertokens.storage.postgresql.Error")
+                : createLoggerForFile(start, errorLogPath, "io.supertokens.storage.postgresql.Error");
     }
 
     private static Logging getInstance(Start start) {
@@ -140,7 +138,7 @@ public static void stopLogging(Start start) {
 
     private Logger createLoggerForFile(Start start, String file, String name) {
         LoggerContext lc = (LoggerContext) LoggerFactory.getILoggerFactory();
-        LayoutWrappingEncoder ple = new LayoutWrappingEncoder(start);
+        LayoutWrappingEncoder ple = new LayoutWrappingEncoder(start.getProcessId());
         ple.setContext(lc);
         ple.start();
         FileAppender<ILoggingEvent> fileAppender = new FileAppender<>();
@@ -158,7 +156,7 @@ private Logger createLoggerForFile(Start start, String file, String name) {
 
     private Logger createLoggerForConsole(Start start, String name) {
         LoggerContext lc = (LoggerContext) LoggerFactory.getILoggerFactory();
-        LayoutWrappingEncoder ple = new LayoutWrappingEncoder(start);
+        LayoutWrappingEncoder ple = new LayoutWrappingEncoder(start.getProcessId());
         ple.setContext(lc);
         ple.start();
         ConsoleAppender<ILoggingEvent> logConsoleAppender = new ConsoleAppender<>();

src/test/java/io/supertokens/storage/postgresql/test/DeadlockTest.java

@@ -160,7 +160,7 @@ public void testCodeCreationRapidly() throws Exception {
         TestingProcessManager.TestingProcess process = TestingProcessManager.start(args);
         assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STARTED));
 
-        ExecutorService es = Executors.newCachedThreadPool();
+        ExecutorService es = Executors.newFixedThreadPool(1000);
 
         AtomicBoolean pass = new AtomicBoolean(true);
 
@@ -196,7 +196,7 @@ public void testCodeCreationRapidlyWithDifferentEmails() throws Exception {
         TestingProcessManager.TestingProcess process = TestingProcessManager.start(args);
         assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STARTED));
 
-        ExecutorService es = Executors.newCachedThreadPool();
+        ExecutorService es = Executors.newFixedThreadPool(1000);
 
         AtomicBoolean pass = new AtomicBoolean(true);
 

src/test/java/io/supertokens/storage/postgresql/test/InMemoryDBTest.java

@@ -92,6 +92,7 @@ public void checkThatInMemDVWorksEvenIfWrongConfig()
 
         {
             String[] args = { "../" };
+            StorageLayer.close();
             TestingProcessManager.TestingProcess process = TestingProcessManager.start(args);
             assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STARTED));
 

src/test/java/io/supertokens/storage/postgresql/test/LoggingTest.java

@@ -54,6 +54,7 @@ public void beforeEach() {
     @Test
     public void defaultLogging() throws Exception {
         String[] args = { "../" };
+        StorageLayer.close();
         TestingProcessManager.TestingProcess process = TestingProcessManager.start(args);
 
         assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STARTED));
@@ -148,14 +149,15 @@ public void customLogging() throws Exception {
     @Test
     public void confirmLoggerClosed() throws Exception {
         String[] args = { "../" };
+        StorageLayer.close();
         TestingProcessManager.TestingProcess process = TestingProcessManager.start(args);
 
         assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STARTED));
 
         ch.qos.logback.classic.Logger postgresqlInfo = (ch.qos.logback.classic.Logger) LoggerFactory
-                .getLogger("io.supertokens.storage.postgresql.Info." + process.getProcess().getProcessId());
+                .getLogger("io.supertokens.storage.postgresql.Info");
         ch.qos.logback.classic.Logger postgresqlError = (ch.qos.logback.classic.Logger) LoggerFactory
-                .getLogger("io.supertokens.storage.postgresql.Error." + process.getProcess().getProcessId());
+                .getLogger("io.supertokens.storage.postgresql.Error");
 
         ch.qos.logback.classic.Logger hikariLogger = (Logger) LoggerFactory.getLogger("com.zaxxer.hikari");
 

src/test/java/io/supertokens/storage/postgresql/test/TableCreationTest.java

@@ -18,6 +18,7 @@
 package io.supertokens.storage.postgresql.test;
 
 import io.supertokens.ProcessState;
+import io.supertokens.storageLayer.StorageLayer;
 import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.Rule;
@@ -55,6 +56,8 @@ public void checkingCreationOfNewTable() throws InterruptedException {
         process.kill();
         assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STOPPED));
 
+        StorageLayer.close();
+
         process = TestingProcessManager.start(args);
         assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STARTED));
 

src/test/java/io/supertokens/storage/postgresql/test/Utils.java

@@ -20,6 +20,7 @@
 import io.supertokens.Main;
 import io.supertokens.pluginInterface.PluginInterfaceTesting;
 import io.supertokens.storage.postgresql.Start;
+import io.supertokens.storageLayer.StorageLayer;
 import org.apache.tomcat.util.http.fileupload.FileUtils;
 import org.junit.rules.TestRule;
 import org.junit.rules.TestWatcher;
@@ -31,6 +32,7 @@
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Paths;
+import java.util.Arrays;
 import java.util.List;
 
 abstract class Utils extends Mockito {
@@ -78,15 +80,20 @@ static void reset() {
         Main.makeConsolePrintSilent = true;
         String installDir = "../";
         try {
-            // move from temp folder to installDir
-            ProcessBuilder pb = new ProcessBuilder("cp", "temp/licenseKey", "./licenseKey");
-            pb.directory(new File(installDir));
-            Process process = pb.start();
-            process.waitFor();
+            // if the default config is not the same as the current config, we must reset the storage layer
+            File ogConfig = new File("../temp/config.yaml");
+            File currentConfig = new File("../config.yaml");
+            if (currentConfig.isFile()) {
+                byte[] ogConfigContent = Files.readAllBytes(ogConfig.toPath());
+                byte[] currentConfigContent = Files.readAllBytes(currentConfig.toPath());
+                if (!Arrays.equals(ogConfigContent, currentConfigContent)) {
+                    StorageLayer.close();
+                }
+            }
 
-            pb = new ProcessBuilder("cp", "temp/config.yaml", "./config.yaml");
+            ProcessBuilder pb = new ProcessBuilder("cp", "temp/config.yaml", "./config.yaml");
             pb.directory(new File(installDir));
-            process = pb.start();
+            Process process = pb.start();
             process.waitFor();
 
             TestingProcessManager.killAll();
@@ -98,6 +105,7 @@ static void reset() {
         } catch (Exception e) {
             e.printStackTrace();
         }
+        System.gc();
     }
 
     static void stopLicenseKeyFromUpdatingToLatest(TestingProcessManager.TestingProcess process) {
@@ -121,6 +129,8 @@ public void write(int b) {
     }
 
     public static void setValueInConfig(String key, String value) throws FileNotFoundException, IOException {
+        // we close the storage layer since there might be a change in the db related config.
+        StorageLayer.close();
 
         String oldStr = "((#\\s)?)" + key + "(:|((:\\s).+))\n";
         String newStr = key + ": " + value + "\n";
@@ -139,6 +149,9 @@ public static void setValueInConfig(String key, String value) throws FileNotFoun
     }
 
     public static void commentConfigValue(String key) throws IOException {
+        // we close the storage layer since there might be a change in the db related config.
+        StorageLayer.close();
+
         String oldStr = "((#\\s)?)" + key + "(:|((:\\s).+))\n";
         String newStr = "# " + key + ":";