Skip to content

Fix a race condition during the computation of uriToWorkspaceCache #1170

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Apr 19, 2024
Merged

Fix a race condition during the computation of uriToWorkspaceCache #1170

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions Sources/SKSupport/AsyncUtils.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,21 @@ public extension Task {
}
}

public extension Task where Failure == Never {
/// Awaits the value of the result.
///
/// If the current task is cancelled, this will cancel the subtask as well.
var valuePropagatingCancellation: Success {
get async {
await withTaskCancellationHandler {
return await self.value
} onCancel: {
self.cancel()
}
}
}
}

/// Allows the execution of a cancellable operation that returns the results
/// via a completion handler.
///
Expand Down
203 changes: 109 additions & 94 deletions Sources/SourceKitLSP/SourceKitLSPServer.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -408,6 +408,10 @@ public actor SourceKitLSPServer {
/// request's task before handling any cancellation request for it.
private let cancellationMessageHandlingQueue = AsyncQueue<Serial>()

/// The queue on which all modifications of `uriToWorkspaceCache` happen. This means that the value of
/// `workspacesAndIsImplicit` and `uriToWorkspaceCache` can't change while executing a closure on `workspaceQueue`.
private let workspaceQueue = AsyncQueue<Serial>()

/// The connection to the editor.
public let client: Connection

Expand All @@ -432,14 +436,19 @@ public actor SourceKitLSPServer {
}

/// Caches which workspace a document with the given URI should be opened in.
/// Must only be accessed from `queue`.
///
/// - Important: Must only be modified from `workspaceQueue`. This means that the value of `uriToWorkspaceCache`
/// can't change while executing an operation on `workspaceQueue`.
private var uriToWorkspaceCache: [DocumentURI: WeakWorkspace] = [:]

/// The open workspaces.
///
/// Implicit workspaces are workspaces that weren't actually specified by the client during initialization or by a
/// `didChangeWorkspaceFolders` request. Instead, they were opened by sourcekit-lsp because a file could not be
/// handled by any of the open workspaces but one of the file's parent directories had handling capabilities for it.
///
/// - Important: Must only be modified from `workspaceQueue`. This means that the value of `workspacesAndIsImplicit`
/// can't change while executing an operation on `workspaceQueue`.
private var workspacesAndIsImplicit: [(workspace: Workspace, isImplicit: Bool)] = [] {
didSet {
uriToWorkspaceCache = [:]
Expand All @@ -452,7 +461,9 @@ public actor SourceKitLSPServer {

@_spi(Testing)
public func setWorkspaces(_ newValue: [(workspace: Workspace, isImplicit: Bool)]) {
self.workspacesAndIsImplicit = newValue
workspaceQueue.async {
self.workspacesAndIsImplicit = newValue
}
}

/// The requests that we are currently handling.
Expand Down Expand Up @@ -517,54 +528,58 @@ public actor SourceKitLSPServer {
}

public func workspaceForDocument(uri: DocumentURI) async -> Workspace? {
if let cachedWorkspace = uriToWorkspaceCache[uri]?.value {
if let cachedWorkspace = self.uriToWorkspaceCache[uri]?.value {
return cachedWorkspace
}

// Pick the workspace with the best FileHandlingCapability for this file.
// If there is a tie, use the workspace that occurred first in the list.
var bestWorkspace: (workspace: Workspace?, fileHandlingCapability: FileHandlingCapability) = (nil, .unhandled)
for workspace in workspaces {
let fileHandlingCapability = await workspace.buildSystemManager.fileHandlingCapability(for: uri)
if fileHandlingCapability > bestWorkspace.fileHandlingCapability {
bestWorkspace = (workspace, fileHandlingCapability)
// Execute the computation of the workspace on `workspaceQueue` to ensure that the file handling capabilities of the
// workspaces don't change during the computation. Otherwise, we could run into a race condition like the following:
// 1. We don't have an entry for file `a.swift` in `uriToWorkspaceCache` and start the computation
// 2. We find that the first workspace in `self.workspaces` can handle this file.
// 3. During the `await ... .fileHandlingCapability` for a second workspace the file handling capabilities for the
// first workspace change, meaning it can no longer handle the document. This resets `uriToWorkspaceCache`
// assuming that the URI to workspace relation will get re-computed.
// 4. But we then set `uriToWorkspaceCache[uri]` to the workspace found in step (2), caching an out-of-date result.
//
// Furthermore, the computation of the workspace for a URI can create a new implicit workspace, which modifies
// `workspacesAndIsImplicit` and which must only be modified on `workspaceQueue`.
return await self.workspaceQueue.async {
// Pick the workspace with the best FileHandlingCapability for this file.
// If there is a tie, use the workspace that occurred first in the list.
var bestWorkspace: (workspace: Workspace?, fileHandlingCapability: FileHandlingCapability) = (nil, .unhandled)
for workspace in self.workspaces {
let fileHandlingCapability = await workspace.buildSystemManager.fileHandlingCapability(for: uri)
if fileHandlingCapability > bestWorkspace.fileHandlingCapability {
bestWorkspace = (workspace, fileHandlingCapability)
}
}
}
if bestWorkspace.fileHandlingCapability < .handled {
// We weren't able to handle the document with any of the known workspaces. See if any of the document's parent
// directories contain a workspace that can handle the document.
let rootUris = workspaces.map(\.rootUri)
if let workspace = await findWorkspaceCapableOfHandlingDocument(at: uri) {
if workspaces.map(\.rootUri) != rootUris {
// Workspaces was modified while running `findWorkspaceCapableOfHandlingDocument`, so we raced.
// This is unlikely to happen unless the user opens many files that in sub-workspaces simultaneously.
// Try again based on the new data. Very likely the workspace that can handle this document is now in
// `workspaces` and we will be able to return it without having to search again.
logger.debug("findWorkspaceCapableOfHandlingDocument raced with another workspace creation. Trying again.")
return await workspaceForDocument(uri: uri)
if bestWorkspace.fileHandlingCapability < .handled {
// We weren't able to handle the document with any of the known workspaces. See if any of the document's parent
// directories contain a workspace that can handle the document.
if let workspace = await self.findWorkspaceCapableOfHandlingDocument(at: uri) {
// Appending a workspace is fine and doesn't require checking if we need to re-open any documents because:
// - Any currently open documents that have FileHandlingCapability `.handled` will continue to be opened in
// their current workspace because it occurs further in front inside the workspace list
// - Any currently open documents that have FileHandlingCapability < `.handled` also went through this check
// and didn't find any parent workspace that was able to handle them. We assume that a workspace can only
// properly handle files within its root directory, so those files now also can't be handled by the new
// workspace.
logger.log("Opening implicit workspace at \(workspace.rootUri.forLogging) to handle \(uri.forLogging)")
self.workspacesAndIsImplicit.append((workspace: workspace, isImplicit: true))
bestWorkspace = (workspace, .handled)
}
// Appending a workspace is fine and doesn't require checking if we need to re-open any documents because:
// - Any currently open documents that have FileHandlingCapability `.handled` will continue to be opened in
// their current workspace because it occurs further in front inside the workspace list
// - Any currently open documents that have FileHandlingCapability < `.handled` also went through this check
// and didn't find any parent workspace that was able to handle them. We assume that a workspace can only
// properly handle files within its root directory, so those files now also can't be handled by the new
// workspace.
logger.log("Opening implicit workspace at \(workspace.rootUri.forLogging) to handle \(uri.forLogging)")
workspacesAndIsImplicit.append((workspace: workspace, isImplicit: true))
bestWorkspace = (workspace, .handled)
}
}
uriToWorkspaceCache[uri] = WeakWorkspace(bestWorkspace.workspace)
if let workspace = bestWorkspace.workspace {
return workspace
}
if let workspace = workspaces.only {
// Special handling: If there is only one workspace, open all files in it, even it it cannot handle the document.
// This retains the behavior of SourceKit-LSP before it supported multiple workspaces.
return workspace
}
return nil
self.uriToWorkspaceCache[uri] = WeakWorkspace(bestWorkspace.workspace)
if let workspace = bestWorkspace.workspace {
return workspace
}
if let workspace = self.workspaces.only {
// Special handling: If there is only one workspace, open all files in it, even it it cannot handle the document.
// This retains the behavior of SourceKit-LSP before it supported multiple workspaces.
return workspace
}
return nil
}.valuePropagatingCancellation
}

/// Execute `notificationHandler` with the request as well as the workspace
Expand Down Expand Up @@ -1077,8 +1092,10 @@ extension SourceKitLSPServer: BuildSystemDelegate {
}

public func fileHandlingCapabilityChanged() {
logger.log("Resetting URI to workspace cache because file handling capability of a workspace changed")
self.uriToWorkspaceCache = [:]
workspaceQueue.async {
logger.log("Resetting URI to workspace cache because file handling capability of a workspace changed")
self.uriToWorkspaceCache = [:]
}
}
}

Expand Down Expand Up @@ -1184,47 +1201,43 @@ extension SourceKitLSPServer {

capabilityRegistry = CapabilityRegistry(clientCapabilities: req.capabilities)

if let workspaceFolders = req.workspaceFolders {
self.workspacesAndIsImplicit += await workspaceFolders.asyncCompactMap {
guard let workspace = await self.createWorkspace($0) else {
return nil
await workspaceQueue.async {
if let workspaceFolders = req.workspaceFolders {
self.workspacesAndIsImplicit += await workspaceFolders.asyncCompactMap {
guard let workspace = await self.createWorkspace($0) else {
return nil
}
return (workspace: workspace, isImplicit: false)
}
} else if let uri = req.rootURI {
let workspaceFolder = WorkspaceFolder(uri: uri)
if let workspace = await self.createWorkspace(workspaceFolder) {
self.workspacesAndIsImplicit.append((workspace: workspace, isImplicit: false))
}
} else if let path = req.rootPath {
let workspaceFolder = WorkspaceFolder(uri: DocumentURI(URL(fileURLWithPath: path)))
if let workspace = await self.createWorkspace(workspaceFolder) {
self.workspacesAndIsImplicit.append((workspace: workspace, isImplicit: false))
}
return (workspace: workspace, isImplicit: false)
}
} else if let uri = req.rootURI {
let workspaceFolder = WorkspaceFolder(uri: uri)
if let workspace = await self.createWorkspace(workspaceFolder) {
self.workspacesAndIsImplicit.append((workspace: workspace, isImplicit: false))
}
} else if let path = req.rootPath {
let workspaceFolder = WorkspaceFolder(uri: DocumentURI(URL(fileURLWithPath: path)))
if let workspace = await self.createWorkspace(workspaceFolder) {
self.workspacesAndIsImplicit.append((workspace: workspace, isImplicit: false))
}
}

if self.workspaces.isEmpty {
logger.error("no workspace found")

let workspace = await Workspace(
documentManager: self.documentManager,
rootUri: req.rootURI,
capabilityRegistry: self.capabilityRegistry!,
toolchainRegistry: self.toolchainRegistry,
buildSetup: self.options.buildSetup,
underlyingBuildSystem: nil,
index: nil,
indexDelegate: nil
)

// Another workspace might have been added while we awaited the
// construction of the workspace above. If that race happened, just
// discard the workspace we created here since `workspaces` now isn't
// empty anymore.
if self.workspaces.isEmpty {
logger.error("no workspace found")

let workspace = await Workspace(
documentManager: self.documentManager,
rootUri: req.rootURI,
capabilityRegistry: self.capabilityRegistry!,
toolchainRegistry: self.toolchainRegistry,
buildSetup: self.options.buildSetup,
underlyingBuildSystem: nil,
index: nil,
indexDelegate: nil
)

self.workspacesAndIsImplicit.append((workspace: workspace, isImplicit: false))
}
}
}.value

assert(!self.workspaces.isEmpty)
for workspace in self.workspaces {
Expand Down Expand Up @@ -1552,21 +1565,23 @@ extension SourceKitLSPServer {
for docUri in self.documentManager.openDocuments {
preChangeWorkspaces[docUri] = await self.workspaceForDocument(uri: docUri)
}
if let removed = notification.event.removed {
self.workspacesAndIsImplicit.removeAll { workspace in
// Close all implicit workspaces as well because we could have opened a new explicit workspace that now contains
// files from a previous implicit workspace.
return workspace.isImplicit
|| removed.contains(where: { workspaceFolder in workspace.workspace.rootUri == workspaceFolder.uri })
await workspaceQueue.async {
if let removed = notification.event.removed {
self.workspacesAndIsImplicit.removeAll { workspace in
// Close all implicit workspaces as well because we could have opened a new explicit workspace that now contains
// files from a previous implicit workspace.
return workspace.isImplicit
|| removed.contains(where: { workspaceFolder in workspace.workspace.rootUri == workspaceFolder.uri })
}
}
}
if let added = notification.event.added {
let newWorkspaces = await added.asyncCompactMap { await self.createWorkspace($0) }
for workspace in newWorkspaces {
await workspace.buildSystemManager.setDelegate(self)
if let added = notification.event.added {
let newWorkspaces = await added.asyncCompactMap { await self.createWorkspace($0) }
for workspace in newWorkspaces {
await workspace.buildSystemManager.setDelegate(self)
}
self.workspacesAndIsImplicit += newWorkspaces.map { (workspace: $0, isImplicit: false) }
}
self.workspacesAndIsImplicit += newWorkspaces.map { (workspace: $0, isImplicit: false) }
}
}.value

// For each document that has moved to a different workspace, close it in
// the old workspace and open it in the new workspace.
Expand Down