Unregister SQLite JDBC driver to avoid memory leaks.

The SQLite library used by Room registers the JDBC driver in a static block and thus conveniently has the driver available but never unregister it. In a normal application, this would be fine since they tend to have a single class loader, but in Gradle builds, multiple class loaders are used and the driver is registered multiple times without ever being unregistered, this can lead to memory leaks.

To avoid leaks, Room tries to manage the life of the driver by unregistering it once processing is done. In the case that the same instance of Room is used to do new processing, Room re-registers the driver since using the same instance of Room would mean the same class loader and JDBC's static block and driver registering logic would not execute. Re-register a driver that is already present is a no-op.

See: google/ksp#1063

Test: Verified manually with sample project in linked bug.
Change-Id: I8916fd0bcb42337314feebef9afa0b54d4f479bc

room/room-compiler/src/main/kotlin/androidx/room/RoomKspProcessor.kt

@@ -17,9 +17,12 @@
 package androidx.room
 
 import androidx.room.DatabaseProcessingStep.Companion.ENV_CONFIG
+import androidx.room.compiler.processing.XProcessingEnv
+import androidx.room.compiler.processing.XRoundEnv
 import androidx.room.compiler.processing.ksp.KspBasicAnnotationProcessor
 import androidx.room.processor.Context.BooleanProcessorOptions.USE_NULL_AWARE_CONVERTER
 import androidx.room.processor.ProcessorErrors
+import androidx.room.verifier.DatabaseVerifier
 import com.google.devtools.ksp.processing.SymbolProcessor
 import com.google.devtools.ksp.processing.SymbolProcessorEnvironment
 import com.google.devtools.ksp.processing.SymbolProcessorProvider
@@ -51,6 +54,12 @@ class RoomKspProcessor(
         DatabaseProcessingStep()
     )
 
+    override fun postRound(env: XProcessingEnv, round: XRoundEnv) {
+        if (round.isProcessingOver) {
+            DatabaseVerifier.cleanup()
+        }
+    }
+
     class Provider : SymbolProcessorProvider {
         override fun create(environment: SymbolProcessorEnvironment): SymbolProcessor {
             return RoomKspProcessor(environment)

room/room-compiler/src/main/kotlin/androidx/room/RoomProcessor.kt

@@ -17,10 +17,13 @@
 package androidx.room
 
 import androidx.room.DatabaseProcessingStep.Companion.ENV_CONFIG
+import androidx.room.compiler.processing.XProcessingEnv
+import androidx.room.compiler.processing.XRoundEnv
 import androidx.room.compiler.processing.javac.JavacBasicAnnotationProcessor
 import androidx.room.processor.Context
 import androidx.room.processor.ProcessorErrors
 import androidx.room.util.SimpleJavaVersion
+import androidx.room.verifier.DatabaseVerifier
 import androidx.room.vo.Warning
 import javax.lang.model.SourceVersion
 
@@ -113,4 +116,10 @@ class RoomProcessor : JavacBasicAnnotationProcessor({
     override fun getSupportedSourceVersion(): SourceVersion {
         return SourceVersion.latest()
     }
+
+    override fun postRound(env: XProcessingEnv, round: XRoundEnv) {
+        if (round.isProcessingOver) {
+            DatabaseVerifier.cleanup()
+        }
+    }
 }

room/room-compiler/src/main/kotlin/androidx/room/verifier/DatabaseVerifier.kt

@@ -24,12 +24,14 @@ import androidx.room.vo.EntityOrView
 import androidx.room.vo.FtsEntity
 import androidx.room.vo.FtsOptions
 import androidx.room.vo.Warning
-import org.sqlite.JDBC
-import org.sqlite.SQLiteJDBCLoader
 import java.io.File
 import java.sql.Connection
+import java.sql.Driver
+import java.sql.DriverManager
 import java.sql.SQLException
 import java.util.regex.Pattern
+import org.sqlite.JDBC
+import org.sqlite.SQLiteJDBCLoader
 
 /**
  * Builds an in-memory version of the database and verifies the queries against it.
@@ -59,6 +61,13 @@ class DatabaseVerifier private constructor(
             "\\s+COLLATE\\s+(LOCALIZED|UNICODE)", Pattern.CASE_INSENSITIVE
         )
 
+        /**
+         * Keep a reference to the SQLite JDBC driver so we can re-register it in the case that Room
+         * finishes processing, cleans up and unregisters the driver but is started again within the
+         * same class loader such that JDBC's static block and driver registration does not occur.
+         */
+        private val DRIVER: Driver
+
         init {
             verifyTempDir()
             // Synchronize on a bootstrap loaded class so that parallel runs of Room in the same JVM
@@ -69,6 +78,10 @@ class DatabaseVerifier private constructor(
             synchronized(System::class.java) {
                 SQLiteJDBCLoader.initialize() // extract and loads native library
                 JDBC.isValidURL(CONNECTION_URL) // call to register driver
+                DRIVER = DriverManager.getDriver("jdbc:sqlite:") // get registered driver
+                check(DRIVER is JDBC) {
+                    "Expected driver to be a '${JDBC::class.java}' but was '${DRIVER::class.java}'"
+                }
             }
         }
 
@@ -103,6 +116,8 @@ class DatabaseVerifier private constructor(
             views: List<DatabaseView>
         ): DatabaseVerifier? {
             try {
+                // Re-register driver in case it was unregistered, this is a no-op is already there.
+                DriverManager.registerDriver(DRIVER)
                 val connection = JDBC.createConnection(CONNECTION_URL, java.util.Properties())
                 return DatabaseVerifier(connection, context, entities, views)
             } catch (ex: Exception) {
@@ -113,6 +128,20 @@ class DatabaseVerifier private constructor(
                 return null
             }
         }
+
+        /**
+         * Unregisters the SQLite JDBC driver used by the verifier.
+         *
+         * This is necessary since the driver is statically registered and never unregistered and
+         * can cause class loader leaks. See https://github.com/google/ksp/issues/1063.
+         */
+        fun cleanup() {
+            try {
+                DriverManager.deregisterDriver(DRIVER)
+            } catch (ignored: SQLException) {
+                // Driver was not found
+            }
+        }
     }
 
     init {