[mlir-lsp] Guard writing output to JSONTransport with mutex

[rriddle]

Summary

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

mlir/include/mlir/Tools/lsp-server-support/Transport.h

notificationHandlers[method] = [method, handler,
       (thisPtr->*handler)(*param);
     };
   }
 
   /// Create an OutgoingNotification object used for the given method.
   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));
     };
   }
 
 private:
   template <typename HandlerT>
   using HandlerMap = llvm::StringMap<llvm::unique_function<HandlerT>>;
 
   HandlerMap<void(llvm::json::Value)> notificationHandlers;
   HandlerMap<void(llvm::json::Value, Callback<llvm::json::Value>)>
       methodHandlers;
 
   JSONTransport &transport;
+
+  /// Mutex to guard sending output messages to the transport.
+  std::mutex transportOutputMutex;
 };
 
 } // namespace lsp
 } // namespace mlir
 
 #endif

mlir/lib/Tools/lsp-server-support/Transport.cpp

 
 namespace {
 /// Function object to reply to an LSP call.
 /// Each instance must be called exactly once, otherwise:
 ///  - if there was no reply, an error reply is sent
 ///  - 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;
   Reply &operator=(const Reply &) = delete;
 
   void operator()(llvm::Expected<llvm::json::Value> reply);
 
 private:
   StringRef method;
   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;
 }
 
 void Reply::operator()(llvm::Expected<llvm::json::Value> reply) {
   if (replied.exchange(true)) {
     Logger::error("Replied twice to message {0}({1})", method, id);
     assert(false && "must reply to each call only once!");
     return;
   }
   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));
   } else {
     llvm::Error error = reply.takeError();
     Logger::info("--> reply:{0}({1})", method, id, error);
     transport->reply(std::move(id), std::move(error));
   }
bool MessageHandler::onNotify(llvm::StringRef method, llvm::json::Value value) {
   }
   return true;
 }
 
 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()) {
     it->second(std::move(params), std::move(reply));
   } else {
     reply(llvm::make_error<LSPError>("method not found: " + method.str(),
                                      ErrorCode::MethodNotFound));
   }