[lldb] Update JSONTransport to use MainLoop for reading. (#152367)

Reapply "[lldb] Update JSONTransport to use MainLoop for reading."
(#152155)

This reverts commit cd40281.

This also includes some updates to try to address the platforms with
failing tests.

I updated the JSONTransport and tests to use std::function instead of
llvm:unique_function. I think the tests were failing due to the
unique_function not being moved correctly in the loop on some platforms.

Author: ashgti

lldb/include/lldb/Host/JSONTransport.h

@@ -13,13 +13,17 @@
 #ifndef LLDB_HOST_JSONTRANSPORT_H
 #define LLDB_HOST_JSONTRANSPORT_H
 
+#include "lldb/Host/MainLoopBase.h"
+#include "lldb/Utility/IOObject.h"
+#include "lldb/Utility/Status.h"
 #include "lldb/lldb-forward.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/JSON.h"
-#include <chrono>
+#include <string>
 #include <system_error>
+#include <vector>
 
 namespace lldb_private {
 
@@ -28,27 +32,26 @@ class TransportEOFError : public llvm::ErrorInfo<TransportEOFError> {
   static char ID;
 
   TransportEOFError() = default;
-
-  void log(llvm::raw_ostream &OS) const override {
-    OS << "transport end of file reached";
-  }
-  std::error_code convertToErrorCode() const override {
-    return llvm::inconvertibleErrorCode();
-  }
+  void log(llvm::raw_ostream &OS) const override;
+  std::error_code convertToErrorCode() const override;
 };
 
-class TransportTimeoutError : public llvm::ErrorInfo<TransportTimeoutError> {
+class TransportUnhandledContentsError
+    : public llvm::ErrorInfo<TransportUnhandledContentsError> {
 public:
   static char ID;
 
-  TransportTimeoutError() = default;
+  explicit TransportUnhandledContentsError(std::string unhandled_contents);
 
-  void log(llvm::raw_ostream &OS) const override {
-    OS << "transport operation timed out";
-  }
-  std::error_code convertToErrorCode() const override {
-    return std::make_error_code(std::errc::timed_out);
+  void log(llvm::raw_ostream &OS) const override;
+  std::error_code convertToErrorCode() const override;
+
+  const std::string &getUnhandledContents() const {
+    return m_unhandled_contents;
   }
+
+private:
+  std::string m_unhandled_contents;
 };
 
 class TransportInvalidError : public llvm::ErrorInfo<TransportInvalidError> {
@@ -57,17 +60,17 @@ class TransportInvalidError : public llvm::ErrorInfo<TransportInvalidError> {
 
   TransportInvalidError() = default;
 
-  void log(llvm::raw_ostream &OS) const override {
-    OS << "transport IO object invalid";
-  }
-  std::error_code convertToErrorCode() const override {
-    return std::make_error_code(std::errc::not_connected);
-  }
+  void log(llvm::raw_ostream &OS) const override;
+  std::error_code convertToErrorCode() const override;
 };
 
 /// A transport class that uses JSON for communication.
 class JSONTransport {
 public:
+  using ReadHandleUP = MainLoopBase::ReadHandleUP;
+  template <typename T>
+  using Callback = std::function<void(MainLoopBase &, const llvm::Expected<T>)>;
+
   JSONTransport(lldb::IOObjectSP input, lldb::IOObjectSP output);
   virtual ~JSONTransport() = default;
 
@@ -83,24 +86,68 @@ class JSONTransport {
     return WriteImpl(message);
   }
 
-  /// Reads the next message from the input stream.
+  /// Registers the transport with the MainLoop.
   template <typename T>
-  llvm::Expected<T> Read(const std::chrono::microseconds &timeout) {
-    llvm::Expected<std::string> message = ReadImpl(timeout);
-    if (!message)
-      return message.takeError();
-    return llvm::json::parse<T>(/*JSON=*/*message);
+  llvm::Expected<ReadHandleUP> RegisterReadObject(MainLoopBase &loop,
+                                                  Callback<T> read_cb) {
+    Status error;
+    ReadHandleUP handle = loop.RegisterReadObject(
+        m_input,
+        [read_cb, this](MainLoopBase &loop) {
+          char buf[kReadBufferSize];
+          size_t num_bytes = sizeof(buf);
+          if (llvm::Error error = m_input->Read(buf, num_bytes).takeError()) {
+            read_cb(loop, std::move(error));
+            return;
+          }
+          if (num_bytes)
+            m_buffer.append(std::string(buf, num_bytes));
+
+          // If the buffer has contents, try parsing any pending messages.
+          if (!m_buffer.empty()) {
+            llvm::Expected<std::vector<std::string>> messages = Parse();
+            if (llvm::Error error = messages.takeError()) {
+              read_cb(loop, std::move(error));
+              return;
+            }
+
+            for (const auto &message : *messages)
+              if constexpr (std::is_same<T, std::string>::value)
+                read_cb(loop, message);
+              else
+                read_cb(loop, llvm::json::parse<T>(message));
+          }
+
+          // On EOF, notify the callback after the remaining messages were
+          // handled.
+          if (num_bytes == 0) {
+            if (m_buffer.empty())
+              read_cb(loop, llvm::make_error<TransportEOFError>());
+            else
+              read_cb(loop, llvm::make_error<TransportUnhandledContentsError>(
+                                std::string(m_buffer)));
+          }
+        },
+        error);
+    if (error.Fail())
+      return error.takeError();
+    return handle;
   }
 
 protected:
+  template <typename... Ts> inline auto Logv(const char *Fmt, Ts &&...Vals) {
+    Log(llvm::formatv(Fmt, std::forward<Ts>(Vals)...).str());
+  }
   virtual void Log(llvm::StringRef message);
 
   virtual llvm::Error WriteImpl(const std::string &message) = 0;
-  virtual llvm::Expected<std::string>
-  ReadImpl(const std::chrono::microseconds &timeout) = 0;
+  virtual llvm::Expected<std::vector<std::string>> Parse() = 0;
+
+  static constexpr size_t kReadBufferSize = 1024;
 
   lldb::IOObjectSP m_input;
   lldb::IOObjectSP m_output;
+  llvm::SmallString<kReadBufferSize> m_buffer;
 };
 
 /// A transport class for JSON with a HTTP header.
@@ -111,14 +158,13 @@ class HTTPDelimitedJSONTransport : public JSONTransport {
   virtual ~HTTPDelimitedJSONTransport() = default;
 
 protected:
-  virtual llvm::Error WriteImpl(const std::string &message) override;
-  virtual llvm::Expected<std::string>
-  ReadImpl(const std::chrono::microseconds &timeout) override;
-
-  // FIXME: Support any header.
-  static constexpr llvm::StringLiteral kHeaderContentLength =
-      "Content-Length: ";
-  static constexpr llvm::StringLiteral kHeaderSeparator = "\r\n\r\n";
+  llvm::Error WriteImpl(const std::string &message) override;
+  llvm::Expected<std::vector<std::string>> Parse() override;
+
+  static constexpr llvm::StringLiteral kHeaderContentLength = "Content-Length";
+  static constexpr llvm::StringLiteral kHeaderFieldSeparator = ":";
+  static constexpr llvm::StringLiteral kHeaderSeparator = "\r\n";
+  static constexpr llvm::StringLiteral kEndOfHeader = "\r\n\r\n";
 };
 
 /// A transport class for JSON RPC.
@@ -129,9 +175,8 @@ class JSONRPCTransport : public JSONTransport {
   virtual ~JSONRPCTransport() = default;
 
 protected:
-  virtual llvm::Error WriteImpl(const std::string &message) override;
-  virtual llvm::Expected<std::string>
-  ReadImpl(const std::chrono::microseconds &timeout) override;
+  llvm::Error WriteImpl(const std::string &message) override;
+  llvm::Expected<std::vector<std::string>> Parse() override;
 
   static constexpr llvm::StringLiteral kMessageSeparator = "\n";
 };