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

This updates JSONTransport to use a MainLoop for reading messages.

This also allows us to read in larger chunks than we did previously.
With the event driven reading operations we can read in chunks and store
the contents in an internal buffer. Separately we can parse the buffer
and split the contents up into messages.

Our previous version approach would read a byte at a time, which is less
efficient.

Author: ashgti

lldb/include/lldb/Host/JSONTransport.h

@@ -13,13 +13,16 @@
 #ifndef LLDB_HOST_JSONTRANSPORT_H
 #define LLDB_HOST_JSONTRANSPORT_H
 
+#include "lldb/Host/MainLoopBase.h"
 #include "lldb/lldb-forward.h"
+#include "llvm/ADT/FunctionExtras.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 +31,33 @@ 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";
-  }
+  void log(llvm::raw_ostream &OS) const override { OS << "transport EOF"; }
   std::error_code convertToErrorCode() const override {
-    return llvm::inconvertibleErrorCode();
+    return std::make_error_code(std::errc::io_error);
   }
 };
 
-class TransportTimeoutError : public llvm::ErrorInfo<TransportTimeoutError> {
+class TransportUnhandledContentsError
+    : public llvm::ErrorInfo<TransportUnhandledContentsError> {
 public:
   static char ID;
 
-  TransportTimeoutError() = default;
+  explicit TransportUnhandledContentsError(std::string unhandled_contents)
+      : m_unhandled_contents(unhandled_contents) {}
 
   void log(llvm::raw_ostream &OS) const override {
-    OS << "transport operation timed out";
+    OS << "transport EOF with unhandled contents " << m_unhandled_contents;
   }
   std::error_code convertToErrorCode() const override {
-    return std::make_error_code(std::errc::timed_out);
+    return std::make_error_code(std::errc::bad_message);
   }
+
+  const std::string &getUnhandledContents() const {
+    return m_unhandled_contents;
+  }
+
+private:
+  std::string m_unhandled_contents;
 };
 
 class TransportInvalidError : public llvm::ErrorInfo<TransportInvalidError> {
@@ -68,6 +77,11 @@ class TransportInvalidError : public llvm::ErrorInfo<TransportInvalidError> {
 /// A transport class that uses JSON for communication.
 class JSONTransport {
 public:
+  using ReadHandleUP = MainLoopBase::ReadHandleUP;
+  template <typename T>
+  using Callback =
+      llvm::unique_function<void(MainLoopBase &, const llvm::Expected<T>)>;
+
   JSONTransport(lldb::IOObjectSP input, lldb::IOObjectSP output);
   virtual ~JSONTransport() = default;
 
@@ -83,24 +97,69 @@ 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> callback) {
+    Status error;
+    ReadHandleUP handle = loop.RegisterReadObject(
+        m_input,
+        [&](MainLoopBase &loop) {
+          char buffer[kReadBufferSize];
+          size_t len = sizeof(buffer);
+          if (llvm::Error error = m_input->Read(buffer, len).takeError()) {
+            callback(loop, std::move(error));
+            return;
+          }
+
+          if (len)
+            m_buffer.append(std::string(buffer, len));
+
+          // 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()) {
+              callback(loop, std::move(error));
+              return;
+            }
+
+            for (const auto &message : *messages)
+              if constexpr (std::is_same<T, std::string>::value)
+                callback(loop, message);
+              else
+                callback(loop, llvm::json::parse<T>(message));
+          }
+
+          // On EOF, notify the callback after the remaining messages were
+          // handled.
+          if (len == 0) {
+            if (m_buffer.empty())
+              callback(loop, llvm::make_error<TransportEOFError>());
+            else
+              callback(loop, llvm::make_error<TransportUnhandledContentsError>(
+                                 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;
 
   lldb::IOObjectSP m_input;
   lldb::IOObjectSP m_output;
+  std::string m_buffer;
+
+  static constexpr size_t kReadBufferSize = 1024;
 };
 
 /// A transport class for JSON with a HTTP header.
@@ -111,14 +170,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 +187,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";
 };