[sourcekit] Fix non-deterministic failure in CompileNotifications tests

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

@@ -101,11 +101,13 @@ static SourceKitRequest ActiveRequest = SourceKitRequest::None;
 static sourcekitd_uid_t SemaDiagnosticStage;
 
 static sourcekitd_uid_t NoteDocUpdate;
-
 static SourceKit::Semaphore semaSemaphore(0);
 static sourcekitd_response_t semaResponse;
 static const char *semaName;
 
+static sourcekitd_uid_t NoteTest;
+static SourceKit::Semaphore noteSyncSemaphore(0);
+
 namespace {
 struct AsyncResponseInfo {
   SourceKit::Semaphore semaphore{0};
@@ -139,7 +141,28 @@ struct NotificationBuffer {
 };
 static NotificationBuffer notificationBuffer;
 
-static void printBufferedNotifications() {
+static void syncNotificationsWithService() {
+  // Send TestNotification request, then wait for the notification. This ensures
+  // that all notifications previously posted on the service side have been
+  // passed to our notification handler.
+  sourcekitd_object_t req = sourcekitd_request_dictionary_create(nullptr, nullptr, 0);
+  sourcekitd_request_dictionary_set_uid(req, KeyRequest, RequestTestNotification);
+  auto resp = sourcekitd_send_request_sync(req);
+  if (sourcekitd_response_is_error(resp)) {
+    sourcekitd_response_description_dump(resp);
+    exit(1);
+  }
+  sourcekitd_response_dispose(resp);
+  sourcekitd_request_release(req);
+  if (noteSyncSemaphore.wait(60 * 1000)) {
+    llvm::report_fatal_error("Test notification not received");
+  }
+}
+
+static void printBufferedNotifications(bool syncWithService = true) {
+  if (syncWithService) {
+    syncNotificationsWithService();
+  }
   notificationBuffer.handleNotifications([](sourcekitd_response_t note) {
     sourcekitd_response_description_dump_filedesc(note, STDOUT_FILENO);
   });
@@ -171,6 +194,7 @@ static int skt_main(int argc, const char **argv) {
   SemaDiagnosticStage = sourcekitd_uid_get_from_cstr("source.diagnostic.stage.swift.sema");
 
   NoteDocUpdate = sourcekitd_uid_get_from_cstr("source.notification.editor.documentupdate");
+  NoteTest = sourcekitd_uid_get_from_cstr("source.notification.test");
 
 #define REQUEST(NAME, CONTENT) Request##NAME = sourcekitd_uid_get_from_cstr(CONTENT);
 #define KIND(NAME, CONTENT) Kind##NAME = sourcekitd_uid_get_from_cstr(CONTENT);
@@ -342,11 +366,12 @@ static int handleTestInvocation(ArrayRef<const char *> Args,
 
   assert(Opts.repeatRequest >= 1);
   for (unsigned i = 0; i < Opts.repeatRequest; ++i) {
-    int ret = handleTestInvocation(Opts, InitOpts);
-    printBufferedNotifications();
-    if (ret) {
+    if (int ret = handleTestInvocation(Opts, InitOpts)) {
+      printBufferedNotifications(/*syncWithService=*/true);
       return ret;
     }
+    // We will sync with the service before exiting; don't do so here.
+    printBufferedNotifications(/*syncWithService=*/false);
   }
   return 0;
 }
@@ -1169,6 +1194,8 @@ static void notification_receiver(sourcekitd_response_t resp) {
     semaResponse = sourcekitd_send_request_sync(edReq);
     sourcekitd_request_release(edReq);
     semaSemaphore.signal();
+  } else if (note == NoteTest) {
+    noteSyncSemaphore.signal();
   } else {
     notificationBuffer.add(resp);
   }

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

@@ -154,15 +154,12 @@ sourcekitd_set_notification_handler(sourcekitd_response_receiver_t receiver) {
 }
 
 void sourcekitd::postNotification(sourcekitd_response_t Notification) {
-  sourcekitd_response_receiver_t receiver = Block_copy(NotificationReceiver);
-  if (!receiver) {
-    sourcekitd_response_dispose(Notification);
-    return;
-  }
-
   WorkQueue::dispatchOnMain([=]{
+    if (!NotificationReceiver) {
+      sourcekitd_response_dispose(Notification);
+      return;
+    }
     // The receiver accepts ownership of the notification object.
-    receiver(Notification);
-    Block_release(receiver);
+    NotificationReceiver(Notification);
   });
 }

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

@@ -322,6 +322,15 @@ void handleRequestImpl(sourcekitd_object_t ReqObj, ResponseReceiver Rec) {
     ::exit(1);
   }
 
+  if (ReqUID == RequestTestNotification) {
+    static UIdent TestNotification("source.notification.test");
+    ResponseBuilder RespBuilder;
+    auto Dict = RespBuilder.getDictionary();
+    Dict.set(KeyNotification, TestNotification);
+    sourcekitd::postNotification(RespBuilder.createResponse());
+    return Rec(ResponseBuilder().createResponse());
+  }
+
   if (ReqUID == RequestDemangle) {
     SmallVector<const char *, 8> MangledNames;
     bool Failed = Req.getStringArray(KeyNames, MangledNames, /*isOptional=*/true);

utils/gyb_sourcekit_support/UIDs.py

@@ -207,6 +207,7 @@ def __init__(self, internal_name, external_name):
     REQUEST('SemanticRefactoring', 'source.request.semantic.refactoring'),
     REQUEST('EnableCompileNotifications',
             'source.request.enable-compile-notifications'),
+    REQUEST('TestNotification', 'source.request.test_notification'),
 ]