Clean up ownership of the connection loss handler

2 files changed, 14 insertions, 10 deletions

diff --git a/src/objective-c/GRPCClient/private/GRPCConnectivityMonitor.h b/src/objective-c/GRPCClient/private/GRPCConnectivityMonitor.h
index d9d6c22ae9..2fae410331 100644
--- a/src/objective-c/GRPCClient/private/GRPCConnectivityMonitor.h
+++ b/src/objective-c/GRPCClient/private/GRPCConnectivityMonitor.h
@@ -68,7 +68,7 @@
 @property(nonatomic, strong, null_resettable) dispatch_queue_t queue;
 
 /**
- * Calls handler the next time the connectivity to this instance's host is lost. If this instance is
+ * Calls handler every time the connectivity to this instance's host is lost. If this instance is
  * released before that happens, the handler won't be called.
  * Only one handler is active at a time, so if this method is called again before the previous
  * handler has been called, it might never be called at all (or yes, if it has already been queued).
diff --git a/src/objective-c/GRPCClient/private/GRPCConnectivityMonitor.m b/src/objective-c/GRPCClient/private/GRPCConnectivityMonitor.m
index f66a77c9b2..b4061bd5ef 100644
--- a/src/objective-c/GRPCClient/private/GRPCConnectivityMonitor.m
+++ b/src/objective-c/GRPCClient/private/GRPCConnectivityMonitor.m
@@ -106,13 +106,16 @@
 
 #pragma mark Connectivity Monitor
 
-// Assumes the third argument is a block that accepts SCNetworkReachabilityFlags, and passes the
+// Assumes the third argument is a block that accepts a GRPCReachabilityFlags object, and passes the
 // received ones to it.
 static void PassFlagsToContextInfoBlock(SCNetworkReachabilityRef target,
                                         SCNetworkReachabilityFlags flags,
                                         void *info) {
   #pragma unused (target)
-  ((__bridge void (^)(SCNetworkReachabilityFlags))info)(flags);
+  // This can be called many times with the same info. The info is retained by SCNetworkReachability
+  // while this function is being executed.
+  void (^handler)(GRPCReachabilityFlags *) = (__bridge void (^)(GRPCReachabilityFlags *))info;
+  handler([[GRPCReachabilityFlags alloc] initWithFlags:flags]);
 }
 
 @implementation GRPCConnectivityMonitor {
@@ -147,29 +150,30 @@ static void PassFlagsToContextInfoBlock(SCNetworkReachabilityRef target,
 }
 
 - (void)handleLossWithHandler:(void (^)())handler {
-  __weak typeof(self) weakSelf = self;
   [self startListeningWithHandler:^(GRPCReachabilityFlags *flags) {
     if (!flags.isHostReachable) {
-      [weakSelf stopListening];
       handler();
     }
   }];
 }
 
 - (void)startListeningWithHandler:(void (^)(GRPCReachabilityFlags *))handler {
+  // Copy to ensure the handler block is in the heap (and so can't be deallocated when this method
+  // returns).
+  void (^copiedHandler)(GRPCReachabilityFlags *) = [handler copy];
   SCNetworkReachabilityContext context = {
     .version = 0,
-    .info = (__bridge_retained void *)^(SCNetworkReachabilityFlags flags){
-      // Pass the flags as as Objective-C wrapper.
-      handler([[GRPCReachabilityFlags alloc] initWithFlags:flags]);
-    },
-    .release = CFRelease
+    .info = (__bridge void *)copiedHandler,
+    .retain = CFRetain,
+    .release = CFRelease,
   };
+  // The following will retain context.info, and release it when the callback is set to NULL.
   SCNetworkReachabilitySetCallback(_reachabilityRef, PassFlagsToContextInfoBlock, &context);
   SCNetworkReachabilitySetDispatchQueue(_reachabilityRef, _queue);
 }
 
 - (void)stopListening {
+  // This releases the block on context.info.
   SCNetworkReachabilitySetCallback(_reachabilityRef, NULL, NULL);
   SCNetworkReachabilitySetDispatchQueue(_reachabilityRef, NULL);
 }