From 1fa20876912495328cd4e89ef5d3d730c8222749 Mon Sep 17 00:00:00 2001
From: Romain Fliedel <romain.fliedel@gmail.com>
Date: Mon, 19 Mar 2018 18:57:23 +0100
Subject: [PATCH] add ability to limit the size of the send queue buffer.

By default _outputBuffer size is not limited and can grow over to the
point that iOS decides to kill the app.
By setting maxTxQueueSize it's now possible to ensure that the _outputBuffer
size will not consume too much memory.
This is done by blocking the send: methods until _outputBuffer size decreases
(once data has actually been sent over the websocket connedtion).
---
 SocketRocket/SRWebSocket.h |  8 +++++++
 SocketRocket/SRWebSocket.m | 49 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/SocketRocket/SRWebSocket.h b/SocketRocket/SRWebSocket.h
index a3806f1a0..18e221840 100644
--- a/SocketRocket/SRWebSocket.h
+++ b/SocketRocket/SRWebSocket.h
@@ -122,6 +122,14 @@ extern NSString *const SRHTTPResponseErrorKey;
  */
 @property (nonatomic, assign, readonly) BOOL allowsUntrustedSSLCertificates;
 
+/**
+ Allow limiting the internal tx queue size to a specific bytes count limit
+ Once the tx queue size is reached the following calls to send: will block until data is sent
+ over the websocket connection.
+ Setting this to 0 does not limit the tx queue size (default)
+ */
+@property (nonatomic, assign) NSUInteger maxTxQueueSize;
+
 ///--------------------------------------
 #pragma mark - Constructors
 ///--------------------------------------
diff --git a/SocketRocket/SRWebSocket.m b/SocketRocket/SRWebSocket.m
index 4e30aef54..539f72b1c 100644
--- a/SocketRocket/SRWebSocket.m
+++ b/SocketRocket/SRWebSocket.m
@@ -138,6 +138,8 @@ @implementation SRWebSocket {
 
     // proxy support
     SRProxyConnect *_proxyConnect;
+
+    NSCondition *__txQueueSizeCond;
 }
 
 @synthesize readyState = _readyState;
@@ -179,6 +181,8 @@ - (instancetype)initWithURLRequest:(NSURLRequest *)request protocols:(NSArray<NS
     _consumerPool = [[SRIOConsumerPool alloc] init];
 
     _scheduledRunloops = [[NSMutableSet alloc] init];
+    _maxTxQueueSize = 0;
+    __txQueueSizeCond = [[NSCondition alloc] init];
 
     return self;
 }
@@ -250,6 +254,11 @@ - (void)assertOnWorkQueue;
     assert(dispatch_get_specific((__bridge void *)self) == (__bridge void *)_workQueue);
 }
 
+- (void)assertNotOnWorkQueue;
+{
+    assert(dispatch_get_specific((__bridge void *)self) != (__bridge void *)_workQueue);
+}
+
 ///--------------------------------------
 #pragma mark - Dealloc
 ///--------------------------------------
@@ -606,6 +615,39 @@ - (void)send:(nullable id)message
     }
 }
 
+- (void)setMaxTxQueueSize:(NSUInteger)maxTxQueueSize
+{
+    if (maxTxQueueSize > 0)
+        assert(maxTxQueueSize > SRDefaultBufferSize());
+    _maxTxQueueSize = maxTxQueueSize;
+    [__txQueueSizeCond lock];
+    [__txQueueSizeCond signal];
+    [__txQueueSizeCond unlock];
+}
+
+- (void)checkTxQueue
+{
+    if (!_maxTxQueueSize) {
+        /* no limit on the tx queue size */
+        return;
+    }
+
+    [self assertNotOnWorkQueue];
+
+    /* check internal queue size */
+    [__txQueueSizeCond lock];
+    while (true) {
+        /* check queue size */
+        NSUInteger txQueueSize = dispatch_data_get_size(_outputBuffer);
+        if (txQueueSize < _maxTxQueueSize)
+            break;
+
+        /* need to block until data is effectively sent over ws connection */
+        [__txQueueSizeCond wait];
+    }
+    [__txQueueSizeCond unlock];
+}
+
 - (BOOL)sendString:(NSString *)string error:(NSError **)error
 {
     if (self.readyState != SR_OPEN) {
@@ -617,6 +659,8 @@ - (BOOL)sendString:(NSString *)string error:(NSError **)error
         return NO;
     }
 
+    [self checkTxQueue];
+
     string = [string copy];
     dispatch_async(_workQueue, ^{
         [self _sendFrameWithOpcode:SROpCodeTextFrame data:[string dataUsingEncoding:NSUTF8StringEncoding]];
@@ -641,6 +685,8 @@ - (BOOL)sendDataNoCopy:(nullable NSData *)data error:(NSError **)error
         return NO;
     }
 
+    [self checkTxQueue];
+
     dispatch_async(_workQueue, ^{
         if (data) {
             [self _sendFrameWithOpcode:SROpCodeBinaryFrame data:data];
@@ -1074,6 +1120,9 @@ - (void)_pumpWriting;
         if (_outputBufferOffset > SRDefaultBufferSize() && _outputBufferOffset > dataLength / 2) {
             _outputBuffer = dispatch_data_create_subrange(_outputBuffer, _outputBufferOffset, dataLength - _outputBufferOffset);
             _outputBufferOffset = 0;
+            [__txQueueSizeCond lock];
+            [__txQueueSizeCond signal];
+            [__txQueueSizeCond unlock];
         }
     }