新增支持外部设置UDP Socket接收缓冲区 (#289)

Author: alexliyu7352

src/Network/BufferSock.cpp

@@ -9,6 +9,7 @@
  */
 
 #include <assert.h>
+#include <limits>
 #include "BufferSock.h"
 #include "Util/logger.h"
 #include "Util/uv_errno.h"
@@ -609,14 +610,37 @@ class SocketRecvFromBuffer : public SocketRecvBuffer {
 
 static constexpr auto kPacketCount = 32;
 static constexpr auto kBufferCapacity = 4 * 1024u;
+static constexpr auto kMaxTotalBufferBytes = 64 * 1024u * 1024u;
 
 SocketRecvBuffer::Ptr SocketRecvBuffer::create(bool is_udp) {
+    return create(is_udp, kPacketCount, kBufferCapacity);
+}
+
+SocketRecvBuffer::Ptr SocketRecvBuffer::create(bool is_udp, size_t packet_count, size_t buffer_capacity) {
+    packet_count = packet_count ? packet_count : kPacketCount;
+    buffer_capacity = buffer_capacity ? buffer_capacity : kBufferCapacity;
+
+    auto use_default = false;
+    if (packet_count < 1 || buffer_capacity < 2) {
+        use_default = true;
+    } else if (packet_count > (std::numeric_limits<size_t>::max)() / buffer_capacity) {
+        use_default = true;
+    } else if (packet_count * buffer_capacity > kMaxTotalBufferBytes) {
+        use_default = true;
+    }
+
+    if (use_default) {
+        WarnL << "Invalid recv buffer config, fallback to defaults: packet_count="
+              << packet_count << ", buffer_capacity=" << buffer_capacity;
+        packet_count = kPacketCount;
+        buffer_capacity = kBufferCapacity;
+    }
 #if defined(__linux) || defined(__linux__)
     if (is_udp) {
-        return std::make_shared<SocketRecvmmsgBuffer>(kPacketCount, kBufferCapacity);
+        return std::make_shared<SocketRecvmmsgBuffer>(packet_count, buffer_capacity);
     }
 #endif
-    return std::make_shared<SocketRecvFromBuffer>(kPacketCount * kBufferCapacity);
+    return std::make_shared<SocketRecvFromBuffer>(packet_count * buffer_capacity);
 }
 
 } //toolkit

src/Network/BufferSock.h

@@ -81,6 +81,7 @@ class SocketRecvBuffer {
     virtual struct sockaddr_storage &getAddress(size_t index) = 0;
 
     static Ptr create(bool is_udp);
+    static Ptr create(bool is_udp, size_t packet_count, size_t buffer_capacity);
 };
 
 }

src/Network/Socket.cpp

@@ -277,6 +277,13 @@ bool Socket::attachEvent(const SockNum::Ptr &sock) {
     // tcp客户端或udp  [AUTO-TRANSLATED:00c16e7f]
     //TCP client or UDP
     auto read_buffer = _poller->getSharedBuffer(sock->type() == SockNum::Sock_UDP);
+    if (sock->type() == SockNum::Sock_UDP) {
+        LOCK_GUARD(_mtx_sock_fd);
+        _udp_recv_buffer_frozen = true;
+        if (_read_buffer) {
+            read_buffer = _read_buffer;
+        }
+    }
     auto result = _poller->addEvent(sock->rawFd(), EventPoller::Event_Read | EventPoller::Event_Error | EventPoller::Event_Write, [weak_self, sock, read_buffer](int event) {
         auto strong_self = weak_self.lock();
         if (!strong_self) {
@@ -298,6 +305,10 @@ bool Socket::attachEvent(const SockNum::Ptr &sock) {
         }
     });
 
+    if (result == -1 && sock->type() == SockNum::Sock_UDP) {
+        LOCK_GUARD(_mtx_sock_fd);
+        _udp_recv_buffer_frozen = false;
+    }
     return -1 != result;
 }
 
@@ -321,7 +332,9 @@ ssize_t Socket::onRead(const SockNum::Ptr &sock, const SocketRecvBuffer::Ptr &bu
                 if (sock->type() == SockNum::Sock_TCP) {
                     emitErr(toSockException(err));
                 } else {
-                    WarnL << "Recv err on udp socket[" << sock->rawFd() << "]: " << uv_strerror(err);
+                    if (!(err == UV_ECONNREFUSED && _ignore_udp_conn_refused.load(std::memory_order_acquire))) {
+                        WarnL << "Recv err on udp socket[" << sock->rawFd() << "]: " << uv_strerror(err);
+                    }
                 }
             }
             return ret;
@@ -485,6 +498,7 @@ void Socket::closeSock(bool close_fd) {
         if (close_fd) {
             _err_emit = false;
             _sock_fd = nullptr;
+            _udp_recv_buffer_frozen = false;
         } else if (_sock_fd) {
             _sock_fd->delEvent();
         }
@@ -703,6 +717,7 @@ void Socket::setSock(SockNum::Ptr sock) {
         SockUtil::get_sock_peer_addr(_sock_fd->rawFd(), _peer_addr);
     } else {
         _sock_fd = nullptr;
+        _udp_recv_buffer_frozen = false;
     }
 }
 
@@ -1021,6 +1036,25 @@ void Socket::setSendFlags(int flags) {
     _sock_flags = flags;
 }
 
+bool Socket::setUdpRecvBuffer(const SocketRecvBuffer::Ptr &buffer) {
+    // This hook is setup-time only. UdpServer creation callbacks may run
+    // before the owner poller starts processing the fd, so the hard
+    // requirement here is "before fd creation/attach", not "already on
+    // poller thread". The customization itself is only honored for UDP
+    // sockets.
+    LOCK_GUARD(_mtx_sock_fd);
+    if (_sock_fd || _udp_recv_buffer_frozen) {
+        WarnL << "setUdpRecvBuffer must be called before the socket fd is created and UDP IO is attached";
+        return false;
+    }
+    _read_buffer = buffer;
+    return true;
+}
+
+void Socket::setIgnoreUdpConnRefused(bool ignore) {
+    _ignore_udp_conn_refused.store(ignore, std::memory_order_release);
+}
+
 ///////////////SockSender///////////////////
 
 SockSender &SockSender::operator<<(const char *buf) {

src/Network/Socket.h

@@ -45,7 +45,11 @@ namespace toolkit {
 
 //默认的socket flags:不触发SIGPIPE,非阻塞发送  [AUTO-TRANSLATED:fefc4946]
 //Default socket flags: do not trigger SIGPIPE, non-blocking send
-#define SOCKET_DEFAULE_FLAGS (FLAG_NOSIGNAL | FLAG_DONTWAIT )
+#define SOCKET_DEFAULT_FLAGS (FLAG_NOSIGNAL | FLAG_DONTWAIT )
+// Backward compatibility alias for the legacy misspelled macro name.
+#ifndef SOCKET_DEFAULE_FLAGS
+#define SOCKET_DEFAULE_FLAGS SOCKET_DEFAULT_FLAGS
+#endif
 
 //发送超时时间，如果在规定时间内一直没有发送数据成功，那么将触发onErr事件  [AUTO-TRANSLATED:9c5d8d87]
 //Send timeout time, if no data is sent successfully within the specified time, the onErr event will be triggered
@@ -649,7 +653,36 @@ class Socket : public std::enable_shared_from_this<Socket>, public noncopyable,
      
      * [AUTO-TRANSLATED:2b11445c]
      */
-    void setSendFlags(int flags = SOCKET_DEFAULE_FLAGS);
+    void setSendFlags(int flags = SOCKET_DEFAULT_FLAGS);
+
+    // Install a UDP-specific recv buffer before the socket starts receiving.
+    // This is intended for setup-time tuning, not runtime reconfiguration
+    // after IO callbacks are active.
+    /**
+     * Replace the UDP recv buffer before the socket fd is created/attached.
+     * This is intended for setup-time customization of special UDP transports
+     * and must not be used as a runtime reconfiguration hook.
+     *
+     * IMPORTANT: custom SocketRecvBuffer implementations must remain
+     * compatible with the batched UDP receive path used by Socket::onRead().
+     * For the active receive batch, both &buffer->getBuffer(0) and
+     * &buffer->getAddress(0) are treated as pointers to contiguous arrays
+     * that can be indexed up to count - 1, and the referenced storage must
+     * remain valid for the duration of the receive operation.
+     *
+     * Passing a SocketRecvBuffer that does not satisfy this layout/lifetime
+     * contract is unsupported and may lead to undefined behavior.
+     *
+     * @return Whether the configuration was accepted.
+     */
+    bool setUdpRecvBuffer(const SocketRecvBuffer::Ptr &buffer);
+
+    // Suppress the UDP ECONNREFUSED read warning on sockets that intentionally
+    // communicate with transient peers, such as QUIC sessions that may receive
+    // a late ICMP port-unreachable after the peer has already closed. This is
+    // a narrow transport-specific knob and should not be enabled casually by
+    // ordinary upper-layer business code.
+    void setIgnoreUdpConnRefused(bool ignore);
 
     /**
      * 关闭套接字
@@ -734,11 +767,10 @@ class Socket : public std::enable_shared_from_this<Socket>, public noncopyable,
     ssize_t send_l(Buffer::Ptr buf, bool is_buf_sock, bool try_flush = true);
     void connect_l(const std::string &url, uint16_t port, const onErrCB &con_cb_in, float timeout_sec, const std::string &local_ip, uint16_t local_port);
     bool fromSock_l(SockNum::Ptr sock);
-
 private:
     // send socket时的flag  [AUTO-TRANSLATED:e364a1bf]
     //Flag for sending socket
-    int _sock_flags = SOCKET_DEFAULE_FLAGS;
+    int _sock_flags = SOCKET_DEFAULT_FLAGS;
     // 最大发送缓存，单位毫秒，距上次发送缓存清空时间不能超过该参数  [AUTO-TRANSLATED:3bd6dba3]
     //Maximum send buffer, in milliseconds, the time since the last send buffer was cleared cannot exceed this parameter
     uint32_t _max_send_buffer_ms = SEND_TIME_OUT_SEC * 1000;
@@ -823,6 +855,13 @@ class Socket : public std::enable_shared_from_this<Socket>, public noncopyable,
     //Object count statistics
     ObjectStatistic<Socket> _statistic;
 
+    // Optional per-socket recv path used by a small number of UDP transports.
+    // This must be configured before the socket fd is created or any IO
+    // callbacks are attached.
+    SocketRecvBuffer::Ptr _read_buffer;
+    bool _udp_recv_buffer_frozen = false;
+    std::atomic<bool> _ignore_udp_conn_refused{false};
+
     // 链接缓存地址,防止tcp reset 导致无法获取对端的地址  [AUTO-TRANSLATED:f8847463]
     //Connection cache address, to prevent TCP reset from causing the inability to obtain the peer's address
     struct sockaddr_storage _local_addr;

tests/test_udpSocketBufferConfig.cpp

@@ -0,0 +1,87 @@
+/*
+ * Copyright (c) 2016 The ZLToolKit project authors. All Rights Reserved.
+ *
+ * This file is part of ZLToolKit(https://github.com/ZLMediaKit/ZLToolKit).
+ *
+ * Use of this source code is governed by MIT license that can be found in the
+ * LICENSE file in the root of the source tree. All contributing project authors
+ * may be found in the AUTHORS file in the root of the source tree.
+ */
+
+#include <csignal>
+#include <iostream>
+#include <string>
+
+#include "Network/Socket.h"
+#include "Thread/semaphore.h"
+#include "Util/logger.h"
+#include "Util/util.h"
+
+using namespace std;
+using namespace toolkit;
+
+int main() {
+    Logger::Instance().add(std::make_shared<ConsoleChannel>());
+    Logger::Instance().setWriter(std::make_shared<AsyncLogWriter>());
+
+    auto recv_sock = Socket::createSocket();
+    auto send_sock = Socket::createSocket();
+    if (!recv_sock || !send_sock) {
+        cerr << "create socket failed" << endl;
+        return 1;
+    }
+
+    auto invalid_cfg_buffer = SocketRecvBuffer::create(true, 0, 1);
+    if (!invalid_cfg_buffer) {
+        cerr << "create invalid_cfg_buffer failed" << endl;
+        return 2;
+    }
+
+    if (!recv_sock->setUdpRecvBuffer(invalid_cfg_buffer)) {
+        cerr << "setUdpRecvBuffer should succeed before fd creation" << endl;
+        return 3;
+    }
+
+    recv_sock->setIgnoreUdpConnRefused(true);
+
+    if (!recv_sock->bindUdpSock(0, "127.0.0.1")) {
+        cerr << "bind recv socket failed" << endl;
+        return 4;
+    }
+    if (!send_sock->bindUdpSock(0, "127.0.0.1")) {
+        cerr << "bind send socket failed" << endl;
+        return 5;
+    }
+
+    if (recv_sock->setUdpRecvBuffer(SocketRecvBuffer::create(true, 1, 4096))) {
+        cerr << "setUdpRecvBuffer should fail after fd creation" << endl;
+        return 6;
+    }
+
+    semaphore sem;
+    string received;
+    recv_sock->setOnRead([&](const Buffer::Ptr &buf, struct sockaddr *, int) {
+        received.assign(buf ? buf->data() : "", buf ? buf->size() : 0);
+        sem.post();
+    });
+
+    auto dst = SockUtil::make_sockaddr("127.0.0.1", recv_sock->get_local_port());
+    const string payload = "udp-buffer-config-smoke";
+    if (send_sock->send(payload, reinterpret_cast<struct sockaddr *>(&dst)) <= 0) {
+        cerr << "send payload failed" << endl;
+        return 7;
+    }
+
+    if (!sem.wait(3000)) {
+        cerr << "recv timeout" << endl;
+        return 8;
+    }
+
+    if (received != payload) {
+        cerr << "unexpected payload: " << received << endl;
+        return 9;
+    }
+
+    cout << "udp socket buffer config regression passed" << endl;
+    return 0;
+}