RUM-717 feat!: Add network instrumentation for async/await URLSession…

… APIs

CHANGELOG.md

@@ -1,5 +1,7 @@
 # Unreleased
 
+- [FEATURE] Add network instrumentation for async/await URLSession APIs. See [#1394][]
+
 # 2.3.0 / 02-10-2023
 
 - [IMPROVEMENT] Add UIBackgroundTask for uploading jobs. See [#1412][]
@@ -537,6 +539,7 @@ Release `2.0` introduces breaking changes. Follow the [Migration Guide](MIGRATIO
 [#1502]: https://github.com/DataDog/dd-sdk-ios/pull/1502
 [#1465]: https://github.com/DataDog/dd-sdk-ios/pull/1465
 [#1498]: https://github.com/DataDog/dd-sdk-ios/pull/1498
+[#1394]: https://github.com/DataDog/dd-sdk-ios/pull/1394
 [@00fa9a]: https://github.com/00FA9A
 [@britton-earnin]: https://github.com/Britton-Earnin
 [@hengyu]: https://github.com/Hengyu

Datadog/Datadog.xcodeproj/xcshareddata/xcschemes/DatadogCore iOS.xcscheme

@@ -106,7 +106,7 @@
          </EnvironmentVariable>
          <EnvironmentVariable
             key = "DD_DISABLE_NETWORK_INSTRUMENTATION"
-            value = "$(DD_DISABLE_NETWORK_INSTRUMENTATION)"
+            value = "1"
             isEnabled = "YES">
          </EnvironmentVariable>
          <EnvironmentVariable

Datadog/Datadog.xcodeproj/xcshareddata/xcschemes/DatadogCore tvOS.xcscheme

@@ -92,7 +92,7 @@
          </EnvironmentVariable>
          <EnvironmentVariable
             key = "DD_DISABLE_NETWORK_INSTRUMENTATION"
-            value = "$(DD_DISABLE_NETWORK_INSTRUMENTATION)"
+            value = "1"
             isEnabled = "YES">
          </EnvironmentVariable>
          <EnvironmentVariable

DatadogCore/Tests/DatadogObjc/DDNSURLSessionDelegateTests.swift

@@ -29,6 +29,27 @@ private class DDURLSessionDelegateMock: DDURLSessionDelegate {
 }
 
 class DDNSURLSessionDelegateTests: XCTestCase {
+    private var core: FeatureRegistrationCoreMock! // swiftlint:disable:this implicitly_unwrapped_optional
+
+    override func setUp() {
+        super.setUp()
+
+        core = FeatureRegistrationCoreMock()
+        CoreRegistry.register(default: core)
+
+        let config = DDRUMConfiguration(applicationID: "fake-id")
+        config.setURLSessionTracking(.init())
+        DDRUM.enable(with: config)
+    }
+
+    override func tearDown() {
+        DDURLSessionInstrumentation.disable(delegateClass: DDNSURLSessionDelegate.self)
+        CoreRegistry.unregisterDefault()
+        core = nil
+
+        super.tearDown()
+    }
+
     func testInit() {
         let delegate = DDNSURLSessionDelegate()
         let url = URL(string: "foo.com")

DatadogCore/Tests/DatadogObjc/DDURLSessionInstrumentationConfigurationTests.swift

@@ -0,0 +1,30 @@
+/*
+ * Unless explicitly stated otherwise all files in this repository are licensed under the Apache License Version 2.0.
+ * This product includes software developed at Datadog (https://www.datadoghq.com/).
+ * Copyright 2019-Present Datadog, Inc.
+ */
+
+import XCTest
+import TestUtilities
+import DatadogInternal
+@testable import DatadogObjc
+
+final class DDURLSessionInstrumentationConfigurationTests: XCTestCase {
+    private var objc = DDURLSessionInstrumentationConfiguration(delegateClass: MockDelegate.self)
+    private var swift: URLSessionInstrumentation.Configuration { objc.swiftConfig }
+
+    func testDelegateClass() {
+        XCTAssertTrue(objc.delegateClass === MockDelegate.self)
+    }
+
+    func testFirstPartyHostsTracing() {
+        objc.setFirstPartyHostsTracing(.init(hosts: ["example.com", "example.org"]))
+        DDAssertReflectionEqual(swift.firstPartyHostsTracing, .trace(hosts: ["example.com", "example.org"]))
+
+        objc.setFirstPartyHostsTracing(.init(hostsWithHeaderTypes: ["example.com": [.b3, .datadog]]))
+        DDAssertReflectionEqual(swift.firstPartyHostsTracing, .traceWithHeaders(hostsWithHeaders: ["example.com": [.b3, .datadog]]))
+    }
+
+    class MockDelegate: NSObject, URLSessionDataDelegate {
+    }
+}

DatadogCore/Tests/DatadogObjc/ObjcAPITests/DDNSURLSessionDelegate+apiTests.m

@@ -6,6 +6,7 @@
 
 #import <XCTest/XCTest.h>
 @import DatadogObjc;
+@import DatadogTrace;
 
 @interface DDNSURLSessionDelegate_apiTests : XCTestCase
 @end
@@ -18,6 +19,27 @@ @implementation DDNSURLSessionDelegate_apiTests
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wunused-value"
 
+- (void)setUp {
+    [super setUp];
+
+    DDConfiguration *configuration = [[DDConfiguration alloc] initWithClientToken:@"abc" env:@"def"];
+    [DDDatadog initializeWithConfiguration:configuration trackingConsent:[DDTrackingConsent notGranted]];
+
+    DDTraceConfiguration *config = [[DDTraceConfiguration alloc] init];
+    DDTraceFirstPartyHostsTracing *tracing = [[DDTraceFirstPartyHostsTracing alloc] initWithHosts:[NSSet new] sampleRate:20];
+    DDTraceURLSessionTracking *urlSessionTracking = [[DDTraceURLSessionTracking alloc] initWithFirstPartyHostsTracing:tracing];
+    [config setURLSessionTracking:urlSessionTracking];
+    [DDTrace enableWith:config];
+}
+
+- (void)tearDown {
+    [super tearDown];
+
+    [DDURLSessionInstrumentation disableWithDelegateClass:[DDNSURLSessionDelegate class]];
+    [DDDatadog clearAllData];
+    [DDDatadog flushAndDeinitialize];
+}
+
 - (void)testDDNSURLSessionDelegateAPI {
     [[DDNSURLSessionDelegate alloc] init];
     [[DDNSURLSessionDelegate alloc] initWithAdditionalFirstPartyHosts:[NSSet setWithArray:@[]]];

DatadogCore/Tests/DatadogObjc/ObjcAPITests/DDURLSessionInstrumentationTests+apiTests.m

@@ -0,0 +1,68 @@
+/*
+* Unless explicitly stated otherwise all files in this repository are licensed under the Apache License Version 2.0.
+* This product includes software developed at Datadog (https://www.datadoghq.com/).
+* Copyright 2019-Present Datadog, Inc.
+*/
+
+#import <XCTest/XCTest.h>
+#include <sys/wait.h>
+@import DatadogObjc;
+@import DatadogTrace;
+
+#import <Foundation/Foundation.h>
+
+@interface MockDelegate : NSObject <NSURLSessionDataDelegate>
+@end
+
+@implementation MockDelegate
+@end
+
+@interface DDURLSessionInstrumentationTests_apiTests : XCTestCase
+@end
+
+@implementation DDURLSessionInstrumentationTests_apiTests
+
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wunused-value"
+
+- (void)setUp {
+    [super setUp];
+
+    DDConfiguration *configuration = [[DDConfiguration alloc] initWithClientToken:@"abc" env:@"def"];
+    [DDDatadog initializeWithConfiguration:configuration trackingConsent:[DDTrackingConsent notGranted]];
+
+    DDTraceConfiguration *config = [[DDTraceConfiguration alloc] init];
+    DDTraceFirstPartyHostsTracing *tracing = [[DDTraceFirstPartyHostsTracing alloc] initWithHosts:[NSSet new] sampleRate:20];
+    DDTraceURLSessionTracking *urlSessionTracking = [[DDTraceURLSessionTracking alloc] initWithFirstPartyHostsTracing:tracing];
+    [config setURLSessionTracking:urlSessionTracking];
+    [DDTrace enableWith:config];
+}
+
+- (void)tearDown {
+    [super tearDown];
+
+    [DDDatadog clearAllData];
+    [DDDatadog flushAndDeinitialize];
+}
+
+- (void)testWorkflow {
+    XCTestExpectation *expectation = [self expectationWithDescription:@"task completed"];
+    DDURLSessionInstrumentationConfiguration *config = [[DDURLSessionInstrumentationConfiguration alloc] initWithDelegateClass:[MockDelegate class]];
+    [DDURLSessionInstrumentation enableWithConfiguration:config];
+
+    NSURLSession *session = [NSURLSession sessionWithConfiguration:[NSURLSessionConfiguration defaultSessionConfiguration]
+                                                          delegate:[MockDelegate new] delegateQueue:nil];
+    NSURLSessionTask *task = [session dataTaskWithURL:[NSURL URLWithString:@"https://status.datadoghq.com"]
+                                    completionHandler:^(NSData * _Nullable data, NSURLResponse * _Nullable response, NSError * _Nullable error) {
+        [expectation fulfill];
+    }];
+    [task resume];
+
+    [self waitForExpectationsWithTimeout:10 handler:nil];
+
+    [DDURLSessionInstrumentation disableWithDelegateClass:[MockDelegate class]];
+}
+
+#pragma clang diagnostic pop
+
+@end

DatadogCore/Tests/TestsObserver/DatadogTestsObserver.swift

@@ -33,22 +33,13 @@ internal class DatadogTestsObserver: NSObject, XCTestObservation {
             """
         ),
         .init(
-            assert: { activeSwizzlingNames.isEmpty },
+            assert: { Swizzling.activeSwizzlingNames.isEmpty },
             problem: "No swizzling must be applied.",
             solution: """
             Make sure all applied swizzling are reset by the end of test with `unswizzle()`.
 
-            `DatadogTestsObserver` found \(activeSwizzlingNames.count) leaked swizzlings:
-            \(activeSwizzlingNames.joined(separator: ", "))
-            """
-        ),
-        .init(
-            assert: { URLSessionSwizzler.bindingsCount == 0 },
-            problem: "No `URLSessionSwizzler` must be bonded.",
-            solution: """
-            Make sure all applied `URLSessionSwizzler.bind()` are reset by the end of test with `URLSessionSwizzler.unbind()`.
-
-            `DatadogTestsObserver` found \(URLSessionSwizzler.bindingsCount) bindings left.
+            `DatadogTestsObserver` found \(Swizzling.activeSwizzlingNames.count) leaked swizzlings:
+            \(Swizzling.activeSwizzlingNames.joined(separator: ", "))
             """
         ),
         .init(
@@ -141,6 +132,34 @@ internal class DatadogTestsObserver: NSObject, XCTestObservation {
 
             If all above conditions are met, this failure might indicate a memory leak in the implementation.
             """
+        ),
+        .init(
+            assert: { URLSessionTaskDelegateSwizzler.isBinded == false },
+            problem: "No URLSessionTaskDelegate swizzling must be applied.",
+            solution: """
+            Make sure all the binded delegates are unbinded by the end of test with `URLSessionTaskDelegateSwizzler.unbind(delegate:)`.
+            """
+        ),
+        .init(
+            assert: { URLSessionDataDelegateSwizzler.isBinded == false },
+            problem: "No URLSessionDataDelegate swizzling must be applied.",
+            solution: """
+            Make sure all the binded delegates are unbinded by the end of test with `URLSessionDataDelegateSwizzler.unbind(delegate:)`.
+            """
+        ),
+        .init(
+            assert: { URLSessionTaskSwizzler.isBinded == false },
+            problem: "No URLSessionTask swizzling must be applied.",
+            solution: """
+            Make sure all the binded delegates are unbinded by the end of test with `URLSessionTaskSwizzler.unbind()`.
+            """
+        ),
+        .init(
+            assert: { URLSessionSwizzler.isBinded == false },
+            problem: "No URLSession swizzling must be applied.",
+            solution: """
+            Make sure all the binded delegates are unbinded by the end of test with `URLSessionSwizzler.unbind()`.
+            """
         )
     ]
 

DatadogInternal/Sources/NetworkInstrumentation/DatadogURLSessionHandler.swift

@@ -46,7 +46,7 @@ extension DatadogCoreProtocol {
     ///
     /// - Parameter urlSessionHandler: The `URLSession` handlers to register.
     public func register(urlSessionHandler: DatadogURLSessionHandler) throws {
-        let feature = try get(feature: NetworkInstrumentationFeature.self) ?? .init()
+        let feature = get(feature: NetworkInstrumentationFeature.self) ?? .init()
         feature.handlers.append(urlSessionHandler)
         try register(feature: feature)
     }

DatadogInternal/Sources/NetworkInstrumentation/FirstPartyHosts.swift

@@ -8,7 +8,7 @@ import Foundation
 
 /// A struct that represents a dictionary of host names and tracing header types.
 public struct FirstPartyHosts: Equatable {
-    fileprivate var hostsWithTracingHeaderTypes: [String: Set<TracingHeaderType>]
+    internal var hostsWithTracingHeaderTypes: [String: Set<TracingHeaderType>]
 
     public var hosts: Set<String> {
         return Set(hostsWithTracingHeaderTypes.keys)
@@ -37,6 +37,17 @@ public struct FirstPartyHosts: Equatable {
         self.init(hostsWithTracingHeaderTypes: [:])
     }
 
+    internal init?(firstPartyHosts: URLSessionInstrumentation.FirstPartyHostsTracing?) {
+        switch firstPartyHosts {
+        case .trace(let hosts):
+            self.init(hosts)
+        case .traceWithHeaders(let hostsWithHeaders):
+            self.init(hostsWithTracingHeaderTypes: hostsWithHeaders)
+        default:
+            return nil
+        }
+    }
+
     internal init(
         hostsWithTracingHeaderTypes: [String: Set<TracingHeaderType>],
         hostsSanitizer: HostsSanitizing = HostsSanitizer()