From 1dc7e6f67a1f50234f82fb2141c17fe6f0b4c9d1 Mon Sep 17 00:00:00 2001
From: Elijah Quartey <equartey19@gmail.com>
Date: Thu, 8 Aug 2024 09:04:31 -0500
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 |  4 ++++
 6 files changed, 45 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 71c7a0db1d6..0f82a922a53 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 7e780985dfa..03fa96a26c8 100644
--- a/packages/amplify_datastore/ios/Classes/SwiftAmplifyDataStorePlugin.swift
+++ b/packages/amplify_datastore/ios/Classes/SwiftAmplifyDataStorePlugin.swift
@@ -612,6 +612,9 @@ public class SwiftAmplifyDataStorePlugin: NSObject, FlutterPlugin, NativeAmplify
                         error: error,
                         flutterResult: flutterResult)
                 case .success():
+                    DispatchQueue.main.async {
+                        self.nativeApiPlugin.onStop {}
+                    }
                     print("Successfully stopped datastore cloud syncing")
                     flutterResult(nil)
                 }
diff --git a/packages/amplify_datastore/ios/Classes/pigeons/NativePluginBindings.swift b/packages/amplify_datastore/ios/Classes/pigeons/NativePluginBindings.swift
index 580f662bb9f..9f4386ced43 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 bf3f424961a..4ffd264e926 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 d2ffdb76db3..49bb3c66b7e 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 d9c2b69043d..fc7b93ab5e3 100644
--- a/packages/amplify_datastore/pigeons/native_plugin.dart
+++ b/packages/amplify_datastore/pigeons/native_plugin.dart
@@ -41,6 +41,10 @@ abstract class NativeApiPlugin {
 
   @async
   void unsubscribe(String subscriptionId);
+
+
+  @async
+  void onStop();
 }
 
 /// Bridge for calling Amplify from Flutter into Native