RUMM-731 Refactor URLSession instrumentation to support RUM and Tracing

Datadog/Example/Scenarios/Tracing/NSURLSessionAutoInstrumentation/TracingNSURLSessionViewController.m

@@ -6,6 +6,7 @@
 
 #import "Example-Swift.h"
 #import "TracingNSURLSessionViewController.h"
+@import DatadogObjc;
 
 @interface TracingNSURLSessionViewController()
 @property TracingNSURLSessionScenario *testScenario;
@@ -17,24 +18,28 @@ @implementation TracingNSURLSessionViewController
 -(void)awakeFromNib {
     [super awakeFromNib];
     self.testScenario = SwiftGlobals.currentTestScenario;
-    self.session = [NSURLSession sharedSession];
+    self.session = [NSURLSession sessionWithConfiguration:[NSURLSessionConfiguration defaultSessionConfiguration]
+                                                 delegate:[DDNSURLSessionDelegate new]
+                                            delegateQueue:nil];
 
     assert(self.testScenario != nil);
 }
 
 - (void)viewDidLoad {
     [super viewDidLoad];
 
-    [self callSuccessfullURLWithCompletionHandler: ^{
-        [self callSuccessfullURLRequestWithCompletionHandler: ^{
-            [self callBadURLWithCompletionHandler: ^{
-                [self useNotInstrumentedAPIs];
-            }];
+    [self callSuccessfullFirstPartyURLWithCompletionHandler: ^{
+        [self callSuccessfullFirstPartyURLRequestWithCompletionHandler: ^{
+            [self callBadFirstPartyURL];
         }];
     }];
+
+    [self callThirdPartyURL];
+    [self callThirdPartyURLRequest];
 }
 
-- (void)callSuccessfullURLWithCompletionHandler:(void (^)(void))completionHandler {
+- (void)callSuccessfullFirstPartyURLWithCompletionHandler:(void (^)(void))completionHandler {
+    // This request is instrumented. It sends the `Span`.
     NSURLSessionTask *task = [self.session dataTaskWithURL:self.testScenario.customGETResourceURL
                                          completionHandler:^(NSData * _Nullable data, NSURLResponse * _Nullable response, NSError * _Nullable error) {
         assert(error == nil);
@@ -43,7 +48,8 @@ - (void)callSuccessfullURLWithCompletionHandler:(void (^)(void))completionHandle
     [task resume];
 }
 
-- (void)callSuccessfullURLRequestWithCompletionHandler:(void (^)(void))completionHandler {
+- (void)callSuccessfullFirstPartyURLRequestWithCompletionHandler:(void (^)(void))completionHandler {
+    // This request is instrumented. It sends the `Span`.
     NSURLSessionTask *task = [self.session dataTaskWithRequest:self.testScenario.customPOSTRequest
                                              completionHandler:^(NSData * _Nullable data, NSURLResponse * _Nullable response, NSError * _Nullable error) {
         assert(error == nil);
@@ -52,28 +58,22 @@ - (void)callSuccessfullURLRequestWithCompletionHandler:(void (^)(void))completio
     [task resume];
 }
 
-- (void)callBadURLWithCompletionHandler:(void (^)(void))completionHandler {
-    NSURLSessionTask *task = [self.session dataTaskWithURL:self.testScenario.badResourceURL
-                                         completionHandler:^(NSData * _Nullable data, NSURLResponse * _Nullable response, NSError * _Nullable error) {
-        assert(error != nil);
-        completionHandler();
-    }];
+- (void)callBadFirstPartyURL {
+    // This request is instrumented. It sends the `Span`.
+    NSURLSessionTask *task = [self.session dataTaskWithURL:self.testScenario.badResourceURL];
     [task resume];
 }
 
-/// Calls `NSURLSession` APIs which are currently not auto instrumented.
-/// This is just a sanity check to make sure the `URLSession` swizzling works fine in different edge case usages of the `NSURLSession`.
-- (void)useNotInstrumentedAPIs {
-    NSURLRequest *badResourceRequest = [[NSURLRequest alloc] initWithURL:self.testScenario.badResourceURL];
-    // Use APIs with no completion block:
-    [[self.session dataTaskWithRequest:badResourceRequest] resume];
-    [[self.session dataTaskWithURL:self.testScenario.badResourceURL] resume];
-    // Nullify the completion handlers:
-#pragma clang diagnostic push
-#pragma clang diagnostic ignored "-Wnonnull"
-    [[self.session dataTaskWithRequest:badResourceRequest completionHandler:nil] resume];
-    [[self.session dataTaskWithURL:self.testScenario.badResourceURL completionHandler:nil] resume];
-#pragma clang diagnostic pop
+- (void)callThirdPartyURL {
+    // This request is NOT instrumented. We test that it does not send the `Span`.
+    NSURLSessionTask *task = [self.session dataTaskWithURL:self.testScenario.thirdPartyURL];
+    [task resume];
+}
+
+- (void)callThirdPartyURLRequest {
+    // This request is NOT instrumented. We test that it does not send the `Span`.
+    NSURLSessionTask *task = [self.session dataTaskWithRequest:self.testScenario.thirdPartyRequest];
+    [task resume];
 }
 
 @end

Datadog/Example/Scenarios/Tracing/URLSessionAutoInstrumentation/TracingURLSessionViewController.swift

@@ -6,10 +6,15 @@
 
 import Foundation
 import UIKit
+import Datadog
 
 internal class TracingURLSessionViewController: UIViewController {
     private var testScenario: TracingURLSessionScenario!
-    private let session = URLSession.shared
+    private lazy var session = URLSession(
+        configuration: .default,
+        delegate: DDURLSessionDelegate(),
+        delegateQueue: nil
+    )
 
     override func awakeFromNib() {
         super.awakeFromNib()
@@ -19,45 +24,49 @@ internal class TracingURLSessionViewController: UIViewController {
     override func viewDidLoad() {
         super.viewDidLoad()
 
-        callSuccessfullURL {
-            self.callSuccessfullURLRequest {
-                self.callBadURL {
-                    self.useNotInstrumentedAPIs()
-                }
+        callSuccessfullFirstPartyURL {
+            self.callSuccessfullFirstPartyURLRequest {
+                self.callBadFirstPartyURL()
             }
         }
+
+        callThirdPartyURL()
+        callThirdPartyURLRequest()
     }
 
-    private func callSuccessfullURL(completion: @escaping () -> Void) {
+    private func callSuccessfullFirstPartyURL(completion: @escaping () -> Void) {
+        // This request is instrumented. It sends the `Span`.
         let task = session.dataTask(with: testScenario.customGETResourceURL) { _, _, error in
             assert(error == nil)
             completion()
         }
         task.resume()
     }
 
-    private func callSuccessfullURLRequest(completion: @escaping () -> Void) {
+    private func callSuccessfullFirstPartyURLRequest(completion: @escaping () -> Void) {
+        // This request is instrumented. It sends the `Span`.
         let task = session.dataTask(with: testScenario.customPOSTRequest) { _, _, error in
             assert(error == nil)
             completion()
         }
         task.resume()
     }
 
-    private func callBadURL(completion: @escaping () -> Void) {
-        let task = session.dataTask(with: testScenario.badResourceURL) { _, _, error in
-            assert(error != nil)
-            completion()
-        }
+    private func callBadFirstPartyURL() {
+        // This request is instrumented. It sends the `Span`.
+        let task = session.dataTask(with: testScenario.badResourceURL)
         task.resume()
     }
 
-    /// Calls `URLSession` APIs which are currently not auto instrumented.
-    /// This is just a sanity check to make sure the `URLSession` swizzling works fine.
-    private func useNotInstrumentedAPIs() {
-        let badResourceRequest = URLRequest(url: testScenario.badResourceURL)
-        // Use APIs with no completion block:
-        session.dataTask(with: badResourceRequest).resume()
-        session.dataTask(with: testScenario.badResourceURL).resume()
+    private func callThirdPartyURL() {
+        // This request is NOT instrumented. We test that it does not send the `Span`.
+        let task = session.dataTask(with: testScenario.thirdPartyURL)
+        task.resume()
+    }
+
+    private func callThirdPartyURLRequest() {
+        // This request is NOT instrumented. We test that it does not send the `Span`.
+        let task = session.dataTask(with: testScenario.thirdPartyRequest)
+        task.resume()
     }
 }

Datadog/Example/TestScenarios.swift

@@ -97,13 +97,23 @@ class _TracingURLSessionBaseScenario: NSObject {
     /// The URL to custom GET resource, observed by Tracing auto instrumentation.
     @objc
     let customGETResourceURL: URL
+
     /// The `URLRequest` to custom POST resource,  observed by Tracing auto instrumentation.
     @objc
     let customPOSTRequest: URLRequest
+
     /// An unresolvable URL to fake resource DNS resolution error,  observed by Tracing auto instrumentation.
     @objc
     let badResourceURL: URL
 
+    /// The `URLRequest` to fake 3rd party resource. As it's 3rd party, it won't be observed by Tracing auto instrumentation.
+    @objc
+    let thirdPartyRequest: URLRequest
+
+    /// The `URL` to fake 3rd party resource. As it's 3rd party, it won't be observed by Tracing auto instrumentation.
+    @objc
+    let thirdPartyURL: URL
+
     override init() {
         if ProcessInfo.processInfo.arguments.contains("IS_RUNNING_UI_TESTS") {
             let customURL = Environment.customEndpointURL()!
@@ -125,11 +135,13 @@ class _TracingURLSessionBaseScenario: NSObject {
             }()
             badResourceURL = URL(string: "https://foo.bar")!
         }
+        thirdPartyURL = URL(string: "https://www.bitrise.io")!
+        thirdPartyRequest = URLRequest(url: thirdPartyURL)
     }
 
     func configureSDK(builder: Datadog.Configuration.Builder) {
         _ = builder
-            .set(tracedHosts: [customGETResourceURL.host!, customPOSTRequest.url!.host!, badResourceURL.host!])
+            .track(firstPartyHosts: [customGETResourceURL.host!, customPOSTRequest.url!.host!, badResourceURL.host!])
     }
 }