Use dispatch barriers for open/edit/close on the message handling queue

Author: ahoppen

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

@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "sourcekitd/DictionaryKeys.h"
 #include "sourcekitd/Internal.h"
 #include "sourcekitd/Service.h"
 
@@ -34,6 +35,8 @@
 
 using namespace SourceKit;
 
+WorkQueue *msgHandlingQueue = nullptr;
+
 static void postNotification(sourcekitd_response_t Notification);
 
 static void getToolchainPrefixPath(llvm::SmallVectorImpl<char> &Path) {
@@ -96,6 +99,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(),
@@ -152,9 +158,12 @@ void sourcekitd_send_request(sourcekitd_object_t req,
     *out_handle = request_handle;
   }
 
+  sourcekitd_uid_t ReqUID =
+      sourcekitd::RequestDict(req).getUID(sourcekitd::KeyRequest);
+
   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 +172,13 @@ void sourcekitd_send_request(sourcekitd_object_t req,
                                 Block_release(receiver);
                               });
     sourcekitd_request_release(req);
-  });
+  };
+
+  if (sourcekitd::requestIsBarrier(ReqUID)) {
+    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

@@ -10,7 +10,9 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "sourcekitd/DictionaryKeys.h"
 #include "sourcekitd/Internal-XPC.h"
+#include "sourcekitd/Internal.h"
 #include "sourcekitd/Logging.h"
 #include "sourcekitd/Service.h"
 
@@ -255,43 +257,57 @@ 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);
+
+      sourcekitd_uid_t ReqUID = RequestDict(req).getUID(KeyRequest);
+
+      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 (requestIsBarrier(ReqUID)) {
+        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))) {
+      dispatch_async(msgHandlingQueue, ^{
         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/Internal.h

@@ -229,6 +229,9 @@ struct VariantFunctions {
   const void *(*data_get_ptr)(sourcekitd_variant_t obj);
 };
 
-}
+/// Whether a request with the given UID should be dispatched as a barrier onto
+/// the message handling queue.
+bool requestIsBarrier(sourcekitd_uid_t reqUID);
 
+} // namespace sourcekitd
 #endif

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

@@ -66,3 +66,16 @@ UIdent sourcekitd::UIdentFromSKDUID(sourcekitd_uid_t uid) {
 
   return UIdent::getFromOpaqueValue(uid);
 }
+
+static sourcekitd_uid_t RequestEditorOpen =
+    sourcekitd::SKDUIDFromUIdent(UIdent("source.request.editor.open"));
+static sourcekitd_uid_t RequestEditorClose =
+    sourcekitd::SKDUIDFromUIdent(UIdent("source.request.editor.close"));
+static sourcekitd_uid_t RequestEditorReplaceText =
+    sourcekitd::SKDUIDFromUIdent(UIdent("source.request.editor.replacetext"));
+
+bool sourcekitd::requestIsBarrier(sourcekitd_uid_t reqUID) {
+  // TODO: Do we also need to add barriers for interface open etc?
+  return reqUID == RequestEditorOpen || reqUID == RequestEditorReplaceText ||
+         reqUID == RequestEditorClose;
+}