[Performance] Fix missing network traces for Swift async/await URLSession requests on iOS 16+ (#11861)

FirebasePerformance/CHANGELOG.md

@@ -1,4 +1,5 @@
 # Unreleased
+- [Performance] Fix missing network traces for Swift async/await URLSession requests on iOS 16+. (#11861)
 - [fixed] Fixed NSURLSession delegate instrumentation for NSProxy delegates. (#14478)
 - [fixed] Address crash by deferring class disposal in FPRObjectSwizzler. (#14473)
 - [fixed] Prevent race condition crashes in FPRScreenTraceTracker by replacing NSMapTable with

FirebasePerformance/Sources/Instrumentation/Network/Delegates/FPRNSURLSessionDelegate.h

@@ -25,3 +25,24 @@ NS_EXTENSION_UNAVAILABLE("Firebase Performance is not supported for extensions."
                                                NSURLSessionDownloadDelegate>
 
 @end
+
+@class FPRNetworkTrace;
+
+/** Attaches a new FPRNetworkTrace to task if one has not already been attached.
+ *  Called from both FPRNSURLSessionDelegate and FPRNSURLSessionDelegateInstrument.
+ *
+ *  @param task The task that was created.
+ */
+FOUNDATION_EXTERN
+NS_EXTENSION_UNAVAILABLE("Firebase Performance is not supported for extensions.")
+void FPRHandleDidCreateTask(NSURLSessionTask *task);
+
+/** Completes and removes the FPRNetworkTrace attached to task using data from metrics.
+ *  Called from both FPRNSURLSessionDelegate and FPRNSURLSessionDelegateInstrument.
+ *
+ *  @param task    The task that finished.
+ *  @param metrics The metrics collected for the task.
+ */
+FOUNDATION_EXTERN
+NS_EXTENSION_UNAVAILABLE("Firebase Performance is not supported for extensions.")
+void FPRHandleDidFinishCollectingMetrics(NSURLSessionTask *task, NSURLSessionTaskMetrics *metrics);

FirebasePerformance/Sources/Instrumentation/Network/Delegates/FPRNSURLSessionDelegate.m

@@ -17,8 +17,77 @@
 #import "FirebasePerformance/Sources/FPRConsoleLogger.h"
 #import "FirebasePerformance/Sources/Instrumentation/FPRNetworkTrace.h"
 
+#pragma mark - Async/await support helpers
+
+void FPRHandleDidCreateTask(NSURLSessionTask *task) {
+  @try {
+    // Skip if trace was already attached by the ObjC task creation swizzle path.
+    if ([FPRNetworkTrace networkTraceFromObject:task] != nil) {
+      return;
+    }
+    // Guard against nil request.
+    if (!task.originalRequest) {
+      return;
+    }
+    FPRNetworkTrace *trace = [[FPRNetworkTrace alloc] initWithURLRequest:task.originalRequest];
+    [trace start];
+    [FPRNetworkTrace addNetworkTrace:trace toObject:task];
+  } @catch (NSException *exception) {
+    FPRLogWarning(kFPRNetworkTraceNotTrackable, @"Unable to track network request.");
+  }
+}
+
+void FPRHandleDidFinishCollectingMetrics(NSURLSessionTask *task, NSURLSessionTaskMetrics *metrics) {
+  @try {
+    FPRNetworkTrace *trace = [FPRNetworkTrace networkTraceFromObject:task];
+    if (!trace) {
+      return;
+    }
+    // Use task metrics for accurate byte counts (more reliable than incremental callbacks).
+    NSURLSessionTaskTransactionMetrics *transactionMetrics = metrics.transactionMetrics.lastObject;
+    if (transactionMetrics) {
+      trace.responseSize = transactionMetrics.countOfResponseBodyBytesReceived;
+      trace.requestSize = transactionMetrics.countOfRequestBodyBytesSent;
+    }
+    // Ensure intermediate checkpoints are recorded before completion.
+    [trace checkpointState:FPRNetworkTraceCheckpointStateRequestCompleted];
+    [trace checkpointState:FPRNetworkTraceCheckpointStateResponseReceived];
+    // Complete the trace. For ObjC delegate tasks this fires before didCompleteWithError:,
+    // making that subsequent call a no op via the traceCompleted guard.
+    [trace didCompleteRequestWithResponse:task.response error:task.error];
+    [FPRNetworkTrace removeNetworkTraceFromObject:task];
+  } @catch (NSException *exception) {
+    FPRLogWarning(kFPRNetworkTraceNotTrackable, @"Unable to track network request.");
+  }
+}
+
 @implementation FPRNSURLSessionDelegate
 
+#pragma mark - Async/await support
+
+/** Fires for every task created on the session, including tasks created by Swift async/await
+ *  methods (iOS 16+). This is the only reliable hook that fires for async created tasks
+ *  where delegate data transfer callbacks are suppressed by the system.
+ *  Only creates a trace when one has not already been attached by the ObjC task creation swizzle.
+ */
+- (void)URLSession:(NSURLSession *)session
+     didCreateTask:(NSURLSessionTask *)task API_AVAILABLE(ios(16.0), tvos(16.0)) {
+  FPRHandleDidCreateTask(task);
+}
+
+/** Fires after every task completes (iOS 10+), including async/await created tasks where
+ *  didCompleteWithError: is suppressed by the system. Using metrics also provides more accurate
+ *  byte counts than incremental delegate callbacks.
+ *  The traceCompleted guard inside FPRNetworkTrace ensures this is safe to call more than once.
+ */
+- (void)URLSession:(NSURLSession *)session
+                          task:(NSURLSessionTask *)task
+    didFinishCollectingMetrics:(NSURLSessionTaskMetrics *)metrics {
+  FPRHandleDidFinishCollectingMetrics(task, metrics);
+}
+
+#pragma mark - NSURLSessionTaskDelegate
+
 - (void)URLSession:(NSURLSession *)session
                     task:(NSURLSessionTask *)task
     didCompleteWithError:(NSError *)error {

FirebasePerformance/Sources/Instrumentation/Network/Delegates/FPRNSURLSessionDelegateInstrument.m

@@ -125,6 +125,57 @@ void InstrumentURLSessionDataTaskDidReceiveData(FPRClassInstrumentor *instrument
   }
 }
 
+#pragma mark - Async/await support (NSURLSessionTaskDelegate)
+
+/** Instruments URLSession:didCreateTask: (iOS 16+ / tvOS 16+).
+ *  Fires for every task created on the session, including tasks created by Swift async/await
+ *  methods, which bypass the ObjC task creation swizzle entirely on iOS 16+.
+ *
+ *  @param instrumentor The FPRClassInstrumentor to add the selector instrumentor to.
+ */
+FOUNDATION_STATIC_INLINE
+NS_EXTENSION_UNAVAILABLE("Firebase Performance is not supported for extensions.")
+void InstrumentURLSessionDidCreateTask(FPRClassInstrumentor *instrumentor) {
+  SEL selector = @selector(URLSession:didCreateTask:);
+  FPRSelectorInstrumentor *selectorInstrumentor =
+      [instrumentor instrumentorForInstanceSelector:selector];
+  if (selectorInstrumentor) {
+    IMP currentIMP = selectorInstrumentor.currentIMP;
+    [selectorInstrumentor
+        setReplacingBlock:^(id object, NSURLSession *session, NSURLSessionTask *task) {
+          FPRHandleDidCreateTask(task);
+          typedef void (*OriginalImp)(id, SEL, NSURLSession *, NSURLSessionTask *);
+          ((OriginalImp)currentIMP)(object, selector, session, task);
+        }];
+  }
+}
+
+/** Instruments URLSession:task:didFinishCollectingMetrics: (iOS 10+).
+ *  Fires after every task completes, including async/await-created tasks where
+ *  didCompleteWithError: is suppressed by the system. The traceCompleted guard inside
+ *  FPRNetworkTrace makes this idempotent with the existing ObjC completion paths.
+ *
+ *  @param instrumentor The FPRClassInstrumentor to add the selector instrumentor to.
+ */
+FOUNDATION_STATIC_INLINE
+NS_EXTENSION_UNAVAILABLE("Firebase Performance is not supported for extensions.")
+void InstrumentURLSessionTaskDidFinishCollectingMetrics(FPRClassInstrumentor *instrumentor) {
+  SEL selector = @selector(URLSession:task:didFinishCollectingMetrics:);
+  FPRSelectorInstrumentor *selectorInstrumentor =
+      [instrumentor instrumentorForInstanceSelector:selector];
+  if (selectorInstrumentor) {
+    IMP currentIMP = selectorInstrumentor.currentIMP;
+    [selectorInstrumentor
+        setReplacingBlock:^(id object, NSURLSession *session, NSURLSessionTask *task,
+                            NSURLSessionTaskMetrics *metrics) {
+          FPRHandleDidFinishCollectingMetrics(task, metrics);
+          typedef void (*OriginalImp)(id, SEL, NSURLSession *, NSURLSessionTask *,
+                                      NSURLSessionTaskMetrics *);
+          ((OriginalImp)currentIMP)(object, selector, session, task, metrics);
+        }];
+  }
+}
+
 #pragma mark - NSURLSessionDownloadDelegate methods.
 
 /** Instruments URLSession:downloadTask:didFinishDownloadingToURL:.
@@ -231,6 +282,14 @@ - (void)registerClass:(Class)aClass {
     InstrumentURLSessionDownloadTaskDidWriteDataTotalBytesWrittenTotalBytesExpectedToWrite(
         instrumentor);
 
+    // Async/await support: captures tasks created by Swift async/await methods (iOS 16+).
+    if (@available(iOS 16.0, tvOS 16.0, *)) {
+      InstrumentURLSessionDidCreateTask(instrumentor);
+    }
+    // Task metrics (iOS 10+): reliable completion hook for async tasks and more accurate byte
+    // counts for all task types. Idempotent with ObjC completion paths via traceCompleted guard.
+    InstrumentURLSessionTaskDidFinishCollectingMetrics(instrumentor);
+
     [instrumentor swizzle];
   });
 }
@@ -261,6 +320,12 @@ - (void)registerObject:(id)object {
             URLSession:downloadTask:didWriteData:totalBytesWritten:totalBytesExpectedToWrite:),
         instrumentor);
 
+    // Async/await support.
+    if (@available(iOS 16.0, tvOS 16.0, *)) {
+      CopySelector(@selector(URLSession:didCreateTask:), instrumentor);
+    }
+    CopySelector(@selector(URLSession:task:didFinishCollectingMetrics:), instrumentor);
+
     [instrumentor swizzle];
   });
 }

FirebasePerformance/Tests/Unit/Instruments/FPRNSURLConnectionInstrumentTest.m

@@ -284,15 +284,13 @@ - (void)testConnectionDidFailWithError {
   self.appFake.fakeIsDataCollectionDefaultEnabled = YES;
   FPRNSURLConnectionInstrument *instrument = [[FPRNSURLConnectionInstrument alloc] init];
   [instrument registerInstrumentors];
-  NSURLRequest *request = [NSURLRequest requestWithURL:[NSURL URLWithString:@"http://nonurl/"]];
   [self.testServer stop];
-  [[NSRunLoop currentRunLoop] runUntilDate:[NSDate dateWithTimeIntervalSinceNow:0.2]];
+  NSURLRequest *request = [NSURLRequest requestWithURL:self.testServer.serverURL];
   FPRNSURLConnectionCompleteTestDelegate *delegate =
       [[FPRNSURLConnectionCompleteTestDelegate alloc] init];
   NSURLConnection *connection = [[NSURLConnection alloc] initWithRequest:request delegate:delegate];
   [connection start];
-  XCTAssertNotNil([FPRNetworkTrace networkTraceFromObject:connection]);
-  [[NSRunLoop currentRunLoop] runUntilDate:[NSDate dateWithTimeIntervalSinceNow:2.0]];
+  [[NSRunLoop currentRunLoop] runUntilDate:[NSDate dateWithTimeIntervalSinceNow:1.0]];
   XCTAssertTrue(delegate.connectionDidFailWithErrorCalled);
   XCTAssertNil([FPRNetworkTrace networkTraceFromObject:connection]);
   [self.testServer start];

FirebasePerformance/Tests/Unit/Instruments/FPRNSURLSessionInstrumentTest.m

@@ -385,8 +385,8 @@
   [instrument registerInstrumentors];
   FPRNSURLSessionCompleteTestDelegate *delegate =
       [[FPRNSURLSessionCompleteTestDelegate alloc] init];
-  // This request needs to fail.
-  NSURLRequest *request = [NSURLRequest requestWithURL:[NSURL URLWithString:@"http://nonurl"]];
+  // This request needs to fail. Point at the stopped local server for deterministic refusal.
+  NSURLRequest *request = [NSURLRequest requestWithURL:self.testServer.serverURL];
   NSURLSessionConfiguration *configuration =
       [NSURLSessionConfiguration defaultSessionConfiguration];
   NSURLSession *session = [NSURLSession sessionWithConfiguration:configuration
@@ -396,8 +396,7 @@
   @autoreleasepool {
     task = [session dataTaskWithRequest:request];
     [task resume];
-    XCTAssertNotNil([FPRNetworkTrace networkTraceFromObject:task]);
-    [[NSRunLoop currentRunLoop] runUntilDate:[NSDate dateWithTimeIntervalSinceNow:5.0]];
+    [[NSRunLoop currentRunLoop] runUntilDate:[NSDate dateWithTimeIntervalSinceNow:1.0]];
   }
   XCTAssertNil([FPRNetworkTrace networkTraceFromObject:task]);
   XCTAssertTrue(delegate.URLSessionTaskDidCompleteWithErrorCalled);
@@ -623,8 +622,8 @@
       [[FPRNSURLSessionDelegateProxy alloc] initWithDelegate:delegate];
   NSURLSessionConfiguration *configuration =
       [NSURLSessionConfiguration defaultSessionConfiguration];
   [NSURLSession sessionWithConfiguration:configuration delegate:proxyDelegate delegateQueue:nil];
   [NSURLSession sessionWithConfiguration:configuration delegate:proxyDelegate delegateQueue:nil];
   XCTAssertEqual(instrument.delegateInstrument.classInstrumentors.count, 1);
   XCTAssertEqual(instrument.delegateInstrument.instrumentedClasses.count, 1);
   XCTAssertTrue(
@@ -645,7 +644,7 @@
   NSURLSessionConfiguration *configuration =
       [NSURLSessionConfiguration defaultSessionConfiguration];
   NSURLSession *session = [NSURLSession sessionWithConfiguration:configuration
                                                         delegate:proxyDelegate
                                                    delegateQueue:nil];
   NSURLSessionTask *task;
   @autoreleasepool {
@@ -671,7 +670,7 @@
   NSURLSessionConfiguration *configuration =
       [NSURLSessionConfiguration defaultSessionConfiguration];
   NSURLSession *session = [NSURLSession sessionWithConfiguration:configuration
                                                         delegate:proxyDelegate
                                                    delegateQueue:nil];
   NSURL *URL = [self.testServer.serverURL URLByAppendingPathComponent:@"testUpload"];
   NSMutableURLRequest *request = [NSMutableURLRequest requestWithURL:URL];
@@ -708,7 +707,7 @@
   NSURL *URL = [self.testServer.serverURL URLByAppendingPathComponent:@"testRedirect"];
   NSURLRequest *request = [NSURLRequest requestWithURL:URL];
   NSURLSession *session = [NSURLSession sessionWithConfiguration:configuration
                                                         delegate:proxyDelegate
                                                    delegateQueue:nil];
   NSURLSessionTask *task = [session dataTaskWithRequest:request];
   [task resume];
@@ -735,7 +734,7 @@
     NSURLSessionConfiguration *configuration =
         [NSURLSessionConfiguration defaultSessionConfiguration];
     NSURLSession *session = [NSURLSession sessionWithConfiguration:configuration
                                                           delegate:proxyDelegate
                                                      delegateQueue:nil];
     NSURL *URL = [self.testServer.serverURL URLByAppendingPathComponent:@"testBigDownload"];
     dataTask = [session dataTaskWithURL:URL];
@@ -761,7 +760,7 @@
   NSURLSessionConfiguration *configuration =
       [NSURLSessionConfiguration defaultSessionConfiguration];
   NSURLSession *session = [NSURLSession sessionWithConfiguration:configuration
                                                         delegate:proxyDelegate
                                                    delegateQueue:nil];
   NSURL *URL = [self.testServer.serverURL URLByAppendingPathComponent:@"testDownload"];
   NSURLSessionDownloadTask *downloadTask = [session downloadTaskWithURL:URL];
@@ -1142,6 +1141,106 @@
   [instrument deregisterInstrumentors];
 }
 
+#pragma mark - Testing async/await support (didCreateTask: + didFinishCollectingMetrics:)
+
+/** Tests that URLSession:task:didFinishCollectingMetrics: is called and wrapped for a
+ *  delegate-based task, and that the network trace is completed and removed after it fires.
+ */
+- (void)testDelegateURLSessionTaskDidFinishCollectingMetrics {
+  FPRNSURLSessionInstrument *instrument = [[FPRNSURLSessionInstrument alloc] init];
+  [instrument registerInstrumentors];
+  FPRNSURLSessionCompleteTestDelegate *delegate =
+      [[FPRNSURLSessionCompleteTestDelegate alloc] init];
+  NSURLSessionConfiguration *configuration =
+      [NSURLSessionConfiguration defaultSessionConfiguration];
+  NSURLSession *session = [NSURLSession sessionWithConfiguration:configuration
+                                                        delegate:delegate
+                                                   delegateQueue:nil];
+  NSURLSessionDataTask *task = [session dataTaskWithURL:self.testServer.serverURL];
+  [task resume];
+  [self waitAndRunBlockAfterResponse:^(id self, GCDWebServerRequest *_Nonnull request,
+                                       GCDWebServerResponse *_Nonnull response) {
+    XCTAssertNil([FPRNetworkTrace networkTraceFromObject:task]);
+  }];
+  XCTAssertTrue(delegate.URLSessionTaskDidFinishCollectingMetricsCalled);
+  [instrument deregisterInstrumentors];
+}
+
+/** Tests that URLSession:task:didFinishCollectingMetrics: is instrumented on user-provided
+ *  delegate classes that don't originally implement it (via CopySelector).
+ */
+- (void)testDelegateURLSessionTaskDidFinishCollectingMetricsCopiedForNonImplementingDelegate {
+  FPRNSURLSessionInstrument *instrument = [[FPRNSURLSessionInstrument alloc] init];
+  [instrument registerInstrumentors];
+  // FPRNSURLSessionTestDelegate only implements NSURLSessionDelegate, not the task delegate.
+  FPRNSURLSessionTestDelegate *delegate = [[FPRNSURLSessionTestDelegate alloc] init];
+  NSURLSessionConfiguration *configuration =
+      [NSURLSessionConfiguration defaultSessionConfiguration];
+  NSURLSession *session = [NSURLSession sessionWithConfiguration:configuration
+                                                        delegate:delegate
+                                                   delegateQueue:nil];
+  // The delegate should now respond to didFinishCollectingMetrics: after CopySelector.
+  XCTAssertTrue(
+      [delegate respondsToSelector:@selector(URLSession:task:didFinishCollectingMetrics:)]);
+  (void)session;
+  [instrument deregisterInstrumentors];
+}
+
+/** Tests that a network trace attached by the ObjC task-creation swizzle is not duplicated
+ *  when URLSession:didCreateTask: also fires on iOS 16+. The trace from the ObjC path takes
+ *  precedence and didCreateTask: should be a no-op for that task.
+ */
+- (void)testDidCreateTaskDoesNotDoubleTraceObjCCreatedTask {
+  if (@available(iOS 16.0, tvOS 16.0, *)) {
+    FPRNSURLSessionInstrument *instrument = [[FPRNSURLSessionInstrument alloc] init];
+    [instrument registerInstrumentors];
+    FPRNSURLSessionCompleteTestDelegate *delegate =
+        [[FPRNSURLSessionCompleteTestDelegate alloc] init];
+    NSURLSessionConfiguration *configuration =
+        [NSURLSessionConfiguration defaultSessionConfiguration];
+    NSURLSession *session = [NSURLSession sessionWithConfiguration:configuration
+                                                          delegate:delegate
+                                                     delegateQueue:nil];
+    NSURL *URL = [self.testServer.serverURL URLByAppendingPathComponent:@"test"];
+    NSURLRequest *request = [NSURLRequest requestWithURL:URL];
+    // Task is created via ObjC swizzled method, trace attached immediately.
+    NSURLSessionTask *task = [session dataTaskWithRequest:request];
+    FPRNetworkTrace *traceFromObjCPath = [FPRNetworkTrace networkTraceFromObject:task];
+    XCTAssertNotNil(traceFromObjCPath);
+    // didCreateTask: fires here (iOS 16+). The trace already exists, so it must be a no-op.
+    // Verify the same trace object is still on the task (not replaced).
+    XCTAssertEqual([FPRNetworkTrace networkTraceFromObject:task], traceFromObjCPath);
+    [task resume];
+    [self waitAndRunBlockAfterResponse:^(id self, GCDWebServerRequest *_Nonnull req,
+                                         GCDWebServerResponse *_Nonnull resp) {
+      XCTAssertNil([FPRNetworkTrace networkTraceFromObject:task]);
+    }];
+    [instrument deregisterInstrumentors];
+  }
+}
+
+/** Tests that URLSession:didCreateTask: is called and wrapped for a user-provided delegate
+ *  on iOS 16+.
+ */
+- (void)testDelegateURLSessionDidCreateTaskCalledOnIOS16 {
+  if (@available(iOS 16.0, tvOS 16.0, *)) {
+    FPRNSURLSessionInstrument *instrument = [[FPRNSURLSessionInstrument alloc] init];
+    [instrument registerInstrumentors];
+    FPRNSURLSessionCompleteTestDelegate *delegate =
+        [[FPRNSURLSessionCompleteTestDelegate alloc] init];
+    NSURLSessionConfiguration *configuration =
+        [NSURLSessionConfiguration defaultSessionConfiguration];
+    NSURLSession *session = [NSURLSession sessionWithConfiguration:configuration
+                                                          delegate:delegate
+                                                     delegateQueue:nil];
+    NSURL *URL = [self.testServer.serverURL URLByAppendingPathComponent:@"test"];
+    NSURLSessionTask *task = [session dataTaskWithRequest:[NSURLRequest requestWithURL:URL]];
+    (void)task;
+    XCTAssertTrue(delegate.URLSessionDidCreateTaskCalled);
+    [instrument deregisterInstrumentors];
+  }
+}
+
 @end
 
 #endif  // SWIFT_PACKAGE

FirebasePerformance/Tests/Unit/Instruments/FPRNSURLSessionInstrumentTestDelegates.h

@@ -73,6 +73,12 @@ NS_ASSUME_NONNULL_BEGIN
  */
 @property(nonatomic) BOOL URLSessionDownloadTaskDidWriteDataTotalBytesWrittenTotalBytesCalled;
 
+/** Set to YES when URLSession:didCreateTask: is called (iOS 16+), used for testing. */
+@property(nonatomic) BOOL URLSessionDidCreateTaskCalled;
+
+/** Set to YES when URLSession:task:didFinishCollectingMetrics: is called, used for testing. */
+@property(nonatomic) BOOL URLSessionTaskDidFinishCollectingMetricsCalled;
+
 @end
 
 /** This class implements the methods necessary to cancel and resume a download. */

FirebasePerformance/Tests/Unit/Instruments/FPRNSURLSessionInstrumentTestDelegates.m

@@ -103,6 +103,19 @@ - (void)URLSession:(NSURLSession *)session
   self.URLSessionDownloadTaskDidWriteDataTotalBytesWrittenTotalBytesCalled = YES;
 }
 
+#pragma mark - Async/await support
+
+- (void)URLSession:(NSURLSession *)session
+     didCreateTask:(NSURLSessionTask *)task API_AVAILABLE(ios(16.0), tvos(16.0)) {
+  self.URLSessionDidCreateTaskCalled = YES;
+}
+
+- (void)URLSession:(NSURLSession *)session
+                          task:(NSURLSessionTask *)task
+    didFinishCollectingMetrics:(NSURLSessionTaskMetrics *)metrics {
+  self.URLSessionTaskDidFinishCollectingMetricsCalled = YES;
+}
+
 @end
 
 @interface FPRNSURLSessionTestDownloadDelegate ()