From fc209e4158ccbe383baf7cd6bcba1605e4d3d01b Mon Sep 17 00:00:00 2001
From: River Riddle <riddleriver@gmail.com>
Date: Wed, 26 Jul 2023 00:08:25 -0700
Subject: [PATCH] [mlir-lsp] Guard writing output to JSONTransport with mutex

This allows for users of the lsp transport libraries to process replies
in parallel, without overlapping/clobbering the output.

Differential Revision: https://reviews.llvm.org/D156295
---
 mlir/include/mlir/Tools/lsp-server-support/Transport.h |  4 ++++
 mlir/lib/Tools/lsp-server-support/Transport.cpp        | 16 ++++++++++------
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/mlir/include/mlir/Tools/lsp-server-support/Transport.h b/mlir/include/mlir/Tools/lsp-server-support/Transport.h
index 0deb7e3..ce742be 100644
--- a/mlir/include/mlir/Tools/lsp-server-support/Transport.h
+++ b/mlir/include/mlir/Tools/lsp-server-support/Transport.h
@@ -158,6 +158,7 @@ public:
   template <typename T>
   OutgoingNotification<T> outgoingNotification(llvm::StringLiteral method) {
     return [&, method](const T &params) {
+      std::lock_guard<std::mutex> transportLock(transportOutputMutex);
       Logger::info("--> {0}", method);
       transport.notify(method, llvm::json::Value(params));
     };
@@ -172,6 +173,9 @@ private:
       methodHandlers;
 
   JSONTransport &transport;
+
+  /// Mutex to guard sending output messages to the transport.
+  std::mutex transportOutputMutex;
 };
 
 } // namespace lsp
diff --git a/mlir/lib/Tools/lsp-server-support/Transport.cpp b/mlir/lib/Tools/lsp-server-support/Transport.cpp
index 92171f1..3915146 100644
--- a/mlir/lib/Tools/lsp-server-support/Transport.cpp
+++ b/mlir/lib/Tools/lsp-server-support/Transport.cpp
@@ -30,8 +30,8 @@ namespace {
 ///  - if there were multiple replies, only the first is sent
 class Reply {
 public:
-  Reply(const llvm::json::Value &id, StringRef method,
-        JSONTransport &transport);
+  Reply(const llvm::json::Value &id, StringRef method, JSONTransport &transport,
+        std::mutex &transportOutputMutex);
   Reply(Reply &&other);
   Reply &operator=(Reply &&) = delete;
   Reply(const Reply &) = delete;
@@ -44,16 +44,19 @@ private:
   std::atomic<bool> replied = {false};
   llvm::json::Value id;
   JSONTransport *transport;
+  std::mutex &transportOutputMutex;
 };
 } // namespace
 
 Reply::Reply(const llvm::json::Value &id, llvm::StringRef method,
-             JSONTransport &transport)
-    : id(id), transport(&transport) {}
+             JSONTransport &transport, std::mutex &transportOutputMutex)
+    : id(id), transport(&transport),
+      transportOutputMutex(transportOutputMutex) {}
 
 Reply::Reply(Reply &&other)
     : replied(other.replied.load()), id(std::move(other.id)),
-      transport(other.transport) {
+      transport(other.transport),
+      transportOutputMutex(other.transportOutputMutex) {
   other.transport = nullptr;
 }
 
@@ -65,6 +68,7 @@ void Reply::operator()(llvm::Expected<llvm::json::Value> reply) {
   }
   assert(transport && "expected valid transport to reply to");
 
+  std::lock_guard<std::mutex> transportLock(transportOutputMutex);
   if (reply) {
     Logger::info("--> reply:{0}({1})", method, id);
     transport->reply(std::move(id), std::move(reply));
@@ -98,7 +102,7 @@ bool MessageHandler::onCall(llvm::StringRef method, llvm::json::Value params,
                             llvm::json::Value id) {
   Logger::info("--> {0}({1})", method, id);
 
-  Reply reply(id, method, transport);
+  Reply reply(id, method, transport, transportOutputMutex);
 
   auto it = methodHandlers.find(method);
   if (it != methodHandlers.end()) {
-- 
2.7.4