From d9fe71a52b59696b71c7c9416aa5c25bad218d26 Mon Sep 17 00:00:00 2001
From: Elijah Quartey <equartey19@gmail.com>
Date: Fri, 9 Aug 2024 13:57:04 -0600
Subject: [PATCH] fix(datastore): Clear subscriptions on Stop

---
 .../pigeons/NativePluginBindings.kt              |  6 ++++++
 .../Classes/SwiftAmplifyDataStorePlugin.swift    |  3 +++
 .../Classes/pigeons/NativePluginBindings.swift   |  6 ++++++
 .../amplify_datastore/lib/amplify_datastore.dart | 10 ++++++++++
 .../lib/src/native_plugin.g.dart                 | 16 ++++++++++++++++
 .../amplify_datastore/pigeons/native_plugin.dart |  3 +++
 6 files changed, 44 insertions(+)

diff --git a/packages/amplify_datastore/android/src/main/kotlin/com/amazonaws/amplify/amplify_datastore/pigeons/NativePluginBindings.kt b/packages/amplify_datastore/android/src/main/kotlin/com/amazonaws/amplify/amplify_datastore/pigeons/NativePluginBindings.kt
index 71c7a0db1d..0f82a922a5 100644
--- a/packages/amplify_datastore/android/src/main/kotlin/com/amazonaws/amplify/amplify_datastore/pigeons/NativePluginBindings.kt
+++ b/packages/amplify_datastore/android/src/main/kotlin/com/amazonaws/amplify/amplify_datastore/pigeons/NativePluginBindings.kt
@@ -387,6 +387,12 @@ class NativeApiPlugin(private val binaryMessenger: BinaryMessenger) {
       callback()
     }
   }
+  fun onStop(callback: () -> Unit) {
+    val channel = BasicMessageChannel<Any?>(binaryMessenger, "dev.flutter.pigeon.amplify_datastore.NativeApiPlugin.onStop", codec)
+    channel.send(null) {
+      callback()
+    }
+  }
 }
 /**
  * Bridge for calling Amplify from Flutter into Native
diff --git a/packages/amplify_datastore/ios/Classes/SwiftAmplifyDataStorePlugin.swift b/packages/amplify_datastore/ios/Classes/SwiftAmplifyDataStorePlugin.swift
index 7e780985df..7ee5db422f 100644
--- a/packages/amplify_datastore/ios/Classes/SwiftAmplifyDataStorePlugin.swift
+++ b/packages/amplify_datastore/ios/Classes/SwiftAmplifyDataStorePlugin.swift
@@ -234,6 +234,9 @@ public class SwiftAmplifyDataStorePlugin: NSObject, FlutterPlugin, NativeAmplify
             onStart(flutterResult: result)
         case "stop":
             onStop(flutterResult: result)
+            DispatchQueue.main.async {
+                self.nativeApiPlugin.onStop {}
+            }
         default:
             result(FlutterMethodNotImplemented)
         }
diff --git a/packages/amplify_datastore/ios/Classes/pigeons/NativePluginBindings.swift b/packages/amplify_datastore/ios/Classes/pigeons/NativePluginBindings.swift
index 580f662bb9..9f4386ced4 100644
--- a/packages/amplify_datastore/ios/Classes/pigeons/NativePluginBindings.swift
+++ b/packages/amplify_datastore/ios/Classes/pigeons/NativePluginBindings.swift
@@ -401,6 +401,12 @@ class NativeApiPlugin {
       completion()
     }
   }
+  func onStop(completion: @escaping () -> Void) {
+    let channel = FlutterBasicMessageChannel(name: "dev.flutter.pigeon.amplify_datastore.NativeApiPlugin.onStop", binaryMessenger: binaryMessenger, codec: codec)
+    channel.sendMessage(nil) { _ in
+      completion()
+    }
+  }
 }
 /// Bridge for calling Amplify from Flutter into Native
 ///
diff --git a/packages/amplify_datastore/lib/amplify_datastore.dart b/packages/amplify_datastore/lib/amplify_datastore.dart
index bf3f424961..3e4fc8ec4b 100644
--- a/packages/amplify_datastore/lib/amplify_datastore.dart
+++ b/packages/amplify_datastore/lib/amplify_datastore.dart
@@ -375,4 +375,14 @@ class NativeAmplifyApi
       _subscriptionsCache.remove(subscriptionId);
     }
   }
+
+  /// Amplify.DataStore.Stop() callback
+  ///
+  /// Clean up subscriptions on stop.
+  @override
+  Future<void> onStop() async {
+    _subscriptionsCache.forEach((subId, _) async {
+      await unsubscribe(subId);
+    });
+  }
 }
diff --git a/packages/amplify_datastore/lib/src/native_plugin.g.dart b/packages/amplify_datastore/lib/src/native_plugin.g.dart
index d2ffdb76db..49bb3c66b7 100644
--- a/packages/amplify_datastore/lib/src/native_plugin.g.dart
+++ b/packages/amplify_datastore/lib/src/native_plugin.g.dart
@@ -358,6 +358,8 @@ abstract class NativeApiPlugin {
 
   Future<void> unsubscribe(String subscriptionId);
 
+  Future<void> onStop();
+
   static void setup(NativeApiPlugin? api, {BinaryMessenger? binaryMessenger}) {
     {
       final BasicMessageChannel<Object?> channel = BasicMessageChannel<Object?>(
@@ -462,6 +464,20 @@ abstract class NativeApiPlugin {
         });
       }
     }
+    {
+      final BasicMessageChannel<Object?> channel = BasicMessageChannel<Object?>(
+          'dev.flutter.pigeon.amplify_datastore.NativeApiPlugin.onStop', codec,
+          binaryMessenger: binaryMessenger);
+      if (api == null) {
+        channel.setMessageHandler(null);
+      } else {
+        channel.setMessageHandler((Object? message) async {
+          // ignore message
+          await api.onStop();
+          return;
+        });
+      }
+    }
   }
 }
 
diff --git a/packages/amplify_datastore/pigeons/native_plugin.dart b/packages/amplify_datastore/pigeons/native_plugin.dart
index d9c2b69043..8c49e0d3cb 100644
--- a/packages/amplify_datastore/pigeons/native_plugin.dart
+++ b/packages/amplify_datastore/pigeons/native_plugin.dart
@@ -41,6 +41,9 @@ abstract class NativeApiPlugin {
 
   @async
   void unsubscribe(String subscriptionId);
+
+  @async
+  void onStop();
 }
 
 /// Bridge for calling Amplify from Flutter into Native