From cda9f27a55f14daa95e7bc6146c95df492bc873d Mon Sep 17 00:00:00 2001
From: Simon Binder <oss@simonbinder.eu>
Date: Sat, 4 Jan 2025 14:33:59 +0100
Subject: [PATCH] Fix async lock with fake_async zones

---
 drift/CHANGELOG.md                      |  2 ++
 drift/lib/src/utils/synchronized.dart   | 26 ++++++++++++++++++++--
 drift/pubspec.yaml                      |  1 +
 drift/test/utils/synchronized_test.dart | 29 +++++++++++++++++++++++--
 4 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/drift/CHANGELOG.md b/drift/CHANGELOG.md
index 03ad272ec..67ec3f8dd 100644
--- a/drift/CHANGELOG.md
+++ b/drift/CHANGELOG.md
@@ -4,6 +4,8 @@
   SQLite databases.
 - Don't attempt to roll-back transactions that failed to begin.
 - Fix unhandled exception when cancelling transactions.
+- Fix deadlock when drift databases are used in a `fake_async` Zone and then
+  closed outside that zone.
 
 ## 2.23.0
 
diff --git a/drift/lib/src/utils/synchronized.dart b/drift/lib/src/utils/synchronized.dart
index fda1a0dc4..3a4f5f7c4 100644
--- a/drift/lib/src/utils/synchronized.dart
+++ b/drift/lib/src/utils/synchronized.dart
@@ -11,10 +11,32 @@ class Lock {
     // This completer may not be sync: It must complete just after
     // callBlockAndComplete completes.
     final blockCompleted = Completer<void>();
-    _last = blockCompleted.future;
+    final blockReleasedLock = blockCompleted.future;
+    _last = blockReleasedLock;
 
     Future<T> callBlockAndComplete() {
-      return Future.sync(block).whenComplete(blockCompleted.complete);
+      return Future.sync(block).whenComplete(() {
+        blockCompleted.complete();
+
+        if (identical(_last, blockReleasedLock)) {
+          // There's no subsequent waiter entering the lock now, so we can reset
+          // the entire state.
+          _last = null;
+
+          // This doesn't affect the correctness of the lock, but is helpful
+          // when drift is used in `fake_async` scenarios but then cleaned up
+          // outside of that `fake_async` scope (a very common pattern in
+          // Flutter widget tests).
+          // Waiting on `previous.then` on a completed `previous` future will
+          // schedule a microtask, so if we call synchronized in a zone outside
+          // of fake_async and the lock was previously locked in a fake_async
+          // zone, that microtask might not run if no one completes the pending
+          // fake_async microtasks.
+          // Since the lock is idle anyway, the next waiter can just call
+          // callBlockAndComplete() directly without calling `.then()` on a
+          // future that will no longer notify listeners.
+        }
+      });
     }
 
     if (previous != null) {
diff --git a/drift/pubspec.yaml b/drift/pubspec.yaml
index 990676407..83cf0dc24 100644
--- a/drift/pubspec.yaml
+++ b/drift/pubspec.yaml
@@ -48,3 +48,4 @@ dev_dependencies:
   vm_service: ^15.0.0
   rxdart: ^0.28.0
   build_daemon: ^4.0.3
+  fake_async: ^1.3.0
diff --git a/drift/test/utils/synchronized_test.dart b/drift/test/utils/synchronized_test.dart
index ce09b7dbf..dd3419977 100644
--- a/drift/test/utils/synchronized_test.dart
+++ b/drift/test/utils/synchronized_test.dart
@@ -1,18 +1,43 @@
+import 'dart:async';
+
 import 'package:drift/src/utils/synchronized.dart';
+import 'package:fake_async/fake_async.dart';
 import 'package:test/test.dart';
 
 void main() {
   test('synchronized runs code in sequence', () async {
     final lock = Lock();
     var i = 0;
+    var inSynchronizedBlock = 0;
     final completionOrder = <int>[];
     final futures = List.generate(
         100,
-        (index) => lock.synchronized(() => i++)
-          ..whenComplete(() => completionOrder.add(index)));
+        (index) => lock.synchronized(() async {
+              expect(inSynchronizedBlock, 0);
+              inSynchronizedBlock = 1;
+              await pumpEventQueue();
+              inSynchronizedBlock--;
+              return i++;
+            })
+              ..whenComplete(() => completionOrder.add(index)));
     final results = await Future.wait(futures);
 
     expect(results, List.generate(100, (index) => index));
     expect(completionOrder, List.generate(100, (index) => index));
   });
+
+  test('can wait on lock used in fakeAsync zone', () async {
+    final lock = Lock();
+    final completer = Completer<void>();
+
+    fakeAsync((async) {
+      lock
+          .synchronized(expectAsync0(() async {}))
+          .then((_) => completer.complete());
+      async.flushTimers();
+    });
+
+    await completer.future;
+    await lock.synchronized(() async {});
+  });
 }