Skip to content

Commit b0678ca

Browse files
ahoppenahoppen
authored
Merge pull request #39494 from ahoppen/pr/sourcekit-cancellation-tokens
[SourceKit] Allow explicit cancellation of requests with a cancellation token
2 parents 6e3d9a7 + fb26447 commit b0678ca

23 files changed

+602
-304
lines changed

‎test/SourceKit/Misc/cancellation.swift

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
// Check that we can cancel requests.
2+
// We need to wait a little bit after request scheduling and cancellation to make sure we are not cancelling the request before it got scheduled.
3+
4+
// RUN: not %sourcekitd-test -req=cursor -id=slow -async -pos=10:3 %s -- %s == \
5+
// RUN: -shell -- sleep 1 == \
6+
// RUN: -cancel=slow 2>&1 \
7+
// RUN: | %FileCheck %s
8+
9+
func foo(x: Invalid1, y: Invalid2) {
10+
x / y / x / y / x / y / x / y
11+
}
12+
13+
// CHECK: error response (Request Cancelled)

‎tools/SourceKit/include/SourceKit/Core/LangSupport.h

Lines changed: 46 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,14 @@
1414
#define LLVM_SOURCEKIT_CORE_LANGSUPPORT_H
1515

1616
#include "SourceKit/Core/LLVM.h"
17+
#include "SourceKit/Support/CancellationToken.h"
1718
#include "SourceKit/Support/UIdent.h"
18-
#include "llvm/Support/VersionTuple.h"
19+
#include "swift/AST/Type.h"
1920
#include "llvm/ADT/ArrayRef.h"
2021
#include "llvm/ADT/IntrusiveRefCntPtr.h"
2122
#include "llvm/ADT/Optional.h"
2223
#include "llvm/ADT/SmallString.h"
23-
#include "swift/AST/Type.h"
24+
#include "llvm/Support/VersionTuple.h"
2425
#include "llvm/Support/VirtualFileSystem.h"
2526
#include <functional>
2627
#include <memory>
@@ -737,6 +738,8 @@ class LangSupport {
737738

738739
virtual void dependencyUpdated() {}
739740

741+
virtual void cancelRequest(SourceKitCancellationToken CancellationToken) = 0;
742+
740743
virtual void indexSource(StringRef Filename,
741744
IndexingConsumer &Consumer,
742745
ArrayRef<const char *> Args) = 0;
@@ -794,10 +797,11 @@ class LangSupport {
794797
bool SynthesizedExtensions,
795798
StringRef swiftVersion) = 0;
796799

797-
virtual void editorOpenSwiftSourceInterface(StringRef Name,
798-
StringRef SourceName,
799-
ArrayRef<const char *> Args,
800-
std::shared_ptr<EditorConsumer> Consumer) = 0;
800+
virtual void
801+
editorOpenSwiftSourceInterface(StringRef Name, StringRef SourceName,
802+
ArrayRef<const char *> Args,
803+
SourceKitCancellationToken CancellationToken,
804+
std::shared_ptr<EditorConsumer> Consumer) = 0;
801805

802806
virtual void editorClose(StringRef Name, bool RemoveCache) = 0;
803807

@@ -821,33 +825,37 @@ class LangSupport {
821825
unsigned Length,
822826
EditorConsumer &Consumer) = 0;
823827

828+
virtual void getCursorInfo(
829+
StringRef Filename, unsigned Offset, unsigned Length, bool Actionables,
830+
bool SymbolGraph, bool CancelOnSubsequentRequest,
831+
ArrayRef<const char *> Args, Optional<VFSOptions> vfsOptions,
832+
SourceKitCancellationToken CancellationToken,
833+
std::function<void(const RequestResult<CursorInfoData> &)> Receiver) = 0;
834+
824835
virtual void
825-
getCursorInfo(StringRef Filename, unsigned Offset, unsigned Length,
826-
bool Actionables, bool SymbolGraph,
827-
bool CancelOnSubsequentRequest, ArrayRef<const char *> Args,
828-
Optional<VFSOptions> vfsOptions,
829-
std::function<void(const RequestResult<CursorInfoData> &)> Receiver) = 0;
830-
831-
virtual void getNameInfo(StringRef Filename, unsigned Offset,
832-
NameTranslatingInfo &Input,
833-
ArrayRef<const char *> Args,
834-
std::function<void(const RequestResult<NameTranslatingInfo> &)> Receiver) = 0;
835-
836-
virtual void getRangeInfo(StringRef Filename, unsigned Offset, unsigned Length,
837-
bool CancelOnSubsequentRequest,
838-
ArrayRef<const char *> Args,
839-
std::function<void(const RequestResult<RangeInfo> &)> Receiver) = 0;
836+
getNameInfo(StringRef Filename, unsigned Offset, NameTranslatingInfo &Input,
837+
ArrayRef<const char *> Args,
838+
SourceKitCancellationToken CancellationToken,
839+
std::function<void(const RequestResult<NameTranslatingInfo> &)>
840+
Receiver) = 0;
841+
842+
virtual void getRangeInfo(
843+
StringRef Filename, unsigned Offset, unsigned Length,
844+
bool CancelOnSubsequentRequest, ArrayRef<const char *> Args,
845+
SourceKitCancellationToken CancellationToken,
846+
std::function<void(const RequestResult<RangeInfo> &)> Receiver) = 0;
840847

841848
virtual void getCursorInfoFromUSR(
842849
StringRef Filename, StringRef USR, bool CancelOnSubsequentRequest,
843850
ArrayRef<const char *> Args, Optional<VFSOptions> vfsOptions,
851+
SourceKitCancellationToken CancellationToken,
844852
std::function<void(const RequestResult<CursorInfoData> &)> Receiver) = 0;
845853

846-
virtual void findRelatedIdentifiersInFile(StringRef Filename,
847-
unsigned Offset,
848-
bool CancelOnSubsequentRequest,
849-
ArrayRef<const char *> Args,
850-
std::function<void(const RequestResult<RelatedIdentsInfo> &)> Receiver) = 0;
854+
virtual void findRelatedIdentifiersInFile(
855+
StringRef Filename, unsigned Offset, bool CancelOnSubsequentRequest,
856+
ArrayRef<const char *> Args, SourceKitCancellationToken CancellationToken,
857+
std::function<void(const RequestResult<RelatedIdentsInfo> &)>
858+
Receiver) = 0;
851859

852860
virtual llvm::Optional<std::pair<unsigned, unsigned>>
853861
findUSRRange(StringRef DocumentName, StringRef USR) = 0;
@@ -872,25 +880,29 @@ class LangSupport {
872880
virtual void
873881
findLocalRenameRanges(StringRef Filename, unsigned Line, unsigned Column,
874882
unsigned Length, ArrayRef<const char *> Args,
883+
SourceKitCancellationToken CancellationToken,
875884
CategorizedRenameRangesReceiver Receiver) = 0;
876885

877-
virtual void semanticRefactoring(StringRef Filename, SemanticRefactoringInfo Info,
878-
ArrayRef<const char*> Args,
886+
virtual void semanticRefactoring(StringRef Filename,
887+
SemanticRefactoringInfo Info,
888+
ArrayRef<const char *> Args,
889+
SourceKitCancellationToken CancellationToken,
879890
CategorizedEditsReceiver Receiver) = 0;
880891

881-
virtual void collectExpressionTypes(StringRef FileName,
882-
ArrayRef<const char *> Args,
883-
ArrayRef<const char *> ExpectedProtocols,
884-
bool CanonicalType,
885-
std::function<void(const
886-
RequestResult<ExpressionTypesInFile> &)> Receiver) = 0;
892+
virtual void collectExpressionTypes(
893+
StringRef FileName, ArrayRef<const char *> Args,
894+
ArrayRef<const char *> ExpectedProtocols, bool CanonicalType,
895+
SourceKitCancellationToken CancellationToken,
896+
std::function<void(const RequestResult<ExpressionTypesInFile> &)>
897+
Receiver) = 0;
887898

888899
/// Collects variable types for a range defined by `Offset` and `Length` in
889900
/// the source file. If `Offset` or `Length` are empty, variable types for
890901
/// the entire document are collected.
891902
virtual void collectVariableTypes(
892903
StringRef FileName, ArrayRef<const char *> Args,
893904
Optional<unsigned> Offset, Optional<unsigned> Length,
905+
SourceKitCancellationToken CancellationToken,
894906
std::function<void(const RequestResult<VariableTypesInFile> &)>
895907
Receiver) = 0;
896908

‎tools/SourceKit/include/SourceKit/Support/CancellationToken.h

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
//===--- CancellationToken.h - ----------------------------------*- C++ -*-===//
2+
//
3+
// This source file is part of the Swift.org open source project
4+
//
5+
// Copyright (c) 2014 - 2021 Apple Inc. and the Swift project authors
6+
// Licensed under Apache License v2.0 with Runtime Library Exception
7+
//
8+
// See https://swift.org/LICENSE.txt for license information
9+
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
#ifndef LLVM_SOURCEKIT_SUPPORT_CANCELLATION_TOKEN_H
14+
#define LLVM_SOURCEKIT_SUPPORT_CANCELLATION_TOKEN_H
15+
16+
namespace SourceKit {
17+
18+
/// A token that uniquely identifies a SourceKit request that's served by a
19+
/// \c SwiftASTConsumer. Used to cancel the request.
20+
/// If the cancellation token is \c nullptr, it means that cancellation is not
21+
/// supported.
22+
typedef const void *SourceKitCancellationToken;
23+
24+
} // namespace SourceKit
25+
26+
#endif

‎tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -580,9 +580,17 @@ struct SwiftASTManager::Implementation {
580580
/// Since we only keep a reference to the consumers to cancel them, the
581581
/// reference to the consumer itself is weak - if it's already deallocated,
582582
/// there is no need to cancel it anymore.
583+
/// The \c CancellationToken that allows cancellation of this consumer.
584+
/// Multiple consumers might share the same \c CancellationToken if they were
585+
/// created from the same SourceKit request. E.g. a \c CursorInfoConsumer
586+
/// might schedule a second \c CursorInfoConsumer if it discovers that the AST
587+
/// that was used to serve the first request is not up-to-date enough.
588+
/// If \c CancellationToken is \c nullptr, the consumer can't be cancelled
589+
/// using a cancellation token.
583590
struct ScheduledConsumer {
584591
SwiftASTConsumerWeakRef Consumer;
585592
const void *OncePerASTToken;
593+
SourceKitCancellationToken CancellationToken;
586594
};
587595

588596
llvm::sys::Mutex ScheduledConsumersMtx;
@@ -739,7 +747,7 @@ SwiftInvocationRef SwiftASTManager::getInvocation(
739747

740748
void SwiftASTManager::processASTAsync(
741749
SwiftInvocationRef InvokRef, SwiftASTConsumerRef ASTConsumer,
742-
const void *OncePerASTToken,
750+
const void *OncePerASTToken, SourceKitCancellationToken CancellationToken,
743751
llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> fileSystem,
744752
ArrayRef<ImmutableTextSnapshotRef> Snapshots) {
745753
assert(fileSystem);
@@ -758,13 +766,32 @@ void SwiftASTManager::processASTAsync(
758766
}
759767
}
760768
}
761-
Impl.ScheduledConsumers.push_back({ASTConsumer, OncePerASTToken});
769+
Impl.ScheduledConsumers.push_back(
770+
{ASTConsumer, OncePerASTToken, CancellationToken});
762771
}
763772

764773
Producer->enqueueConsumer(ASTConsumer, fileSystem, Snapshots,
765774
shared_from_this());
766775
}
767776

777+
void SwiftASTManager::cancelASTConsumer(
778+
SourceKitCancellationToken CancellationToken) {
779+
if (!CancellationToken) {
780+
return;
781+
}
782+
Impl.cleanDeletedConsumers();
783+
llvm::sys::ScopedLock L(Impl.ScheduledConsumersMtx);
784+
for (auto ScheduledConsumer : Impl.ScheduledConsumers) {
785+
if (ScheduledConsumer.CancellationToken == CancellationToken) {
786+
if (auto Consumer = ScheduledConsumer.Consumer.lock()) {
787+
Consumer->requestCancellation();
788+
// Multiple consumers might share the same cancellation token (see
789+
// documentation on ScheduledConsumer), so we can't break here.
790+
}
791+
}
792+
}
793+
}
794+
768795
void SwiftASTManager::removeCachedAST(SwiftInvocationRef Invok) {
769796
Impl.ASTCache.remove(Invok->Impl.Key);
770797
}

‎tools/SourceKit/lib/SwiftLang/SwiftASTManager.h

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@
7575
#define LLVM_SOURCEKIT_LIB_SWIFTLANG_SWIFTASTMANAGER_H
7676

7777
#include "SourceKit/Core/LLVM.h"
78+
#include "SourceKit/Support/CancellationToken.h"
7879
#include "SwiftInvocation.h"
7980
#include "llvm/ADT/ArrayRef.h"
8081
#include "llvm/ADT/IntrusiveRefCntPtr.h"
@@ -244,10 +245,17 @@ class SwiftASTManager : public std::enable_shared_from_this<SwiftASTManager> {
244245
void
245246
processASTAsync(SwiftInvocationRef Invok, SwiftASTConsumerRef ASTConsumer,
246247
const void *OncePerASTToken,
248+
SourceKitCancellationToken CancellationToken,
247249
llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> fileSystem,
248250
ArrayRef<ImmutableTextSnapshotRef> Snapshots =
249251
ArrayRef<ImmutableTextSnapshotRef>());
250252

253+
/// Request the \c SwiftASTConsumer with the given \p CancellationToken to be
254+
/// cancelled. If \p CancellationToken is \c nullptr or no consumer with the
255+
/// given cancellation token exists (e.g. because the consumer already
256+
/// finished), this is a no-op.
257+
void cancelASTConsumer(SourceKitCancellationToken CancellationToken);
258+
251259
std::unique_ptr<llvm::MemoryBuffer> getMemoryBuffer(StringRef Filename,
252260
std::string &Error);
253261

‎tools/SourceKit/lib/SwiftLang/SwiftDocSupport.cpp

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1407,7 +1407,8 @@ void SwiftLangSupport::findRenameRanges(
14071407

14081408
void SwiftLangSupport::findLocalRenameRanges(
14091409
StringRef Filename, unsigned Line, unsigned Column, unsigned Length,
1410-
ArrayRef<const char *> Args, CategorizedRenameRangesReceiver Receiver) {
1410+
ArrayRef<const char *> Args, SourceKitCancellationToken CancellationToken,
1411+
CategorizedRenameRangesReceiver Receiver) {
14111412
std::string Error;
14121413
SwiftInvocationRef Invok = ASTMgr->getInvocation(Args, Filename, Error);
14131414
if (!Invok) {
@@ -1447,6 +1448,7 @@ void SwiftLangSupport::findLocalRenameRanges(
14471448
/// don't use 'OncePerASTToken'.
14481449
static const char OncePerASTToken = 0;
14491450
getASTManager()->processASTAsync(Invok, ASTConsumer, &OncePerASTToken,
1451+
CancellationToken,
14501452
llvm::vfs::getRealFileSystem());
14511453
}
14521454

0 commit comments

Comments
 (0)