Skip to content

Commit

Permalink
NR-359272: Fix setState signature and add NRMAHandledRequest associat…
Browse files Browse the repository at this point in the history 
…ed object to track handled URLSessionTasks
  • Loading branch information
@cdillard-NewRelic
cdillard-NewRelic committed Jan 24, 2025
1 parent 86cca04 commit 8c3d711
Show file tree
Hide file tree
Showing 10 changed files with 203 additions and 30 deletions.
41 changes: 26 additions & 15 deletions Agent/Instrumentation/NSURLSession/NRMAURLSessionOverride.m
View file
Original file line number Diff line number Diff line change
Expand Up @@ -216,10 +216,11 @@ + (void)swizzleURLSessionTask

NSMutableURLRequest* mutableRequest = [NRMAHTTPUtilities addCrossProcessIdentifier:request];
NSURLSessionTask* task = ((id(*)(id,SEL,NSURLRequest*))originalImp)(self,_cmd,request);
objc_setAssociatedObject(task, NRMAHandledRequestKey, @YES, OBJC_ASSOCIATION_RETAIN_NONATOMIC);

NRMAPayloadContainer* payload = [NRMAHTTPUtilities addConnectivityHeader:mutableRequest];
[NRMAHTTPUtilities attachPayload:payload
to:task.originalRequest];
NRMAPayloadContainer* payload = [NRMAHTTPUtilities addConnectivityHeader:mutableRequest];
[NRMAHTTPUtilities attachPayload:payload
to:task.originalRequest];

// Try to override the methods of the private class that is returned by this method.
[NRMAURLSessionTaskOverride instrumentConcreteClass:[task class]];
Expand Down Expand Up @@ -256,6 +257,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);

[NRMAHTTPUtilities attachPayload:payloadHolder.cppPayload to:task.originalRequest];

Expand All @@ -267,11 +269,14 @@ + (void)swizzleURLSessionTask

[NRMAHTTPUtilities attachPayload:payloadHolder.cppPayload to:task.originalRequest];

// NSLog(@"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]];
Expand All @@ -291,7 +296,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;
Expand All @@ -312,10 +318,11 @@ + (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);

NRMAPayloadContainer* payload = [NRMAHTTPUtilities addConnectivityHeader:mutableRequest];
[NRMAHTTPUtilities attachPayload:payload
to:task.originalRequest];
NRMAPayloadContainer* payload = [NRMAHTTPUtilities addConnectivityHeader:mutableRequest];
[NRMAHTTPUtilities attachPayload:payload
to:task.originalRequest];

[NRMAURLSessionTaskOverride instrumentConcreteClass:[task class]];

Expand All @@ -335,6 +342,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);

NRMAPayloadContainer* payload = [NRMAHTTPUtilities addConnectivityHeader:mutableRequest];
[NRMAHTTPUtilities attachPayload:payload to:task.originalRequest];
Expand All @@ -356,7 +364,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;
Expand All @@ -381,8 +390,7 @@ + (void)swizzleURLSessionTask
if (completionHandler == nil) {
task = ((NSURLSessionUploadTask*(*)(id,SEL,NSURLRequest*,NSURL*,void(^)(NSData*,NSURLResponse*,NSError*)))originalIMP)(self,_cmd,mutableRequest,fileURL,completionHandler);


[NRMAHTTPUtilities attachPayload:payloadHolder.cppPayload to:task.originalRequest];
[NRMAHTTPUtilities attachPayload:payloadHolder.cppPayload to:task.originalRequest];

[NRMAURLSessionTaskOverride instrumentConcreteClass:[task class]];
return task;
Expand All @@ -391,13 +399,15 @@ + (void)swizzleURLSessionTask
task = ((NSURLSessionUploadTask*(*)(id,SEL,NSURLRequest*,NSURL*,void(^)(NSData*,NSURLResponse*,NSError*)))originalIMP)(self,_cmd,mutableRequest,fileURL,^(NSData* data,
NSURLResponse* response,
NSError* error){
[NRMAHTTPUtilities attachPayload:payloadHolder.cppPayload to:task.originalRequest];

[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]];
Expand Down Expand Up @@ -425,22 +435,23 @@ + (void)swizzleURLSessionTask
if (completionHandler == nil) {
task = ((NSURLSessionUploadTask*(*)(id,SEL,NSURLRequest*,NSData*,void(^)(NSData*,NSURLResponse*,NSError*)))originalIMP)(self,_cmd,mutableRequest,bodyData,completionHandler);


[NRMAHTTPUtilities attachPayload:payloadHolder.cppPayload to:task.originalRequest];
[NRMAHTTPUtilities attachPayload:payloadHolder.cppPayload to:task.originalRequest];

[NRMAURLSessionTaskOverride instrumentConcreteClass:[task class]];
return task;
}

task = ((NSURLSessionUploadTask*(*)(id,SEL,NSURLRequest*,NSData*,void(^)(NSData*,NSURLResponse*,NSError*)))originalIMP)(self,_cmd,mutableRequest,bodyData,^(NSData* data, NSURLResponse* response, NSError* error){


[NRMAHTTPUtilities attachPayload:payloadHolder.cppPayload to:task.originalRequest];
[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]];
Expand Down
2 changes: 1 addition & 1 deletion Agent/Instrumentation/NSURLSession/NRMAURLSessionTaskOverride.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
37 changes: 27 additions & 10 deletions Agent/Instrumentation/NSURLSession/NRMAURLSessionTaskOverride.m
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -108,40 +109,56 @@ 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
NSMutableURLRequest* mutableRequest = [NRMAHTTPUtilities addCrossProcessIdentifier:currentRequest];
mutableRequest = [NRMAHTTPUtilities addConnectivityHeaderAndPayload:mutableRequest];

NRMAPayloadContainer* payload = [NRMAHTTPUtilities addConnectivityHeader:mutableRequest];
[NRMAHTTPUtilities attachPayload:payload
to:task.originalRequest];
[NRMAHTTPUtilities attachPayload:payload to:task.originalRequest];

NSData *data = NRMA__getDataForSessionTask(task);
NRMA__recordTask(task, data, task.response, task.error);

// 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);
}
}
}
}
}
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);
}
}

Expand Down
2 changes: 2 additions & 0 deletions Agent/Public/NRConstants.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,8 @@ typedef NSString NRMetricUnit;
#define NRMA_METRIC_APP_LAUNCH_COLD @"AppLaunch/Cold"
#define NRMA_METRIC_APP_LAUNCH_RESUME @"AppLaunch/Hot"

#define NRMAHandledRequestKey @"NRMAHandledRequest"

// Network Failure Codes
enum NRNetworkFailureCode {
NRURLErrorUnknown = -1,
Expand Down
3 changes: 3 additions & 0 deletions Test Harness/NRTestApp/NRTestApp/AppDelegate.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,9 @@ class AppDelegate: UIResponder, UIApplicationDelegate {
NewRelic.enableFeatures(NRMAFeatureFlags.NRFeatureFlag_SwiftInteractionTracing)
#endif

NewRelic.enableFeatures(NRMAFeatureFlags.NRFeatureFlag_SwiftAsyncURLSessionSupport)


// Generate your own api key to see data get sent to your app's New Relic web services. Also be sure to put your key in the `Run New Relic dSYM Upload Tool` build phase.
guard let apiKey = plistHelper.objectFor(key: "NRAPIKey", plist: "NRAPI-Info") as? String else {return true}

Expand Down
9 changes: 9 additions & 0 deletions Test Harness/NRTestApp/NRTestApp/Helpers/ApodURL.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)"
}
}
18 changes: 16 additions & 2 deletions Test Harness/NRTestApp/NRTestApp/ViewControllers/ViewController.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,11 @@ 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 refreshActionAsync()}))


}

func utilitiesAction() {
Expand All @@ -131,13 +136,22 @@ class ViewController: UIViewController {
func refreshAction() {
viewModel.loadApodData()
}
func brokeRefreshAction() {
viewModel.loadApodDataBrokeData()
}

func refreshActionAsync() {
Task {
await viewModel.loadApodDataAsync()
await viewModel.loadApodDataAsyncBrokeData()
}
}


func brokeRefreshActionAsync() {
Task {
await viewModel.loadApodDataAsyncBrokeData()
}
}

func makeButton(title: String) -> UIButton {
let button = UIButton(type: .system)
button.setTitle(title, for: .normal)
Expand Down
47 changes: 47 additions & 0 deletions Test Harness/NRTestApp/NRTestApp/ViewModels/ApodViewModel.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,4 +58,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
}
}
}
16 changes: 16 additions & 0 deletions Test Harness/NRTestApp/NRTestApp/ViewModels/UtilViewModel.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,11 @@ class UtilViewModel {
try await doAsyncDataTask()
}
}))
options.append(UtilOption(title: "Async URLSession dataTask Fail", handler: { [self] in
Task {
try await doAsyncDataTaskFail()
}
}))

options.append(UtilOption(title: "Shut down New Relic Agent", handler: { [self] in shutDown()}))
}
Expand Down Expand Up @@ -175,6 +180,17 @@ class UtilViewModel {
print("Data: \(data)")
}

func doAsyncDataTaskFail() async throws {
let urlSession = URLSession(configuration: URLSession.shared.configuration, delegate: taskProcessor, delegateQueue: nil)

guard let url = URL(string: "https://www.goo3gle.c3om") else { return }

let request = URLRequest(url: url)
let (data, _) = try await urlSession.data(for: request)

print("Data: \(data)")
}

func shutDown() {
NewRelic.shutdown()
}
Expand Down
Loading

0 comments on commit 8c3d711

Please sign in to comment.