Skip to content

Commit 5999e62

Browse files
MustafaHaddaraMustafaHaddara
authored
URLInstrumentation: only set a Task Delegate if there is no Session Delegate (#747)
* remove DelegateProxy, only set FakeDelegate if there is no task delegate or session delegate * add task delegate tests for upload and download tasks
1 parent 93fa564 commit 5999e62

File tree

3 files changed

+152
-7
lines changed

3 files changed

+152
-7
lines changed

‎Examples/Network Sample/main.swift‎

Lines changed: 88 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,14 +28,54 @@ func simpleNetworkCall() {
2828

2929
class SessionDelegate: NSObject, URLSessionDataDelegate, URLSessionTaskDelegate {
3030
let semaphore = DispatchSemaphore(value: 0)
31+
var callCount = 0
3132

3233
func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
3334
semaphore.signal()
3435
}
36+
37+
func urlSession(
38+
_ session: URLSession,
39+
task: URLSessionTask,
40+
didFinishCollecting metrics: URLSessionTaskMetrics
41+
) {
42+
semaphore.signal()
43+
callCount += 1
44+
print("delegate called")
45+
}
3546
}
3647

3748
let delegate = SessionDelegate()
3849

50+
enum TimeoutError: Error {
51+
case timeout
52+
}
53+
54+
func waitForSemaphore(withTimeoutSecs: Int) async {
55+
do {
56+
let _ = try await withThrowingTaskGroup(of: Bool.self) { group in
57+
group.addTask {
58+
try await Task.sleep(nanoseconds: UInt64(withTimeoutSecs) * NSEC_PER_SEC)
59+
throw TimeoutError.timeout
60+
}
61+
group.addTask {
62+
let semaphoreTask = Task {
63+
DispatchQueue.global().async {
64+
delegate.semaphore.wait()
65+
}
66+
}
67+
await semaphoreTask.value
68+
try Task.checkCancellation()
69+
return true
70+
}
71+
72+
return try await group.next()!
73+
}
74+
} catch {
75+
print("timed out waiting for semaphore")
76+
}
77+
}
78+
3979
func simpleNetworkCallWithDelegate() {
4080
let session = URLSession(configuration: .default, delegate: delegate, delegateQueue: nil)
4181

@@ -48,6 +88,38 @@ func simpleNetworkCallWithDelegate() {
4888
delegate.semaphore.wait()
4989
}
5090

91+
@available(macOS 10.15, iOS 15.0, watchOS 8.0, tvOS 15.0, *)
92+
func asyncNetworkCallWithTaskDelegate() async {
93+
let session = URLSession(configuration: .default)
94+
95+
let url = URL(string: "http://httpbin.org/get")!
96+
let request = URLRequest(url: url)
97+
98+
do {
99+
_ = try await session.data(for: request, delegate: delegate)
100+
} catch {
101+
return
102+
}
103+
104+
await waitForSemaphore(withTimeoutSecs: 3)
105+
}
106+
107+
@available(macOS 10.15, iOS 15.0, tvOS 13.0, *)
108+
func asyncNetworkCallWithSessionDelegate() async {
109+
let session = URLSession(configuration: .default, delegate: delegate, delegateQueue: nil)
110+
111+
let url = URL(string: "http://httpbin.org/get")!
112+
let request = URLRequest(url: url)
113+
114+
do {
115+
_ = try await session.data(for: request)
116+
} catch {
117+
return
118+
}
119+
120+
await waitForSemaphore(withTimeoutSecs: 3)
121+
}
122+
51123
let spanProcessor = SimpleSpanProcessor(spanExporter: StdoutSpanExporter(isDebug: true))
52124
OpenTelemetry.registerTracerProvider(tracerProvider:
53125
TracerProviderBuilder()
@@ -57,6 +129,21 @@ OpenTelemetry.registerTracerProvider(tracerProvider:
57129

58130
let networkInstrumentation = URLSessionInstrumentation(configuration: URLSessionInstrumentationConfiguration())
59131

60-
simpleNetworkCall()
132+
print("making simple call")
133+
var callCount = delegate.callCount
61134
simpleNetworkCallWithDelegate()
135+
assert(delegate.callCount == callCount + 1)
136+
137+
if #available(macOS 10.15, iOS 15.0, watchOS 8.0, tvOS 15.0, *) {
138+
print("making simple call with task delegate")
139+
callCount = delegate.callCount
140+
await asyncNetworkCallWithTaskDelegate()
141+
assert(delegate.callCount == callCount + 1, "async task delegate not called")
142+
143+
print("making simple call with session delegate")
144+
callCount = delegate.callCount
145+
await asyncNetworkCallWithSessionDelegate()
146+
assert(delegate.callCount == callCount + 1, "async session delegate not called")
147+
}
148+
62149
sleep(1)

‎Sources/Instrumentation/URLSession/URLSessionInstrumentation.swift‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -733,8 +733,8 @@ public class URLSessionInstrumentation {
733733
task.setValue(instrumentedRequest, forKey: "currentRequest")
734734
}
735735
self.setIdKey(value: taskId, for: task)
736-
737-
if task.delegate == nil, task.state != .running {
736+
737+
if task.delegate == nil, task.state != .running, (task.value(forKey: "session") as? URLSession)?.delegate == nil {
738738
task.delegate = FakeDelegate()
739739
}
740740
}

‎Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift‎

Lines changed: 62 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,23 @@ class URLSessionInstrumentationTests: XCTestCase {
3535
semaphore.signal()
3636
}
3737
}
38+
39+
class CountingSessionDelegate: NSObject, URLSessionDelegate, URLSessionDataDelegate {
40+
var callCount: Int = 0
41+
42+
func urlSession(_ session: URLSession, didBecomeInvalidWithError error: Error?) {
43+
callCount += 1
44+
45+
}
46+
47+
func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
48+
callCount += 1
49+
}
50+
51+
func urlSession(_ session: URLSession, task: URLSessionTask, didFinishCollecting metrics: URLSessionTaskMetrics) {
52+
callCount += 1
53+
}
54+
}
3855

3956
static var requestCopy: URLRequest!
4057
static var responseCopy: HTTPURLResponse!
@@ -594,9 +611,27 @@ class URLSessionInstrumentationTests: XCTestCase {
594611
testDataTaskWithRequestDelegateAsync() async throws {
595612
let request = URLRequest(url: URL(string: "http://localhost:33333/success")!)
596613

597-
let session = URLSession(configuration: URLSessionConfiguration.default, delegate: sessionDelegate, delegateQueue: nil)
598-
_ = try await session.data(for: request)
614+
let delegate = CountingSessionDelegate()
615+
616+
let session = URLSession(configuration: URLSessionConfiguration.default, delegate: delegate, delegateQueue: nil)
617+
let _ = try await session.data(for: request)
618+
619+
XCTAssertEqual(1, delegate.callCount)
620+
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
621+
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
622+
}
623+
624+
@available(macOS 10.15, iOS 15.0, tvOS 15.0, watchOS 8.0, *)
625+
public func
626+
testDataTaskWithTaskDelegateAsync() async throws {
627+
let request = URLRequest(url: URL(string: "http://localhost:33333/success")!)
628+
629+
let delegate = CountingSessionDelegate()
630+
631+
let session = URLSession(configuration: URLSessionConfiguration.default)
632+
let _ = try await session.data(for: request, delegate: delegate)
599633

634+
XCTAssertEqual(1, delegate.callCount)
600635
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
601636
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
602637
}
@@ -623,7 +658,7 @@ class URLSessionInstrumentationTests: XCTestCase {
623658
}
624659

625660
@available(macOS 12, iOS 15.0, tvOS 15.0, watchOS 8.0, *)
626-
public func testDownloadTaskWithRequestDelegateAsync() async throws {
661+
public func testDownloadTaskWithSessionDelegateAsync() async throws {
627662
let url = URL(string: "http://localhost:33333/success")!
628663
let request = URLRequest(url: url)
629664

@@ -634,9 +669,21 @@ class URLSessionInstrumentationTests: XCTestCase {
634669
XCTAssertTrue(URLSessionInstrumentationTests.checker.receivedResponseCalled)
635670
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
636671
}
672+
673+
@available(macOS 12, iOS 15.0, tvOS 15.0, watchOS 8.0, *)
674+
public func testDownloadTaskWithRequestDelegateAsync() async throws {
675+
let url = URL(string: "http://localhost:33333/success")!
676+
let request = URLRequest(url: url)
677+
678+
let session = URLSession(configuration: URLSessionConfiguration.default)
679+
_ = try await session.download(for: request, delegate: sessionDelegate)
680+
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
681+
XCTAssertTrue(URLSessionInstrumentationTests.checker.receivedResponseCalled)
682+
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
683+
}
637684

638685
@available(macOS 10.15, iOS 13.0, tvOS 13.0, *)
639-
public func testUploadTaskWithRequestDelegateAsync() async throws {
686+
public func testUploadTaskWithSessionDelegateAsync() async throws {
640687
let url = URL(string: "http://localhost:33333/success")!
641688
let request = URLRequest(url: url)
642689
let session = URLSession(configuration: URLSessionConfiguration.default, delegate: sessionDelegate, delegateQueue: nil)
@@ -645,6 +692,17 @@ class URLSessionInstrumentationTests: XCTestCase {
645692
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
646693
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
647694
}
695+
696+
@available(macOS 10.15, iOS 15.0, tvOS 15.0, watchOS 8.0, *)
697+
public func testUploadTaskWithRequestDelegateAsync() async throws {
698+
let url = URL(string: "http://localhost:33333/success")!
699+
let request = URLRequest(url: url)
700+
let session = URLSession(configuration: URLSessionConfiguration.default)
701+
_ = try await session.upload(for: request, from: Data(), delegate: sessionDelegate)
702+
703+
XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
704+
XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
705+
}
648706

649707
public func testNonInstrumentedRequestCompletes() {
650708
let request = URLRequest(url: URL(string: "http://localhost:33333/dontinstrument")!)

0 commit comments

Comments
 (0)