Merge pull request #33 from dumganhar/fix/SRWebSocket-memory-leak

Memory leak fix and cancel connection timeout timer in scheduleCleanup
cocos · Jul 20, 2017 · b43a5b3 · b43a5b3
2 parents 87575ef + 4eb591a
commit b43a5b3
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 22 deletions.
diff --git a/sources/SocketRocket/Internal/Delegate/SRDelegateController.h b/sources/SocketRocket/Internal/Delegate/SRDelegateController.h
@@ -33,11 +33,7 @@ typedef void(^SRDelegateBlock)(id<SRWebSocketDelegate> _Nullable delegate, SRDel
 @property (nonatomic, weak) id<SRWebSocketDelegate> delegate;
 @property (atomic, readonly) SRDelegateAvailableMethods availableDelegateMethods;
 
-#if OS_OBJECT_USE_OBJC
 @property (nullable, nonatomic, strong) dispatch_queue_t dispatchQueue;
-#else
-@property (nullable, nonatomic, assign) dispatch_queue_t dispatchQueue;
-#endif
 @property (nullable, nonatomic, strong) NSOperationQueue *operationQueue;
 
 ///--------------------------------------

diff --git a/sources/SocketRocket/Internal/IOConsumer/SRIOConsumerPool.h b/sources/SocketRocket/Internal/IOConsumer/SRIOConsumerPool.h
@@ -25,4 +25,6 @@
                           unmaskBytes:(BOOL)unmaskBytes;
 - (void)returnConsumer:(SRIOConsumer *)consumer;
 
+- (void)clear;
+
 @end
diff --git a/sources/SocketRocket/Internal/IOConsumer/SRIOConsumerPool.m b/sources/SocketRocket/Internal/IOConsumer/SRIOConsumerPool.m
@@ -61,4 +61,10 @@ - (void)returnConsumer:(SRIOConsumer *)consumer;
     }
 }
 
+-(void)clear
+{
+    _poolSize = 0;
+    _bufferedConsumers = nil;
+}
+
 @end
diff --git a/sources/SocketRocket/Internal/Proxy/SRProxyConnect.m b/sources/SocketRocket/Internal/Proxy/SRProxyConnect.m
@@ -234,11 +234,11 @@ - (void)_fetchPAC:(NSURL *)PACurl withProxySettings:(NSDictionary *)proxySetting
         [self _openConnection];
         return;
     }
-    __weak __typeof__(self) wself = self;
+    __weak __typeof(self) wself = self;
     NSURLRequest *request = [NSURLRequest requestWithURL:PACurl];
     NSURLSession *session = [NSURLSession sharedSession];
     [[session dataTaskWithRequest:request completionHandler:^(NSData *data, NSURLResponse *response, NSError *error) {
-        __strong __typeof__(wself) sself = wself;
+        __strong __typeof(wself) sself = wself;
         if (!error) {
             NSString *script = [[NSString alloc] initWithData:data encoding:NSUTF8StringEncoding];
             [sself _runPACScript:script withProxySettings:proxySettings];
@@ -454,9 +454,9 @@ - (void)_writeData:(NSData *)data
 {
     const uint8_t * bytes = data.bytes;
     __block NSInteger timeout = (NSInteger)(SRProxyConnectWriteTimeout * 1000000); // wait timeout before giving up
-    __weak __typeof__(self) wself = self;
+    __weak __typeof(self) wself = self;
     dispatch_async(_writeQueue, ^{
-        __strong __typeof__(wself) sself = self;
+        __strong __typeof(wself) sself = self;
         if (!sself) {
             return;
         }

diff --git a/sources/SocketRocket/SRWebSocket.m b/sources/SocketRocket/SRWebSocket.m
@@ -295,6 +295,14 @@ + (BOOL)automaticallyNotifiesObserversOfReadyState {
     return NO;
 }
 
+-(void)_onTimeout
+{
+    if (self.readyState == SR_CONNECTING) {
+        NSError *error = SRErrorWithDomainCodeDescription(NSURLErrorDomain, NSURLErrorTimedOut, @"Timed out connecting to server.");
+        [self _failWithError:error];
+    }
+}
+
 ///--------------------------------------
 #pragma mark - Open / Close
 ///--------------------------------------
@@ -307,18 +315,12 @@ - (void)open
     _selfRetain = self;
 
     if (_urlRequest.timeoutInterval > 0) {
-        dispatch_time_t popTime = dispatch_time(DISPATCH_TIME_NOW, (int64_t)(_urlRequest.timeoutInterval * NSEC_PER_SEC));
-        dispatch_after(popTime, dispatch_get_main_queue(), ^{
-            if (self.readyState == SR_CONNECTING) {
-                NSError *error = SRErrorWithDomainCodeDescription(NSURLErrorDomain, NSURLErrorTimedOut, @"Timed out connecting to server.");
-                [self _failWithError:error];
-            }
-        });
+        [self performSelector:@selector(_onTimeout) withObject:nil afterDelay:_urlRequest.timeoutInterval];
     }
 
     _proxyConnect = [[SRProxyConnect alloc] initWithURL:_url];
 
-    __weak __typeof__(self) wself = self;
+    __weak __typeof(self) wself = self;
     [_proxyConnect openNetworkStreamWithCompletion:^(NSError *error, NSInputStream *readStream, NSOutputStream *writeStream) {
         [wself _connectionDoneWithError:error readStream:readStream writeStream:writeStream];
     }];
@@ -419,14 +421,23 @@ - (void)_readHTTPHeader;
         _receivedHTTPHeaders = CFHTTPMessageCreateEmpty(NULL, NO);
     }
 
+    // Uses weak self object in the block, otherwise Consumers will retain SRWebSocket instance,
+    // and SRWebSocket instance also hold consumers, cycle reference will occur.
+    __weak __typeof(self) wself = self; 
+
     [self _readUntilHeaderCompleteWithCallback:^(SRWebSocket *socket,  NSData *data) {
-        CFHTTPMessageAppendBytes(_receivedHTTPHeaders, (const UInt8 *)data.bytes, data.length);
 
-        if (CFHTTPMessageIsHeaderComplete(_receivedHTTPHeaders)) {
-            SRDebugLog(@"Finished reading headers %@", CFBridgingRelease(CFHTTPMessageCopyAllHeaderFields(_receivedHTTPHeaders)));
-            [self _HTTPHeadersDidFinish];
+        __strong __typeof(wself) sself = wself;
+        if (sself == nil)
+            return;
+
+        CFHTTPMessageAppendBytes(sself.receivedHTTPHeaders, (const UInt8 *)data.bytes, data.length);
+
+        if (CFHTTPMessageIsHeaderComplete(sself.receivedHTTPHeaders)) {
+            SRDebugLog(@"Finished reading headers %@", CFBridgingRelease(CFHTTPMessageCopyAllHeaderFields(sself.receivedHTTPHeaders)));
+            [sself _HTTPHeadersDidFinish];
         } else {
-            [self _readHTTPHeader];
+            [sself _readHTTPHeader];
         }
     }];
 }
@@ -1127,6 +1138,16 @@ - (void)_scheduleCleanup
 
         _cleanupScheduled = YES;
 
+        // _consumers retain SRWebSocket instance by block copy, if there are consumers here, clear them.
+        [_consumers removeAllObjects];
+        [_consumerPool clear];
+
+        // Cancel the timer which retains SRWebSocket instance.
+        // If we don't cancel the timer, the 'dealloc' method will be invoked only after the time (default: 60s) have come, which may cause memory increase.
+        dispatch_async(dispatch_get_main_queue(), ^(){
+            [NSObject cancelPreviousPerformRequestsWithTarget:self selector:@selector(_onTimeout) object:nil];
+        });
+
         // Cleanup NSStream delegate's in the same RunLoop used by the streams themselves:
         // This way we'll prevent race conditions between handleEvent and SRWebsocket's dealloc
         NSTimer *timer = [NSTimer timerWithTimeInterval:(0.0f) target:self selector:@selector(_cleanupSelfReference:) userInfo:nil repeats:NO];
@@ -1392,7 +1413,7 @@ - (void)_sendFrameWithOpcode:(SROpCode)opCode data:(NSData *)data
 
 - (void)stream:(NSStream *)aStream handleEvent:(NSStreamEvent)eventCode
 {
-    __weak __typeof__(self) wself = self;
+    __weak __typeof(self) wself = self;
 
     if (_requestRequiresSSL && !_streamSecurityValidated &&
         (eventCode == NSStreamEventHasBytesAvailable || eventCode == NSStreamEventHasSpaceAvailable)) {