[sourcekitd] Allow client to enable use of dispatch barriers for open/edit/close on the message handling queue

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

@@ -90,16 +90,6 @@ class RequestTracker {
   }
 
 public:
-  /// Returns \c true if the request with the given \p CancellationToken has
-  /// already been cancelled.
-  bool isCancelled(SourceKitCancellationToken CancellationToken) {
-    if (!CancellationToken) {
-      return false;
-    }
-    llvm::sys::ScopedLock L(RequestsMtx);
-    return isCancelledImpl(CancellationToken);
-  }
-
   /// Adds a \p CancellationHandler that will be called when the request
   /// associated with the \p CancellationToken is cancelled.
   /// If that request has already been cancelled when this method is called,

tools/SourceKit/tools/sourcekitd/bin/InProc/sourcekitdInProc.cpp

@@ -34,6 +34,15 @@
 
 using namespace SourceKit;
 
+/// The queue on which the all incoming requests will be handled. If barriers
+/// are enabled, open/edit/close requests will be dispatched as barriers on this
+/// queue.
+WorkQueue *msgHandlingQueue = nullptr;
+
+/// Whether request barriers have been enabled, i.e. whether open/edit/close
+/// requests should be dispatched as barriers.
+static bool RequestBarriersEnabled = false;
+
 static void postNotification(sourcekitd_response_t Notification);
 
 static void getToolchainPrefixPath(llvm::SmallVectorImpl<char> &Path) {
@@ -96,6 +105,9 @@ static std::string getDiagnosticDocumentationPath() {
 }
 
 void sourcekitd_initialize(void) {
+  assert(msgHandlingQueue == nullptr && "Cannot initialize service twice");
+  msgHandlingQueue = new WorkQueue(WorkQueue::Dequeuing::Concurrent,
+                                   "sourcekitdInProc.msgHandlingQueue");
   if (sourcekitd::initializeClient()) {
     LOG_INFO_FUNC(High, "initializing");
     sourcekitd::initializeService(getSwiftExecutablePath(), getRuntimeLibPath(),
@@ -154,7 +166,7 @@ void sourcekitd_send_request(sourcekitd_object_t req,
 
   sourcekitd_request_retain(req);
   receiver = Block_copy(receiver);
-  WorkQueue::dispatchConcurrent([=] {
+  auto handler = [=] {
     sourcekitd::handleRequest(req, /*CancellationToken=*/request_handle,
                               [=](sourcekitd_response_t resp) {
                                 // The receiver accepts ownership of the
@@ -163,7 +175,20 @@ void sourcekitd_send_request(sourcekitd_object_t req,
                                 Block_release(receiver);
                               });
     sourcekitd_request_release(req);
-  });
+  };
+
+  if (sourcekitd::requestIsEnableBarriers(req)) {
+    RequestBarriersEnabled = true;
+    sourcekitd::sendBarriersEnabledResponse([=](sourcekitd_response_t resp) {
+      // The receiver accepts ownership of the response.
+      receiver(resp);
+      Block_release(receiver);
+    });
+  } else if (RequestBarriersEnabled && sourcekitd::requestIsBarrier(req)) {
+    msgHandlingQueue->dispatchBarrier(handler);
+  } else {
+    msgHandlingQueue->dispatchConcurrent(handler);
+  }
 }
 
 void sourcekitd_cancel_request(sourcekitd_request_handle_t handle) {

tools/SourceKit/tools/sourcekitd/bin/XPC/Service/XPCService.cpp

@@ -32,6 +32,7 @@ using namespace SourceKit;
 using namespace sourcekitd;
 
 static xpc_connection_t MainConnection = nullptr;
+static bool RequestBarriersEnabled = false;
 
 static void postNotification(sourcekitd_response_t Notification) {
   xpc_connection_t peer = MainConnection;
@@ -255,43 +256,66 @@ static void sourcekitdServer_peer_event_handler(xpc_connection_t peer,
     assert(type == XPC_TYPE_DICTIONARY);
     // Handle the message
     xpc_retain(event);
-    dispatch_async(msgHandlingQueue, ^{
-      if (xpc_object_t contents =
-              xpc_dictionary_get_value(event, xpc::KeyMsg)) {
+    if (xpc_object_t contents = xpc_dictionary_get_value(event, xpc::KeyMsg)) {
+      assert(xpc_get_type(contents) == XPC_TYPE_ARRAY);
+      sourcekitd_object_t req = xpc_array_get_value(contents, 0);
+
+      void (^handler)(void) = ^{
         SourceKitCancellationToken cancelToken =
             reinterpret_cast<SourceKitCancellationToken>(
                 xpc_dictionary_get_uint64(event, xpc::KeyCancelToken));
         auto Responder = std::make_shared<XPCResponder>(event, peer);
         xpc_release(event);
 
-        assert(xpc_get_type(contents) == XPC_TYPE_ARRAY);
-        sourcekitd_object_t req = xpc_array_get_value(contents, 0);
         sourcekitd::handleRequest(req, /*CancellationToken=*/cancelToken,
                                   [Responder](sourcekitd_response_t response) {
                                     Responder->sendReply(response);
                                   });
-      } else if (xpc_object_t contents =
-                     xpc_dictionary_get_value(event, "ping")) {
+      };
+
+      if (sourcekitd::requestIsEnableBarriers(req)) {
+        dispatch_barrier_async(msgHandlingQueue, ^{
+          auto Responder = std::make_shared<XPCResponder>(event, peer);
+          xpc_release(event);
+          RequestBarriersEnabled = true;
+          sourcekitd::sendBarriersEnabledResponse([Responder](sourcekitd_response_t response) {
+            Responder->sendReply(response);
+          });
+        });
+      } else if (RequestBarriersEnabled && sourcekitd::requestIsBarrier(req)) {
+        dispatch_barrier_async(msgHandlingQueue, handler);
+      } else {
+        dispatch_async(msgHandlingQueue, handler);
+      }
+    } else if (xpc_object_t contents =
+                   xpc_dictionary_get_value(event, "ping")) {
+      dispatch_async(msgHandlingQueue, ^{
         // Ping back.
         xpc_object_t reply = xpc_dictionary_create_reply(event);
         xpc_release(event);
         assert(reply);
         xpc_connection_send_message(peer, reply);
         xpc_release(reply);
-      } else if (SourceKitCancellationToken cancelToken =
-                     reinterpret_cast<SourceKitCancellationToken>(
-                         xpc_dictionary_get_uint64(event,
-                                                   xpc::KeyCancelRequest))) {
+      });
+    } else if (SourceKitCancellationToken cancelToken =
+                   reinterpret_cast<SourceKitCancellationToken>(
+                       xpc_dictionary_get_uint64(event,
+                                                 xpc::KeyCancelRequest))) {
+      // Execute cancellation on a queue other than `msgHandling` so that we
+      // don’t block the cancellation of a request with a barrier
+      dispatch_async(dispatch_get_global_queue(QOS_CLASS_USER_INITIATED, 0), ^{
         sourcekitd::cancelRequest(/*CancellationToken=*/cancelToken);
-      } else if (SourceKitCancellationToken cancelToken =
-                     reinterpret_cast<SourceKitCancellationToken>(
-                         xpc_dictionary_get_uint64(
-                             event, xpc::KeyDisposeRequestHandle))) {
+      });
+    } else if (SourceKitCancellationToken cancelToken =
+                   reinterpret_cast<SourceKitCancellationToken>(
+                       xpc_dictionary_get_uint64(
+                           event, xpc::KeyDisposeRequestHandle))) {
+      dispatch_async(msgHandlingQueue, ^{
         sourcekitd::disposeCancellationToken(/*CancellationToken=*/cancelToken);
-      } else {
-        assert(false && "unexpected message");
-      }
-    });
+      });
+    } else {
+      assert(false && "unexpected message");
+    }
   }
 }
 

tools/SourceKit/tools/sourcekitd/include/sourcekitd/Service.h

@@ -45,6 +45,21 @@ void cancelRequest(SourceKitCancellationToken CancellationToken);
 
 void disposeCancellationToken(SourceKitCancellationToken CancellationToken);
 
+/// Returns \c true if \p Request is of a request kind that should be issued as
+/// a dispatch barrier of the message handling queue. In practice, this returns
+/// \c true for open, edit and close requets.
+///
+/// This does not check if dispatch barriers have been enabled by the sourckitd
+/// client.
+bool requestIsBarrier(sourcekitd_object_t Request);
+
+/// Returns \c true if this is a request to enable dispatch barriers in
+/// sourcekitd.
+bool requestIsEnableBarriers(sourcekitd_object_t Request);
+
+/// Send the response that request barriers have been enabled to \p Receiver.
+void sendBarriersEnabledResponse(ResponseReceiver Receiver);
+
 } // namespace sourcekitd
 
 #endif // LLVM_SOURCEKITD_SERVICE_H

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

@@ -405,6 +405,26 @@ void sourcekitd::disposeCancellationToken(
   getGlobalContext().getRequestTracker()->stopTracking(CancellationToken);
 }
 
+bool sourcekitd::requestIsBarrier(sourcekitd_object_t ReqObj) {
+  RequestDict Req(ReqObj);
+  sourcekitd_uid_t ReqUID = Req.getUID(KeyRequest);
+  return ReqUID == RequestEditorOpen || ReqUID == RequestEditorReplaceText ||
+         ReqUID == RequestEditorClose;
+}
+
+bool sourcekitd::requestIsEnableBarriers(sourcekitd_object_t ReqObj) {
+  RequestDict Req(ReqObj);
+  sourcekitd_uid_t ReqUID = Req.getUID(KeyRequest);
+  return ReqUID == RequestEnableRequestBarriers;
+}
+
+void sourcekitd::sendBarriersEnabledResponse(ResponseReceiver Receiver) {
+  ResponseBuilder RespBuilder;
+  auto Elem = RespBuilder.getDictionary();
+  Elem.setBool(KeyBarriersEnabled, true);
+  Receiver(RespBuilder.createResponse());
+}
+
 static std::unique_ptr<llvm::MemoryBuffer> getInputBufForRequest(
     Optional<StringRef> SourceFile, Optional<StringRef> SourceText,
     const Optional<VFSOptions> &vfsOptions, llvm::SmallString<64> &ErrBuf) {

utils/gyb_sourcekit_support/UIDs.py

@@ -207,7 +207,8 @@ def __init__(self, internal_name, external_name):
     # in this time. For cancellation testing purposes.
     KEY('SimulateLongRequest', 'key.simulate_long_request'),
     KEY('IsSynthesized', 'key.is_synthesized'),
-    KEY('BufferName', 'key.buffer_name')
+    KEY('BufferName', 'key.buffer_name'),
+    KEY('BarriersEnabled', 'key.barriers_enabled'),
 ]
 
 
@@ -272,6 +273,7 @@ def __init__(self, internal_name, external_name):
     REQUEST('Diagnostics', 'source.request.diagnostics'),
     REQUEST('Compile', 'source.request.compile'),
     REQUEST('CompileClose', 'source.request.compile.close'),
+    REQUEST('EnableRequestBarriers', 'source.request.enable_request_barriers'),
 ]