Older/ToolKit/Network/BufferSock.cpp

591 lines
16 KiB
C++
Raw Normal View History

2024-09-28 23:55:00 +08:00
/*
* 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 <assert.h>
#include "BufferSock.h"
#include "Util/logger.h"
#include "Util/uv_errno.h"
#if defined(__linux__) || defined(__linux)
#ifndef _GNU_SOURCE
#define _GNU_SOURCE
#endif
#ifndef MSG_WAITFORONE
#define MSG_WAITFORONE 0x10000
#endif
#ifndef HAVE_SENDMMSG_API
#include <unistd.h>
#include <sys/syscall.h>
static inline int sendmmsg(int fd, struct mmsghdr *mmsg,
unsigned vlen, unsigned flags)
{
return syscall(__NR_sendmmsg, fd, mmsg, vlen, flags);
}
#endif
#ifndef HAVE_RECVMMSG_API
#include <unistd.h>
#include <sys/syscall.h>
static inline int recvmmsg(int fd, struct mmsghdr *mmsg,
unsigned vlen, unsigned flags, struct timespec *timeout)
{
return syscall(__NR_recvmmsg, fd, mmsg, vlen, flags, timeout);
}
#endif
#endif// defined(__linux__) || defined(__linux)
namespace toolkit {
StatisticImp(BufferList)
/////////////////////////////////////// BufferSock ///////////////////////////////////////
BufferSock::BufferSock(Buffer::Ptr buffer, struct sockaddr *addr, int addr_len) {
if (addr) {
_addr_len = addr_len ? addr_len : SockUtil::get_sock_len(addr);
memcpy(&_addr, addr, _addr_len);
}
assert(buffer);
_buffer = std::move(buffer);
}
char *BufferSock::data() const {
return _buffer->data();
}
size_t BufferSock::size() const {
return _buffer->size();
}
const struct sockaddr *BufferSock::sockaddr() const {
return (struct sockaddr *)&_addr;
}
socklen_t BufferSock::socklen() const {
return _addr_len;
}
/////////////////////////////////////// BufferCallBack ///////////////////////////////////////
class BufferCallBack {
public:
BufferCallBack(List<std::pair<Buffer::Ptr, bool> > list, BufferList::SendResult cb)
: _cb(std::move(cb))
, _pkt_list(std::move(list)) {}
~BufferCallBack() {
sendCompleted(false);
}
void sendCompleted(bool flag) {
if (_cb) {
//全部发送成功或失败回调
while (!_pkt_list.empty()) {
_cb(_pkt_list.front().first, flag);
_pkt_list.pop_front();
}
} else {
_pkt_list.clear();
}
}
void sendFrontSuccess() {
if (_cb) {
//发送成功回调
_cb(_pkt_list.front().first, true);
}
_pkt_list.pop_front();
}
protected:
BufferList::SendResult _cb;
List<std::pair<Buffer::Ptr, bool> > _pkt_list;
};
/////////////////////////////////////// BufferSendMsg ///////////////////////////////////////
#if defined(_WIN32)
using SocketBuf = WSABUF;
#else
using SocketBuf = iovec;
#endif
class BufferSendMsg final : public BufferList, public BufferCallBack {
public:
using SocketBufVec = std::vector<SocketBuf>;
BufferSendMsg(List<std::pair<Buffer::Ptr, bool> > list, SendResult cb);
~BufferSendMsg() override = default;
bool empty() override;
size_t count() override;
ssize_t send(int fd, int flags) override;
private:
void reOffset(size_t n);
ssize_t send_l(int fd, int flags);
private:
size_t _iovec_off = 0;
size_t _remain_size = 0;
SocketBufVec _iovec;
};
bool BufferSendMsg::empty() {
return _remain_size == 0;
}
size_t BufferSendMsg::count() {
return _iovec.size() - _iovec_off;
}
ssize_t BufferSendMsg::send_l(int fd, int flags) {
ssize_t n;
#if !defined(_WIN32)
do {
struct msghdr msg;
msg.msg_name = nullptr;
msg.msg_namelen = 0;
msg.msg_iov = &(_iovec[_iovec_off]);
msg.msg_iovlen = _iovec.size() - _iovec_off;
if (msg.msg_iovlen > IOV_MAX) {
msg.msg_iovlen = IOV_MAX;
}
msg.msg_control = nullptr;
msg.msg_controllen = 0;
msg.msg_flags = flags;
n = sendmsg(fd, &msg, flags);
} while (-1 == n && UV_EINTR == get_uv_error(true));
#else
do {
DWORD sent = 0;
n = WSASend(fd, const_cast<LPWSABUF>(&_iovec[0]), static_cast<DWORD>(_iovec.size()), &sent, static_cast<DWORD>(flags), 0, 0);
if (n == SOCKET_ERROR) {return -1;}
n = sent;
} while (n < 0 && UV_ECANCELED == get_uv_error(true));
#endif
if (n >= (ssize_t)_remain_size) {
//全部写完了
_remain_size = 0;
sendCompleted(true);
return n;
}
if (n > 0) {
//部分发送成功
reOffset(n);
return n;
}
//一个字节都未发送
return n;
}
ssize_t BufferSendMsg::send(int fd, int flags) {
auto remain_size = _remain_size;
while (_remain_size && send_l(fd, flags) != -1);
ssize_t sent = remain_size - _remain_size;
if (sent > 0) {
//部分或全部发送成功
return sent;
}
//一个字节都未发送成功
return -1;
}
void BufferSendMsg::reOffset(size_t n) {
_remain_size -= n;
size_t offset = 0;
for (auto i = _iovec_off; i != _iovec.size(); ++i) {
auto &ref = _iovec[i];
#if !defined(_WIN32)
offset += ref.iov_len;
#else
offset += ref.len;
#endif
if (offset < n) {
//此包发送完毕
sendFrontSuccess();
continue;
}
_iovec_off = i;
if (offset == n) {
//这是末尾发送完毕的一个包
++_iovec_off;
sendFrontSuccess();
break;
}
//这是末尾发送部分成功的一个包
size_t remain = offset - n;
#if !defined(_WIN32)
ref.iov_base = (char *)ref.iov_base + ref.iov_len - remain;
ref.iov_len = remain;
#else
ref.buf = (CHAR *)ref.buf + ref.len - remain;
ref.len = remain;
#endif
break;
}
}
BufferSendMsg::BufferSendMsg(List<std::pair<Buffer::Ptr, bool>> list, SendResult cb)
: BufferCallBack(std::move(list), std::move(cb))
, _iovec(_pkt_list.size()) {
auto it = _iovec.begin();
_pkt_list.for_each([&](std::pair<Buffer::Ptr, bool> &pr) {
#if !defined(_WIN32)
it->iov_base = pr.first->data();
it->iov_len = pr.first->size();
_remain_size += it->iov_len;
#else
it->buf = pr.first->data();
it->len = pr.first->size();
_remain_size += it->len;
#endif
++it;
});
}
/////////////////////////////////////// BufferSendTo ///////////////////////////////////////
class BufferSendTo final: public BufferList, public BufferCallBack {
public:
BufferSendTo(List<std::pair<Buffer::Ptr, bool> > list, SendResult cb, bool is_udp);
~BufferSendTo() override = default;
bool empty() override;
size_t count() override;
ssize_t send(int fd, int flags) override;
private:
bool _is_udp;
size_t _offset = 0;
};
BufferSendTo::BufferSendTo(List<std::pair<Buffer::Ptr, bool>> list, BufferList::SendResult cb, bool is_udp)
: BufferCallBack(std::move(list), std::move(cb))
, _is_udp(is_udp) {}
bool BufferSendTo::empty() {
return _pkt_list.empty();
}
size_t BufferSendTo::count() {
return _pkt_list.size();
}
static inline BufferSock *getBufferSockPtr(std::pair<Buffer::Ptr, bool> &pr) {
if (!pr.second) {
return nullptr;
}
return static_cast<BufferSock *>(pr.first.get());
}
ssize_t BufferSendTo::send(int fd, int flags) {
size_t sent = 0;
ssize_t n;
while (!_pkt_list.empty()) {
auto &front = _pkt_list.front();
auto &buffer = front.first;
if (_is_udp) {
auto ptr = getBufferSockPtr(front);
n = ::sendto(fd, buffer->data() + _offset, buffer->size() - _offset, flags, ptr ? ptr->sockaddr() : nullptr, ptr ? ptr->socklen() : 0);
} else {
n = ::send(fd, buffer->data() + _offset, buffer->size() - _offset, flags);
}
if (n >= 0) {
assert(n);
_offset += n;
if (_offset == buffer->size()) {
sendFrontSuccess();
_offset = 0;
}
sent += n;
continue;
}
//n == -1的情况
if (get_uv_error(true) == UV_EINTR) {
//被打断,需要继续发送
continue;
}
//其他原因导致的send返回-1
break;
}
return sent ? sent : -1;
}
/////////////////////////////////////// BufferSendMmsg ///////////////////////////////////////
#if defined(__linux__) || defined(__linux)
class BufferSendMMsg : public BufferList, public BufferCallBack {
public:
BufferSendMMsg(List<std::pair<Buffer::Ptr, bool> > list, SendResult cb);
~BufferSendMMsg() override = default;
bool empty() override;
size_t count() override;
ssize_t send(int fd, int flags) override;
private:
void reOffset(size_t n);
ssize_t send_l(int fd, int flags);
private:
size_t _remain_size = 0;
std::vector<struct iovec> _iovec;
std::vector<struct mmsghdr> _hdrvec;
};
bool BufferSendMMsg::empty() {
return _remain_size == 0;
}
size_t BufferSendMMsg::count() {
return _hdrvec.size();
}
ssize_t BufferSendMMsg::send_l(int fd, int flags) {
ssize_t n;
do {
n = sendmmsg(fd, &_hdrvec[0], _hdrvec.size(), flags);
} while (-1 == n && UV_EINTR == get_uv_error(true));
if (n > 0) {
//部分或全部发送成功
reOffset(n);
return n;
}
//一个字节都未发送
return n;
}
ssize_t BufferSendMMsg::send(int fd, int flags) {
auto remain_size = _remain_size;
while (_remain_size && send_l(fd, flags) != -1);
ssize_t sent = remain_size - _remain_size;
if (sent > 0) {
//部分或全部发送成功
return sent;
}
//一个字节都未发送成功
return -1;
}
void BufferSendMMsg::reOffset(size_t n) {
for (auto it = _hdrvec.begin(); it != _hdrvec.end();) {
auto &hdr = *it;
auto &io = *(hdr.msg_hdr.msg_iov);
assert(hdr.msg_len <= io.iov_len);
_remain_size -= hdr.msg_len;
if (hdr.msg_len == io.iov_len) {
//这个udp包全部发送成功
it = _hdrvec.erase(it);
sendFrontSuccess();
continue;
}
//部分发送成功
io.iov_base = (char *)io.iov_base + hdr.msg_len;
io.iov_len -= hdr.msg_len;
break;
}
}
BufferSendMMsg::BufferSendMMsg(List<std::pair<Buffer::Ptr, bool>> list, SendResult cb)
: BufferCallBack(std::move(list), std::move(cb))
, _iovec(_pkt_list.size())
, _hdrvec(_pkt_list.size()) {
auto i = 0U;
_pkt_list.for_each([&](std::pair<Buffer::Ptr, bool> &pr) {
auto &io = _iovec[i];
io.iov_base = pr.first->data();
io.iov_len = pr.first->size();
_remain_size += io.iov_len;
auto ptr = getBufferSockPtr(pr);
auto &mmsg = _hdrvec[i];
auto &msg = mmsg.msg_hdr;
mmsg.msg_len = 0;
msg.msg_name = ptr ? (void *)ptr->sockaddr() : nullptr;
msg.msg_namelen = ptr ? ptr->socklen() : 0;
msg.msg_iov = &io;
msg.msg_iovlen = 1;
msg.msg_control = nullptr;
msg.msg_controllen = 0;
msg.msg_flags = 0;
++i;
});
}
#endif //defined(__linux__) || defined(__linux)
BufferList::Ptr BufferList::create(List<std::pair<Buffer::Ptr, bool> > list, SendResult cb, bool is_udp) {
#if defined(_WIN32)
if (is_udp) {
// sendto/send 方案,待优化
return std::make_shared<BufferSendTo>(std::move(list), std::move(cb), is_udp);
}
// WSASend方案
return std::make_shared<BufferSendMsg>(std::move(list), std::move(cb));
#elif defined(__linux__) || defined(__linux)
if (is_udp) {
// sendmmsg方案
return std::make_shared<BufferSendMMsg>(std::move(list), std::move(cb));
}
// sendmsg方案
return std::make_shared<BufferSendMsg>(std::move(list), std::move(cb));
#else
if (is_udp) {
// sendto/send 方案, 可优化?
return std::make_shared<BufferSendTo>(std::move(list), std::move(cb), is_udp);
}
// sendmsg方案
return std::make_shared<BufferSendMsg>(std::move(list), std::move(cb));
#endif
}
#if defined(__linux) || defined(__linux__)
class SocketRecvmmsgBuffer : public SocketRecvBuffer {
public:
SocketRecvmmsgBuffer(size_t count, size_t size)
: _size(size)
, _iovec(count)
, _mmsgs(count)
, _buffers(count)
, _address(count) {
for (auto i = 0u; i < count; ++i) {
auto buf = BufferRaw::create();
buf->setCapacity(size);
_buffers[i] = buf;
auto &mmsg = _mmsgs[i];
auto &addr = _address[i];
mmsg.msg_len = 0;
mmsg.msg_hdr.msg_name = &addr;
mmsg.msg_hdr.msg_namelen = sizeof(addr);
mmsg.msg_hdr.msg_iov = &_iovec[i];
mmsg.msg_hdr.msg_iov->iov_base = buf->data();
mmsg.msg_hdr.msg_iov->iov_len = buf->getCapacity() - 1;
mmsg.msg_hdr.msg_iovlen = 1;
mmsg.msg_hdr.msg_control = nullptr;
mmsg.msg_hdr.msg_controllen = 0;
mmsg.msg_hdr.msg_flags = 0;
}
}
ssize_t recvFromSocket(int fd, ssize_t &count) override {
for (auto i = 0; i < _last_count; ++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);
} while (-1 == count && UV_EINTR == get_uv_error(true));
_last_count = count;
if (count <= 0) {
return count;
}
ssize_t nread = 0;
for (auto i = 0; i < count; ++i) {
auto &mmsg = _mmsgs[i];
nread += mmsg.msg_len;
auto buf = std::static_pointer_cast<BufferRaw>(_buffers[i]);
buf->setSize(mmsg.msg_len);
buf->data()[mmsg.msg_len] = '\0';
}
return nread;
}
Buffer::Ptr &getBuffer(size_t index) override { return _buffers[index]; }
struct sockaddr_storage &getAddress(size_t index) override { return _address[index]; }
private:
size_t _size;
ssize_t _last_count { 0 };
std::vector<struct iovec> _iovec;
std::vector<struct mmsghdr> _mmsgs;
std::vector<Buffer::Ptr> _buffers;
std::vector<struct sockaddr_storage> _address;
};
#endif
class SocketRecvFromBuffer : public SocketRecvBuffer {
public:
SocketRecvFromBuffer(size_t size): _size(size) {}
ssize_t recvFromSocket(int fd, ssize_t &count) override {
ssize_t nread;
socklen_t len = sizeof(_address);
if (!_buffer) {
allocBuffer();
}
do {
nread = recvfrom(fd, _buffer->data(), _buffer->getCapacity() - 1, 0, (struct sockaddr *)&_address, &len);
} while (-1 == nread && UV_EINTR == get_uv_error(true));
if (nread > 0) {
count = 1;
_buffer->data()[nread] = '\0';
std::static_pointer_cast<BufferRaw>(_buffer)->setSize(nread);
}
return nread;
}
Buffer::Ptr &getBuffer(size_t index) override { return _buffer; }
struct sockaddr_storage &getAddress(size_t index) override { return _address; }
private:
void allocBuffer() {
auto buf = BufferRaw::create();
buf->setCapacity(_size);
_buffer = std::move(buf);
}
private:
size_t _size;
Buffer::Ptr _buffer;
struct sockaddr_storage _address;
};
static constexpr auto kPacketCount = 32;
static constexpr auto kBufferCapacity = 4 * 1024u;
SocketRecvBuffer::Ptr SocketRecvBuffer::create(bool is_udp) {
#if defined(__linux) || defined(__linux__)
if (is_udp) {
return std::make_shared<SocketRecvmmsgBuffer>(kPacketCount, kBufferCapacity);
}
#endif
return std::make_shared<SocketRecvFromBuffer>(kPacketCount * kBufferCapacity);
}
} //toolkit