[sourcekit] Fix non-deterministic failure in CompileNotifications tests

This is a test-only change except for the introduction of a new request
that is only used by tests.

Our notifications are dispatched in the XPC event handler, which is not
synchronized with replies to explicit XPC send_message_with_reply calls.
This is fine for most users of sourcekitd, since the notifications are
already enqueued on the client side, but for testing we need a way to
guarantee that all notifications are passed to the client-side handler
before we exit. This commit introduces a new request for testing that
triggers a notification, allowing the client to wait on that
notification to ensure all previously posted notifications have been
handled.

Note: the non-deterministic test failures can be triggered by adding a
sleep of ~100 ms in the event handler before the notification is
dispatched to the main queue.

rdar://40311995

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);
+    if (int ret = handleTestInvocation(Opts, InitOpts)) {
-    printBufferedNotifications();
+      printBufferedNotifications(/*syncWithService=*/true);
-    if (ret) {
       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/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 @@ UID_REQUESTS = [
     REQUEST('SemanticRefactoring', 'source.request.semantic.refactoring'),
     REQUEST('EnableCompileNotifications',
             'source.request.enable-compile-notifications'),
+    REQUEST('TestNotification', 'source.request.test_notification'),
 ]