Fix infinite loop due to keepAlive

Fix ref.exists not checking ancestor containers.

fixes #2177
fixes #2044

packages/riverpod/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 - Fix "pending timer" issue inside tests when using `ref.keepAlive()`.
 - Fix `Ref.invalidate`/`Ref.refresh` not throwing on circular dependency.
+- Fix an infinite loop caused by `ref.keepAlive` if the `KeepAliveLink` is immediately closed.
+- Fix `container.exists(provider)` on nested containers not checking their
+  parent containers.
 
 ## 2.4.8 - 2023-11-20
 

packages/riverpod/lib/src/framework/container.dart

@@ -243,7 +243,7 @@ class ProviderContainer implements Node {
 
   /// {@macro riverpod.exists}
   bool exists(ProviderBase<Object?> provider) {
-    final element = _stateReaders[provider]?._element;
+    final element = _getOrNull(provider)?._element;
 
     return element != null;
   }
@@ -306,9 +306,9 @@ class ProviderContainer implements Node {
   /// {@macro riverpod.invalidate}
   void invalidate(ProviderOrFamily provider) {
     if (provider is ProviderBase) {
-      final reader = _getStateReader(provider._origin);
+      final reader = _getOrNull(provider);
 
-      reader._element?.invalidateSelf();
+      reader?._element?.invalidateSelf();
     } else {
       provider as Family;
 
@@ -329,18 +329,21 @@ class ProviderContainer implements Node {
   }
 
   void _disposeProvider(ProviderBase<Object?> provider) {
-    final element = readProviderElement(provider);
-    element.dispose();
+    final reader = _getOrNull(provider);
+    // The provider is already disposed, so we don't need to do anything
+    if (reader == null) return;
 
-    final reader = _stateReaders[element._origin]!;
+    reader._element?.dispose();
 
     if (reader.isDynamicallyCreated) {
       // Since the StateReader is implicitly created, we don't keep it
       // on provider dispose, to avoid memory leak
 
       void removeStateReaderFrom(ProviderContainer container) {
-        if (container._stateReaders[element._origin] == reader) {
-          container._stateReaders.remove(element._origin);
+        /// Checking if the reader is the same instance is important,
+        /// as it is possible that the provider was overridden.
+        if (container._stateReaders[provider] == reader) {
+          container._stateReaders.remove(provider);
         }
         container._children.forEach(removeStateReaderFrom);
       }
@@ -439,7 +442,7 @@ class ProviderContainer implements Node {
       );
     }
 
-    final reader = _getStateReader(provider);
+    final reader = _putIfAbsent(provider);
 
     assert(
       () {
@@ -493,7 +496,20 @@ final b = Provider((ref) => ref.watch(a), dependencies: [a]);
     return reader.getElement() as ProviderElementBase<State>;
   }
 
-  _StateReader _getStateReader(ProviderBase<Object?> provider) {
+  /// Obtains a [_StateReader] for a provider, but do not create it if it does
+  /// not exist.
+  _StateReader? _getOrNull(ProviderBase<Object?> provider) {
+    return _stateReaders[provider] ??
+
+        /// No need to check "parent". We can directly check "root", because
+        /// if the provider is not in the root, it must have been overridden.
+        /// In which case, it is guaranteed to be in the current container already.
+        _root?._getOrNull(provider);
+  }
+
+  /// Create a [_StateReader] for a provider if it does not exist.
+  /// If one already exists, returns it.
+  _StateReader _putIfAbsent(ProviderBase<Object?> provider) {
     final currentReader = _stateReaders[provider];
     if (currentReader != null) return currentReader;
 

packages/riverpod/test/framework/auto_dispose_test.dart

@@ -49,6 +49,38 @@ Future<void> main() async {
   });
 
   group('ref.keepAlive', () {
+    test('Does not cause an infinite loop if aborted directly in the callback',
+        () async {
+      final container = createContainer();
+      var buildCount = 0;
+      var disposeCount = 0;
+      final provider = Provider.autoDispose<String>((ref) {
+        buildCount++;
+        ref.onDispose(() => disposeCount++);
+        final link = ref.keepAlive();
+        link.close();
+        return 'value';
+      });
+
+      container.read(provider);
+
+      expect(buildCount, 1);
+      expect(disposeCount, 0);
+      expect(
+        container.getAllProviderElements().map((e) => e.provider),
+        [provider],
+      );
+
+      await container.pump();
+
+      expect(buildCount, 1);
+      expect(disposeCount, 1);
+      expect(
+        container.getAllProviderElements().map((e) => e.provider),
+        isEmpty,
+      );
+    });
+
     test('when the provider rebuilds, links are cleared', () async {
       final container = createContainer();
       final dep = StateProvider((ref) => 0);
@@ -354,15 +386,15 @@ final alwaysAlive = Provider((ref) {
         return 0;
       },
     );
-    final isDependendingOnDependency = StateProvider(
-      name: 'isDependendingOnDependency',
+    final isDependingOnDependency = StateProvider(
+      name: 'isDependingOnDependency',
       (ref) => true,
     );
     final provider = Provider.autoDispose(
       name: 'provider',
       (ref) {
         ref.maintainState = true;
-        if (ref.watch(isDependendingOnDependency)) {
+        if (ref.watch(isDependingOnDependency)) {
           ref.watch(dependency);
         }
       },
@@ -376,19 +408,19 @@ final alwaysAlive = Provider((ref) {
       unorderedEquals(<Object>[
         dependency,
         provider,
-        isDependendingOnDependency,
+        isDependingOnDependency,
       ]),
     );
 
-    container.read(isDependendingOnDependency.notifier).state = false;
+    container.read(isDependingOnDependency.notifier).state = false;
     await container.pump();
 
     expect(dependencyDisposeCount, 1);
     expect(
       container.getAllProviderElements().map((e) => e.provider),
       unorderedEquals(<Object>[
         provider,
-        isDependendingOnDependency,
+        isDependingOnDependency,
       ]),
     );
   });

packages/riverpod/test/framework/provider_element_test.dart

@@ -8,6 +8,17 @@ import '../utils.dart';
 
 void main() {
   group('Ref.exists', () {
+    test('Returns true if available on ancestor container', () {
+      final root = createContainer();
+      final container = createContainer(parent: root);
+      final provider = Provider((ref) => 0);
+
+      root.read(provider);
+
+      expect(container.exists(provider), true);
+      expect(root.exists(provider), true);
+    });
+
     test('simple use-case', () {
       final container = createContainer();
       final provider = Provider((ref) => 0);

packages/riverpod/test/framework/provider_reference_test.dart → packages/riverpod/test/framework/ref_test.dart

@@ -66,8 +66,6 @@ void main() {
       });
     });
 
-    test('ref.read should keep providers alive', () {}, skip: true);
-
     group('listen', () {
       test('ref.listen on outdated provider causes it to rebuild', () {
         final dep = StateProvider((ref) => 0);