NR-363895; port setState swift async request changes from bespoke build to 7.5.x

Agent/Instrumentation/NSURLSession/NRMAURLSessionOverride.m

@@ -225,6 +225,8 @@ + (void)swizzleURLSessionTask
     }
 
     NSURLSessionTask* task = ((id(*)(id,SEL,NSURLRequest*))originalImp)(self,_cmd,mutableRequest);
+    objc_setAssociatedObject(task, NRMAHandledRequestKey, @YES, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
+
     if([NRMAFlags shouldEnableNewEventSystem]){
         [NRMAHTTPUtilities attachNRMAPayload:payloadHolder.objcPayload
                                       to:task.originalRequest];
@@ -272,6 +274,7 @@ + (void)swizzleURLSessionTask
 
     if (completionHandler == nil) {
         task  = ((id(*)(id,SEL,NSURLRequest*,void(^)(NSData*,NSURLResponse*,NSError*)))originalImp)(self,_cmd,mutableRequest,completionHandler);
+        objc_setAssociatedObject(task, NRMAHandledRequestKey, @YES, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
 
         if([NRMAFlags shouldEnableNewEventSystem]) {
             [NRMAHTTPUtilities attachNRMAPayload:payloadHolder.objcPayload to:task.originalRequest];
@@ -290,11 +293,14 @@ + (void)swizzleURLSessionTask
         } else {
             [NRMAHTTPUtilities attachPayload:payloadHolder.cppPayload to:task.originalRequest];
         }
+
+        // NRLOG_AGENT_VERBOSE(@"NRMA__recordTask called from NRMAOverride__dataTaskWithRequest_completionHandler");
 
         NRMA__recordTask(task,data,response,error);
 
         completionHandler(data,response,error);
     });
+    objc_setAssociatedObject(task, NRMAHandledRequestKey, @YES, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
 
     // Try to override the methods of the private class that is returned by this method.
     [NRMAURLSessionTaskOverride instrumentConcreteClass:[task class]];
@@ -314,7 +320,8 @@ + (void)swizzleURLSessionTask
     }
 
     NSURLSessionTask* task = ((id(*)(id,SEL,NSURL*))originalImp)(self,_cmd,url);
-
+    objc_setAssociatedObject(task, NRMAHandledRequestKey, @YES, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
+
     // Try to override the methods of the private class that is returned by this method.
     [NRMAURLSessionTaskOverride instrumentConcreteClass:[task class]];
     return task;
@@ -335,6 +342,7 @@ + (void)swizzleURLSessionTask
 
     NSMutableURLRequest* mutableRequest = [NRMAHTTPUtilities addCrossProcessIdentifier:request];
     NSURLSessionTask* task = ((NSURLSessionTask*(*)(id,SEL,NSURLRequest*,NSURL*))originalImp)(self,_cmd,mutableRequest,fileURL);
+    objc_setAssociatedObject(task, NRMAHandledRequestKey, @YES, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
     if([NRMAFlags shouldEnableNewEventSystem]){
         NRMAPayload* payload = [NRMAHTTPUtilities addConnectivityHeaderNRMAPayload:mutableRequest];
         [NRMAHTTPUtilities attachNRMAPayload:payload
@@ -363,6 +371,7 @@ + (void)swizzleURLSessionTask
 
     NSMutableURLRequest* mutableRequest = [NRMAHTTPUtilities addCrossProcessIdentifier:request];
     NSURLSessionTask* task = ((NSURLSessionTask*(*)(id,SEL,NSURLRequest*,NSData*))originalImp)(self, _cmd, mutableRequest, data);
+    objc_setAssociatedObject(task, NRMAHandledRequestKey, @YES, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
     if([NRMAFlags shouldEnableNewEventSystem]){
         NRMAPayload* payload = [NRMAHTTPUtilities addConnectivityHeaderNRMAPayload:mutableRequest];
         [NRMAHTTPUtilities attachNRMAPayload:payload to:task.originalRequest];
@@ -388,7 +397,8 @@ + (void)swizzleURLSessionTask
     }
 
     NSURLSessionTask* task = ((NSURLSessionTask*(*)(id,SEL,NSURLRequest*))originalImp)(self, _cmd,request);
-
+    objc_setAssociatedObject(task, NRMAHandledRequestKey, @YES, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
+
     [NRMAURLSessionTaskOverride instrumentConcreteClass:[task class]];
 
     return task;
@@ -415,7 +425,8 @@ + (void)swizzleURLSessionTask
 
     if (completionHandler == nil) {
         task = ((NSURLSessionUploadTask*(*)(id,SEL,NSURLRequest*,NSURL*,void(^)(NSData*,NSURLResponse*,NSError*)))originalIMP)(self,_cmd,mutableRequest,fileURL,completionHandler);
-
+        objc_setAssociatedObject(task, NRMAHandledRequestKey, @YES, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
+
         if([NRMAFlags shouldEnableNewEventSystem]) {
             [NRMAHTTPUtilities attachNRMAPayload:payloadHolder.objcPayload to:task.originalRequest];
         } else {
@@ -435,11 +446,14 @@ + (void)swizzleURLSessionTask
             [NRMAHTTPUtilities attachPayload:payloadHolder.cppPayload to:task.originalRequest];
         }
 
+        //  NSLog(@"NRMA__recordTask called from NRMAOverride__uploadTaskWithRequest_fromFile_completionHandler");
+
         NRMA__recordTask(task,data,response,error);
 
         completionHandler(data,response,error);
     });
-
+    objc_setAssociatedObject(task, NRMAHandledRequestKey, @YES, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
+
     // Try to override the methods of the private class that is returned by this method.
     [NRMAURLSessionTaskOverride instrumentConcreteClass:[task class]];
     return task;
@@ -468,6 +482,7 @@ + (void)swizzleURLSessionTask
 
     if (completionHandler == nil) {
         task = ((NSURLSessionUploadTask*(*)(id,SEL,NSURLRequest*,NSData*,void(^)(NSData*,NSURLResponse*,NSError*)))originalIMP)(self,_cmd,mutableRequest,bodyData,completionHandler);
+        objc_setAssociatedObject(task, NRMAHandledRequestKey, @YES, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
 
         if([NRMAFlags shouldEnableNewEventSystem]) {
             [NRMAHTTPUtilities attachNRMAPayload:payloadHolder.objcPayload to:task.originalRequest];
@@ -491,7 +506,8 @@ + (void)swizzleURLSessionTask
 
         completionHandler(data,response,error);
     });
-
+    objc_setAssociatedObject(task, NRMAHandledRequestKey, @YES, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
+
     // Try to override the methods of the private class that is returned by this method.
     [NRMAURLSessionTaskOverride instrumentConcreteClass:[task class]];
     return task;

Agent/Instrumentation/NSURLSession/NRMAURLSessionTaskOverride.h

@@ -9,7 +9,7 @@
 #import <Foundation/Foundation.h>
 #import "NRTimer.h"
 void NRMAOverride__resume(id self, SEL _cmd);
-void NRMAOverride__urlSessionTask_SetState(NSURLSessionTask *sessionTask, SEL _cmd, NSURLSessionTaskState *newState);
+void NRMAOverride__urlSessionTask_SetState(NSURLSessionTask *sessionTask, SEL _cmd, NSURLSessionTaskState newState);
 
 @interface NRMAURLSessionTaskOverride : NSObject
 

Agent/Instrumentation/NSURLSession/NRMAURLSessionTaskOverride.m

@@ -15,6 +15,7 @@
 #import "NRMAHTTPUtilities.h"
 #import "NRMANetworkFacade.h"
 #import "NRMAFlags.h"
+#import "NRConstants.h"
 
 static IMP NRMAOriginal__resume;
 static IMP NRMAOriginal__urlSessionTask_SetState;
@@ -85,10 +86,6 @@ + (void) deinstrument
     }
 }
 
-+ (NSInteger) statusCode:(NSURLResponse*)response {
-    return [response isKindOfClass:[NSHTTPURLResponse class]] ? [((NSHTTPURLResponse*)response) statusCode] : -1;
-}
-
 // Currently we support NSURLSessionDataTask, NSURLSessionDownloadTask, and NSURLSessionUploadTask.
 + (bool) isSupportedTaskType:(NSURLSessionTask*) task {
     return [task isKindOfClass:[NSURLSessionDataTask class]] || [task isKindOfClass:[NSURLSessionDownloadTask class]] || [task isKindOfClass:[NSURLSessionUploadTask class]];
@@ -112,21 +109,30 @@ void NRMAOverride__resume(id self, SEL _cmd)
 }
 
 // This is the only way we have right now to record an swift async await web request.
-void NRMAOverride__urlSessionTask_SetState(NSURLSessionTask* task, SEL _cmd, NSURLSessionTaskState *newState)
+void NRMAOverride__urlSessionTask_SetState(NSURLSessionTask* task, SEL _cmd, NSURLSessionTaskState newState)
 {
     @synchronized(lock) {
         @synchronized(task) {
             if ([NRMAURLSessionTaskOverride isSupportedTaskType: task]) {
-                // Checking for NEW_RELIC_CROSS_PROCESS_ID_HEADER_KEY in the headers here. The data usually isn't link to the task yet here so, if that header exists we are handling the task elsewhere and have a better chance of getting the data so we don't need to record it here.
+
+                NSNumber *isHandled = objc_getAssociatedObject(task, NRMAHandledRequestKey);
+
+                if (isHandled != nil && [isHandled boolValue]) {
+                    if (NRMAOriginal__urlSessionTask_SetState!= nil) {
+                        // Call original setState function.
+                        ((void(*)(NSURLSessionTask *,SEL,NSURLSessionTaskState))NRMAOriginal__urlSessionTask_SetState)(task, _cmd, newState);
+                    }
+                    return;
+                }
+
                 NSURLRequest  *currentRequest = task.currentRequest;
-                
-                if(currentRequest != nil && [currentRequest valueForHTTPHeaderField:NEW_RELIC_CROSS_PROCESS_ID_HEADER_KEY] != nil) {
+
+                if(currentRequest == nil) {
                     return;
                 }
-                
+
                 NSURL *url = [currentRequest URL];
-                if (url != nil &&
-                    task.state == NSURLSessionTaskStateRunning) {
+                if (url != nil) {
 
                     // Added this section to add Distributed Tracing traceId\trace.id, guid,id and payload.
                     //1
@@ -143,10 +149,15 @@ void NRMAOverride__urlSessionTask_SetState(NSURLSessionTask* task, SEL _cmd, NSU
                                                       to:task.originalRequest];
                     }
 
-                    // get response code
-                    NSUInteger responseCode = [NRMAURLSessionTaskOverride statusCode:task.response];
-                    if (responseCode != -1) {
-                        NSData *data = NRMA__getDataForSessionTask(task);
+
+                    NSData *data = NRMA__getDataForSessionTask(task);
+
+                    // log the task and data that we will record
+                    //NSLog(@"NRMAOverride__urlSessionTask_SetState newState: %ld, taskState:%ld  task: %@ data: %@", (long) newState, (long)task.state, task, data);
+
+                    if (newState == NSURLSessionTaskStateCompleted) {
+                        // NSLog(@"NRMAOverride NRMA__recordTask called because newState  == NSURLSessionTaskStateCompleted  newState: %ld, taskState:%ld  task: %@ data: %@", (long) newState, (long)task.state, task, data);
+
                         NRMA__recordTask(task, data, task.response, task.error);
                     }
                 }
@@ -155,7 +166,7 @@ void NRMAOverride__urlSessionTask_SetState(NSURLSessionTask* task, SEL _cmd, NSU
     }
     if (NRMAOriginal__urlSessionTask_SetState!= nil) {
         // Call original setState function.
-        ((void(*)(NSURLSessionTask *,SEL,NSURLSessionTaskState *))NRMAOriginal__urlSessionTask_SetState)(task, _cmd, newState);
+        ((void(*)(NSURLSessionTask *,SEL,NSURLSessionTaskState))NRMAOriginal__urlSessionTask_SetState)(task, _cmd, newState);
     }
 }
 

Agent/Public/NRConstants.h

@@ -170,6 +170,8 @@ typedef NSString NRMetricUnit;
 #define kNRMALoggingMetricFailedUpload    kNRMALoggingMetric @"/FailedUpload"
 #define kNRMALoggingMetricSuccessfulSize    kNRMALoggingMetric @"/Size/Uncompressed"
 
+#define NRMAHandledRequestKey @"NRMAHandledRequest"
+
 // Network Failure Codes
 enum NRNetworkFailureCode {
         NRURLErrorUnknown = -1,

Test Harness/NRTestApp/NRTestApp/Helpers/ApodURL.swift

@@ -15,3 +15,12 @@ struct ApodURL {
         self.url = "https://api.nasa.gov/planetary/apod?api_key=L9fVBfet3ldADKiogWO5EZyOOOHczSE45du4FhXT&date=\(date)"
     }
 }
+
+struct ApodURLBroke {
+
+    let url: String
+
+    init(date:String) {
+        self.url = "https://api.nasa.gov/planetary/apod?date=\(date)"
+    }
+}

Test Harness/NRTestApp/NRTestApp/ViewControllers/ViewController.swift

@@ -118,6 +118,9 @@ class ViewController: UIViewController {
 
         options.append(UtilOption(title: "Change Image (Async)", handler: { [self] in refreshActionAsync()}))
 
+        options.append(UtilOption(title: "Change Image Error", handler: { [self] in brokeRefreshAction()}))
+
+        options.append(UtilOption(title: "Change Image Error (Async)", handler: { [self] in brokeRefreshActionAsync()}))
     }
 
     func utilitiesAction() {
@@ -131,13 +134,22 @@ class ViewController: UIViewController {
     func refreshAction() {
         viewModel.loadApodData()
     }
+    func brokeRefreshAction() {
+        viewModel.loadApodDataBrokeData()
+    }
 
     func refreshActionAsync() {
         Task {
             await viewModel.loadApodDataAsync()
         }
     }
-
+
+    func brokeRefreshActionAsync() {
+         Task {
+             await viewModel.loadApodDataAsyncBrokeData()
+         }
+     }
+
     func makeButton(title: String) -> UIButton {
         let button = UIButton(type: .system)
         button.setTitle(title, for: .normal)

Test Harness/NRTestApp/NRTestApp/ViewModels/ApodViewModel.swift

@@ -60,4 +60,51 @@ class ApodViewModel {
             self.error.value = error
         }
     }
+
+    // Broke Data
+
+    func loadApodDataBrokeData() {
+        let nasaUrl = ApodURLBroke(date: Date.randomBetween(start: "2015-10-31", end: Date().dateString()))
+        service.getApod(nasaURL: URL(string: nasaUrl.url)!, completion: { [weak self] result in
+            switch result {
+            case .success(let response):
+                // We do not want a video, so if we get one try again
+                if response.media_type == "video"{
+                    self?.loadApodData()
+                    return
+                }
+                NewRelic.logInfo("ApodViewModel loadApodData finished.")
+
+                self?.apodResponse.value = response
+            case .failure(let error):
+                NewRelic.logError("ApodViewModel loadApodData encountered error=error=\(error.localizedDescription).")
+
+                self?.error.value = error
+            }
+        })
+    }
+
+    func loadApodDataAsyncBrokeData() async {
+        do {
+            let nasaUrl = ApodURLBroke(date: Date.randomBetween(start: "2015-10-31", end: Date().dateString()))
+            guard let url = URL(string: nasaUrl.url) else { return }
+
+            let request = URLRequest(url: url)
+            let (data, _) = try await URLSession.shared.data(for: request)
+
+            let decoded = try JSONDecoder().decode(ApodResult.self, from: data)
+
+            if decoded.media_type == "video" {
+                return await loadApodDataAsync()
+            }
+            NewRelic.logInfo("ApodViewModel loadApodDataAsync finished.")
+
+            self.apodResponse.value = decoded
+        } catch {
+
+            NewRelic.logError("ApodViewModel loadApodDataAsync encountered error=\(error.localizedDescription).")
+
+            self.error.value = error
+        }
+    }
 }