Skip to content

Commit 36c39d6

Browse files
lynnswaplynnswap
authored
Merge pull request #51 from lynnswap:codex/refactor/deterministic-test-timing
refactor(test): make timing-sensitive tests deterministic
2 parents 67ab6fc + 2839002 commit 36c39d6

16 files changed

Lines changed: 999 additions & 455 deletions

‎Package.swift‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ let package = Package(
8080
"ProxyCore",
8181
"ProxyRuntime",
8282
.product(name: "Logging", package: "swift-log"),
83+
.product(name: "NIOConcurrencyHelpers", package: "swift-nio"),
8384
.product(name: "NIO", package: "swift-nio"),
8485
],
8586
path: "Sources/ProxyFeatureXcode",
@@ -138,7 +139,8 @@ let package = Package(
138139
.target(
139140
name: "XcodeMCPTestSupport",
140141
dependencies: [
141-
.product(name: "NIO", package: "swift-nio")
142+
.product(name: "NIO", package: "swift-nio"),
143+
.product(name: "NIOConcurrencyHelpers", package: "swift-nio"),
142144
],
143145
swiftSettings: strictSwiftSettings
144146
),

‎Sources/ProxyFeatureXcode/RefreshCodeIssuesCoordinator.swift‎

Lines changed: 54 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import Foundation
2+
import NIOConcurrencyHelpers
23

34
package actor RefreshCodeIssuesCoordinator {
45
package struct Permit: Sendable {
@@ -20,7 +21,8 @@ package actor RefreshCodeIssuesCoordinator {
2021

2122
package nonisolated let maxPendingPerKey: Int
2223
package nonisolated let maxPendingTotal: Int
23-
package nonisolated let queueWaitTimeoutNanoseconds: UInt64
24+
package nonisolated let queueWaitTimeout: Duration
25+
package nonisolated let queueWaitClock: any Clock<Duration> & Sendable
2426
private var nextWaiterID: UInt64 = 0
2527
private var busyKeys: Set<String> = []
2628
private var pendingWaiterCount = 0
@@ -40,14 +42,29 @@ package actor RefreshCodeIssuesCoordinator {
4042
maxPendingPerKey: Int = 4,
4143
maxPendingTotal: Int = 32,
4244
queueWaitTimeout: TimeInterval = 30
45+
) {
46+
self.init(
47+
maxPendingPerKey: maxPendingPerKey,
48+
maxPendingTotal: maxPendingTotal,
49+
queueWaitTimeout: Self.duration(from: queueWaitTimeout),
50+
queueWaitClock: ContinuousClock()
51+
)
52+
}
53+
54+
package init(
55+
maxPendingPerKey: Int = 4,
56+
maxPendingTotal: Int = 32,
57+
queueWaitTimeout: Duration,
58+
queueWaitClock: any Clock<Duration> & Sendable = ContinuousClock()
4359
) {
4460
self.maxPendingPerKey = max(0, maxPendingPerKey)
4561
self.maxPendingTotal = max(0, maxPendingTotal)
46-
self.queueWaitTimeoutNanoseconds = Self.nanoseconds(from: queueWaitTimeout)
62+
self.queueWaitTimeout = queueWaitTimeout
63+
self.queueWaitClock = queueWaitClock
4764
}
4865

4966
package nonisolated var queueWaitTimeoutSeconds: Double {
50-
Double(queueWaitTimeoutNanoseconds) / 1_000_000_000
67+
Self.seconds(from: queueWaitTimeout)
5168
}
5269

5370
package func withPermit<T: Sendable>(
@@ -88,18 +105,16 @@ package actor RefreshCodeIssuesCoordinator {
88105
pendingTotal: pendingWaiterCount + 1
89106
)
90107

91-
let timeoutTask = Task { [queueWaitTimeoutNanoseconds] in
92-
do {
93-
try await Task.sleep(nanoseconds: queueWaitTimeoutNanoseconds)
94-
self.timeoutWaiter(key: key, waiterID: waiterID)
95-
} catch {
96-
return
97-
}
98-
}
108+
let timeoutTaskBox = NIOLockedValueBox<Task<Void, Never>?>(nil)
99109

100110
return try await withTaskCancellationHandler(
101111
operation: {
102-
defer { timeoutTask.cancel() }
112+
defer {
113+
timeoutTaskBox.withLockedValue { task in
114+
task?.cancel()
115+
task = nil
116+
}
117+
}
103118

104119
return try await withCheckedThrowingContinuation { continuation in
105120
waitersByKey[key, default: []].append(
@@ -111,6 +126,18 @@ package actor RefreshCodeIssuesCoordinator {
111126
)
112127
pendingWaiterCount += 1
113128

129+
let timeoutTask = Task { [queueWaitClock, queueWaitTimeout] in
130+
do {
131+
try await queueWaitClock.sleep(for: queueWaitTimeout)
132+
self.timeoutWaiter(key: key, waiterID: waiterID)
133+
} catch {
134+
return
135+
}
136+
}
137+
timeoutTaskBox.withLockedValue { task in
138+
task = timeoutTask
139+
}
140+
114141
if Task.isCancelled {
115142
failWaiter(
116143
key: key,
@@ -121,7 +148,10 @@ package actor RefreshCodeIssuesCoordinator {
121148
}
122149
},
123150
onCancel: {
124-
timeoutTask.cancel()
151+
timeoutTaskBox.withLockedValue { task in
152+
task?.cancel()
153+
task = nil
154+
}
125155
Task {
126156
await self.cancelWaiter(key: key, waiterID: waiterID)
127157
}
@@ -183,4 +213,15 @@ package actor RefreshCodeIssuesCoordinator {
183213
}
184214
return UInt64(nanoseconds.rounded(.up))
185215
}
216+
217+
private static func seconds(from duration: Duration) -> Double {
218+
let components = duration.components
219+
return Double(components.seconds)
220+
+ (Double(components.attoseconds) / 1_000_000_000_000_000_000)
221+
}
222+
223+
private static func duration(from interval: TimeInterval) -> Duration {
224+
let clampedNanoseconds = min(nanoseconds(from: interval), UInt64(Int64.max))
225+
return .nanoseconds(Int64(clampedNanoseconds))
226+
}
186227
}

‎Sources/ProxyRuntime/InitializeGate.swift‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,21 +21,21 @@ package final class InitializeGate: Sendable {
2121
}
2222

2323
package struct SuccessPreparation: Sendable {
24-
package let timeout: Scheduled<Void>?
24+
package let timeout: RuntimeScheduledTimeout?
2525
package let shouldWarmSecondary: Bool
2626
package let cachedResult: JSONValue?
2727
}
2828

2929
package struct FailureResult: Sendable {
3030
package let pending: [PendingInitialize]
31-
package let timeout: Scheduled<Void>?
31+
package let timeout: RuntimeScheduledTimeout?
3232
package let upstreamID: Int64?
3333
package let shouldRetryEagerInitialize: Bool
3434
}
3535

3636
package struct ExitResult: Sendable {
3737
package let pending: [PendingInitialize]
38-
package let timeout: Scheduled<Void>?
38+
package let timeout: RuntimeScheduledTimeout?
3939
package let hadGlobalInit: Bool
4040
package let wasInFlight: Bool
4141
package let primaryInitUpstreamID: Int64?
@@ -53,7 +53,7 @@ package final class InitializeGate: Sendable {
5353
var initResult: JSONValue?
5454
var initPending: [PendingInitialize] = []
5555
var initInFlight = false
56-
var initTimeout: Scheduled<Void>?
56+
var initTimeout: RuntimeScheduledTimeout?
5757
var isShuttingDown = false
5858
var didWarmSecondary = false
5959
var primaryInitUpstreamID: Int64?
@@ -64,7 +64,7 @@ package final class InitializeGate: Sendable {
6464

6565
package init() {}
6666

67-
package func beginShutdown() -> (pending: [PendingInitialize], timeout: Scheduled<Void>?) {
67+
package func beginShutdown() -> (pending: [PendingInitialize], timeout: RuntimeScheduledTimeout?) {
6868
state.withLockedValue { state in
6969
state.isShuttingDown = true
7070
state.initInFlight = false
@@ -253,7 +253,7 @@ package final class InitializeGate: Sendable {
253253
}
254254
}
255255

256-
package func replaceInitTimeout(_ timeout: Scheduled<Void>) -> Scheduled<Void>? {
256+
package func replaceInitTimeout(_ timeout: RuntimeScheduledTimeout) -> RuntimeScheduledTimeout? {
257257
state.withLockedValue { state in
258258
let existing = state.initTimeout
259259
state.initTimeout = timeout
@@ -296,7 +296,7 @@ package final class InitializeGate: Sendable {
296296
}
297297
}
298298

299-
package func resetForDebug() -> (pending: [PendingInitialize], timeout: Scheduled<Void>?) {
299+
package func resetForDebug() -> (pending: [PendingInitialize], timeout: RuntimeScheduledTimeout?) {
300300
state.withLockedValue { state in
301301
let result = (pending: state.initPending, timeout: state.initTimeout)
302302
state.initResult = nil

‎Sources/ProxyRuntime/RuntimeCoordinator+Health.swift‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ extension RuntimeCoordinator {
1313
}
1414

1515
func markRequestTimedOut(upstreamIndex: Int) {
16-
let nowUptimeNs = DispatchTime.now().uptimeNanoseconds
16+
let nowUptimeNs = nowUptimeNanoseconds()
1717
let result = upstreamSelectionPolicy.markRequestTimedOut(
1818
upstreamIndex: upstreamIndex,
1919
nowUptimeNs: nowUptimeNs
@@ -111,7 +111,7 @@ extension RuntimeCoordinator {
111111
success: Bool,
112112
reason: String
113113
) {
114-
let nowUptimeNs = DispatchTime.now().uptimeNanoseconds
114+
let nowUptimeNs = nowUptimeNanoseconds()
115115
upstreamSelectionPolicy.finishHealthProbe(
116116
upstreamIndex: upstreamIndex,
117117
probeGeneration: probeGeneration,
@@ -164,7 +164,7 @@ extension RuntimeCoordinator {
164164
}
165165

166166
let refreshTimeout: TimeAmount = .seconds(5)
167-
let nowUptimeNs = DispatchTime.now().uptimeNanoseconds
167+
let nowUptimeNs = nowUptimeNanoseconds()
168168
let internalSessionID = toolsListInternalSessionID()
169169
_ = session(id: internalSessionID)
170170

@@ -277,7 +277,7 @@ extension RuntimeCoordinator {
277277
else {
278278
return
279279
}
280-
let timeout = eventLoop.scheduleTask(in: timeoutAmount) { [weak self] in
280+
let timeout = scheduleRuntimeTimeout(timeoutAmount) { [weak self] in
281281
guard let self else { return }
282282
self.handleUpstreamInitTimeout(upstreamIndex: upstreamIndex, upstreamID: upstreamID)
283283
}

‎Sources/ProxyRuntime/RuntimeCoordinator+Initialization.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -226,7 +226,7 @@ extension RuntimeCoordinator {
226226
else {
227227
return
228228
}
229-
let timeout = eventLoop.scheduleTask(in: timeoutAmount) { [weak self] in
229+
let timeout = scheduleRuntimeTimeout(timeoutAmount) { [weak self] in
230230
guard let self else { return }
231231
self.failInitPending(error: TimeoutError())
232232
}

‎Sources/ProxyRuntime/RuntimeCoordinator+UpstreamRouting.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -576,7 +576,7 @@ extension RuntimeCoordinator {
576576
upstreamIndex: Int
577577
) {
578578
debugRecorder.recordProtocolViolation(protocolViolation, upstreamIndex: upstreamIndex)
579-
let nowUptimeNs = DispatchTime.now().uptimeNanoseconds
579+
let nowUptimeNs = nowUptimeNanoseconds()
580580
let initSnapshot = initializeGate.snapshot()
581581
let transition = upstreamSelectionPolicy.markProtocolViolation(
582582
upstreamIndex: upstreamIndex,

‎Sources/ProxyRuntime/RuntimeCoordinator.swift‎

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,9 @@ package final class RuntimeCoordinator: Sendable, RuntimeCoordinating {
149149

150150
package let upstreamSelectionPolicy: UpstreamSelectionPolicy
151151
package let upstreamSlotScheduler: UpstreamSlotScheduler
152+
package let nowUptimeNanoseconds: @Sendable () -> UInt64
153+
package let scheduleRuntimeTimeout: @Sendable (TimeAmount, @escaping @Sendable () -> Void) ->
154+
RuntimeScheduledTimeout
152155

153156
package convenience init(config: ProxyConfig, eventLoop: EventLoop) {
154157
let count = max(1, min(config.upstreamProcessCount, 10))
@@ -157,9 +160,27 @@ package final class RuntimeCoordinator: Sendable, RuntimeCoordinating {
157160
self.init(config: config, eventLoop: eventLoop, upstreams: upstreams)
158161
}
159162

160-
package init(config: ProxyConfig, eventLoop: EventLoop, upstreams: [any UpstreamSlotControlling]) {
163+
package init(
164+
config: ProxyConfig,
165+
eventLoop: EventLoop,
166+
upstreams: [any UpstreamSlotControlling],
167+
nowUptimeNanoseconds: @escaping @Sendable () -> UInt64 = {
168+
DispatchTime.now().uptimeNanoseconds
169+
},
170+
scheduleRuntimeTimeout: (@Sendable (TimeAmount, @escaping @Sendable () -> Void) ->
171+
RuntimeScheduledTimeout)? = nil
172+
) {
161173
precondition(!upstreams.isEmpty, "upstreams must not be empty")
162174
let schedulerProbeStarter = NIOLockedValueBox<(@Sendable ([HealthProbeRequest]) -> Void)?>(nil)
175+
let uptimeProvider = nowUptimeNanoseconds
176+
let timeoutScheduler = scheduleRuntimeTimeout
177+
?? { delay, operation in
178+
RuntimeScheduledTimeout.schedule(
179+
on: eventLoop,
180+
in: delay,
181+
operation: operation
182+
)
183+
}
163184
self.config = config
164185
self.eventLoop = eventLoop
165186
self.upstreams = upstreams
@@ -168,11 +189,13 @@ package final class RuntimeCoordinator: Sendable, RuntimeCoordinating {
168189
self.requestLeaseRegistry = RequestLeaseRegistry()
169190
self.responseCorrelationStore = ResponseCorrelationStore(upstreamCount: upstreams.count)
170191
self.upstreamSelectionPolicy = UpstreamSelectionPolicy(upstreamCount: upstreams.count)
192+
self.nowUptimeNanoseconds = nowUptimeNanoseconds
193+
self.scheduleRuntimeTimeout = timeoutScheduler
171194
self.upstreamSlotScheduler = UpstreamSlotScheduler(
172195
upstreamCount: upstreams.count,
173196
defaultCapacity: 1,
174197
canUseUpstream: { [weak upstreamSelectionPolicy = self.upstreamSelectionPolicy] upstreamIndex in
175-
let nowUptimeNs = DispatchTime.now().uptimeNanoseconds
198+
let nowUptimeNs = uptimeProvider()
176199
guard let upstreamSelectionPolicy else { return false }
177200
let evaluation = upstreamSelectionPolicy.evaluateUsableInitialized(
178201
index: upstreamIndex,
@@ -186,7 +209,7 @@ package final class RuntimeCoordinator: Sendable, RuntimeCoordinating {
186209
return evaluation.0
187210
},
188211
selectUpstream: { [weak upstreamSelectionPolicy = self.upstreamSelectionPolicy] occupied in
189-
let nowUptimeNs = DispatchTime.now().uptimeNanoseconds
212+
let nowUptimeNs = uptimeProvider()
190213
let chooseResult = upstreamSelectionPolicy?.chooseBestInitializedUpstream(
191214
nowUptimeNs: nowUptimeNs,
192215
occupiedUpstreams: occupied
@@ -340,7 +363,7 @@ package final class RuntimeCoordinator: Sendable, RuntimeCoordinating {
340363
}
341364

342365
package func chooseUpstreamIndex() -> Int? {
343-
let nowUptimeNs = DispatchTime.now().uptimeNanoseconds
366+
let nowUptimeNs = nowUptimeNanoseconds()
344367
let occupiedUpstreams = upstreamSlotScheduler.occupiedUpstreamIndices()
345368

346369
let chooseResult = upstreamSelectionPolicy.chooseBestInitializedUpstream(

‎Sources/ProxyRuntime/RuntimeScheduledTimeout.swift‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
import NIO
2+
3+
package final class RuntimeScheduledTimeout: @unchecked Sendable {
4+
private let cancelImpl: @Sendable () -> Void
5+
6+
package init(cancel: @escaping @Sendable () -> Void) {
7+
self.cancelImpl = cancel
8+
}
9+
10+
package func cancel() {
11+
cancelImpl()
12+
}
13+
14+
package static func schedule(
15+
on eventLoop: EventLoop,
16+
in delay: TimeAmount,
17+
operation: @escaping @Sendable () -> Void
18+
) -> RuntimeScheduledTimeout {
19+
let task = eventLoop.scheduleTask(in: delay) {
20+
operation()
21+
}
22+
return RuntimeScheduledTimeout {
23+
task.cancel()
24+
}
25+
}
26+
}

0 commit comments

Comments
 (0)