fix: Avoid keeping strong reference to self instance in delegates

Author: mykola-mokhnach

WebDriverAgentLib/Routing/FBTCPSocket.h

@@ -38,7 +38,11 @@ NS_ASSUME_NONNULL_BEGIN
 
 @interface FBTCPSocket : NSObject
 
-@property (nullable, nonatomic) id<FBTCPSocketDelegate> delegate;
+#if __has_feature(objc_arc_weak)
+@property (nullable, nonatomic, weak) id<FBTCPSocketDelegate> delegate;
+#else
+@property (nullable, nonatomic, assign) id<FBTCPSocketDelegate> delegate;
+#endif
 
 /**
  Creates TCP socket isntance which is going to be started on the specified port

WebDriverAgentLib/Routing/FBTCPSocket.m

@@ -48,11 +48,13 @@ - (BOOL)startWithError:(NSError **)error
 - (void)stop
 {
   @synchronized(self.connectedClients) {
-    for (NSUInteger i = 0; i < [self.connectedClients count]; i++) {
-      [[self.connectedClients objectAtIndex:i] disconnect];
+    NSArray *clients = self.connectedClients.copy;
+    for (GCDAsyncSocket *client in clients) {
+      [client disconnect];
     }
   }
 
+  self.delegate = nil;
   [self.listeningSocket disconnect];
 }
 
@@ -66,20 +68,29 @@ - (void)socket:(GCDAsyncSocket *)sock didAcceptNewSocket:(GCDAsyncSocket *)newSo
   @synchronized(self.connectedClients) {
     [self.connectedClients addObject:newSocket];
   }
-  [self.delegate didClientConnect:newSocket];
+  id<FBTCPSocketDelegate> delegate = self.delegate;
+  if (nil != delegate) {
+    [delegate didClientConnect:newSocket];
+  }
 }
 
 - (void)socket:(GCDAsyncSocket *)sock didReadData:(NSData *)data withTag:(long)tag
 {
-  [self.delegate didClientSendData:sock];
+  id<FBTCPSocketDelegate> delegate = self.delegate;
+  if (nil != delegate) {
+    [delegate didClientSendData:sock];
+  }
 }
 
 - (void)socketDidDisconnect:(GCDAsyncSocket *)sock withError:(NSError *)err
 {
   @synchronized(self.connectedClients) {
     [self.connectedClients removeObject:sock];
   }
-  [self.delegate didClientDisconnect:sock];
+  id<FBTCPSocketDelegate> delegate = self.delegate;
+  if (nil != delegate) {
+    [delegate didClientDisconnect:sock];
+  }
 }
 
 @end

WebDriverAgentLib/Routing/FBWebServer.m

@@ -51,6 +51,11 @@ @interface FBWebServer ()
 
 @implementation FBWebServer
 
+- (void)dealloc
+{
+  [self stopScreenshotsBroadcaster];
+}
+
 + (NSArray<Class<FBCommandHandler>> *)collectCommandHandlerClasses
 {
   NSArray *handlersClasses = FBClassesThatConformsToProtocol(@protocol(FBCommandHandler));
@@ -125,12 +130,14 @@ - (void)startHTTPServer
 - (void)initScreenshotsBroadcaster
 {
   [self readMjpegSettingsFromEnv];
+  FBMjpegServer *mjpegServer = [[FBMjpegServer alloc] init];
   self.screenshotsBroadcaster = [[FBTCPSocket alloc]
                                  initWithPort:(uint16_t)FBConfiguration.mjpegServerPort];
-  self.screenshotsBroadcaster.delegate = [[FBMjpegServer alloc] init];
+  self.screenshotsBroadcaster.delegate = mjpegServer;
   NSError *error;
   if (![self.screenshotsBroadcaster startWithError:&error]) {
     [FBLogger logFmt:@"Cannot init screenshots broadcaster service on port %@. Original error: %@", @(FBConfiguration.mjpegServerPort), error.description];
+    [mjpegServer stopStreaming];
     self.screenshotsBroadcaster = nil;
   }
 }
@@ -141,7 +148,13 @@ - (void)stopScreenshotsBroadcaster
     return;
   }
 
+  id<FBTCPSocketDelegate> delegate = self.screenshotsBroadcaster.delegate;
+  if ([delegate respondsToSelector:@selector(stopStreaming)]) {
+    [(FBMjpegServer *)delegate stopStreaming];
+  }
+  self.screenshotsBroadcaster.delegate = nil;
   [self.screenshotsBroadcaster stop];
+  self.screenshotsBroadcaster = nil;
 }
 
 - (void)readMjpegSettingsFromEnv
@@ -164,6 +177,8 @@ - (void)stopServing
   if (self.server.isRunning) {
     [self.server stop:NO];
   }
+  self.server = nil;
+  self.exceptionHandler = nil;
   self.keepAlive = NO;
 }
 
@@ -192,10 +207,15 @@ - (BOOL)attemptToStartServer:(RoutingHTTPServer *)server onPort:(NSInteger)port
 
 - (void)registerRouteHandlers:(NSArray *)commandHandlerClasses
 {
+  __weak typeof(self) weakSelf = self;
   for (Class<FBCommandHandler> commandHandler in commandHandlerClasses) {
     NSArray *routes = [commandHandler routes];
     for (FBRoute *route in routes) {
       [self.server handleMethod:route.verb withPath:route.path block:^(RouteRequest *request, RouteResponse *response) {
+        __strong typeof(weakSelf) strongSelf = weakSelf;
+        if (nil == strongSelf) {
+          return;
+        }
         NSDictionary *arguments = [NSJSONSerialization JSONObjectWithData:request.body options:NSJSONReadingMutableContainers error:NULL];
         FBRouteRequest *routeParams = [FBRouteRequest
           routeRequestWithURL:request.url
@@ -209,7 +229,7 @@ - (void)registerRouteHandlers:(NSArray *)commandHandlerClasses
           [route mountRequest:routeParams intoResponse:response];
         }
         @catch (NSException *exception) {
-          [self handleException:exception forResponse:response];
+          [strongSelf handleException:exception forResponse:response];
         }
       }];
     }
@@ -237,9 +257,14 @@ - (void)registerServerKeyRouteHandlers
     [response respondWithString:calibrationPage];
   }];
 
+  __weak typeof(self) weakSelf = self;
   [self.server get:@"/wda/shutdown" withBlock:^(RouteRequest *request, RouteResponse *response) {
+    __strong typeof(weakSelf) strongSelf = weakSelf;
+    if (nil == strongSelf) {
+      return;
+    }
     [response respondWithString:@"Shutting down"];
-    [self.delegate webServerDidRequestShutdown:self];
+    [strongSelf.delegate webServerDidRequestShutdown:strongSelf];
   }];
 
   [self registerRouteHandlers:@[FBUnknownCommands.class]];

WebDriverAgentLib/Utilities/FBImageProcessor.m

@@ -27,6 +27,7 @@ @interface FBImageProcessor ()
 @property (nonatomic) NSData *nextImage;
 @property (nonatomic, readonly) NSLock *nextImageLock;
 @property (nonatomic, readonly) dispatch_queue_t scalingQueue;
+@property (atomic, assign) BOOL isScalingScheduled;
 
 @end
 
@@ -38,6 +39,7 @@ - (id)init
   if (self) {
     _nextImageLock = [[NSLock alloc] init];
     _scalingQueue = dispatch_queue_create("image.scaling.queue", NULL);
+    _isScalingScheduled = NO;
   }
   return self;
 }
@@ -51,32 +53,43 @@ - (void)submitImageData:(NSData *)image
     [FBLogger verboseLog:@"Discarding screenshot"];
   }
   self.nextImage = image;
+  BOOL shouldSchedule = !self.isScalingScheduled;
+  if (shouldSchedule) {
+    self.isScalingScheduled = YES;
+  }
   [self.nextImageLock unlock];
+  if (!shouldSchedule) {
+    return;
+  }
 
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wcompletion-handler"
   dispatch_async(self.scalingQueue, ^{
-    [self.nextImageLock lock];
-    NSData *nextImageData = self.nextImage;
-    self.nextImage = nil;
-    [self.nextImageLock unlock];
-    if (nextImageData == nil) {
-      return;
-    }
+    while (YES) {
+      [self.nextImageLock lock];
+      NSData *nextImageData = self.nextImage;
+      self.nextImage = nil;
+      if (nextImageData == nil) {
+        self.isScalingScheduled = NO;
+        [self.nextImageLock unlock];
+        return;
+      }
+      [self.nextImageLock unlock];
 
-    // We do not want this value to be too high because then we get images larger in size than original ones
-    // Although, we also don't want to lose too much of the quality on recompression
-    CGFloat recompressionQuality = MAX(0.9,
-                                       MIN(FBMaxCompressionQuality, FBConfiguration.mjpegServerScreenshotQuality / 100.0));
-    NSData *thumbnailData = [self.class fixedImageDataWithImageData:nextImageData
-                                                      scalingFactor:scalingFactor
-                                                                uti:UTTypeJPEG
-                                                 compressionQuality:recompressionQuality
-    // iOS always returns screnshots in portrait orientation, but puts the real value into the metadata
-    // Use it with care. See https://github.com/appium/WebDriverAgent/pull/812
-                                                     fixOrientation:FBConfiguration.mjpegShouldFixOrientation
-                                                 desiredOrientation:nil];
-    completionHandler(thumbnailData ?: nextImageData);
+      // We do not want this value to be too high because then we get images larger in size than original ones
+      // Although, we also don't want to lose too much of the quality on recompression
+      CGFloat recompressionQuality = MAX(0.9,
+                                         MIN(FBMaxCompressionQuality, FBConfiguration.mjpegServerScreenshotQuality / 100.0));
+      NSData *thumbnailData = [self.class fixedImageDataWithImageData:nextImageData
+                                                        scalingFactor:scalingFactor
+                                                                  uti:UTTypeJPEG
+                                                   compressionQuality:recompressionQuality
+      // iOS always returns screnshots in portrait orientation, but puts the real value into the metadata
+      // Use it with care. See https://github.com/appium/WebDriverAgent/pull/812
+                                                       fixOrientation:FBConfiguration.mjpegShouldFixOrientation
+                                                   desiredOrientation:nil];
+      completionHandler(thumbnailData ?: nextImageData);
+    }
   });
 #pragma clang diagnostic pop
 }

WebDriverAgentLib/Utilities/FBMjpegServer.h

@@ -19,6 +19,11 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (instancetype)init;
 
+/**
+ Stops screenshot broadcasting and prevents future scheduling.
+ */
+- (void)stopStreaming;
+
 @end
 
 NS_ASSUME_NONNULL_END

WebDriverAgentLib/Utilities/FBMjpegServer.m

@@ -35,6 +35,9 @@ @interface FBMjpegServer()
 @property (nonatomic, readonly) FBImageProcessor *imageProcessor;
 @property (nonatomic, readonly) long long mainScreenID;
 @property (nonatomic, assign) NSUInteger consecutiveScreenshotFailures;
+@property (atomic, assign) BOOL isStreaming;
+@property (nonatomic, assign) NSUInteger sentFramesCount;
+@property (nonatomic, assign) NSUInteger sentBytesCount;
 
 @end
 
@@ -45,11 +48,15 @@ - (instancetype)init
 {
   if ((self = [super init])) {
     _consecutiveScreenshotFailures = 0;
+    _isStreaming = YES;
+    _sentFramesCount = 0;
+    _sentBytesCount = 0;
     _listeningClients = [NSMutableArray array];
     dispatch_queue_attr_t queueAttributes = dispatch_queue_attr_make_with_qos_class(DISPATCH_QUEUE_SERIAL, QOS_CLASS_UTILITY, 0);
     _backgroundQueue = dispatch_queue_create(QUEUE_NAME, queueAttributes);
+    __weak typeof(self) weakSelf = self;
     dispatch_async(_backgroundQueue, ^{
-      [self streamScreenshot];
+      [weakSelf streamScreenshot];
     });
     _imageProcessor = [[FBImageProcessor alloc] init];
     _mainScreenID = [XCUIScreen.mainScreen displayID];
@@ -59,22 +66,29 @@ - (instancetype)init
 
 - (void)scheduleNextScreenshotWithInterval:(uint64_t)timerInterval timeStarted:(uint64_t)timeStarted
 {
+  if (!self.isStreaming) {
+    return;
+  }
   uint64_t timeElapsed = clock_gettime_nsec_np(CLOCK_MONOTONIC_RAW) - timeStarted;
   int64_t nextTickDelta = timerInterval - timeElapsed;
+  __weak typeof(self) weakSelf = self;
   if (nextTickDelta > 0) {
     dispatch_after(dispatch_time(DISPATCH_TIME_NOW, nextTickDelta), self.backgroundQueue, ^{
-      [self streamScreenshot];
+      [weakSelf streamScreenshot];
     });
   } else {
     // Try to do our best to keep the FPS at a decent level
     dispatch_async(self.backgroundQueue, ^{
-      [self streamScreenshot];
+      [weakSelf streamScreenshot];
     });
   }
 }
 
 - (void)streamScreenshot
 {
+  if (!self.isStreaming) {
+    return;
+  }
   NSUInteger framerate = FBConfiguration.mjpegServerFramerate;
   uint64_t timerInterval = (uint64_t)(1.0 / ((0 == framerate || framerate > MAX_FPS) ? MAX_FPS : framerate) * NSEC_PER_SEC);
   uint64_t timeStarted = clock_gettime_nsec_np(CLOCK_MONOTONIC_RAW);
@@ -106,10 +120,11 @@ - (void)streamScreenshot
   self.consecutiveScreenshotFailures = 0;
 
   CGFloat scalingFactor = FBConfiguration.mjpegScalingFactor / 100.0;
+  __weak typeof(self) weakSelf = self;
   [self.imageProcessor submitImageData:screenshotData
                          scalingFactor:scalingFactor
                      completionHandler:^(NSData * _Nonnull scaled) {
-    [self sendScreenshot:scaled];
+    [weakSelf sendScreenshot:scaled];
   }];
 
   [self scheduleNextScreenshotWithInterval:timerInterval timeStarted:timeStarted];
@@ -122,7 +137,17 @@ - (void)sendScreenshot:(NSData *)screenshotData {
   [chunk appendData:(id)[@"\r\n\r\n" dataUsingEncoding:NSUTF8StringEncoding]];
   @synchronized (self.listeningClients) {
     for (GCDAsyncSocket *client in self.listeningClients) {
-      [client writeData:chunk withTimeout:-1 tag:0];
+      // Slow clients should fail/close instead of buffering indefinitely.
+      [client writeData:chunk withTimeout:FRAME_TIMEOUT tag:0];
+    }
+    self.sentFramesCount++;
+    self.sentBytesCount += chunk.length * self.listeningClients.count;
+    NSUInteger framerate = MAX(1, MIN(MAX_FPS, FBConfiguration.mjpegServerFramerate));
+    if (0 == self.sentFramesCount % framerate) {
+      [FBLogger verboseLog:[NSString stringWithFormat:@"MJPEG stats: clients=%@ sentFrames=%@ sentBytes=%@",
+                            @(self.listeningClients.count),
+                            @(self.sentFramesCount),
+                            @(self.sentBytesCount)]];
     }
   }
 }
@@ -158,4 +183,22 @@ - (void)didClientDisconnect:(GCDAsyncSocket *)client
   [FBLogger log:@"Disconnected a client from screenshots broadcast"];
 }
 
+- (void)stopStreaming
+{
+  self.isStreaming = NO;
+  @synchronized (self.listeningClients) {
+    NSArray<GCDAsyncSocket *> *clients = self.listeningClients.copy;
+    [self.listeningClients removeAllObjects];
+    for (GCDAsyncSocket *client in clients) {
+      [client disconnect];
+    }
+  }
+}
+
+- (void)dealloc
+{
+  [self stopStreaming];
+  [FBLogger verboseLog:@"FBMjpegServer deallocated"];
+}
+
 @end