refactor: 整理 IOStream read 相关接口

MaaXYZ · Dec 5, 2023 · d80f09d · d80f09d
1 parent a964114
commit d80f09d
Show file tree

Hide file tree

Showing 6 changed files with 76 additions and 26 deletions.
diff --git a/source/MaaAdbControlUnit/Screencap/Minicap/MinicapStream.cpp b/source/MaaAdbControlUnit/Screencap/Minicap/MinicapStream.cpp
@@ -150,7 +150,7 @@ std::optional<std::string> MinicapStream::read(size_t count)
     }
 
     using namespace std::chrono_literals;
-    return stream_handle_->read(5s, count);
+    return stream_handle_->read_some(count, 5s);
 }
 
 void MinicapStream::working_thread()

diff --git a/source/MaaUtils/IOStream/ChildPipeIOStream.cpp b/source/MaaUtils/IOStream/ChildPipeIOStream.cpp
@@ -1,10 +1,12 @@
 #include "Utils/IOStream/ChildPipeIOStream.h"
 
+#include "Utils/Time.hpp"
+
 MAA_NS_BEGIN
 
 ChildPipeIOStream::ChildPipeIOStream(const std::filesystem::path& exec, const std::vector<std::string>& args)
     : child_( //
-          exec, args, boost::process::std_out > pout_, boost::process::std_err > pout_, boost::process::std_in < pin_
+          exec, args, boost::process::std_out > pin_, boost::process::std_err > pin_, boost::process::std_in < pout_
 #ifdef _WIN32
           ,
           boost::process::windows::create_no_window
@@ -24,25 +26,43 @@ bool ChildPipeIOStream::write(std::string_view data)
         return false;
     }
 
-    pin_ << data << std::endl;
+    pout_ << data << std::endl;
     return true;
 }
 
-std::string ChildPipeIOStream::read(std::chrono::seconds timeout, size_t count)
+std::string ChildPipeIOStream::read(duration_t timeout)
 {
-    auto start_time = std::chrono::steady_clock::now();
-    auto check_timeout = [&](const auto& start_time) -> bool {
-        return timeout < std::chrono::seconds(0) || std::chrono::steady_clock::now() - start_time < timeout;
-    };
+    return read_some(std::numeric_limits<size_t>::max(), timeout);
+}
 
+std::string ChildPipeIOStream::read_some(size_t count, duration_t timeout)
+{
+    auto start_time = std::chrono::steady_clock::now();
     std::string result;
 
-    while (check_timeout(start_time) && count > result.size() && child_.running()) {
+    while (child_.running() && result.size() < count && duration_since(start_time) < timeout) {
         auto read_size = std::min(kBufferSize, count - result.size());
-        auto read_num = pout_.readsome(buffer_.get(), read_size);
-        if (read_num > 0) {
-            result.append(buffer_.get(), read_num);
+        auto read_num = pin_.readsome(buffer_.get(), read_size);
+        result.append(buffer_.get(), read_num);
+    }
+
+    return result;
+}
+
+std::string ChildPipeIOStream::read_until(std::string_view delimiter, duration_t timeout)
+{
+    auto start_time = std::chrono::steady_clock::now();
+
+    std::string result;
+
+    while (!result.ends_with(delimiter)) {
+        auto sub_timeout = timeout - duration_since<duration_t>(start_time);
+        if (sub_timeout < duration_t::zero()) {
+            break;
         }
+
+        auto sub_str = read_some(1, sub_timeout);
+        result.append(std::move(sub_str));
     }
 
     return result;

diff --git a/source/MaaUtils/IOStream/SockIOStream.cpp b/source/MaaUtils/IOStream/SockIOStream.cpp
@@ -83,21 +83,39 @@ bool SockIOStream::write(std::string_view data)
     return true;
 }
 
-std::string SockIOStream::read(std::chrono::seconds timeout, size_t count)
+std::string SockIOStream::read(duration_t timeout)
 {
-    auto start_time = std::chrono::steady_clock::now();
-    auto check_timeout = [&](const auto& start_time) -> bool {
-        return timeout < std::chrono::seconds(0) || std::chrono::steady_clock::now() - start_time < timeout;
-    };
+    return read_some(std::numeric_limits<size_t>::max(), timeout);
+}
 
+std::string SockIOStream::read_some(size_t count, duration_t timeout)
+{
+    auto start_time = std::chrono::steady_clock::now();
     std::string result;
 
-    while (check_timeout(start_time) && count > result.size() && sock_.is_open()) {
+    while (sock_.is_open() && result.size() < count && duration_since(start_time) < timeout) {
         auto read_size = std::min(kBufferSize, count - result.size());
         auto read_num = sock_.read_some(boost::asio::mutable_buffer(buffer_.get(), read_size));
-        if (read_num > 0) {
-            result.append(buffer_.get(), read_num);
+        result.append(buffer_.get(), read_num);
+    }
+
+    return result;
+}
+
+std::string SockIOStream::read_until(std::string_view delimiter, duration_t timeout)
+{
+    auto start_time = std::chrono::steady_clock::now();
+
+    std::string result;
+
+    while (!result.ends_with(delimiter)) {
+        auto sub_timeout = timeout - duration_since<duration_t>(start_time);
+        if (sub_timeout < duration_t::zero()) {
+            break;
         }
+
+        auto sub_str = read_some(1, sub_timeout);
+        result.append(std::move(sub_str));
     }
 
     return result;

diff --git a/source/include/Utils/IOStream/ChildPipeIOStream.h b/source/include/Utils/IOStream/ChildPipeIOStream.h
@@ -10,6 +10,8 @@ MAA_NS_BEGIN
 
 class MAA_UTILS_API ChildPipeIOStream : public NonCopyButMovable
 {
+    using duration_t = std::chrono::milliseconds;
+
     static constexpr size_t kBufferSize = 128 * 1024;
 
 public:
@@ -19,12 +21,16 @@ class MAA_UTILS_API ChildPipeIOStream : public NonCopyButMovable
 
 public:
     bool write(std::string_view data);
-    std::string read(std::chrono::seconds timeout = std::chrono::seconds(-1), size_t count = SIZE_MAX);
+
+    std::string read(duration_t timeout = duration_t::max());
+    std::string read_some(size_t count, duration_t timeout = duration_t::max());
+    std::string read_until(std::string_view delimiter, duration_t timeout = duration_t::max());
+
     int release();
 
 private:
-    boost::process::opstream pin_;
-    boost::process::ipstream pout_;
+    boost::process::ipstream pin_;
+    boost::process::opstream pout_;
     boost::process::child child_;
 
     std::unique_ptr<char[]> buffer_ = nullptr;

diff --git a/source/include/Utils/IOStream/SockIOStream.h b/source/include/Utils/IOStream/SockIOStream.h
@@ -42,6 +42,8 @@ class MAA_UTILS_API ClientSockIOFactory : public NonCopyButMovable
 
 class MAA_UTILS_API SockIOStream : public NonCopyButMovable
 {
+    using duration_t = std::chrono::milliseconds;
+
     static constexpr size_t kBufferSize = 128 * 1024;
 
 public:
@@ -50,7 +52,10 @@ class MAA_UTILS_API SockIOStream : public NonCopyButMovable
 
 public:
     bool write(std::string_view data);
-    std::string read(std::chrono::seconds timeout = std::chrono::seconds(-1), size_t count = SIZE_MAX);
+
+    std::string read(duration_t timeout = duration_t::max());
+    std::string read_some(size_t count, duration_t timeout = duration_t::max());
+    std::string read_until(std::string_view delimiter, duration_t timeout = duration_t::max());
 
 private:
     boost::asio::ip::tcp::socket sock_;

diff --git a/source/include/Utils/Time.hpp b/source/include/Utils/Time.hpp
@@ -43,9 +43,10 @@ inline std::string now_filestem()
 #endif
 }
 
-inline std::chrono::milliseconds duration_since(const std::chrono::steady_clock::time_point& start_time)
+template <typename duration_t = std::chrono::milliseconds>
+inline duration_t duration_since(const std::chrono::steady_clock::time_point& start_time)
 {
-    return std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::steady_clock::now() - start_time);
+    return std::chrono::duration_cast<duration_t>(std::chrono::steady_clock::now() - start_time);
 }
 
 MAA_NS_END