fix(connections): serialize concurrent pool acquire for the same key (#275)

lib/core/database/connection_pool_lock.dart

@@ -0,0 +1,40 @@
+import 'dart:async';
+
+/// Serializes the synchronous check-and-set for creating a pool entry per key.
+///
+/// Multiple callers waiting for the same key all receive the same creation
+/// [Future], so only one underlying connection is produced. Different keys can
+/// still be created concurrently because the lock is only held while the map
+/// of pending futures is inspected/updated.
+class PoolEntryLock<T> {
+  final Map<String, Future<T>> _pending = {};
+  Future<void>? _lock;
+
+  /// Returns a [Future] that resolves to the created value for [key].
+  /// If a creation for [key] is already in progress, the existing future is
+  /// returned. Otherwise, [create] is started and its future is stored.
+  Future<T> createIfAbsent(String key, Future<T> Function() create) async {
+    // Wait for any other caller that is currently updating the pending map.
+    while (_lock != null) {
+      await _lock;
+    }
+    final completer = Completer<void>();
+    _lock = completer.future;
+    try {
+      final existing = _pending[key];
+      if (existing != null) return existing;
+      final future = create();
+      _pending[key] = future;
+      // Ensure the pending entry is removed once the creation finishes, and
+      // swallow errors on the cleanup chain so they don't become unhandled.
+      final guarded = future.whenComplete(() => _pending.remove(key));
+      guarded.then((_) {}, onError: (_) {});
+      return future;
+    } finally {
+      completer.complete();
+      if (_lock == completer.future) {
+        _lock = null;
+      }
+    }
+  }
+}

lib/core/database/mysql_connection_pool.dart

@@ -1,5 +1,6 @@
 import 'dart:async';
 
+import 'package:querya_desktop/core/database/connection_pool_lock.dart';
 import 'package:querya_desktop/core/database/mysql_connection.dart';
 import 'package:querya_desktop/core/storage/local_db.dart';
 
@@ -48,6 +49,7 @@ class MysqlConnectionPool {
   final int maxEntries;
 
   final Map<String, _PoolEntry> _pool = {};
+  final PoolEntryLock<MysqlConnection> _creationLock = PoolEntryLock();
 
   String keyFor(int? id, String database, MysqlSessionMode mode) =>
       '${id ?? 0}::$database::${mode.name}';
@@ -73,12 +75,25 @@ class MysqlConnectionPool {
       return MysqlLease._(this, k, entry.connection);
     }
 
-    _evictIfNeededBeforeNewSlot();
+    await _creationLock.createIfAbsent(k, () async {
+      _evictIfNeededBeforeNewSlot();
+      final conn = await createAndConnect(row, database: database, mode: mode);
+      _pool[k] = _PoolEntry(conn);
+      return conn;
+    });
 
-    final conn = await createAndConnect(row, database: database, mode: mode);
-    entry = _PoolEntry(conn)..refs = 1;
-    _pool[k] = entry;
-    return MysqlLease._(this, k, conn);
+    entry = _pool[k]!;
+    entry.touch();
+    entry.idleTimer?.cancel();
+    entry.idleTimer = null;
+    entry.refs++;
+    if (!entry.connection.isConnected) {
+      await entry.connection.connect();
+      await entry.connection.setSessionReadOnly(
+        mode == MysqlSessionMode.readOnly,
+      );
+    }
+    return MysqlLease._(this, k, entry.connection);
   }
 
   void _evictIfNeededBeforeNewSlot() {

lib/core/database/postgres_connection_pool.dart

@@ -1,5 +1,6 @@
 import 'dart:async';
 
+import 'package:querya_desktop/core/database/connection_pool_lock.dart';
 import 'package:querya_desktop/core/database/postgres_connection.dart';
 import 'package:querya_desktop/core/storage/local_db.dart';
 
@@ -60,6 +61,7 @@ class PostgresConnectionPool {
   final int maxEntries;
 
   final Map<String, _PoolEntry> _pool = {};
+  final PoolEntryLock<PostgresConnection> _creationLock = PoolEntryLock();
 
   String keyFor(int? id, String database, PgSessionMode mode) =>
       '${id ?? 0}::$database::${mode.name}';
@@ -85,13 +87,15 @@ class PostgresConnectionPool {
       return PgLease._(this, k, entry.connection);
     }
 
-    _evictIfNeededBeforeNewSlot();
-
     try {
-      final conn = await createAndConnect(row, database: database, mode: mode);
-      entry = _PoolEntry(conn)..refs = 1;
-      _pool[k] = entry;
-      return PgLease._(this, k, conn);
+      await _creationLock.createIfAbsent(k, () async {
+        _evictIfNeededBeforeNewSlot();
+        final conn = await createAndConnect(row, database: database, mode: mode);
+        _pool[k] = _PoolEntry(conn);
+        return conn;
+      });
+    } on StateError {
+      rethrow;
     } on PostgresConnectionException {
       rethrow;
     } catch (e, st) {
@@ -104,6 +108,18 @@ class PostgresConnectionPool {
         st,
       );
     }
+
+    entry = _pool[k]!;
+    entry.touch();
+    entry.idleTimer?.cancel();
+    entry.idleTimer = null;
+    entry.refs++;
+    if (!entry.connection.isConnected) {
+      await entry.connection.connect();
+      await entry.connection
+          .setSessionReadOnly(mode == PgSessionMode.readOnly);
+    }
+    return PgLease._(this, k, entry.connection);
   }
 
   /// Drops idle LRU slots until there is room for one more key.

lib/core/database/sqlite_connection_pool.dart

@@ -1,4 +1,6 @@
 import 'dart:async';
+
+import 'package:querya_desktop/core/database/connection_pool_lock.dart';
 import 'package:querya_desktop/core/database/sqlite_connection.dart';
 import 'package:querya_desktop/core/storage/local_db.dart';
 
@@ -59,6 +61,7 @@ class SqliteConnectionPool {
   final int maxEntries;
 
   final Map<String, _PoolEntry> _pool = {};
+  final PoolEntryLock<SqliteConnection> _creationLock = PoolEntryLock();
 
   String keyFor(int? id, SqliteSessionMode mode) =>
       '${id ?? 0}::${mode.name}';
@@ -80,12 +83,22 @@ class SqliteConnectionPool {
       return SqliteLease._(this, k, entry.connection);
     }
 
-    _evictIfNeededBeforeNewSlot();
-
-    final conn = await createAndConnect(row, mode: mode);
-    entry = _PoolEntry(conn)..refs = 1;
-    _pool[k] = entry;
-    return SqliteLease._(this, k, conn);
+    await _creationLock.createIfAbsent(k, () async {
+      _evictIfNeededBeforeNewSlot();
+      final conn = await createAndConnect(row, mode: mode);
+      _pool[k] = _PoolEntry(conn);
+      return conn;
+    });
+
+    entry = _pool[k]!;
+    entry.touch();
+    entry.idleTimer?.cancel();
+    entry.idleTimer = null;
+    entry.refs++;
+    if (!entry.connection.isConnected) {
+      await entry.connection.connect();
+    }
+    return SqliteLease._(this, k, entry.connection);
   }
 
   void _evictIfNeededBeforeNewSlot() {

test/core/database/postgres_connection_pool_test.dart

@@ -123,6 +123,61 @@ void main() {
     });
   });
 
+  group('PostgresConnectionPool concurrent acquire', () {
+    test('only one factory call for the same key when racing', () async {
+      int factoryCalls = 0;
+      Future<PostgresConnection> factory(
+        ConnectionRow row, {
+        required String database,
+        required PgSessionMode mode,
+      }) async {
+        factoryCalls++;
+        await Future<void>.delayed(const Duration(milliseconds: 50));
+        final c = FakePostgresConnection();
+        await c.connect();
+        return c;
+      }
+
+      final pool = PostgresConnectionPool(createAndConnect: factory);
+      final r = _row();
+      final f1 = pool.acquire(r, database: 'postgres');
+      final f2 = pool.acquire(r, database: 'postgres');
+      final l1 = await f1;
+      final l2 = await f2;
+
+      expect(identical(l1.connection, l2.connection), isTrue);
+      expect(factoryCalls, 1);
+      l1.release();
+      l2.release();
+    });
+
+    test('different keys are created concurrently', () async {
+      int factoryCalls = 0;
+      Future<PostgresConnection> factory(
+        ConnectionRow row, {
+        required String database,
+        required PgSessionMode mode,
+      }) async {
+        factoryCalls++;
+        await Future<void>.delayed(const Duration(milliseconds: 30));
+        final c = FakePostgresConnection(id: row.id ?? 0);
+        await c.connect();
+        return c;
+      }
+
+      final pool = PostgresConnectionPool(createAndConnect: factory);
+      final f1 = pool.acquire(_row(id: 1), database: 'postgres');
+      final f2 = pool.acquire(_row(id: 2), database: 'postgres');
+      final l1 = await f1;
+      final l2 = await f2;
+
+      expect(identical(l1.connection, l2.connection), isFalse);
+      expect(factoryCalls, 2);
+      l1.release();
+      l2.release();
+    });
+  });
+
   group('PostgresConnectionPool refcount & reuse', () {
     test('second acquire reuses same connection without new factory', () async {
       FakePostgresConnection? sole;