fix(iot): handling iot crashes

harsh62 · harsh62 · commit 3a75df4c7d27 · 2025-09-05T09:34:41.000-04:00
patch AWSStrem dealloc crash
diff --git a/AWSIoT/Internal/AWSIoTMQTTClient.m b/AWSIoT/Internal/AWSIoTMQTTClient.m
@@ -118,8 +118,38 @@ - (instancetype)init {
 
 - (void)dealloc
 {
+    // Invalidate timers first to prevent callbacks
     [self.reconnectTimer invalidate];
+    self.reconnectTimer = nil;
     [self.connectionAgeTimer invalidate];
+    self.connectionAgeTimer = nil;
+    
+    // Cancel stream thread safely
+    if (self.streamsThread) {
+        [self.streamsThread cancelAndDisconnect:YES];
+        self.streamsThread = nil;
+    }
+    
+    // Close websocket if open
+    if (self.webSocket) {
+        [self.webSocket close];
+        self.webSocket = nil;
+    }
+    
+    // Clean up session
+    if (self.session) {
+        self.session.delegate = nil;
+        [self.session close];
+        self.session = nil;
+    }
+    
+    // Clear callbacks and delegates
+    self.connectStatusCallback = nil;
+    self.clientDelegate = nil;
+    
+    // Clear data structures
+    self.ackCallbackDictionary = nil;
+    self.topicListeners = nil;
 }
 
 - (instancetype)initWithDelegate:(id<AWSIoTMQTTClientDelegate>)delegate {
@@ -1262,9 +1292,19 @@ - (void)webSocket:(AWSSRWebSocket *)webSocket didFailWithError:(NSError *)error
     // Also, the webSocket can be set to nil
     [self cleanUpWebsocketOutputStream];
 
-    [self.encoderOutputStream close];
-    [self.webSocket close];
-    self.webSocket = nil;
+    @synchronized(self.encoderOutputStream) {
+        if (self.encoderOutputStream) {
+            [self.encoderOutputStream close];
+            self.encoderOutputStream = nil;
+        }
+    }
+    
+    @synchronized(self.webSocket) {
+        if (self.webSocket) {
+            [self.webSocket close];
+            self.webSocket = nil;
+        }
+    }
     
     // If this is not because of user initated disconnect, setup timer to retry.
     if (!self.userDidIssueDisconnect ) {
@@ -1300,9 +1340,19 @@ - (void)webSocket:(AWSSRWebSocket *)webSocket didCloseWithCode:(NSInteger)code r
     // The WebSocket has closed. The input/output streams can be closed here.
     [self cleanUpWebsocketOutputStream];
 
-    [self.encoderOutputStream close];
-    [self.webSocket close];
-    self.webSocket = nil;
+    @synchronized(self.encoderOutputStream) {
+        if (self.encoderOutputStream) {
+            [self.encoderOutputStream close];
+            self.encoderOutputStream = nil;
+        }
+    }
+    
+    @synchronized(self.webSocket) {
+        if (self.webSocket) {
+            [self.webSocket close];
+            self.webSocket = nil;
+        }
+    }
     
     // If this is not because of user initated disconnect, setup timer to retry.
     if (!self.userDidIssueDisconnect ) {
diff --git a/AWSIoT/Internal/AWSIoTStreamThread.m b/AWSIoT/Internal/AWSIoTStreamThread.m
@@ -22,17 +22,84 @@ @interface AWSIoTStreamThread()
 @property(nonatomic, strong, nullable) NSOutputStream *encoderOutputStream;
 @property(nonatomic, strong, nullable) NSInputStream  *decoderInputStream;
 @property(nonatomic, strong, nullable) NSOutputStream *outputStream;
-@property(nonatomic, strong, nullable) NSTimer *defaultRunLoopTimer;
-@property(nonatomic, strong, nullable) NSRunLoop *runLoopForStreamsThread;
+@property(atomic, strong, nullable) NSTimer *defaultRunLoopTimer;
+@property(atomic, strong, nullable) NSRunLoop *runLoopForStreamsThread;
 @property(nonatomic, assign) NSTimeInterval defaultRunLoopTimeInterval;
-@property(nonatomic, assign) BOOL isRunning;
-@property(nonatomic, assign) BOOL shouldDisconnect;
-@property(nonatomic, strong) dispatch_queue_t serialQueue;
-@property(nonatomic, assign) BOOL didCleanUp;
+@property(atomic, assign) BOOL isRunning;
+@property(atomic, assign) BOOL shouldDisconnect;
+@property(atomic, assign) BOOL didCleanUp;
+@property(atomic, assign) BOOL isDeallocationInProgress;
+
 @end
 
 @implementation AWSIoTStreamThread
 
+- (void)dealloc {
+    _isDeallocationInProgress = YES;
+
+    [[NSNotificationCenter defaultCenter] removeObserver:self];
+    
+    // ALWAYS use minimal cleanup during dealloc to avoid crashes
+    // - Minimal cleanup is safer during object destruction
+    // - Avoids runloop operations that could crash during dealloc
+    // - Ignores shouldDisconnect flag for safety
+    [self performMinimalCleanup];
+}
+
+/**
+ * MINIMAL CLEANUP - Used during dealloc for crash safety
+ * 
+ * What it does:
+ * - Stops timer to prevent callbacks
+ * - Closes all streams and session directly with @synchronized for thread safety
+ * - Clears delegates and blocks
+ * - NO runloop operations (removeFromRunLoop)
+ * - NO conditional logic based on shouldDisconnect
+ * 
+ * Why minimal:
+ * - Safe during object deallocation
+ * - Avoids method calls that could crash
+ * - Simple operations only
+ * - Prevents use-after-free crashes
+ * 
+ * Thread Safety:
+ * - Uses @synchronized per stream object (customer's proven approach)
+ * - Combined with atomic properties for comprehensive protection
+ * - Double-checks stream existence within synchronized blocks
+ */
+- (void)performMinimalCleanup {
+    if (_didCleanUp) {
+        return;
+    }
+    _didCleanUp = YES;
+    
+    // Stop timer to prevent callbacks during/after deallocation
+    if (_defaultRunLoopTimer) {
+        [_defaultRunLoopTimer invalidate];
+        _defaultRunLoopTimer = nil;
+    }
+
+    if (_outputStream) {
+        _outputStream.delegate = nil;
+        _outputStream = nil;
+    }
+    
+    if (_decoderInputStream) {
+        _decoderInputStream = nil;
+    }
+    
+    if (_encoderOutputStream) {
+        _encoderOutputStream = nil;
+    }
+    
+    if (_session) {
+        _session = nil;
+    }
+    
+    // Clear callback to break retain cycles
+    _onStop = nil;
+}
+
 - (nonnull instancetype)initWithSession:(nonnull AWSMQTTSession *)session
                      decoderInputStream:(nonnull NSInputStream *)decoderInputStream
                     encoderOutputStream:(nonnull NSOutputStream *)encoderOutputStream {
@@ -53,8 +120,8 @@ - (instancetype)initWithSession:(nonnull AWSMQTTSession *)session
         _outputStream = outputStream;
         _defaultRunLoopTimeInterval = 10;
         _shouldDisconnect = NO;
-        _serialQueue = dispatch_queue_create("com.amazonaws.iot.streamthread.syncQueue", DISPATCH_QUEUE_SERIAL);
         _didCleanUp = NO;
+        _isDeallocationInProgress = NO;
     }
     return self;
 }
@@ -92,11 +159,16 @@ - (void)main {
     [self.session connectToInputStream:self.decoderInputStream
                           outputStream:self.encoderOutputStream];
 
-    while ([self shouldContinueRunning]) {
-        //This will continue run until the thread is cancelled
-        //Run one cycle of the runloop. This will return after a input source event or timer event is processed
-        [self.runLoopForStreamsThread runMode:NSDefaultRunLoopMode
-                                   beforeDate:[NSDate dateWithTimeIntervalSinceNow:self.defaultRunLoopTimeInterval]];
+    // Add protection against runloop corruption from multiple threads
+    @try {
+        while ([self shouldContinueRunning]) {
+            //This will continue run until the thread is cancelled
+            //Run one cycle of the runloop. This will return after a input source event or timer event is processed
+            [self.runLoopForStreamsThread runMode:NSDefaultRunLoopMode
+                                       beforeDate:[NSDate dateWithTimeIntervalSinceNow:self.defaultRunLoopTimeInterval]];
+        }
+    } @catch (NSException *exception) {
+        AWSDDLogError(@"Exception in runloop execution: %@", exception);
     }
 
     [self cleanUp];
@@ -105,84 +177,126 @@ - (void)main {
 }
 
 - (BOOL)shouldContinueRunning {
-    __block BOOL shouldRun;
-    dispatch_sync(self.serialQueue, ^{
-        shouldRun = self.isRunning && !self.isCancelled && self.defaultRunLoopTimer != nil;
-    });
-    return shouldRun;
+    if (self.isDeallocationInProgress) {
+        return NO;
+    }
+    
+    return self.isRunning && !self.isCancelled && self.defaultRunLoopTimer != nil;
 }
 
 - (void)cancel {
     AWSDDLogVerbose(@"Issued Cancel on thread [%@]", (NSThread *)self);
-    dispatch_sync(self.serialQueue, ^{
-        self.isRunning = NO;
-        [super cancel];
-    });
+    
+    if (self.isDeallocationInProgress) {
+        return;
+    }
+    
+    // Atomic property, no synchronization needed
+    self.isRunning = NO;
+    [super cancel];
 }
 
 - (void)cancelAndDisconnect:(BOOL)shouldDisconnect {
     AWSDDLogVerbose(@"Issued Cancel and Disconnect = [%@] on thread [%@]", shouldDisconnect ? @"YES" : @"NO", (NSThread *)self);
-    dispatch_sync(self.serialQueue, ^{
-        self.shouldDisconnect = shouldDisconnect;
-        self.isRunning = NO;
-        [super cancel];
-    });
+    
+    if (self.isDeallocationInProgress) {
+        return;
+    }
+    
+    // Set flags and cancel - properties are atomic
+    self.shouldDisconnect = shouldDisconnect;
+    self.isRunning = NO;
+    [super cancel];
 }
 
+/**
+ * FULL CLEANUP - Called during normal thread shutdown
+ * 
+ * When used:
+ * - Thread finishes main() execution normally
+ * - Thread is cancelled via cancel() or cancelAndDisconnect()
+ * - App is terminating or going to background
+ * 
+ * What it does:
+ * - Stops timer to prevent callbacks
+ * - Conditionally closes streams based on shouldDisconnect flag
+ * - Properly removes streams from runloop before closing with @synchronized
+ * - Follows original disconnect logic
+ * - Handles onStop callback on main thread
+ * 
+ * Why full cleanup is safe here:
+ * - Object is still valid and not being deallocated
+ * - Safe to perform runloop operations
+ * - Need to respect shouldDisconnect flag for proper behavior
+ * - Can safely dispatch onStop callback to main thread
+ * 
+ * Thread Safety Strategy:
+ * - @synchronized blocks prevent race conditions during stream operations
+ * - Atomic properties ensure consistent state across threads
+ * - Double-checks within sync blocks prevent operating on closed streams
+ * - Combines customer's surgical approach with our comprehensive coverage
+ */
 - (void)cleanUp {
-    dispatch_sync(self.serialQueue, ^{
-        if (self.didCleanUp) {
-            AWSDDLogVerbose(@"Clean up already called for thread: [%@]", (NSThread *)self);
-            return;
-        }
-
-        self.didCleanUp = YES;
-        if (self.defaultRunLoopTimer) {
-            [self.defaultRunLoopTimer invalidate];
-            self.defaultRunLoopTimer = nil;
+    if (self.didCleanUp) {
+        return;
+    }
+    self.didCleanUp = YES;
+    
+    // Stop timer to prevent callbacks
+    if (self.defaultRunLoopTimer) {
+        [self.defaultRunLoopTimer invalidate];
+        self.defaultRunLoopTimer = nil;
+    }
+    
+    // Conditional cleanup based on shouldDisconnect flag
+    if (self.shouldDisconnect) {
+        // Close session first to cleanly terminate MQTT connection
+        if (self.session) {
+            [self.session close];
+            self.session = nil;
         }
-
-        if (self.shouldDisconnect) {
-            // Properly handle session closure first
-            if (self.session) {
-                [self.session close];
-                self.session = nil;
-            }
-
-            // Make sure we handle the streams in a thread-safe way
+        
+        // Properly remove stream from runloop before closing with synchronized access
+        @synchronized(self.outputStream) {
             if (self.outputStream) {
-                // Remove from runLoop first before closing
+                self.outputStream.delegate = nil;
+                // Safe to do runloop operations during normal cleanup
                 if (self.runLoopForStreamsThread) {
-                    [self.outputStream removeFromRunLoop:self.runLoopForStreamsThread
-                                                 forMode:NSDefaultRunLoopMode];
+                    [self.outputStream removeFromRunLoop:self.runLoopForStreamsThread forMode:NSDefaultRunLoopMode];
                 }
-                self.outputStream.delegate = nil;
                 [self.outputStream close];
                 self.outputStream = nil;
             }
-
+        }
+        
+        @synchronized(self.decoderInputStream) {
             if (self.decoderInputStream) {
                 [self.decoderInputStream close];
                 self.decoderInputStream = nil;
             }
-
+        }
+        
+        @synchronized(self.encoderOutputStream) {
             if (self.encoderOutputStream) {
                 [self.encoderOutputStream close];
                 self.encoderOutputStream = nil;
             }
-        } else {
-            AWSDDLogVerbose(@"Skipping disconnect for thread: [%@]", (NSThread *)self);
         }
-
-        // Make sure onStop is called on the main thread to avoid UI issues
+    } else {
+        // Preserve streams/session for potential reuse
+        AWSDDLogVerbose(@"Skipping disconnect for thread: [%@]", (NSThread *)self);
+    }
+    
+    // Handle onStop callback on main thread (skip during deallocation to avoid async operations)
+    if (!self.isDeallocationInProgress) {
         void (^stopBlock)(void) = self.onStop;
         if (stopBlock) {
             self.onStop = nil;
             dispatch_async(dispatch_get_main_queue(), ^{
                 stopBlock();
             });
         }
-    });
+    }
 }
 
 @end
diff --git a/AWSIoTUnitTests/AWSIoTStreamThreadTests.m b/AWSIoTUnitTests/AWSIoTStreamThreadTests.m
