fix(storage, apple): clean up event channel, stream handler and task …

…on completion (#16708)
firebase · Dec 4, 2024 · 14b4a55 · 14b4a55
1 parent 03e85ae
commit 14b4a55
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 10 deletions.
diff --git a/packages/firebase_storage/firebase_storage/ios/Classes/FLTFirebaseStoragePlugin.h b/packages/firebase_storage/firebase_storage/ios/Classes/FLTFirebaseStoragePlugin.h
@@ -19,4 +19,5 @@
 
 + (NSDictionary *)parseTaskSnapshot:(FIRStorageTaskSnapshot *)snapshot;
 + (NSDictionary *)NSDictionaryFromNSError:(NSError *)error;
+- (void)cleanUpTask:(NSString *)channelName handle:(NSNumber *)handle;
 @end
diff --git a/packages/firebase_storage/firebase_storage/ios/Classes/FLTFirebaseStoragePlugin.m b/packages/firebase_storage/firebase_storage/ios/Classes/FLTFirebaseStoragePlugin.m
@@ -412,10 +412,10 @@ - (void)referencePutDataApp:(PigeonStorageFirebaseApp *)app
     self->_tasks[handle] = task;
   }
 
-  completion([self setupTaskListeners:task], nil);
+  completion([self setupTaskListeners:task handle:handle], nil);
 }
 
-- (NSString *)setupTaskListeners:(FIRStorageObservableTask *)task {
+- (NSString *)setupTaskListeners:(FIRStorageObservableTask *)task handle:(NSNumber *)handle {
   // Generate a random UUID to register with
   NSString *uuid = [[NSUUID UUID] UUIDString];
 
@@ -426,7 +426,10 @@ - (NSString *)setupTaskListeners:(FIRStorageObservableTask *)task {
   FlutterEventChannel *channel = [FlutterEventChannel eventChannelWithName:channelName
                                                            binaryMessenger:_binaryMessenger];
   FLTTaskStateChannelStreamHandler *handler =
-      [[FLTTaskStateChannelStreamHandler alloc] initWithTask:task];
+      [[FLTTaskStateChannelStreamHandler alloc] initWithTask:task
+                                               storagePlugin:self
+                                                 channelName:channelName
+                                                      handle:handle];
   [channel setStreamHandler:handler];
 
   [_eventChannels setObject:channel forKey:channelName];
@@ -435,6 +438,23 @@ - (NSString *)setupTaskListeners:(FIRStorageObservableTask *)task {
   return uuid;
 }
 
+- (void)cleanUpTask:(NSString *)channelName handle:(NSNumber *)handle {
+  NSObject<FlutterStreamHandler> *handler = [_streamHandlers objectForKey:channelName];
+  if (handler) {
+    [_streamHandlers removeObjectForKey:channelName];
+  }
+
+  FlutterEventChannel *channel = [_eventChannels objectForKey:channelName];
+  if (channel) {
+    [channel setStreamHandler:nil];
+    [_eventChannels removeObjectForKey:channelName];
+  }
+
+  @synchronized(self->_tasks) {
+    [self->_tasks removeObjectForKey:handle];
+  }
+}
+
 - (void)referencePutStringApp:(PigeonStorageFirebaseApp *)app
                     reference:(PigeonStorageReference *)reference
                          data:(NSString *)data
@@ -456,7 +476,7 @@ - (void)referencePutStringApp:(PigeonStorageFirebaseApp *)app
     self->_tasks[handle] = task;
   }
 
-  completion([self setupTaskListeners:task], nil);
+  completion([self setupTaskListeners:task handle:handle], nil);
 }
 
 - (void)referencePutFileApp:(PigeonStorageFirebaseApp *)app
@@ -481,7 +501,7 @@ - (void)referencePutFileApp:(PigeonStorageFirebaseApp *)app
     self->_tasks[handle] = task;
   }
 
-  completion([self setupTaskListeners:task], nil);
+  completion([self setupTaskListeners:task handle:handle], nil);
 }
 
 - (void)referenceDownloadFileApp:(PigeonStorageFirebaseApp *)app
@@ -501,7 +521,7 @@ - (void)referenceDownloadFileApp:(PigeonStorageFirebaseApp *)app
     self->_tasks[handle] = task;
   }
 
-  completion([self setupTaskListeners:task], nil);
+  completion([self setupTaskListeners:task handle:handle], nil);
 }
 
 - (void)referenceUpdateMetadataApp:(PigeonStorageFirebaseApp *)app

diff --git a/packages/firebase_storage/firebase_storage/ios/Classes/FLTTaskStateChannelStreamHandler.h b/packages/firebase_storage/firebase_storage/ios/Classes/FLTTaskStateChannelStreamHandler.h
@@ -17,8 +17,10 @@
 NS_ASSUME_NONNULL_BEGIN
 
 @interface FLTTaskStateChannelStreamHandler : NSObject <FlutterStreamHandler>
-
-- (instancetype)initWithTask:(FIRStorageObservableTask *)task;
+- (instancetype)initWithTask:(FIRStorageObservableTask *)task
+               storagePlugin:(FLTFirebaseStoragePlugin *)storagePlugin
+                 channelName:(NSString *)channelName
+                      handle:(NSNumber *)handle;
 
 @end
 

diff --git a/packages/firebase_storage/firebase_storage/ios/Classes/FLTTaskStateChannelStreamHandler.m b/packages/firebase_storage/firebase_storage/ios/Classes/FLTTaskStateChannelStreamHandler.m
@@ -9,17 +9,25 @@
 
 @implementation FLTTaskStateChannelStreamHandler {
   FIRStorageObservableTask *_task;
-
+  FLTFirebaseStoragePlugin *_storagePlugin;
+  NSString *_channelName;
+  NSNumber *_handle;
   FIRStorageHandle successHandle;
   FIRStorageHandle failureHandle;
   FIRStorageHandle pausedHandle;
   FIRStorageHandle progressHandle;
 }
 
-- (instancetype)initWithTask:(FIRStorageObservableTask *)task {
+- (instancetype)initWithTask:(FIRStorageObservableTask *)task
+               storagePlugin:(FLTFirebaseStoragePlugin *)storagePlugin
+                 channelName:(NSString *)channelName
+                      handle:(NSNumber *)handle {
   self = [super init];
   if (self) {
     _task = task;
+    _storagePlugin = storagePlugin;
+    _channelName = channelName;
+    _handle = handle;
   }
   return self;
 }
@@ -98,6 +106,10 @@ - (FlutterError *)onCancelWithArguments:(id)arguments {
   }
   progressHandle = nil;
 
+  if (_storagePlugin) {
+    [_storagePlugin cleanUpTask:_channelName handle:_handle];
+  }
+
   return nil;
 }