@@ -54,43 +54,118 @@ class GlobalConfig {
5454 Settings::CompletionOptions getCompletionOpts () const ;
5555};
5656
57- class SlowRequestSimulator {
58- std::map<SourceKitCancellationToken, std::shared_ptr<Semaphore>>
59- InProgressRequests;
57+ // / Keeps track of all requests that are currently in progress and coordinates
58+ // / cancellation. All operations are no-ops if \c nullptr is used as a
59+ // / cancellation token.
60+ // /
61+ // / Tracked requests will be removed from this object when the request returns
62+ // / a response (\c sourcekitd::handleRequest disposes the cancellation token).
63+ // / We leak memory if a request is cancelled after it has returned a result
64+ // / because we add an \c IsCancelled entry in that case but \c
65+ // / sourcekitd::handleRequest won't have a chance to dispose of the token.
66+ class RequestTracker {
67+ struct TrackedRequest {
68+ // / Whether the request has already been cancelled.
69+ bool IsCancelled = false ;
70+ // / A handler that will be called as the request gets cancelled.
71+ std::function<void (void )> CancellationHandler;
72+ };
6073
61- // / Mutex guarding \c InProgressRequests.
62- llvm::sys::Mutex InProgressRequestsMutex;
74+ // / Once we have information about a request (either a cancellation handler
75+ // / or information that it has been cancelled), it is added to this list.
76+ std::map<SourceKitCancellationToken, TrackedRequest> Requests;
77+
78+ // / Guards the \c Requests variable.
79+ llvm::sys::Mutex RequestsMtx;
80+
81+ // / Must only be called if \c RequestsMtx has been claimed.
82+ // / Returns \c true if the request has already been cancelled. Requests that
83+ // / are not tracked are not assumed to be cancelled.
84+ bool isCancelledImpl (SourceKitCancellationToken CancellationToken) {
85+ if (Requests.count (CancellationToken) == 0 ) {
86+ return false ;
87+ } else {
88+ return Requests[CancellationToken].IsCancelled ;
89+ }
90+ }
6391
6492public:
93+ // / Returns \c true if the request with the given \p CancellationToken has
94+ // / already been cancelled.
95+ bool isCancelled (SourceKitCancellationToken CancellationToken) {
96+ if (!CancellationToken) {
97+ return false ;
98+ }
99+ llvm::sys::ScopedLock L (RequestsMtx);
100+ return isCancelledImpl (CancellationToken);
101+ }
102+
103+ // / Adds a \p CancellationHandler that will be called when the request
104+ // / associated with the \p CancellationToken is cancelled.
105+ // / If that request has already been cancelled when this method is called,
106+ // / \p CancellationHandler will be called synchronously.
107+ void setCancellationHandler (SourceKitCancellationToken CancellationToken,
108+ std::function<void (void )> CancellationHandler) {
109+ if (!CancellationToken) {
110+ return ;
111+ }
112+ llvm::sys::ScopedLock L (RequestsMtx);
113+ if (isCancelledImpl (CancellationToken)) {
114+ if (CancellationHandler) {
115+ CancellationHandler ();
116+ }
117+ } else {
118+ Requests[CancellationToken].CancellationHandler = CancellationHandler;
119+ }
120+ }
121+
122+ // / Cancel the request with the given \p CancellationToken. If a cancellation
123+ // / handler is associated with it, call it.
124+ void cancel (SourceKitCancellationToken CancellationToken) {
125+ if (!CancellationToken) {
126+ return ;
127+ }
128+ llvm::sys::ScopedLock L (RequestsMtx);
129+ Requests[CancellationToken].IsCancelled = true ;
130+ if (auto CancellationHandler =
131+ Requests[CancellationToken].CancellationHandler ) {
132+ CancellationHandler ();
133+ }
134+ }
135+
136+ // / Stop tracking the request with the given \p CancellationToken, freeing up
137+ // / any memory needed for the tracking.
138+ void stopTracking (SourceKitCancellationToken CancellationToken) {
139+ if (!CancellationToken) {
140+ return ;
141+ }
142+ llvm::sys::ScopedLock L (RequestsMtx);
143+ Requests.erase (CancellationToken);
144+ }
145+ };
146+
147+ class SlowRequestSimulator {
148+ std::shared_ptr<RequestTracker> ReqTracker;
149+
150+ public:
151+ SlowRequestSimulator (std::shared_ptr<RequestTracker> ReqTracker)
152+ : ReqTracker(ReqTracker) {}
153+
65154 // / Simulate that a request takes \p DurationMs to execute. While waiting that
66155 // / duration, the request can be cancelled using the \p CancellationToken.
67156 // / Returns \c true if the request waited the required duration and \c false
68157 // / if it was cancelled.
69158 bool simulateLongRequest (int64_t DurationMs,
70159 SourceKitCancellationToken CancellationToken) {
71- auto Sema = std::make_shared<Semaphore>(0 );
72- {
73- llvm::sys::ScopedLock L (InProgressRequestsMutex);
74- InProgressRequests[CancellationToken] = Sema;
75- }
76- bool DidTimeOut = Sema->wait (DurationMs);
77- {
78- llvm::sys::ScopedLock L (InProgressRequestsMutex);
79- InProgressRequests[CancellationToken] = nullptr ;
80- }
160+ auto Sema = Semaphore (0 );
161+ ReqTracker->setCancellationHandler (CancellationToken,
162+ [&] { Sema.signal (); });
163+ bool DidTimeOut = Sema.wait (DurationMs);
164+ ReqTracker->setCancellationHandler (CancellationToken, nullptr );
81165 // If we timed out, we waited the required duration. If we didn't time out,
82166 // the semaphore was cancelled.
83167 return DidTimeOut;
84168 }
85-
86- // / Cancel a simulated long request. If the required wait duration already
87- // / elapsed, this is a no-op.
88- void cancel (SourceKitCancellationToken CancellationToken) {
89- llvm::sys::ScopedLock L (InProgressRequestsMutex);
90- if (auto InProgressSema = InProgressRequests[CancellationToken]) {
91- InProgressSema->signal ();
92- }
93- }
94169};
95170
96171class Context {
@@ -99,6 +174,7 @@ class Context {
99174 std::unique_ptr<LangSupport> SwiftLang;
100175 std::shared_ptr<NotificationCenter> NotificationCtr;
101176 std::shared_ptr<GlobalConfig> Config;
177+ std::shared_ptr<RequestTracker> ReqTracker;
102178 std::shared_ptr<SlowRequestSimulator> SlowRequestSim;
103179
104180public:
@@ -122,6 +198,8 @@ class Context {
122198 std::shared_ptr<SlowRequestSimulator> getSlowRequestSimulator () {
123199 return SlowRequestSim;
124200 }
201+
202+ std::shared_ptr<RequestTracker> getRequestTracker () { return ReqTracker; }
125203};
126204
127205} // namespace SourceKit
0 commit comments