From 91f8919a091a045eeb6ce3796f69d5a94e5c64a4 Mon Sep 17 00:00:00 2001
From: xiongziliang <771730766@qq.com>
Date: Mon, 24 Jun 2024 21:59:12 +0800
Subject: [PATCH] =?UTF-8?q?Socket=E6=8E=A5=E6=94=B6=E6=95=B0=E6=8D=AE?=
 =?UTF-8?q?=E5=9B=9E=E8=B0=83=E6=94=AF=E6=8C=81=E6=97=A0=E6=8B=B7=E8=B4=9D?=
 =?UTF-8?q?=E6=95=B0=E6=8D=AE=E8=BD=AC=E7=A7=BB?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/Network/Socket.cpp | 18 ++++++++++++++----
 src/Network/Socket.h   |  4 ++--
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/src/Network/Socket.cpp b/src/Network/Socket.cpp
index 7cb5d721..5e4f611d 100644
--- a/src/Network/Socket.cpp
+++ b/src/Network/Socket.cpp
@@ -282,7 +282,8 @@ class MMsgBuffer {
 #endif
 
     MMsgBuffer(size_t count, size_t size)
-        : _iovec(count)
+        : _size(size)
+        , _iovec(count)
         , _mmsgs(count)
         , _buffers(count)
         , _address(count) {
@@ -307,8 +308,16 @@ class MMsgBuffer {
     }
 
     ssize_t recvFromSocket(int fd, ssize_t &count) {
-        for (auto &mmsg : _mmsgs) {
+        for (auto i = 0u; i < _mmsgs.size(); ++i) {
+            auto &mmsg = _mmsgs[i];
             mmsg.msg_hdr.msg_namelen = sizeof(struct sockaddr_storage);
+            auto &buf = _buffers[i];
+            if (!buf) {
+                auto raw = BufferRaw::create();
+                raw->setCapacity(_size);
+                buf = raw;
+                mmsg.msg_hdr.msg_iov->iov_base = buf->data();
+            }
         }
         do {
             count = recvmmsg(fd, &_mmsgs[0], _mmsgs.size(), 0, nullptr);
@@ -330,9 +339,9 @@ class MMsgBuffer {
         return nread;
     }
 
-    const Buffer::Ptr &getBuffer(size_t index) const { return _buffers[index]; }
+    Buffer::Ptr &getBuffer(size_t index) { return _buffers[index]; }
 
-    const struct sockaddr_storage &getAddress(size_t index) const { return _address[index]; }
+    struct sockaddr_storage &getAddress(size_t index) { return _address[index]; }
 
 private:
 #if !defined(__linux)
@@ -347,6 +356,7 @@ class MMsgBuffer {
 #endif
 
 private:
+    size_t _size;
     std::vector<struct iovec> _iovec;
     std::vector<struct mmsghdr> _mmsgs;
     std::vector<Buffer::Ptr> _buffers;
diff --git a/src/Network/Socket.h b/src/Network/Socket.h
index a2bd4af2..97f970ee 100644
--- a/src/Network/Socket.h
+++ b/src/Network/Socket.h
@@ -282,8 +282,8 @@ class Socket : public std::enable_shared_from_this<Socket>, public noncopyable,
 public:
     using Ptr = std::shared_ptr<Socket>;
     //接收数据回调
-    using onReadCB = std::function<void(const Buffer::Ptr &buf, struct sockaddr *addr, int addr_len)>;
-    using onMultiReadCB = std::function<void(const Buffer::Ptr *buf, const struct sockaddr_storage *addr, size_t count)>;
+    using onReadCB = std::function<void(Buffer::Ptr &buf, struct sockaddr *addr, int addr_len)>;
+    using onMultiReadCB = std::function<void(Buffer::Ptr *buf, struct sockaddr_storage *addr, size_t count)>;
 
     //发生错误回调
     using onErrCB = std::function<void(const SockException &err)>;