fix: reconfigure ReplayKit video on rotation

WebDriverAgentBroadcast/FBExtSessionPipeline.m

@@ -30,12 +30,16 @@ @interface FBExtSessionPipeline () <FBVideoEncoderDelegate> {
 
 @property (nonatomic, weak) id<FBExtMessageSink> sink;
 @property (nonatomic) dispatch_queue_t queue;
-@property (nonatomic, nullable) FBVideoEncoder *encoder;
+@property (atomic, nullable) FBVideoEncoder *encoder;
 @property (nonatomic) VTPixelTransferSessionRef transferSession;
 @property (nonatomic) CVPixelBufferPoolRef bufferPool;
+@property (nonatomic) NSUInteger configuredWidth;
+@property (nonatomic) NSUInteger configuredHeight;
 @property (nonatomic) NSUInteger width;
 @property (nonatomic) NSUInteger height;
 @property (nonatomic) NSUInteger fps;
+@property (nonatomic) NSUInteger bitrate;
+@property (nonatomic) FBVideoCodec codec;
 @property (nonatomic) uint64_t lastSubmitTimeMs;
 /** The next monotonic timestamp at which a frame is due (fps gate accumulator). */
 @property (nonatomic) uint64_t nextDueMs;
@@ -46,6 +50,10 @@ @interface FBExtSessionPipeline () <FBVideoEncoderDelegate> {
 @property (atomic) uint8_t currentOrientation;
 @property (nonatomic) BOOL directSourceEncodingDisabled;
 @property (nonatomic, nullable, copy) NSData *lastSentParameterSets;
+@property (nonatomic) NSUInteger failedReconfigureWidth;
+@property (nonatomic) NSUInteger failedReconfigureHeight;
+@property (nonatomic) NSUInteger reconfigureFailureCount;
+@property (nonatomic) uint64_t nextReconfigureAttemptMs;
 
 // Lightweight diagnostics, exposed via metricsSnapshot/the heartbeat. Increments are not
 // strictly atomic RMW, which is acceptable for metrics.
@@ -70,6 +78,7 @@ - (void)replacePendingSampleBuffer:(CMSampleBufferRef)sampleBuffer
 - (nullable CMSampleBufferRef)copyPendingSampleBufferAtTimeMs:(uint64_t *)timeMs
                                                   orientation:(uint8_t *)orientation;
 - (void)clearPendingSampleBufferLocked;
+- (BOOL)reconfigureForWidth:(NSUInteger)width height:(NSUInteger)height;
 
 @end
 
@@ -108,6 +117,10 @@ - (nullable instancetype)initWithSessionId:(uint32_t)sessionId
       return nil;
     }
     _fps = fps;
+    _bitrate = bitrate > 0 ? bitrate : 6000000;
+    _codec = codec;
+    _configuredWidth = width;
+    _configuredHeight = height;
     _width = width;
     _height = height;
 
@@ -150,7 +163,7 @@ - (nullable instancetype)initWithSessionId:(uint32_t)sessionId
     FBVideoEncoder *encoder = [[FBVideoEncoder alloc] initWithCodec:codec
                                                               width:width
                                                              height:height
-                                                            bitrate:bitrate > 0 ? bitrate : 6000000
+                                                            bitrate:_bitrate
                                                                 fps:fps > 0 ? fps : 30
                                                               error:error];
     if (nil == encoder) {
@@ -372,6 +385,48 @@ - (void)processRetainedSampleBuffer:(CMSampleBufferRef)sampleBuffer
     return;
   }
 
+  FBBroadcastDimensions target = FBBroadcastTargetDimensions(self.configuredWidth,
+                                                              self.configuredHeight,
+                                                              CVPixelBufferGetWidth(sourceBuffer),
+                                                              CVPixelBufferGetHeight(sourceBuffer),
+                                                              orientation);
+  if (target.width == self.width && target.height == self.height) {
+    self.failedReconfigureWidth = 0;
+    self.failedReconfigureHeight = 0;
+    self.reconfigureFailureCount = 0;
+    self.nextReconfigureAttemptMs = 0;
+  } else {
+    BOOL targetChanged = target.width != self.failedReconfigureWidth
+      || target.height != self.failedReconfigureHeight;
+    if (targetChanged) {
+      self.failedReconfigureWidth = target.width;
+      self.failedReconfigureHeight = target.height;
+      self.reconfigureFailureCount = 0;
+      self.nextReconfigureAttemptMs = 0;
+    }
+    if (nowMs >= self.nextReconfigureAttemptMs) {
+      if ([self reconfigureForWidth:target.width height:target.height]) {
+        self.failedReconfigureWidth = 0;
+        self.failedReconfigureHeight = 0;
+        self.reconfigureFailureCount = 0;
+        self.nextReconfigureAttemptMs = 0;
+      } else {
+        self.reconfigureFailureCount += 1;
+        uint64_t retryDelayMs = FBBroadcastReconfigureRetryDelayMs(self.reconfigureFailureCount);
+        uint64_t failureCompletedAtMs = clock_gettime_nsec_np(CLOCK_MONOTONIC_RAW) / NSEC_PER_MSEC;
+        self.nextReconfigureAttemptMs = FBBroadcastReconfigureRetryDeadlineMs(failureCompletedAtMs,
+                                                                              self.reconfigureFailureCount);
+        FBExtLogError("Session %u: cannot reconfigure the encoder for rotation to %lux%lu; keeping %lux%lu and retrying in %llums",
+                      self.sessionId,
+                      (unsigned long)target.width,
+                      (unsigned long)target.height,
+                      (unsigned long)self.width,
+                      (unsigned long)self.height,
+                      (unsigned long long)retryDelayMs);
+      }
+    }
+  }
+
   if (!self.directSourceEncodingDisabled
       && CVPixelBufferGetWidth(sourceBuffer) == self.width
       && CVPixelBufferGetHeight(sourceBuffer) == self.height) {
@@ -416,6 +471,74 @@ - (void)processRetainedSampleBuffer:(CMSampleBufferRef)sampleBuffer
   CVPixelBufferRelease(scaledBuffer);
 }
 
+- (BOOL)reconfigureForWidth:(NSUInteger)width height:(NSUInteger)height
+{
+  VTPixelTransferSessionRef transferSession = NULL;
+  OSStatus transferStatus = VTPixelTransferSessionCreate(kCFAllocatorDefault, &transferSession);
+  if (transferStatus != noErr || NULL == transferSession) {
+    return NO;
+  }
+  VTSessionSetProperty(transferSession, kVTPixelTransferPropertyKey_ScalingMode, kVTScalingMode_Letterbox);
+  VTSessionSetProperty(transferSession, kVTPixelTransferPropertyKey_RealTime, kCFBooleanTrue);
+
+  NSDictionary *pixelBufferAttributes = @{
+    (id)kCVPixelBufferPixelFormatTypeKey: @(kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange),
+    (id)kCVPixelBufferWidthKey: @(width),
+    (id)kCVPixelBufferHeightKey: @(height),
+    (id)kCVPixelBufferIOSurfacePropertiesKey: @{},
+  };
+  NSDictionary *poolAttributes = @{(id)kCVPixelBufferPoolMinimumBufferCountKey: @(POOL_ALLOCATION_THRESHOLD)};
+  CVPixelBufferPoolRef bufferPool = NULL;
+  CVReturn poolStatus = CVPixelBufferPoolCreate(kCFAllocatorDefault,
+                                                 (__bridge CFDictionaryRef)poolAttributes,
+                                                 (__bridge CFDictionaryRef)pixelBufferAttributes,
+                                                 &bufferPool);
+  if (poolStatus != kCVReturnSuccess || NULL == bufferPool) {
+    VTPixelTransferSessionInvalidate(transferSession);
+    CFRelease(transferSession);
+    return NO;
+  }
+
+  NSError *error;
+  FBVideoEncoder *encoder = [[FBVideoEncoder alloc] initWithCodec:self.codec
+                                                           width:width
+                                                          height:height
+                                                         bitrate:self.bitrate
+                                                             fps:self.fps > 0 ? self.fps : 30
+                                                           error:&error];
+  if (nil == encoder) {
+    CVPixelBufferPoolRelease(bufferPool);
+    VTPixelTransferSessionInvalidate(transferSession);
+    CFRelease(transferSession);
+    return NO;
+  }
+
+  FBVideoEncoder *previousEncoder = self.encoder;
+  previousEncoder.delegate = nil;
+  [previousEncoder stop];
+  if (NULL != _repeatBuffer) {
+    CVPixelBufferRelease(_repeatBuffer);
+    _repeatBuffer = NULL;
+  }
+  [self releaseScalerResources];
+
+  self.transferSession = transferSession;
+  self.bufferPool = bufferPool;
+  self.width = width;
+  self.height = height;
+  self.directSourceEncodingDisabled = NO;
+  self.lastSentParameterSets = nil;
+  self.lastEncodeAtMs = 0;
+  self.encoder = encoder;
+  encoder.delegate = self;
+  [encoder requestKeyFrame];
+  FBExtLogInfo("Session %u: reconfigured video encoder for rotation to %lux%lu",
+               self.sessionId,
+               (unsigned long)width,
+               (unsigned long)height);
+  return YES;
+}
+
 - (void)requestKeyFrame
 {
   [self.encoder requestKeyFrame];
@@ -500,6 +623,9 @@ - (void)videoEncoder:(FBVideoEncoder *)encoder
            isKeyFrame:(BOOL)isKeyFrame
    presentationTimeUs:(uint64_t)presentationTimeUs
 {
+  if (encoder != self.encoder) {
+    return;
+  }
   self.encodedCount += 1;
   // The frame pts is the monotonic submit time, so callback-time minus pts is the
   // scale+encode latency.

WebDriverAgentLib/Utilities/FBBroadcastProtocol.h

@@ -83,7 +83,7 @@ typedef struct {
 
 /** VIDEO_FRAME flags bit 0: the frame is a key (IDR) frame. */
 extern const uint8_t FBBroadcastFrameFlagKeyFrame;
-/** VIDEO_FRAME flags bits 1-3: CGImagePropertyOrientation (1-8) of the captured frame. */
+/** VIDEO_FRAME flags bits 1-4: CGImagePropertyOrientation (1-8) of the captured frame. */
 extern const uint8_t FBBroadcastFrameOrientationShift;
 extern const uint8_t FBBroadcastFrameOrientationMask;
 
@@ -118,6 +118,30 @@ extern NSString *const FBBroadcastKeySampleRate;
 /** The FBBroadcastKeyMedia value marking an audio session. */
 extern NSString *const FBBroadcastMediaAudio;
 
+/** An encoder footprint in pixels. */
+typedef struct {
+  NSUInteger width;
+  NSUInteger height;
+} FBBroadcastDimensions;
+
+/**
+ Returns the encoder footprint for a ReplayKit frame while preserving the
+ configured resolution. Orientations 5-8 rotate the source buffer by 90
+ degrees; an unknown orientation falls back to the raw buffer aspect.
+ */
+FBBroadcastDimensions FBBroadcastTargetDimensions(NSUInteger configuredWidth,
+                                                   NSUInteger configuredHeight,
+                                                   NSUInteger sourceBufferWidth,
+                                                   NSUInteger sourceBufferHeight,
+                                                   uint8_t orientation);
+
+/** Milliseconds to wait after a failed encoder reconfiguration attempt. */
+uint64_t FBBroadcastReconfigureRetryDelayMs(NSUInteger failureCount);
+
+/** Monotonic retry deadline measured from the completion of a failed attempt. */
+uint64_t FBBroadcastReconfigureRetryDeadlineMs(uint64_t failureCompletedAtMs,
+                                               NSUInteger failureCount);
+
 /** Codec string values used in SESSION_ADD, matching the HTTP API. */
 extern NSString *const FBBroadcastCodecH264;
 extern NSString *const FBBroadcastCodecH265;

WebDriverAgentLib/Utilities/FBBroadcastProtocol.m

@@ -15,7 +15,7 @@
 
 const uint8_t FBBroadcastFrameFlagKeyFrame = 1 << 0;
 const uint8_t FBBroadcastFrameOrientationShift = 1;
-const uint8_t FBBroadcastFrameOrientationMask = 0x07;
+const uint8_t FBBroadcastFrameOrientationMask = 0x0F;
 
 const uint32_t FBBroadcastAudioSessionIdFlag = 0x80000000u;
 
@@ -40,6 +40,46 @@
 
 NSString *const FBBroadcastMediaAudio = @"audio";
 
+FBBroadcastDimensions FBBroadcastTargetDimensions(NSUInteger configuredWidth,
+                                                   NSUInteger configuredHeight,
+                                                   NSUInteger sourceBufferWidth,
+                                                   NSUInteger sourceBufferHeight,
+                                                   uint8_t orientation)
+{
+  BOOL swapsSourceAxes = orientation >= 5 && orientation <= 8;
+  NSUInteger effectiveWidth = swapsSourceAxes ? sourceBufferHeight : sourceBufferWidth;
+  NSUInteger effectiveHeight = swapsSourceAxes ? sourceBufferWidth : sourceBufferHeight;
+  if (effectiveWidth == 0 || effectiveHeight == 0) {
+    return (FBBroadcastDimensions){configuredWidth, configuredHeight};
+  }
+
+  NSUInteger shortSide = MIN(configuredWidth, configuredHeight);
+  NSUInteger longSide = MAX(configuredWidth, configuredHeight);
+  return effectiveWidth > effectiveHeight
+    ? (FBBroadcastDimensions){longSide, shortSide}
+    : (FBBroadcastDimensions){shortSide, longSide};
+}
+
+uint64_t FBBroadcastReconfigureRetryDelayMs(NSUInteger failureCount)
+{
+  if (failureCount == 0) {
+    return 0;
+  }
+  if (failureCount >= 6) {
+    return 5000;
+  }
+  return MIN((uint64_t)250 << (failureCount - 1), (uint64_t)5000);
+}
+
+uint64_t FBBroadcastReconfigureRetryDeadlineMs(uint64_t failureCompletedAtMs,
+                                               NSUInteger failureCount)
+{
+  uint64_t delayMs = FBBroadcastReconfigureRetryDelayMs(failureCount);
+  return UINT64_MAX - failureCompletedAtMs < delayMs
+    ? UINT64_MAX
+    : failureCompletedAtMs + delayMs;
+}
+
 NSString *const FBBroadcastCodecH264 = @"h264";
 NSString *const FBBroadcastCodecH265 = @"h265";
 NSString *const FBBroadcastCodecOpus = @"opus";

WebDriverAgentTests/UnitTests/FBVideoStreamSessionTests.m

@@ -8,6 +8,7 @@
 
 #import <XCTest/XCTest.h>
 
+#import "FBBroadcastProtocol.h"
 #import "FBScrcpyPacket.h"
 #import "FBVideoStreamSession.h"
 
@@ -97,6 +98,76 @@ - (void)testConfigPacket
   XCTAssertEqualObjects(payload, parameterSets);
 }
 
+- (void)testBroadcastTargetSizeFollowsEffectiveFrameOrientation
+{
+  // ReplayKit keeps a portrait-shaped pixel buffer for this landscape frame
+  // and carries the clockwise rotation in RPVideoSampleOrientationKey.
+  FBBroadcastDimensions landscape = FBBroadcastTargetDimensions(1278, 588, 1170, 2532, 6);
+  XCTAssertEqual(landscape.width, (NSUInteger)1278);
+  XCTAssertEqual(landscape.height, (NSUInteger)588);
+  FBBroadcastDimensions otherLandscape = FBBroadcastTargetDimensions(1278, 588, 1170, 2532, 8);
+  XCTAssertEqual(otherLandscape.width, (NSUInteger)1278);
+  XCTAssertEqual(otherLandscape.height, (NSUInteger)588);
+
+  // Enabling portrait lock changes only the orientation attachment. The
+  // encoder footprint must cross axes or WebRTC keeps reporting landscape.
+  FBBroadcastDimensions portrait = FBBroadcastTargetDimensions(1278, 588, 1170, 2532, 1);
+  XCTAssertEqual(portrait.width, (NSUInteger)588);
+  XCTAssertEqual(portrait.height, (NSUInteger)1278);
+}
+
+- (void)testBroadcastTargetSizeFallsBackToBufferAspectForUnknownOrientation
+{
+  FBBroadcastDimensions landscape = FBBroadcastTargetDimensions(588, 1278, 2532, 1170, 0);
+  XCTAssertEqual(landscape.width, (NSUInteger)1278);
+  XCTAssertEqual(landscape.height, (NSUInteger)588);
+}
+
+- (void)testBroadcastFrameProtocolPreservesExifOrientationEight
+{
+  NSData *picture = [@"frame" dataUsingEncoding:NSUTF8StringEncoding];
+  NSData *message = FBBroadcastEncodeVideoFrameMessage(42, 123, YES, 8, picture);
+
+  FBBroadcastMessageHeader header;
+  NSData *headerData = [message subdataWithRange:NSMakeRange(0, FBBroadcastHeaderLength)];
+  XCTAssertTrue(FBBroadcastParseHeader(headerData, &header));
+  XCTAssertEqual(header.type, (uint8_t)FBBroadcastMessageTypeVideoFrame);
+  XCTAssertEqual(header.sessionId, (uint32_t)42);
+  NSData *payload = [message subdataWithRange:NSMakeRange(FBBroadcastHeaderLength, header.payloadLength)];
+
+  uint64_t pts = 0;
+  BOOL key = NO;
+  uint8_t orientation = 0;
+  NSData *decoded = nil;
+  XCTAssertTrue(FBBroadcastParseVideoFramePayload(payload, &pts, &key, &orientation, &decoded));
+  XCTAssertEqual(pts, (uint64_t)123);
+  XCTAssertTrue(key);
+  XCTAssertEqual(orientation, (uint8_t)8);
+  XCTAssertEqualObjects(decoded, picture);
+}
+
+- (void)testBroadcastEncoderReconfigurationRetryBackoffIsBounded
+{
+  XCTAssertEqual(FBBroadcastReconfigureRetryDelayMs(1), (uint64_t)250);
+  XCTAssertEqual(FBBroadcastReconfigureRetryDelayMs(2), (uint64_t)500);
+  XCTAssertEqual(FBBroadcastReconfigureRetryDelayMs(3), (uint64_t)1000);
+  XCTAssertEqual(FBBroadcastReconfigureRetryDelayMs(4), (uint64_t)2000);
+  XCTAssertEqual(FBBroadcastReconfigureRetryDelayMs(5), (uint64_t)4000);
+  XCTAssertEqual(FBBroadcastReconfigureRetryDelayMs(6), (uint64_t)5000);
+  XCTAssertEqual(FBBroadcastReconfigureRetryDelayMs(100), (uint64_t)5000);
+}
+
+- (void)testBroadcastEncoderReconfigurationRetryStartsAfterSlowFailureCompletes
+{
+  uint64_t attemptStartedAtMs = 1000;
+  uint64_t failureCompletedAtMs = 6000;
+  uint64_t deadlineMs = FBBroadcastReconfigureRetryDeadlineMs(failureCompletedAtMs, 1);
+
+  XCTAssertEqual(deadlineMs, (uint64_t)6250);
+  XCTAssertLessThan(attemptStartedAtMs + FBBroadcastReconfigureRetryDelayMs(1), deadlineMs);
+  XCTAssertLessThan(failureCompletedAtMs, deadlineMs);
+}
+
 // A large timestamp must survive the flag bits untouched (and vice versa).
 - (void)testPtsAndFlagsDoNotCollide
 {