[SourceKit] Add option to simulate a long-running request

This allows us to remove the dependency on a slow-to-typecheck example for a `sourcekitd-test`.

test/SourceKit/Misc/cancellation.swift

@@ -1,13 +1,11 @@
 // Check that we can cancel requests.
 // 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.
 
-// RUN: not %sourcekitd-test -req=cursor -id=slow -async -pos=10:3 %s -- %s == \
+// RUN: not %sourcekitd-test -req=cursor -id=slow -async -pos=9:5 -simulate-long-request=5000 %s -- %s == \
 // RUN: -shell -- sleep 1 == \
 // RUN: -cancel=slow 2>&1 \
 // RUN: | %FileCheck %s
 
-func foo(x: Invalid1, y: Invalid2) {
+let x = "Hello World"
-  x / y / x / y / x / y / x / y
-}
 
 // CHECK: error response (Request Cancelled)

tools/SourceKit/include/SourceKit/Core/Context.h

@@ -14,9 +14,12 @@
 #define LLVM_SOURCEKIT_CORE_CONTEXT_H
 
 #include "SourceKit/Core/LLVM.h"
-#include "llvm/ADT/StringRef.h"
+#include "SourceKit/Support/CancellationToken.h"
+#include "SourceKit/Support/Concurrency.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Mutex.h"
+#include <map>
 #include <memory>
 #include <string>
 
@@ -51,12 +54,52 @@ public:
   Settings::CompletionOptions getCompletionOpts() const;
 };
 
+class SlowRequestSimulator {
+  std::map<SourceKitCancellationToken, std::shared_ptr<Semaphore>>
+      InProgressRequests;
+
+  /// Mutex guarding \c InProgressRequests.
+  llvm::sys::Mutex InProgressRequestsMutex;
+
+public:
+  /// Simulate that a request takes \p DurationMs to execute. While waiting that
+  /// duration, the request can be cancelled using the \p CancellationToken.
+  /// Returns \c true if the request waited the required duration and \c false
+  /// if it was cancelled.
+  bool simulateLongRequest(int64_t DurationMs,
+                           SourceKitCancellationToken CancellationToken) {
+    auto Sema = std::make_shared<Semaphore>(0);
+    {
+      llvm::sys::ScopedLock L(InProgressRequestsMutex);
+      InProgressRequests[CancellationToken] = Sema;
+    }
+    bool DidTimeOut = Sema->wait(DurationMs);
+    {
+      llvm::sys::ScopedLock L(InProgressRequestsMutex);
+      InProgressRequests[CancellationToken] = nullptr;
+    }
+    // If we timed out, we waited the required duration. If we didn't time out,
+    // the semaphore was cancelled.
+    return DidTimeOut;
+  }
+
+  /// Cancel a simulated long request. If the required wait duration already
+  /// elapsed, this is a no-op.
+  void cancel(SourceKitCancellationToken CancellationToken) {
+    llvm::sys::ScopedLock L(InProgressRequestsMutex);
+    if (auto InProgressSema = InProgressRequests[CancellationToken]) {
+      InProgressSema->signal();
+    }
+  }
+};
+
 class Context {
   std::string RuntimeLibPath;
   std::string DiagnosticDocumentationPath;
   std::unique_ptr<LangSupport> SwiftLang;
   std::shared_ptr<NotificationCenter> NotificationCtr;
   std::shared_ptr<GlobalConfig> Config;
+  std::shared_ptr<SlowRequestSimulator> SlowRequestSim;
 
 public:
   Context(StringRef RuntimeLibPath, StringRef DiagnosticDocumentationPath,
@@ -75,6 +118,10 @@ public:
   std::shared_ptr<NotificationCenter> getNotificationCenter() { return NotificationCtr; }
 
   std::shared_ptr<GlobalConfig> getGlobalConfiguration() { return Config; }
+
+  std::shared_ptr<SlowRequestSimulator> getSlowRequestSimulator() {
+    return SlowRequestSim;
+  }
 };
 
 } // namespace SourceKit

tools/SourceKit/include/SourceKit/Support/Concurrency.h

@@ -35,6 +35,8 @@ public:
     return Impl::wait(ImplObj);
   }
 
+  /// Waits for the semaphore, timing out after \p milliseconds.
+  /// Returns \c true if waiting timed out.
   bool wait(long milliseconds) {
     return Impl::wait(ImplObj, milliseconds);
   }

tools/SourceKit/lib/Core/Context.cpp

@@ -44,7 +44,7 @@ SourceKit::Context::Context(
       DiagnosticDocumentationPath(DiagnosticDocumentationPath),
       NotificationCtr(
           new NotificationCenter(shouldDispatchNotificationsOnMain)),
-      Config(new GlobalConfig()) {
+      Config(new GlobalConfig()), SlowRequestSim(new SlowRequestSimulator()) {
   // Should be called last after everything is initialized.
   SwiftLang = LangSupportFactoryFn(*this);
 }

tools/SourceKit/tools/sourcekitd-test/Options.td

@@ -157,6 +157,10 @@ def suppress_config_request : Flag<["-"], "suppress-config-request">,
 def module_cache_path: Separate<["-"], "module-cache-path">, HelpText<"module cache path">;
 def module_cache_path_EQ : Joined<["-"], "module-cache-path=">, Alias<module_cache_path>;
 
+def simulate_long_request : Separate<["-"], "simulate-long-request">,
+  HelpText<"Simulate that the request takes x ms longer to execute. The request can be cancelled while waiting this duration">;
+def simulate_long_request_EQ : Joined<["-"], "simulate-long-request=">, Alias<simulate_long_request>;
+
 def shell: Flag<["-"], "shell">,
   HelpText<"Run shell command">;
 

tools/SourceKit/tools/sourcekitd-test/TestOptions.cpp

@@ -420,6 +420,15 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
       ModuleCachePath = InputArg->getValue();
       break;
 
+    case OPT_simulate_long_request:
+      unsigned SimulatedDuration;
+      if (StringRef(InputArg->getValue()).getAsInteger(10, SimulatedDuration)) {
+        llvm::errs() << "error: expected integer for 'simulate-long-request'\n";
+        return true;
+      }
+      SimulateLongRequest = SimulatedDuration;
+      break;
+
     case OPT_shell:
       ShellExecution = true;
       break;

tools/SourceKit/tools/sourcekitd-test/TestOptions.h

@@ -107,6 +107,9 @@ struct TestOptions {
   std::string RequestId;
   /// If not empty, all requests with this ID should be cancelled.
   std::string CancelRequest;
+  /// If set, simulate that the request takes x ms longer than it actually
+  /// does. The request can be cancelled while waiting this duration.
+  llvm::Optional<uint64_t> SimulateLongRequest;
   bool CheckInterfaceIsASCII = false;
   bool UsedSema = false;
   bool PrintRequest = true;

tools/SourceKit/tools/sourcekitd-test/sourcekitd-test.cpp

@@ -1137,6 +1137,10 @@ static int handleTestInvocation(TestOptions Opts, TestOptions &InitOpts) {
     sourcekitd_request_dictionary_set_int64(Req, KeyCancelOnSubsequentRequest,
                                             *Opts.CancelOnSubsequentRequest);
   }
+  if (Opts.SimulateLongRequest.hasValue()) {
+    sourcekitd_request_dictionary_set_int64(Req, KeySimulateLongRequest,
+                                            *Opts.SimulateLongRequest);
+  }
 
   if (!Opts.SwiftVersion.empty()) {
     if (Opts.PassVersionAsString) {

tools/SourceKit/tools/sourcekitd/lib/API/Requests.cpp

@@ -399,8 +399,8 @@ void sourcekitd::handleRequest(sourcekitd_object_t Req,
 }
 
 void sourcekitd::cancelRequest(SourceKitCancellationToken CancellationToken) {
-  LangSupport &Lang = getGlobalContext().getSwiftLangSupport();
+  getGlobalContext().getSlowRequestSimulator()->cancel(CancellationToken);
-  Lang.cancelRequest(CancellationToken);
+  getGlobalContext().getSwiftLangSupport().cancelRequest(CancellationToken);
 }
 
 static std::unique_ptr<llvm::MemoryBuffer> getInputBufForRequest(
@@ -472,6 +472,16 @@ void handleRequestImpl(sourcekitd_object_t ReqObj,
   ++numRequests;
 
   RequestDict Req(ReqObj);
+
+  if (auto SimulateLongRequestDuration =
+          Req.getOptionalInt64(KeySimulateLongRequest)) {
+    if (!getGlobalContext().getSlowRequestSimulator()->simulateLongRequest(
+            *SimulateLongRequestDuration, CancellationToken)) {
+      Rec(createErrorRequestCancelled());
+      return;
+    }
+  }
+
   sourcekitd_uid_t ReqUID = Req.getUID(KeyRequest);
   if (!ReqUID)
     return Rec(createErrorRequestInvalid("missing 'key.request' with UID value"));

utils/gyb_sourcekit_support/UIDs.py

@@ -199,6 +199,9 @@ UID_KEYS = [
     KEY('EffectiveAccess', 'key.effective_access'),
     KEY('DeclarationLang', 'key.decl_lang'),
     KEY('SecondarySymbols', 'key.secondary_symbols'),
+    # Before executing the actual request wait x ms. The request can be canceled
+    # in this time. For cancellation testing purposes.
+    KEY('SimulateLongRequest', 'key.simulate_long_request'),
 ]