Skip to content

Commit 99b2e2c

Browse files
author
Ignacio Bonafonte
authored
Merge pull request #398 from open-telemetry/fix-network-instrumentation-for-ios16
Fix URLSession instrumentation in iOS 16 iOS 16 broke async/await URLSession instrumentation, add extra code to make it work: Only NSURLSessionTask.resume and URLSessionTaskDelegate.urlSession(_:task:didFinishCollecting:) are being called (contrary to what documentation says), try to move all the logic to these two methods Also for async/await calls without delegate set a task delegate so that the delegate method is called and we can capture it Added many tests exercising async await network calls
2 parents fd2171e + d9909d7 commit 99b2e2c

File tree

4 files changed

+242
-10
lines changed

4 files changed

+242
-10
lines changed

‎Sources/Exporters/DatadogExporter/Utils/Device.swift‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
#if os(macOS)
77
import Foundation
88
import SystemConfiguration
9-
#elseif os(iOS) || targetEnvironment(macCatalyst)
9+
#elseif os(iOS) || targetEnvironment(macCatalyst) || os(tvOS)
1010
import UIKit
1111
#elseif os(watchOS)
1212
import WatchKit
@@ -23,7 +23,8 @@ internal class Device {
2323
init(
2424
model: String,
2525
osName: String,
26-
osVersion: String) {
26+
osVersion: String)
27+
{
2728
self.model = model
2829
self.osName = osName
2930
self.osVersion = osVersion
@@ -36,6 +37,7 @@ internal class Device {
3637
osName: uiDevice.systemName,
3738
osVersion: uiDevice.systemVersion)
3839
}
40+
3941
#elseif os(macOS)
4042
convenience init(processInfo: ProcessInfo) {
4143
self.init(
@@ -64,8 +66,7 @@ internal class Device {
6466
return Device(
6567
model: device.model,
6668
osName: device.systemName,
67-
osVersion: device.systemVersion
68-
)
69+
osVersion: device.systemVersion)
6970
#endif
7071
}
7172
}

‎Sources/Instrumentation/URLSession/InstrumentationUtils.swift‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,11 @@
44
*/
55

66
import Foundation
7+
#if os(iOS) || os(tvOS)
8+
import UIKit
9+
#elseif os(watchOS)
10+
import WatchKit
11+
#endif
712

813
enum InstrumentationUtils {
914
static func objc_getClassList() -> [AnyClass] {
@@ -48,4 +53,28 @@ enum InstrumentationUtils {
4853
ptr = ptr.successor()
4954
}
5055
}
56+
57+
static var usesUndocumentedAsyncAwaitMethods: Bool = {
58+
#if os(macOS)
59+
let os = ProcessInfo.processInfo.operatingSystemVersion
60+
if os.majorVersion >= 13 {
61+
return true
62+
}
63+
#elseif os(watchOS)
64+
let version = WKInterfaceDevice.current().systemVersion
65+
if let versionNumber = Double(version),
66+
versionNumber >= 9.0
67+
{
68+
return true
69+
}
70+
#else
71+
let version = UIDevice.current.systemVersion
72+
if let versionNumber = Double(version),
73+
versionNumber >= 16.0
74+
{
75+
return true
76+
}
77+
#endif
78+
return false
79+
}()
5180
}

‎Sources/Instrumentation/URLSession/URLSessionInstrumentation.swift‎

Lines changed: 39 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,9 @@ public class URLSessionInstrumentation {
105105
injectTaskDidCompleteWithErrorIntoDelegateClass(cls: cls)
106106
injectRespondsToSelectorIntoDelegateClass(cls: cls)
107107
// For future use
108-
// injectTaskDidFinishCollectingMetricsIntoDelegateClass(cls: cls)
108+
if InstrumentationUtils.usesUndocumentedAsyncAwaitMethods {
109+
injectTaskDidFinishCollectingMetricsIntoDelegateClass(cls: cls)
110+
}
109111

110112
// Data tasks
111113
injectDataTaskDidBecomeDownloadTaskIntoDelegateClass(cls: cls)
@@ -549,18 +551,47 @@ public class URLSessionInstrumentation {
549551
let taskId = self.idKeyForTask(task)
550552
if (self.requestMap[taskId]?.request) != nil {
551553
/// Code for instrumenting colletion should be written here
554+
var requestState: NetworkRequestState?
555+
queue.sync {
556+
requestState = requestMap[taskId]
557+
if requestState != nil {
558+
requestMap[taskId] = nil
559+
}
560+
}
561+
if let error = task.error {
562+
let status = (task.response as? HTTPURLResponse)?.statusCode ?? 0
563+
URLSessionLogger.logError(error, dataOrFile: requestState?.dataProcessed, statusCode: status, instrumentation: self, sessionTaskId: taskId)
564+
} else if let response = task.response {
565+
URLSessionLogger.logResponse(response, dataOrFile: requestState?.dataProcessed, instrumentation: self, sessionTaskId: taskId)
566+
}
552567
}
553568
}
554569

555-
private func urlSessionTaskWillResume(_ session: URLSessionTask) {
556-
let taskId = self.idKeyForTask(session)
557-
if let request = session.currentRequest {
570+
private func urlSessionTaskWillResume(_ task: URLSessionTask) {
571+
let taskId = self.idKeyForTask(task)
572+
if let request = task.currentRequest {
558573
queue.sync {
559574
if requestMap[taskId] == nil {
560575
requestMap[taskId] = NetworkRequestState()
561576
}
562577
requestMap[taskId]?.setRequest(request)
563578
}
579+
580+
if InstrumentationUtils.usesUndocumentedAsyncAwaitMethods {
581+
if #available(OSX 13.0, iOS 16.0, watchOS 9.0, tvOS 16.0, *) {
582+
guard Task.basePriority != nil else {
583+
return
584+
}
585+
let instrumentedRequest = URLSessionLogger.processAndLogRequest(request, sessionTaskId: taskId, instrumentation: self, shouldInjectHeaders: true)
586+
task.setValue(instrumentedRequest, forKey: "currentRequest")
587+
self.setIdKey(value: taskId, for: task)
588+
589+
// If not inside a Task basePriority is nil
590+
if task.delegate == nil {
591+
task.delegate = FakeDelegate()
592+
}
593+
}
594+
}
564595
}
565596
}
566597

@@ -586,3 +617,7 @@ public class URLSessionInstrumentation {
586617
}
587618
}
588619
}
620+
621+
class FakeDelegate: NSObject, URLSessionTaskDelegate {
622+
func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {}
623+
}

‎Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift‎

Lines changed: 169 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -195,9 +195,9 @@ class URLSessionInstrumentationTests: XCTestCase {
195195
}
196196

197197
public func testConfigurationCallbacksCalledWhenForbidden() throws {
198-
if #available(watchOS 3.0, *) {
198+
#if os(watchOS)
199199
throw XCTSkip("Implementation needs to be updated for watchOS to make this test pass")
200-
}
200+
#endif
201201

202202
let request = URLRequest(url: URL(string: "http://localhost:33333/forbidden")!)
203203
let task = URLSession.shared.dataTask(with: request) { data, _, _ in
@@ -388,4 +388,171 @@ class URLSessionInstrumentationTests: XCTestCase {
388388
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
389389
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
390390
}
391+
392+
393+
394+
#if swift(>=5.5.2)
395+
@available(macOS 10.15, iOS 13.0, tvOS 13.0, *)
396+
public func testConfigurationCallbacksCalledWhenSuccessAsync() async throws {
397+
let request = URLRequest(url: URL(string: "http://localhost:33333/success")!)
398+
399+
let (data, _) = try await URLSession.shared.data(for: request)
400+
let string = String(decoding: data, as: UTF8.self)
401+
print(string)
402+
403+
XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInstrumentCalled)
404+
XCTAssertTrue(URLSessionInstrumentationTests.checker.nameSpanCalled)
405+
XCTAssertTrue(URLSessionInstrumentationTests.checker.spanCustomizationCalled)
406+
XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInjectTracingHeadersCalled)
407+
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
408+
XCTAssertTrue(URLSessionInstrumentationTests.checker.receivedResponseCalled)
409+
}
410+
411+
@available(macOS 10.15, iOS 13.0, tvOS 13.0, *)
412+
public func testConfigurationCallbacksCalledWhenForbiddenAsync() async throws {
413+
#if os(watchOS)
414+
throw XCTSkip("Implementation needs to be updated for watchOS to make this test pass")
415+
#endif
416+
let request = URLRequest(url: URL(string: "http://localhost:33333/forbidden")!)
417+
let (data, _) = try await URLSession.shared.data(for: request)
418+
let string = String(decoding: data, as: UTF8.self)
419+
print(string)
420+
421+
XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInstrumentCalled)
422+
XCTAssertTrue(URLSessionInstrumentationTests.checker.nameSpanCalled)
423+
XCTAssertTrue(URLSessionInstrumentationTests.checker.spanCustomizationCalled)
424+
XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInjectTracingHeadersCalled)
425+
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
426+
XCTAssertTrue(URLSessionInstrumentationTests.checker.receivedResponseCalled)
427+
XCTAssertFalse(URLSessionInstrumentationTests.checker.receivedErrorCalled)
428+
}
429+
430+
@available(macOS 10.15, iOS 13.0, tvOS 13.0, *)
431+
public func testConfigurationCallbacksCalledWhenErrorAsync() async throws {
432+
let request = URLRequest(url: URL(string: "http://localhost:33333/error")!)
433+
434+
do {
435+
_ = try await URLSession.shared.data(for: request)
436+
} catch {
437+
}
438+
439+
XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInstrumentCalled)
440+
XCTAssertTrue(URLSessionInstrumentationTests.checker.nameSpanCalled)
441+
XCTAssertTrue(URLSessionInstrumentationTests.checker.spanCustomizationCalled)
442+
XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInjectTracingHeadersCalled)
443+
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
444+
XCTAssertFalse(URLSessionInstrumentationTests.checker.receivedResponseCalled)
445+
XCTAssertTrue(URLSessionInstrumentationTests.checker.receivedErrorCalled)
446+
}
447+
448+
@available(macOS 10.15, iOS 13.0, tvOS 13.0, *)
449+
public func testDataTaskWithRequestBlockAsync() async throws {
450+
let request = URLRequest(url: URL(string: "http://localhost:33333/success")!)
451+
452+
_ = try await URLSession.shared.data(for: request)
453+
454+
XCTAssertEqual(0, URLSessionInstrumentationTests.instrumentation.startedRequestSpans.count)
455+
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
456+
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
457+
}
458+
459+
@available(macOS 10.15, iOS 13.0, tvOS 13.0, *)
460+
public func testDataTaskWithUrlBlockAsync() async throws {
461+
let url = URL(string: "http://localhost:33333/success")!
462+
463+
_ = try await URLSession.shared.data(from: url)
464+
465+
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
466+
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
467+
}
468+
469+
@available(macOS 12, iOS 15.0, tvOS 15.0, watchOS 8.0, *)
470+
public func testDownloadTaskWithUrlBlockAsync() async throws {
471+
let url = URL(string: "http://localhost:33333/success")!
472+
473+
_ = try await URLSession.shared.download(from: url)
474+
475+
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
476+
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
477+
}
478+
479+
@available(macOS 12, iOS 15.0, tvOS 15.0, watchOS 8.0, *)
480+
public func testDownloadTaskWithRequestBlockAsync() async throws {
481+
let url = URL(string: "http://localhost:33333/success")!
482+
let request = URLRequest(url: url)
483+
_ = try await URLSession.shared.download(for: request)
484+
485+
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
486+
XCTAssertTrue(URLSessionInstrumentationTests.checker.receivedResponseCalled)
487+
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
488+
}
489+
490+
@available(macOS 10.15, iOS 13.0, tvOS 13.0, *)
491+
public func testUploadTaskWithRequestBlockAsync() async throws {
492+
let url = URL(string: "http://localhost:33333/success")!
493+
let request = URLRequest(url: url)
494+
_ = try await URLSession.shared.upload(for: request, from: Data())
495+
496+
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
497+
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
498+
}
499+
500+
@available(macOS 10.15, iOS 13.0, tvOS 13.0, *)
501+
public func
502+
testDataTaskWithRequestDelegateAsync() async throws {
503+
let request = URLRequest(url: URL(string: "http://localhost:33333/success")!)
504+
505+
let session = URLSession(configuration: URLSessionConfiguration.default, delegate: sessionDelegate, delegateQueue: nil)
506+
_ = try await session.data(for: request)
507+
508+
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
509+
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
510+
}
511+
512+
@available(macOS 10.15, iOS 13.0, tvOS 13.0, *)
513+
public func testDataTaskWithUrlDelegateAsync() async throws {
514+
let url = URL(string: "http://localhost:33333/success")!
515+
516+
let session = URLSession(configuration: URLSessionConfiguration.default, delegate: sessionDelegate, delegateQueue: nil)
517+
_ = try await session.data(from: url)
518+
519+
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
520+
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
521+
}
522+
523+
@available(macOS 12, iOS 15.0, tvOS 15.0, watchOS 8.0, *)
524+
public func testDownloadTaskWithUrlDelegateAsync() async throws {
525+
let url = URL(string: "http://localhost:33333/success")!
526+
527+
_ = try await URLSession.shared.download(from: url, delegate: sessionDelegate)
528+
529+
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
530+
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
531+
}
532+
533+
@available(macOS 12, iOS 15.0, tvOS 15.0, watchOS 8.0, *)
534+
public func testDownloadTaskWithRequestDelegateAsync() async throws {
535+
let url = URL(string: "http://localhost:33333/success")!
536+
let request = URLRequest(url: url)
537+
538+
let session = URLSession(configuration: URLSessionConfiguration.default, delegate: sessionDelegate, delegateQueue: nil)
539+
_ = try await session.download(for: request)
540+
541+
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
542+
XCTAssertTrue(URLSessionInstrumentationTests.checker.receivedResponseCalled)
543+
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
544+
}
545+
546+
@available(macOS 10.15, iOS 13.0, tvOS 13.0, *)
547+
public func testUploadTaskWithRequestDelegateAsync() async throws {
548+
let url = URL(string: "http://localhost:33333/success")!
549+
let request = URLRequest(url: url)
550+
let session = URLSession(configuration: URLSessionConfiguration.default, delegate: sessionDelegate, delegateQueue: nil)
551+
_ = try await session.upload(for: request, from: Data())
552+
553+
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
554+
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
555+
}
556+
557+
#endif
391558
}

0 commit comments

Comments
 (0)