Make the statement cache a thread local

library/src/androidMain/kotlin/com/eygraber/sqldelight/androidx/driver/StatementsCacheThreadLocal.kt

@@ -0,0 +1,24 @@
+package com.eygraber.sqldelight.androidx.driver
+
+import androidx.collection.LruCache
+import androidx.collection.mutableLongObjectMapOf
+
+internal actual class StatementsCacheThreadLocal actual constructor() {
+  private val allCaches = mutableLongObjectMapOf<LruCache<Int, AndroidxStatement>>()
+  private val caches = ThreadLocal<LruCache<Int, AndroidxStatement>>()
+
+  internal actual fun get(): LruCache<Int, AndroidxStatement>? = caches.get()
+
+  internal actual fun set(cache: LruCache<Int, AndroidxStatement>?) {
+    when(cache) {
+      null -> allCaches.remove(Thread.currentThread().id)
+      else -> allCaches[Thread.currentThread().id] = cache
+    }
+
+    caches.set(cache)
+  }
+
+  internal actual fun getAll() = buildList {
+    allCaches.forEachValue { add(it) }
+  }
+}

library/src/commonMain/kotlin/com/eygraber/sqldelight/androidx/driver/AndroidxSqliteDriver.kt

@@ -23,6 +23,13 @@ internal expect class TransactionsThreadLocal() {
   internal fun set(transaction: Transacter.Transaction?)
 }
 
+internal expect class StatementsCacheThreadLocal() {
+  internal fun get(): LruCache<Int, AndroidxStatement>?
+  internal fun set(cache: LruCache<Int, AndroidxStatement>?)
+
+  internal fun getAll(): List<LruCache<Int, AndroidxStatement>>
+}
+
 /**
  * @param databaseType Specifies the type of the database file
  * (see [Sqlite open documentation](https://www.sqlite.org/c3ref/open.html)).
@@ -89,13 +96,11 @@ public class AndroidxSqliteDriver(
 
   private val transactions = TransactionsThreadLocal()
 
-  private val statementsCaches = mutableMapOf<SQLiteConnection, LruCache<Int, AndroidxStatement>>()
+  private val statementsCaches = StatementsCacheThreadLocal()
 
-  private fun getStatementsCache(connection: SQLiteConnection): LruCache<Int, AndroidxStatement> =
-    statementsCaches.getOrPut(
-      connection,
-    ) {
-      object : LruCache<Int, AndroidxStatement>(configuration.cacheSize) {
+  private fun getStatementsCache(): LruCache<Int, AndroidxStatement> =
+    when(val cache = statementsCaches.get()) {
+      null -> object : LruCache<Int, AndroidxStatement>(configuration.cacheSize) {
         override fun entryRemoved(
           evicted: Boolean,
           key: Int,
@@ -104,7 +109,9 @@ public class AndroidxSqliteDriver(
         ) {
           if(evicted) oldValue.close()
         }
-      }
+      }.also { statementsCaches.set(it) }
+
+      else -> cache
     }
 
   private var skipStatementsCache = true
@@ -231,7 +238,7 @@ public class AndroidxSqliteDriver(
     binders: (SqlPreparedStatement.() -> Unit)?,
     result: AndroidxStatement.() -> T,
   ): QueryResult.Value<T> {
-    val statementsCache = if(!skipStatementsCache) getStatementsCache(connection) else null
+    val statementsCache = if(!skipStatementsCache) getStatementsCache() else null
     var statement: AndroidxStatement? = null
     if(identifier != null && statementsCache != null) {
       statement = statementsCache[identifier]
@@ -328,12 +335,15 @@ public class AndroidxSqliteDriver(
     }
   }
 
+  /**
+   * It is the caller's responsibility to ensure that no threads
+   * are using any of the connections starting from when close is invoked
+   */
   override fun close() {
-    statementsCaches.forEach { (_, cache) ->
+    statementsCaches.getAll().forEach { cache ->
       cache.snapshot().values.forEach { it.close() }
       cache.evictAll()
     }
-    statementsCaches.clear()
     connectionPool.close()
   }
 

library/src/jvmMain/kotlin/com/eygraber/sqldelight/androidx/driver/StatementsCacheThreadLocal.kt

@@ -0,0 +1,24 @@
+package com.eygraber.sqldelight.androidx.driver
+
+import androidx.collection.LruCache
+import androidx.collection.mutableLongObjectMapOf
+
+internal actual class StatementsCacheThreadLocal actual constructor() {
+  private val allCaches = mutableLongObjectMapOf<LruCache<Int, AndroidxStatement>>()
+  private val caches = ThreadLocal<LruCache<Int, AndroidxStatement>>()
+
+  internal actual fun get(): LruCache<Int, AndroidxStatement>? = caches.get()
+
+  internal actual fun set(cache: LruCache<Int, AndroidxStatement>?) {
+    when(cache) {
+      null -> allCaches.remove(Thread.currentThread().threadId())
+      else -> allCaches[Thread.currentThread().threadId()] = cache
+    }
+
+    caches.set(cache)
+  }
+
+  internal actual fun getAll() = buildList {
+    allCaches.forEachValue { add(it) }
+  }
+}

library/src/nativeMain/kotlin/com/eygraber/sqldelight/androidx/driver/StatementsCacheThreadLocal.kt

@@ -0,0 +1,27 @@
+package com.eygraber.sqldelight.androidx.driver
+
+import androidx.collection.LruCache
+import androidx.collection.mutableIntObjectMapOf
+import kotlin.native.concurrent.ThreadLocal
+
+@ThreadLocal
+private object ThreadLocalStatementsCache {
+  val threadLocalMap = mutableIntObjectMapOf<LruCache<Int, AndroidxStatement>>()
+}
+
+internal actual class StatementsCacheThreadLocal actual constructor() {
+  private val threadLocalId = ThreadLocalId.next()
+
+  actual fun get() = ThreadLocalStatementsCache.threadLocalMap[threadLocalId]
+
+  actual fun set(cache: LruCache<Int, AndroidxStatement>?) {
+    when(cache) {
+      null -> ThreadLocalStatementsCache.threadLocalMap.remove(threadLocalId)
+      else -> ThreadLocalStatementsCache.threadLocalMap[threadLocalId] = cache
+    }
+  }
+
+  actual fun getAll(): List<LruCache<Int, AndroidxStatement>> = buildList {
+    ThreadLocalStatementsCache.threadLocalMap.forEachValue { add(it) }
+  }
+}

library/src/nativeMain/kotlin/com/eygraber/sqldelight/androidx/driver/ThreadLocalId.kt

@@ -0,0 +1,8 @@
+package com.eygraber.sqldelight.androidx.driver
+
+import kotlin.concurrent.AtomicInt
+
+internal object ThreadLocalId {
+  val id = AtomicInt(0)
+  fun next(): Int = id.incrementAndGet()
+}

library/src/nativeMain/kotlin/com/eygraber/sqldelight/androidx/driver/TransactionsThreadLocal.kt

@@ -2,19 +2,13 @@ package com.eygraber.sqldelight.androidx.driver
 
 import androidx.collection.mutableIntObjectMapOf
 import app.cash.sqldelight.Transacter
-import kotlin.concurrent.AtomicInt
 import kotlin.native.concurrent.ThreadLocal
 
 @ThreadLocal
 private object ThreadLocalTransactions {
   val threadLocalMap = mutableIntObjectMapOf<Transacter.Transaction>()
 }
 
-private object ThreadLocalId {
-  val id = AtomicInt(0)
-  fun next(): Int = id.incrementAndGet()
-}
-
 internal actual class TransactionsThreadLocal actual constructor() {
   private val threadLocalId = ThreadLocalId.next()