diff --git a/src/MeoAssistant/AsstPlatform.h b/src/MeoAssistant/AsstPlatform.h new file mode 100644 index 0000000000..690c15fd4a --- /dev/null +++ b/src/MeoAssistant/AsstPlatform.h @@ -0,0 +1,52 @@ +#pragma once + +#include +#include +#include + +namespace asst::platform +{ + extern const size_t page_size; + + void* aligned_alloc(size_t len, size_t align); + void aligned_free(void* ptr); + + template + requires std::is_trivial_v + class single_page_buffer + { + TElem* _ptr = nullptr; + + public: + single_page_buffer() + { + _ptr = reinterpret_cast(aligned_alloc(page_size, page_size)); + if (!_ptr) throw std::bad_alloc(); + } + + single_page_buffer(nullptr_t) {} + + ~single_page_buffer() + { + if (_ptr) aligned_free(reinterpret_cast(_ptr)); + } + + // disable copy construct + single_page_buffer(const single_page_buffer&) = delete; + single_page_buffer& operator=(const single_page_buffer&) = delete; + + inline single_page_buffer(single_page_buffer&& other) { std::swap(_ptr, other._ptr); } + inline single_page_buffer& operator=(single_page_buffer&& other) + { + if (_ptr) { + aligned_free(reinterpret_cast(_ptr)); + _ptr = nullptr; + } + std::swap(_ptr, other._ptr); + return *this; + } + + inline TElem* get() const { return _ptr; } + inline size_t size() const { return _ptr ? (page_size / sizeof(TElem)) : 0; } + }; +} diff --git a/src/MeoAssistant/AsstPlatformPosix.cpp b/src/MeoAssistant/AsstPlatformPosix.cpp index fe894ebbbf..4a4914346c 100644 --- a/src/MeoAssistant/AsstPlatformPosix.cpp +++ b/src/MeoAssistant/AsstPlatformPosix.cpp @@ -1,6 +1,27 @@ #if __has_include() +#include "AsstPlatform.h" #include "AsstPlatformPosix.h" #include "AsstUtils.hpp" + +#include + +static size_t get_page_size() +{ + return (size_t)sysconf(_SC_PAGESIZE); +} + +const size_t asst::platform::page_size = get_page_size(); + +void* asst::platform::aligned_alloc(size_t len, size_t align) +{ + return aligned_alloc(len, align); +} + +void asst::platform::aligned_free(void* ptr) +{ + free(ptr); +} + std::string asst::utils::callcmd(const std::string& cmdline) { constexpr int PipeBuffSize = 4096; diff --git a/src/MeoAssistant/AsstPlatformWin32.cpp b/src/MeoAssistant/AsstPlatformWin32.cpp index 632a68318f..48344928c1 100644 --- a/src/MeoAssistant/AsstPlatformWin32.cpp +++ b/src/MeoAssistant/AsstPlatformWin32.cpp @@ -1,4 +1,5 @@ #ifdef _WIN32 +#include "AsstPlatform.h" #include "AsstPlatformWin32.h" #include @@ -8,6 +9,25 @@ #include "AsstUtils.hpp" #include "Logger.hpp" +static size_t get_page_size() +{ + SYSTEM_INFO sysInfo {}; + GetSystemInfo(&sysInfo); + return sysInfo.dwPageSize; +} + +const size_t asst::platform::page_size = get_page_size(); + +void* asst::platform::aligned_alloc(size_t len, size_t align) +{ + return _aligned_malloc(len, align); +} + +void asst::platform::aligned_free(void* ptr) +{ + _aligned_free(ptr); +} + bool asst::win32::CreateOverlappablePipe(HANDLE* read, HANDLE* write, SECURITY_ATTRIBUTES* secattr_read, SECURITY_ATTRIBUTES* secattr_write, DWORD bufsize, bool overlapped_read, bool overlapped_write) diff --git a/src/MeoAssistant/Controller.cpp b/src/MeoAssistant/Controller.cpp index d9a7aa398d..4d4b945d8f 100644 --- a/src/MeoAssistant/Controller.cpp +++ b/src/MeoAssistant/Controller.cpp @@ -1,5 +1,6 @@ #include "Controller.h" #include "AsstConf.h" +#include "AsstPlatform.h" #ifdef _WIN32 #include "AsstPlatformWin32.h" @@ -60,17 +61,6 @@ asst::Controller::Controller(AsstCallback callback, void* callback_arg) m_support_socket = false; #endif - m_pipe_buffer = std::make_unique(PipeBuffSize); - if (!m_pipe_buffer) { - throw "controller pipe buffer allocated failed"; - } - if (m_support_socket) { - m_socket_buffer = std::make_unique(SocketBuffSize); - if (!m_socket_buffer) { - throw "controller socket buffer allocated failed"; - } - } - m_cmd_thread = std::thread(&Controller::pipe_working_proc, this); } @@ -204,6 +194,8 @@ std::optional asst::Controller::call_command(const std::string& cmd std::string pipe_data; std::string sock_data; + asst::platform::single_page_buffer pipe_buffer; + std::optional> sock_buffer; auto start_time = steady_clock::now(); @@ -212,7 +204,7 @@ std::optional asst::Controller::call_command(const std::string& cmd DWORD err = 0; HANDLE pipe_parent_read = INVALID_HANDLE_VALUE, pipe_child_write = INVALID_HANDLE_VALUE; SECURITY_ATTRIBUTES sa_inherit { .nLength = sizeof(SECURITY_ATTRIBUTES), .bInheritHandle = TRUE }; - if (!asst::win32::CreateOverlappablePipe(&pipe_parent_read, &pipe_child_write, nullptr, &sa_inherit, PipeBuffSize, + if (!asst::win32::CreateOverlappablePipe(&pipe_parent_read, &pipe_child_write, nullptr, &sa_inherit, (DWORD)pipe_buffer.size(), true, false)) { err = GetLastError(); Log.error("CreateOverlappablePipe failed, err", err); @@ -247,17 +239,21 @@ std::optional asst::Controller::call_command(const std::string& cmd bool socket_eof = false; OVERLAPPED pipeov { .hEvent = CreateEventW(nullptr, TRUE, FALSE, nullptr) }; - (void)ReadFile(pipe_parent_read, m_pipe_buffer.get(), PipeBuffSize, nullptr, &pipeov); + (void)ReadFile(pipe_parent_read, pipe_buffer.get(), (DWORD)pipe_buffer.size(), nullptr, &pipeov); OVERLAPPED sockov {}; SOCKET client_socket = INVALID_SOCKET; + std::unique_lock socket_lock; if (recv_by_socket) { + // acquire socket accept lock + socket_lock = std::unique_lock(m_socket_mutex); + sock_buffer = asst::platform::single_page_buffer(); sockov.hEvent = CreateEventW(nullptr, TRUE, FALSE, nullptr); client_socket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP); DWORD dummy; - if (!m_AcceptEx(m_server_sock, client_socket, m_socket_buffer.get(), - SocketBuffSize - ((sizeof(sockaddr_in) + 16) * 2), sizeof(sockaddr_in) + 16, + if (!m_AcceptEx(m_server_sock, client_socket, sock_buffer.value().get(), + (DWORD)sock_buffer.value().size() - ((sizeof(sockaddr_in) + 16) * 2), sizeof(sockaddr_in) + 16, sizeof(sockaddr_in) + 16, &dummy, &sockov)) { err = WSAGetLastError(); if (err == ERROR_IO_PENDING) { @@ -308,8 +304,8 @@ std::optional asst::Controller::call_command(const std::string& cmd // pipe read DWORD len = 0; if (GetOverlappedResult(pipe_parent_read, &pipeov, &len, FALSE)) { - pipe_data.insert(pipe_data.end(), m_pipe_buffer.get(), m_pipe_buffer.get() + len); - (void)ReadFile(pipe_parent_read, m_pipe_buffer.get(), PipeBuffSize, nullptr, &pipeov); + pipe_data.insert(pipe_data.end(), pipe_buffer.get(), pipe_buffer.get() + len); + (void)ReadFile(pipe_parent_read, pipe_buffer.get(), (DWORD)pipe_buffer.size(), nullptr, &pipeov); } else { err = GetLastError(); @@ -323,9 +319,11 @@ std::optional asst::Controller::call_command(const std::string& cmd // AcceptEx, client_socker is connected and first chunk of data is received DWORD len = 0; if (GetOverlappedResult(reinterpret_cast(m_server_sock), &sockov, &len, FALSE)) { + // unlock after accept + socket_lock.unlock(); accept_pending = false; if (recv_by_socket) - sock_data.insert(sock_data.end(), m_socket_buffer.get(), m_socket_buffer.get() + len); + sock_data.insert(sock_data.end(), sock_buffer.value().get(), sock_buffer.value().get() + len); if (len == 0) { socket_eof = true; @@ -337,24 +335,25 @@ std::optional asst::Controller::call_command(const std::string& cmd sockov = {}; sockov.hEvent = event; - (void)ReadFile(reinterpret_cast(client_socket), m_socket_buffer.get(), PipeBuffSize, - nullptr, &sockov); + (void)ReadFile(reinterpret_cast(client_socket), sock_buffer.value().get(), + (DWORD)sock_buffer.value().size(), nullptr, &sockov); } } + } else { // ReadFile DWORD len = 0; if (GetOverlappedResult(reinterpret_cast(client_socket), &sockov, &len, FALSE)) { if (recv_by_socket) - sock_data.insert(sock_data.end(), m_socket_buffer.get(), m_socket_buffer.get() + len); + sock_data.insert(sock_data.end(), sock_buffer.value().get(), sock_buffer.value().get() + len); if (len == 0) { socket_eof = true; closesocket(client_socket); } else { - (void)ReadFile(reinterpret_cast(client_socket), m_socket_buffer.get(), PipeBuffSize, - nullptr, &sockov); + (void)ReadFile(reinterpret_cast(client_socket), sock_buffer.value().get(), + (DWORD)sock_buffer.value().size(), nullptr, &sockov); } } else { @@ -403,11 +402,11 @@ std::optional asst::Controller::call_command(const std::string& cmd // parent process std::unique_lock pipe_lock(m_pipe_mutex); do { - ssize_t read_num = read(m_pipe_out[PIPE_READ], m_pipe_buffer.get(), PipeBuffSize); + ssize_t read_num = read(m_pipe_out[PIPE_READ], pipe_buffer.get(), PipeBuffSize); while (read_num > 0) { - pipe_data.insert(pipe_data.end(), m_pipe_buffer.get(), m_pipe_buffer.get() + read_num); - read_num = read(m_pipe_out[PIPE_READ], m_pipe_buffer.get(), PipeBuffSize); + pipe_data.insert(pipe_data.end(), pipe_buffer.get(), pipe_buffer.get() + read_num); + read_num = read(m_pipe_out[PIPE_READ], pipe_buffer.get(), PipeBuffSize); }; } while (::waitpid(m_child, &exit_ret, WNOHANG) == 0 && !check_timeout()); } diff --git a/src/MeoAssistant/Controller.h b/src/MeoAssistant/Controller.h index e0b0a79e2e..5c777965bd 100644 --- a/src/MeoAssistant/Controller.h +++ b/src/MeoAssistant/Controller.h @@ -103,11 +103,6 @@ namespace asst std::minstd_rand m_rand_engine; - static constexpr int PipeBuffSize = 4 * 1024 * 1024; // 管道缓冲区大小 - static constexpr int SocketBuffSize = 4 * 1024 * 1024; // socket 缓冲区大小 - std::unique_ptr m_pipe_buffer = nullptr; - std::mutex m_pipe_mutex; - std::unique_ptr m_socket_buffer = nullptr; std::mutex m_socket_mutex; #ifdef _WIN32 diff --git a/src/MeoAssistant/MeoAssistant.vcxproj b/src/MeoAssistant/MeoAssistant.vcxproj index 7cad74e3e0..e3e53cc799 100644 --- a/src/MeoAssistant/MeoAssistant.vcxproj +++ b/src/MeoAssistant/MeoAssistant.vcxproj @@ -23,6 +23,7 @@ + diff --git a/src/MeoAssistant/MeoAssistant.vcxproj.filters b/src/MeoAssistant/MeoAssistant.vcxproj.filters index 2fc6a92906..569c033b2e 100644 --- a/src/MeoAssistant/MeoAssistant.vcxproj.filters +++ b/src/MeoAssistant/MeoAssistant.vcxproj.filters @@ -408,6 +408,9 @@ 头文件\Utils + + 头文件\Utils +