FBMjpegServer.m 5.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152
  1. /**
  2. * Copyright (c) 2015-present, Facebook, Inc.
  3. * All rights reserved.
  4. *
  5. * This source code is licensed under the BSD-style license found in the
  6. * LICENSE file in the root directory of this source tree. An additional grant
  7. * of patent rights can be found in the PATENTS file in the same directory.
  8. */
  9. #import "FBMjpegServer.h"
  10. #import <mach/mach_time.h>
  11. @import UniformTypeIdentifiers;
  12. #import "GCDAsyncSocket.h"
  13. #import "FBConfiguration.h"
  14. #import "FBLogger.h"
  15. #import "FBScreenshot.h"
  16. #import "FBImageProcessor.h"
  17. #import "FBImageUtils.h"
  18. #import "XCUIScreen.h"
  19. static const NSUInteger MAX_FPS = 60;
  20. static const NSTimeInterval FRAME_TIMEOUT = 1.;
  21. static NSString *const SERVER_NAME = @"WDA MJPEG Server";
  22. static const char *QUEUE_NAME = "JPEG Screenshots Provider Queue";
  23. @interface FBMjpegServer()
  24. @property (nonatomic, readonly) dispatch_queue_t backgroundQueue;
  25. @property (nonatomic, readonly) NSMutableArray<GCDAsyncSocket *> *listeningClients;
  26. @property (nonatomic, readonly) FBImageProcessor *imageProcessor;
  27. @property (nonatomic, readonly) long long mainScreenID;
  28. @end
  29. @implementation FBMjpegServer
  30. - (instancetype)init
  31. {
  32. if ((self = [super init])) {
  33. _listeningClients = [NSMutableArray array];
  34. dispatch_queue_attr_t queueAttributes = dispatch_queue_attr_make_with_qos_class(DISPATCH_QUEUE_SERIAL, QOS_CLASS_UTILITY, 0);
  35. _backgroundQueue = dispatch_queue_create(QUEUE_NAME, queueAttributes);
  36. dispatch_async(_backgroundQueue, ^{
  37. [self streamScreenshot];
  38. });
  39. _imageProcessor = [[FBImageProcessor alloc] init];
  40. _mainScreenID = [XCUIScreen.mainScreen displayID];
  41. }
  42. return self;
  43. }
  44. - (void)scheduleNextScreenshotWithInterval:(uint64_t)timerInterval timeStarted:(uint64_t)timeStarted
  45. {
  46. uint64_t timeElapsed = clock_gettime_nsec_np(CLOCK_MONOTONIC_RAW) - timeStarted;
  47. int64_t nextTickDelta = timerInterval - timeElapsed;
  48. if (nextTickDelta > 0) {
  49. dispatch_after(dispatch_time(DISPATCH_TIME_NOW, nextTickDelta), self.backgroundQueue, ^{
  50. [self streamScreenshot];
  51. });
  52. } else {
  53. // Try to do our best to keep the FPS at a decent level
  54. dispatch_async(self.backgroundQueue, ^{
  55. [self streamScreenshot];
  56. });
  57. }
  58. }
  59. - (void)streamScreenshot
  60. {
  61. NSUInteger framerate = FBConfiguration.mjpegServerFramerate;
  62. uint64_t timerInterval = (uint64_t)(1.0 / ((0 == framerate || framerate > MAX_FPS) ? MAX_FPS : framerate) * NSEC_PER_SEC);
  63. uint64_t timeStarted = clock_gettime_nsec_np(CLOCK_MONOTONIC_RAW);
  64. @synchronized (self.listeningClients) {
  65. if (0 == self.listeningClients.count) {
  66. [self scheduleNextScreenshotWithInterval:timerInterval timeStarted:timeStarted];
  67. return;
  68. }
  69. }
  70. NSError *error;
  71. CGFloat compressionQuality = MAX(FBMinCompressionQuality,
  72. MIN(FBMaxCompressionQuality, FBConfiguration.mjpegServerScreenshotQuality / 100.0));
  73. NSData *screenshotData = [FBScreenshot takeInOriginalResolutionWithScreenID:self.mainScreenID
  74. compressionQuality:compressionQuality
  75. uti:UTTypeJPEG
  76. timeout:FRAME_TIMEOUT
  77. error:&error];
  78. if (nil == screenshotData) {
  79. [FBLogger logFmt:@"%@", error.description];
  80. [self scheduleNextScreenshotWithInterval:timerInterval timeStarted:timeStarted];
  81. return;
  82. }
  83. CGFloat scalingFactor = FBConfiguration.mjpegScalingFactor / 100.0;
  84. [self.imageProcessor submitImageData:screenshotData
  85. scalingFactor:scalingFactor
  86. completionHandler:^(NSData * _Nonnull scaled) {
  87. [self sendScreenshot:scaled];
  88. }];
  89. [self scheduleNextScreenshotWithInterval:timerInterval timeStarted:timeStarted];
  90. }
  91. - (void)sendScreenshot:(NSData *)screenshotData {
  92. NSString *chunkHeader = [NSString stringWithFormat:@"--BoundaryString\r\nContent-type: image/jpeg\r\nContent-Length: %@\r\n\r\n", @(screenshotData.length)];
  93. NSMutableData *chunk = [[chunkHeader dataUsingEncoding:NSUTF8StringEncoding] mutableCopy];
  94. [chunk appendData:screenshotData];
  95. [chunk appendData:(id)[@"\r\n\r\n" dataUsingEncoding:NSUTF8StringEncoding]];
  96. @synchronized (self.listeningClients) {
  97. for (GCDAsyncSocket *client in self.listeningClients) {
  98. [client writeData:chunk withTimeout:-1 tag:0];
  99. }
  100. }
  101. }
  102. - (void)didClientConnect:(GCDAsyncSocket *)newClient
  103. {
  104. [FBLogger logFmt:@"Got screenshots broadcast client connection at %@:%d", newClient.connectedHost, newClient.connectedPort];
  105. // Start broadcast only after there is any data from the client
  106. [newClient readDataWithTimeout:-1 tag:0];
  107. }
  108. - (void)didClientSendData:(GCDAsyncSocket *)client
  109. {
  110. @synchronized (self.listeningClients) {
  111. if ([self.listeningClients containsObject:client]) {
  112. return;
  113. }
  114. }
  115. [FBLogger logFmt:@"Starting screenshots broadcast for the client at %@:%d", client.connectedHost, client.connectedPort];
  116. NSString *streamHeader = [NSString stringWithFormat:@"HTTP/1.0 200 OK\r\nServer: %@\r\nConnection: close\r\nMax-Age: 0\r\nExpires: 0\r\nCache-Control: no-cache, private\r\nPragma: no-cache\r\nContent-Type: multipart/x-mixed-replace; boundary=--BoundaryString\r\n\r\n", SERVER_NAME];
  117. [client writeData:(id)[streamHeader dataUsingEncoding:NSUTF8StringEncoding] withTimeout:-1 tag:0];
  118. @synchronized (self.listeningClients) {
  119. [self.listeningClients addObject:client];
  120. }
  121. }
  122. - (void)didClientDisconnect:(GCDAsyncSocket *)client
  123. {
  124. @synchronized (self.listeningClients) {
  125. [self.listeningClients removeObject:client];
  126. }
  127. [FBLogger log:@"Disconnected a client from screenshots broadcast"];
  128. }
  129. @end