Skip to content

Commit 16aa082

Browse files
committed
Fix a race condition in code completion
Each `sourcekitd` uses a single, global code completion session but we were managing code completion sessions on the `SwiftLanguageServer` level. This had two issues: - Code completion is considered non-blocking on `SourceKitServer` and thus we could execute two concurrent code completion requests, which would conflict with each other. - You could have multiple `SwiftLanguageServer`s that each have a connection to `sourcekitd` but share the same `sourcekitd` state. This shouldn't happen in the real world but does happen if we create multiple `SourceKitServer` instances in tests.
1 parent 68a90aa commit 16aa082

File tree

3 files changed

+146
-99
lines changed

3 files changed

+146
-99
lines changed

Sources/SourceKitLSP/Swift/CodeCompletion.swift

Lines changed: 14 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -37,22 +37,6 @@ extension SwiftLanguageServer {
3737

3838
let options = req.sourcekitlspOptions ?? serverOptions.completionOptions
3939

40-
return try await completionWithServerFiltering(
41-
offset: offset,
42-
completionPos: completionPos,
43-
snapshot: snapshot,
44-
request: req,
45-
options: options
46-
)
47-
}
48-
49-
private func completionWithServerFiltering(
50-
offset: Int,
51-
completionPos: Position,
52-
snapshot: DocumentSnapshot,
53-
request req: CompletionRequest,
54-
options: SKCompletionOptions
55-
) async throws -> CompletionList {
5640
guard let start = snapshot.indexOf(utf8Offset: offset),
5741
let end = snapshot.index(of: req.position)
5842
else {
@@ -62,42 +46,20 @@ extension SwiftLanguageServer {
6246

6347
let filterText = String(snapshot.text[start..<end])
6448

65-
let session: CodeCompletionSession
66-
if req.context?.triggerKind == .triggerFromIncompleteCompletions {
67-
guard let currentSession = currentCompletionSession else {
68-
logger.error("triggerFromIncompleteCompletions with no existing completion session")
69-
throw ResponseError.serverCancelled
70-
}
71-
guard currentSession.uri == snapshot.uri, currentSession.utf8StartOffset == offset else {
72-
logger.error(
73-
"""
74-
triggerFromIncompleteCompletions with incompatible completion session; expected \
75-
\(currentSession.uri.forLogging)@\(currentSession.utf8StartOffset), \
76-
but got \(snapshot.uri.forLogging)@\(offset)
77-
"""
78-
)
79-
throw ResponseError.serverCancelled
80-
}
81-
session = currentSession
82-
} else {
83-
let clientSupportsSnippets = capabilityRegistry.clientCapabilities.textDocument?.completion?.completionItem?.snippetSupport ?? false
84-
85-
// FIXME: even if trigger kind is not from incomplete, we could to detect a compatible
86-
// location if we also check that the rest of the snapshot has not changed.
87-
session = CodeCompletionSession(
88-
sourcekitd: sourcekitd,
89-
snapshot: snapshot,
90-
utf8Offset: offset,
91-
position: completionPos,
92-
compileCommand: await buildSettings(for: snapshot.uri),
93-
clientSupportsSnippets: clientSupportsSnippets
94-
)
95-
96-
await currentCompletionSession?.close()
97-
currentCompletionSession = session
98-
}
99-
100-
return try await session.update(filterText: filterText, position: req.position, in: snapshot, options: options)
49+
let clientSupportsSnippets = capabilityRegistry.clientCapabilities.textDocument?.completion?.completionItem?.snippetSupport ?? false
50+
let buildSettings = await buildSettings(for: snapshot.uri)
51+
return try await CodeCompletionSession.completionList(
52+
sourcekitd: sourcekitd,
53+
snapshot: snapshot,
54+
completionPosition: completionPos,
55+
completionUtf8Offset: offset,
56+
cursorPosition: req.position,
57+
compileCommand: buildSettings,
58+
options: options,
59+
clientSupportsSnippets: clientSupportsSnippets,
60+
filterText: filterText,
61+
mustReuse: req.context?.triggerKind == .triggerFromIncompleteCompletions
62+
)
10163
}
10264

10365
/// Adjust completion position to the start of identifier characters.

Sources/SourceKitLSP/Swift/CodeCompletionSession.swift

Lines changed: 132 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,128 @@ import SourceKitD
2525
///
2626
/// At the sourcekitd level, this uses `codecomplete.open`, `codecomplete.update` and
2727
/// `codecomplete.close` requests.
28-
actor CodeCompletionSession {
28+
class CodeCompletionSession {
29+
// MARK: - Public static API
30+
31+
/// The queue on which all code completion requests are executed.
32+
///
33+
/// This is needed because sourcekitd has a single, global code completion
34+
/// session and we need to make sure that multiple code completion requests
35+
/// don't race each other.
36+
///
37+
/// Technically, we would only need one queue for each sourcekitd and different
38+
/// sourcekitd could serve code completion requests simultaneously.
39+
///
40+
/// But it's rare to open multiple sourcekitd instances simultaneously and
41+
/// even rarer to interact with them at the same time, so we have a global
42+
/// queue for now to simplify the implementation.
43+
private static let completionQueue = AsyncQueue(.serial)
44+
45+
/// The code completion session for each sourcekitd instance.
46+
///
47+
/// `sourcekitd` has a global code completion session, that's why we need to
48+
/// have a global mapping from `sourcekitd` to its currently active code
49+
/// completion session.
50+
///
51+
/// Modification of code completion sessions should only happen on
52+
/// `completionQueue`.
53+
private static var completionSessions: [ObjectIdentifier: CodeCompletionSession] = [:]
54+
55+
/// Gets the code completion results for the given parameters.
56+
///
57+
/// If a code completion session that is compatible with the parameters
58+
/// already exists, this just performs an update to the filtering. If it does
59+
/// not, this opens a new code completion session with `sourcekitd` and gets
60+
/// the results.
61+
///
62+
/// - Parameters:
63+
/// - sourcekitd: The `sourcekitd` instance from which to get code
64+
/// completion results
65+
/// - snapshot: The document in which to perform completion.
66+
/// - completionPosition: The position at which to perform completion.
67+
/// This is the position at which the code completion token logically
68+
/// starts. For example when completing `foo.ba|`, then the completion
69+
/// position should be after the `.`.
70+
/// - completionUtf8Offset: Same as `completionPosition` but as a UTF-8
71+
/// offset within the buffer.
72+
/// - cursorPosition: The position at which the cursor is positioned. E.g.
73+
/// when completing `foo.ba|`, this is after the `a` (see
74+
/// `completionPosition` for comparison)
75+
/// - compileCommand: The compiler arguments to use.
76+
/// - options: Further options that can be sent from the editor to control
77+
/// completion.
78+
/// - clientSupportsSnippets: Whether the editor supports LSP snippets.
79+
/// - filterText: The text by which to filter code completion results.
80+
/// - mustReuse: If `true` and there is an active session in this
81+
/// `sourcekitd` instance, cancel the request instead of opening a new
82+
/// session.
83+
/// This is set to `true` when triggering a filter from incomplete results
84+
/// so that clients can rely on results being delivered quickly when
85+
/// getting updated results after updating the filter text.
86+
/// - Returns: The code completion results for those parameters.
87+
static func completionList(
88+
sourcekitd: any SourceKitD,
89+
snapshot: DocumentSnapshot,
90+
completionPosition: Position,
91+
completionUtf8Offset: Int,
92+
cursorPosition: Position,
93+
compileCommand: SwiftCompileCommand?,
94+
options: SKCompletionOptions,
95+
clientSupportsSnippets: Bool,
96+
filterText: String,
97+
mustReuse: Bool
98+
) async throws -> CompletionList {
99+
let task = completionQueue.asyncThrowing {
100+
if let session = completionSessions[ObjectIdentifier(sourcekitd)], session.state == .open {
101+
let isCompatible = session.snapshot.uri == snapshot.uri &&
102+
session.utf8StartOffset == completionUtf8Offset &&
103+
session.position == completionPosition &&
104+
session.compileCommand == compileCommand &&
105+
session.clientSupportsSnippets == clientSupportsSnippets
106+
107+
if isCompatible {
108+
return try await session.update(filterText: filterText, position: cursorPosition, in: snapshot, options: options)
109+
}
110+
111+
if mustReuse {
112+
logger.error(
113+
"""
114+
triggerFromIncompleteCompletions with incompatible completion session; expected \
115+
\(session.uri.forLogging)@\(session.utf8StartOffset), \
116+
but got \(snapshot.uri.forLogging)@\(completionUtf8Offset)
117+
"""
118+
)
119+
throw ResponseError.serverCancelled
120+
}
121+
// The sessions aren't compatible. Close the existing session and open
122+
// a new one below.
123+
session.close()
124+
}
125+
if mustReuse {
126+
logger.error("triggerFromIncompleteCompletions with no existing completion session")
127+
throw ResponseError.serverCancelled
128+
}
129+
let session = CodeCompletionSession(
130+
sourcekitd: sourcekitd,
131+
snapshot: snapshot,
132+
utf8Offset: completionUtf8Offset,
133+
position: completionPosition,
134+
compileCommand: compileCommand,
135+
clientSupportsSnippets: clientSupportsSnippets
136+
)
137+
completionSessions[ObjectIdentifier(sourcekitd)] = session
138+
return try await session.open(filterText: filterText, position: cursorPosition, in: snapshot, options: options)
139+
}
140+
141+
// FIXME: (async) Use valuePropagatingCancellation once we support cancellation
142+
return try await task.value
143+
}
144+
145+
// MARK: - Implementation
146+
29147
private let sourcekitd: any SourceKitD
30148
private let snapshot: DocumentSnapshot
31-
let utf8StartOffset: Int
149+
private let utf8StartOffset: Int
32150
private let position: Position
33151
private let compileCommand: SwiftCompileCommand?
34152
private let clientSupportsSnippets: Bool
@@ -39,10 +157,10 @@ actor CodeCompletionSession {
39157
case open
40158
}
41159

42-
nonisolated var uri: DocumentURI { snapshot.uri }
43-
nonisolated var keys: sourcekitd_keys { return sourcekitd.keys }
160+
private nonisolated var uri: DocumentURI { snapshot.uri }
161+
private nonisolated var keys: sourcekitd_keys { return sourcekitd.keys }
44162

45-
init(
163+
private init(
46164
sourcekitd: any SourceKitD,
47165
snapshot: DocumentSnapshot,
48166
utf8Offset: Int,
@@ -58,30 +176,6 @@ actor CodeCompletionSession {
58176
self.clientSupportsSnippets = clientSupportsSnippets
59177
}
60178

61-
/// Retrieve completions for the given `filterText`, opening or updating the session.
62-
///
63-
/// - parameters:
64-
/// - filterText: The text to use for fuzzy matching the results.
65-
/// - position: The position at the end of the existing text (typically right after the end of
66-
/// `filterText`), which determines the end of the `TextEdit` replacement range
67-
/// in the resulting completions.
68-
/// - snapshot: The current snapshot that the `TextEdit` replacement in results will be in.
69-
/// - options: The completion options, such as the maximum number of results.
70-
func update(
71-
filterText: String,
72-
position: Position,
73-
in snapshot: DocumentSnapshot,
74-
options: SKCompletionOptions
75-
) async throws -> CompletionList {
76-
switch self.state {
77-
case .closed:
78-
self.state = .open
79-
return try await self.open(filterText: filterText, position: position, in: snapshot, options: options)
80-
case .open:
81-
return try await self.updateImpl(filterText: filterText, position: position, in: snapshot, options: options)
82-
}
83-
}
84-
85179
private func open(
86180
filterText: String,
87181
position: Position,
@@ -105,6 +199,7 @@ actor CodeCompletionSession {
105199
}
106200

107201
let dict = try await sourcekitd.send(req)
202+
self.state = .open
108203

109204
guard let completions: SKDResponseArray = dict[keys.results] else {
110205
return CompletionList(isIncomplete: false, items: [])
@@ -121,7 +216,7 @@ actor CodeCompletionSession {
121216
)
122217
}
123218

124-
private func updateImpl(
219+
private func update(
125220
filterText: String,
126221
position: Position,
127222
in snapshot: DocumentSnapshot,
@@ -170,22 +265,18 @@ actor CodeCompletionSession {
170265
return dict
171266
}
172267

173-
private func sendClose() {
174-
let req = SKDRequestDictionary(sourcekitd: sourcekitd)
175-
req[keys.request] = sourcekitd.requests.codecomplete_close
176-
req[keys.offset] = self.utf8StartOffset
177-
req[keys.name] = self.snapshot.uri.pseudoPath
178-
logger.info("Closing code completion session: \(self, privacy: .private)")
179-
_ = try? sourcekitd.sendSync(req)
180-
}
181-
182-
func close() async {
268+
private func close() {
183269
switch self.state {
184270
case .closed:
185271
// Already closed, nothing to do.
186272
break
187273
case .open:
188-
self.sendClose()
274+
let req = SKDRequestDictionary(sourcekitd: sourcekitd)
275+
req[keys.request] = sourcekitd.requests.codecomplete_close
276+
req[keys.offset] = self.utf8StartOffset
277+
req[keys.name] = self.snapshot.uri.pseudoPath
278+
logger.info("Closing code completion session: \(self, privacy: .private)")
279+
_ = try? sourcekitd.sendSync(req)
189280
self.state = .closed
190281
}
191282
}

Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -108,8 +108,6 @@ public actor SwiftLanguageServer: ToolchainLanguageServer {
108108

109109
var currentDiagnostics: [DocumentURI: [CachedDiagnostic]] = [:]
110110

111-
var currentCompletionSession: CodeCompletionSession? = nil
112-
113111
let syntaxTreeManager = SyntaxTreeManager()
114112
let semanticTokensManager = SemanticTokensManager()
115113

@@ -393,10 +391,6 @@ extension SwiftLanguageServer {
393391
}
394392

395393
public func shutdown() async {
396-
if let session = self.currentCompletionSession {
397-
await session.close()
398-
self.currentCompletionSession = nil
399-
}
400394
self.sourcekitd.removeNotificationHandler(self)
401395
}
402396

0 commit comments

Comments
 (0)