fix(database): Run WebSocket keepalive on the connection worker queue (#16767)

Co-authored-by: Andrew Heard <andrewheard@google.com>
Co-authored-by: Paul Beusterien <paulbeusterien@google.com>
Co-authored-by: Nick Cooke <nickcooke@google.com>

FirebaseDatabase/CHANGELOG.md

@@ -2,6 +2,8 @@
 - [fixed] Fixed an issue where the client could stay offline for hours after a
   connection attempt timed out, for example after the device woke from sleep.
   (#9682)
+- [fixed] Keep Realtime Database WebSocket keepalive timers and connection state
+  on the worker queue to avoid cross-thread access. (#16767, #16971)
 
 # 13.1.0
 - [fixed] Fixed a race where calling childByAutoId() concurrently from multiple

FirebaseDatabase/Sources/Realtime/FWebSocketConnection.m

@@ -45,12 +45,14 @@ static NSString *const kGoogleAppIDHeader = @"X-Firebase-GMPID";
     NSMutableString *frame;
     BOOL everConnected;
     BOOL isClosed;
-    NSTimer *keepAlive;
+    dispatch_source_t keepAlive;
+    NSDate *keepAliveFireDate;
 }
 
 - (void)shutdown;
 - (void)onClosed;
 - (void)closeIfNeverConnected;
+- (void)cancelKeepAlive;
 
 #if TARGET_OS_WATCH
 @property(nonatomic, strong) NSURLSessionWebSocketTask *webSocketTask;
@@ -63,13 +65,19 @@ static NSString *const kGoogleAppIDHeader = @"X-Firebase-GMPID";
 @property(nonatomic, readonly) NSString *userAgent;
 @property(nonatomic) dispatch_queue_t dispatchQueue;
 
-- (void)nop:(NSTimer *)timer;
+- (void)nop;
 
 @end
 
 @implementation FWebSocketConnection
 
 @synthesize delegate;
+
+- (void)dealloc {
+    if (keepAlive) {
+        dispatch_source_cancel(keepAlive);
+    }
+}
 #if !TARGET_OS_WATCH
 @synthesize webSocket;
 #endif // !TARGET_OS_WATCH
@@ -112,6 +120,7 @@ static NSString *const kGoogleAppIDHeader = @"X-Firebase-GMPID";
 #if TARGET_OS_WATCH
         // Regular NSURLSession websocket.
         NSOperationQueue *opQueue = [[NSOperationQueue alloc] init];
+        opQueue.maxConcurrentOperationCount = 1;
         opQueue.underlyingQueue = queue;
         NSURLSession *session = [NSURLSession
             sessionWithConfiguration:[NSURLSessionConfiguration
@@ -235,6 +244,7 @@ static NSString *const kGoogleAppIDHeader = @"X-Firebase-GMPID";
     FFLog(@"I-RDB083003", @"(wsc:%@) FWebSocketConnection is being closed.",
           self.connectionId);
     isClosed = YES;
+    [self cancelKeepAlive];
 #if TARGET_OS_WATCH
     [self.webSocketTask
         cancelWithCloseCode:NSURLSessionWebSocketCloseCodeNormalClosure
@@ -276,17 +286,12 @@ static NSString *const kGoogleAppIDHeader = @"X-Firebase-GMPID";
     }
 }
 
-- (void)nop:(NSTimer *)timer {
+- (void)nop {
     if (!isClosed) {
         FFLog(@"I-RDB083004", @"(wsc:%@) nop", self.connectionId);
         // Note: the backend is expecting a string "0" here, not any special
         // ping/pong from build in websocket APIs.
         [self sendStringToWebSocket:@"0"];
-    } else {
-        FFLog(@"I-RDB083005",
-              @"(wsc:%@) No more websocket; invalidating nop timer.",
-              self.connectionId);
-        [timer invalidate];
     }
 }
 
@@ -378,13 +383,13 @@ static NSString *const kGoogleAppIDHeader = @"X-Firebase-GMPID";
                             NSURLSessionWebSocketMessage *_Nullable message,
                             NSError *_Nullable error) {
       __auto_type strongSelf = weakSelf;
-      if (strongSelf == nil) {
+      if (strongSelf == nil || strongSelf->isClosed) {
           return;
       }
 
       if (message) {
           [strongSelf handleIncomingFrame:message.string];
-      } else if (error && !strongSelf->isClosed) {
+      } else if (error) {
           FFWarn(@"I-RDB083020",
                  @"Error received from web socket, closing the connection. %@",
                  error);
@@ -428,16 +433,30 @@ static NSString *const kGoogleAppIDHeader = @"X-Firebase-GMPID";
 
     everConnected = YES;
 
-    dispatch_async(dispatch_get_main_queue(), ^{
-      self->keepAlive =
-          [NSTimer scheduledTimerWithTimeInterval:kWebsocketKeepaliveInterval
-                                           target:self
-                                         selector:@selector(nop:)
-                                         userInfo:nil
-                                          repeats:YES];
-      FFLog(@"I-RDB083009", @"(wsc:%@) nop timer kicked off",
-            self.connectionId);
+    if (isClosed || keepAlive) {
+        return;
+    }
+    keepAlive = dispatch_source_create(DISPATCH_SOURCE_TYPE_TIMER, 0, 0,
+                                       self.dispatchQueue);
+    keepAliveFireDate =
+        [NSDate dateWithTimeIntervalSinceNow:kWebsocketKeepaliveInterval];
+    dispatch_source_set_timer(
+        keepAlive,
+        dispatch_time(DISPATCH_TIME_NOW,
+                      kWebsocketKeepaliveInterval * NSEC_PER_SEC),
+        kWebsocketKeepaliveInterval * NSEC_PER_SEC, NSEC_PER_SEC);
+    __weak FWebSocketConnection *weakSelf = self;
+    dispatch_source_set_event_handler(keepAlive, ^{
+      FWebSocketConnection *strongSelf = weakSelf;
+      if (!strongSelf || strongSelf->isClosed) {
+          return;
+      }
+      strongSelf->keepAliveFireDate =
+          [NSDate dateWithTimeIntervalSinceNow:kWebsocketKeepaliveInterval];
+      [strongSelf nop];
     });
+    dispatch_resume(keepAlive);
+    FFLog(@"I-RDB083009", @"(wsc:%@) nop timer kicked off", self.connectionId);
 }
 
 #pragma mark -
@@ -501,7 +520,11 @@ static NSString *const kGoogleAppIDHeader = @"X-Firebase-GMPID";
 }
 
 - (void)shutdown {
+    if (isClosed) {
+        return;
+    }
     isClosed = YES;
+    [self cancelKeepAlive];
 
     // Call delegate methods
     [self.delegate onDisconnect:self wasEverConnected:everConnected];
@@ -517,23 +540,38 @@ static NSString *const kGoogleAppIDHeader = @"X-Firebase-GMPID";
 #else
     self.webSocket = nil;
 #endif // TARGET_OS_WATCH
-    if (keepAlive.isValid) {
-        [keepAlive invalidate];
-    }
+    [self cancelKeepAlive];
 }
 
 - (void)resetKeepAlive {
+    // Connection callbacks and this timer share the worker queue.
+    if (isClosed || !keepAlive) {
+        return;
+    }
     NSDate *newTime =
         [NSDate dateWithTimeIntervalSinceNow:kWebsocketKeepaliveInterval];
-    // Calling setFireDate is actually kinda' expensive, so wait at least 5
-    // seconds before updating it.
-    if ([newTime timeIntervalSinceDate:keepAlive.fireDate] > 5) {
+    // Calling dispatch_source_set_timer on every frame is unnecessary, so wait
+    // at least 5 seconds before updating it.
+    if ([newTime timeIntervalSinceDate:keepAliveFireDate] > 5) {
         FFLog(@"I-RDB083014", @"(wsc:%@) resetting keepalive, to %@ ; old: %@",
-              self.connectionId, newTime, [keepAlive fireDate]);
-        [keepAlive setFireDate:newTime];
+              self.connectionId, newTime, keepAliveFireDate);
+        dispatch_source_set_timer(
+            keepAlive,
+            dispatch_time(DISPATCH_TIME_NOW,
+                          kWebsocketKeepaliveInterval * NSEC_PER_SEC),
+            kWebsocketKeepaliveInterval * NSEC_PER_SEC, NSEC_PER_SEC);
+        keepAliveFireDate = newTime;
     }
 }
 
+- (void)cancelKeepAlive {
+    if (keepAlive) {
+        dispatch_source_cancel(keepAlive);
+        keepAlive = nil;
+    }
+    keepAliveFireDate = nil;
+}
+
 + (NSURLRequest *)createRequestWithURL:(NSString *)connectionURL
                              userAgent:(NSString *)userAgent
                            googleAppID:(NSString *)googleAppID

FirebaseDatabase/Tests/Unit/FWebSocketConnectionTest.m

@@ -34,6 +34,33 @@
 
 #if !TARGET_OS_WATCH
 
+@interface FWebSocketConnection (KeepAliveTesting)
+- (void)onClosed;
+- (void)resetKeepAlive;
+- (void)webSocketDidOpen;
+- (void)sendStringToWebSocket:(NSString *)string;
+@end
+
+static char workerQueueKey;
+
+@interface FWebSocketKeepAliveSpy : FWebSocketConnection
+@property(nonatomic) NSUInteger keepAliveCount;
+@property(nonatomic) BOOL sentOnWorkerQueue;
+@property(nonatomic, copy) NSString *lastSentString;
+@property(nonatomic, copy) void (^onKeepAlive)(void);
+@end
+
+@implementation FWebSocketKeepAliveSpy
+- (void)sendStringToWebSocket:(NSString *)string {
+  self.lastSentString = string;
+  self.sentOnWorkerQueue = dispatch_get_specific(&workerQueueKey) != NULL;
+  self.keepAliveCount++;
+  if (self.onKeepAlive) {
+    self.onKeepAlive();
+  }
+}
+@end
+
 @interface FWebSocketConnectionTestDelegate : NSObject <FWebSocketDelegate>
 @property(nonatomic) BOOL receivedMessage;
 @property(nonatomic) NSInteger disconnectCount;
@@ -198,6 +225,65 @@
   XCTAssertEqual(delegate.disconnectCount, 1);
 }
 
+- (FWebSocketKeepAliveSpy *)keepAliveConnectionOnQueue:(dispatch_queue_t)queue {
+  FRepoInfo *info = [[FRepoInfo alloc] initWithHost:@"foo.firebaseio.com"
+                                           isSecure:YES
+                                      withNamespace:@"foo"];
+  dispatch_queue_set_specific(queue, &workerQueueKey, &workerQueueKey, NULL);
+  return [[FWebSocketKeepAliveSpy alloc] initWith:info
+                                         andQueue:queue
+                                      googleAppID:@"1:1234:ios:1234"
+                                    lastSessionID:nil
+                                    appCheckToken:nil];
+}
+
+- (void)testKeepAliveRunsOnWorkerQueueAndResetsInline {
+  dispatch_queue_t queue = dispatch_queue_create("keepalive.test", DISPATCH_QUEUE_SERIAL);
+  FWebSocketKeepAliveSpy *connection = [self keepAliveConnectionOnQueue:queue];
+  XCTestExpectation *sent = [self expectationWithDescription:@"worker keepalive"];
+  dispatch_async(queue, ^{
+    [connection webSocketDidOpen];
+    [connection setValue:[NSDate distantPast] forKey:@"keepAliveFireDate"];
+    [connection resetKeepAlive];
+    NSDate *fireDate = [connection valueForKey:@"keepAliveFireDate"];
+    XCTAssertGreaterThan([fireDate timeIntervalSinceNow], 40.0);
+    __weak FWebSocketKeepAliveSpy *weakConnection = connection;
+    connection.onKeepAlive = ^{
+      FWebSocketKeepAliveSpy *activeConnection = weakConnection;
+      XCTAssertNotNil(activeConnection);
+      XCTAssertEqualObjects(activeConnection.lastSentString, @"0");
+      XCTAssertTrue(activeConnection.sentOnWorkerQueue);
+      [activeConnection onClosed];
+      activeConnection.onKeepAlive = nil;
+      [sent fulfill];
+    };
+    dispatch_source_t timer = [connection valueForKey:@"keepAlive"];
+    dispatch_source_set_timer(timer, DISPATCH_TIME_NOW, DISPATCH_TIME_FOREVER, 0);
+  });
+  [self waitForExpectations:@[ sent ] timeout:2];
+}
+
+- (void)testCloseCancelsKeepAliveAndCannotReinstallIt {
+  dispatch_queue_t queue = dispatch_queue_create("keepalive.close.test", DISPATCH_QUEUE_SERIAL);
+  FWebSocketKeepAliveSpy *connection = [self keepAliveConnectionOnQueue:queue];
+  XCTestExpectation *closed = [self expectationWithDescription:@"no keepalive after close"];
+  dispatch_async(queue, ^{
+    [connection webSocketDidOpen];
+    dispatch_source_t timer = [connection valueForKey:@"keepAlive"];
+    dispatch_source_set_timer(timer, dispatch_time(DISPATCH_TIME_NOW, 100 * NSEC_PER_MSEC),
+                              DISPATCH_TIME_FOREVER, 0);
+    [connection close];
+    [connection resetKeepAlive];
+    [connection webSocketDidOpen];
+    XCTAssertNil([connection valueForKey:@"keepAlive"]);
+    dispatch_after(dispatch_time(DISPATCH_TIME_NOW, 200 * NSEC_PER_MSEC), queue, ^{
+      XCTAssertEqual(connection.keepAliveCount, 0U);
+      [closed fulfill];
+    });
+  });
+  [self waitForExpectations:@[ closed ] timeout:2];
+}
+
 @end
 
 #endif  // !TARGET_OS_WATCH