Commit be83edaf authored by Konstantin Tsoy's avatar Konstantin Tsoy Committed by Facebook Github Bot

Add recvmmsg

Summary: Add recvmmsg

Differential Revision: D18927890

fbshipit-source-id: de500dd1c266f21fbfc39c2972df787da50c7a80
parent 3f2526d6
......@@ -455,6 +455,14 @@ ssize_t AsyncUDPSocket::recvmsg(struct msghdr* msg, int flags) {
return netops::recvmsg(fd_, msg, flags);
}
int AsyncUDPSocket::recvmmsg(
struct mmsghdr* msgvec,
unsigned int vlen,
unsigned int flags,
struct timespec* timeout) {
return netops::recvmmsg(fd_, msgvec, vlen, flags, timeout);
}
void AsyncUDPSocket::resumeRead(ReadCallback* cob) {
CHECK(!readCallback_) << "Another read callback already installed";
CHECK_NE(NetworkSocket(), fd_)
......
......@@ -197,6 +197,12 @@ class AsyncUDPSocket : public EventHandler {
virtual ssize_t recvmsg(struct msghdr* msg, int flags);
virtual int recvmmsg(
struct mmsghdr* msgvec,
unsigned int vlen,
unsigned int flags,
struct timespec* timeout);
/**
* Start reading datagrams
*/
......
......@@ -176,7 +176,8 @@ class UDPClient : private AsyncUDPSocket::ReadCallback, private AsyncTimeout {
explicit UDPClient(EventBase* evb) : AsyncTimeout(evb), evb_(evb) {}
void start(const folly::SocketAddress& server, int n) {
void
start(const folly::SocketAddress& server, int n, bool sendClustered = false) {
CHECK(evb_->isInEventBaseThread());
server_ = server;
socket_ = std::make_unique<AsyncUDPSocket>(evb_);
......@@ -196,7 +197,11 @@ class UDPClient : private AsyncUDPSocket::ReadCallback, private AsyncTimeout {
n_ = n;
// Start playing ping pong
sendPing();
if (sendClustered) {
sendPingsClustered();
} else {
sendPing();
}
}
void connect() {
......@@ -227,6 +232,14 @@ class UDPClient : private AsyncUDPSocket::ReadCallback, private AsyncTimeout {
writePing(folly::IOBuf::copyBuffer(folly::to<std::string>("PING ", n_)));
}
void sendPingsClustered() {
scheduleTimeout(5);
while (n_ > 0) {
--n_;
writePing(folly::IOBuf::copyBuffer(folly::to<std::string>("PING ", n_)));
}
}
virtual void writePing(std::unique_ptr<folly::IOBuf> buf) {
auto ret = socket_->write(server_, std::move(buf));
if (ret == -1) {
......@@ -284,6 +297,10 @@ class UDPClient : private AsyncUDPSocket::ReadCallback, private AsyncTimeout {
return error_;
}
void incrementPongCount(int n) {
pongRecvd_ += n;
}
protected:
folly::Optional<folly::SocketAddress> connectAddr_;
EventBase* const evb_{nullptr};
......@@ -303,15 +320,17 @@ class UDPNotifyClient : public UDPClient {
public:
~UDPNotifyClient() override = default;
explicit UDPNotifyClient(EventBase* evb) : UDPClient(evb) {}
explicit UDPNotifyClient(
EventBase* evb,
bool useRecvmmsg = false,
unsigned int numMsgs = 1)
: UDPClient(evb), useRecvmmsg_(useRecvmmsg), numMsgs_(numMsgs) {}
bool shouldOnlyNotify() override {
return true;
}
void onNotifyDataAvailable() noexcept override {
notifyInvoked = true;
void onRecvMsg() {
struct msghdr msg;
memset(&msg, 0, sizeof(msg));
......@@ -347,7 +366,72 @@ class UDPNotifyClient : public UDPClient {
onDataAvailable(addr, size_t(read), false);
}
void onRecvMmsg() {
std::vector<struct mmsghdr> msgs;
msgs.reserve(numMsgs_);
memset(msgs.data(), 0, sizeof(struct mmsghdr) * numMsgs_);
const socklen_t addrLen = sizeof(struct sockaddr_storage);
const size_t dataSize = 1024;
std::vector<char> buf;
buf.reserve(numMsgs_ * dataSize);
memset(buf.data(), 0, numMsgs_ * dataSize);
std::vector<struct sockaddr_storage> addrs;
addrs.reserve(numMsgs_);
memset(addrs.data(), 0, sizeof(struct sockaddr_storage) * numMsgs_);
std::vector<struct iovec> iovecs;
iovecs.reserve(numMsgs_);
memset(iovecs.data(), 0, sizeof(struct iovec) * numMsgs_);
for (unsigned int i = 0; i < numMsgs_; ++i) {
struct msghdr* msg = &msgs[i].msg_hdr;
auto rawAddr = reinterpret_cast<sockaddr*>(&addrs[i]);
rawAddr->sa_family = socket_->address().getFamily();
iovecs[i].iov_base = &buf[i * dataSize];
iovecs[i].iov_len = dataSize;
msg->msg_name = rawAddr;
msg->msg_namelen = addrLen;
msg->msg_iov = &iovecs[i];
msg->msg_iovlen = 1;
}
int ret = socket_->recvmmsg(
msgs.data(), numMsgs_, 0x10000 /* MSG_WAITFORONE */, nullptr);
if (ret < 0) {
if (errno != EAGAIN || errno != EWOULDBLOCK) {
onReadError(folly::AsyncSocketException(
folly::AsyncSocketException::NETWORK_ERROR, "error"));
}
return;
}
incrementPongCount(ret);
if (pongRecvd() == (int)numMsgs_) {
shutdown();
} else {
onRecvMmsg();
}
}
void onNotifyDataAvailable() noexcept override {
notifyInvoked = true;
if (useRecvmmsg_) {
onRecvMmsg();
} else {
onRecvMsg();
}
}
bool notifyInvoked{false};
bool useRecvmmsg_{false};
unsigned int numMsgs_{1};
};
class AsyncSocketIntegrationTest : public Test {
......@@ -389,6 +473,11 @@ class AsyncSocketIntegrationTest : public Test {
folly::SocketAddress writeAddress,
folly::Optional<folly::SocketAddress> connectedAddress);
std::unique_ptr<UDPNotifyClient> performPingPongNotifyMmsgTest(
folly::SocketAddress writeAddress,
unsigned int numMsgs,
folly::Optional<folly::SocketAddress> connectedAddress);
folly::EventBase sevb;
folly::EventBase cevb;
std::unique_ptr<std::thread> serverThread;
......@@ -438,6 +527,30 @@ AsyncSocketIntegrationTest::performPingPongNotifyTest(
return client;
}
std::unique_ptr<UDPNotifyClient>
AsyncSocketIntegrationTest::performPingPongNotifyMmsgTest(
folly::SocketAddress writeAddress,
unsigned int numMsgs,
folly::Optional<folly::SocketAddress> connectedAddress) {
auto client = std::make_unique<UDPNotifyClient>(&cevb, true, numMsgs);
if (connectedAddress) {
client->setShouldConnect(*connectedAddress);
}
// Start event loop in a separate thread
auto clientThread = std::thread([this]() { cevb.loopForever(); });
// Wait for event loop to start
cevb.waitUntilRunning();
// Send ping
cevb.runInEventBaseThread(
[&]() { client->start(writeAddress, numMsgs, true); });
// Wait for client to finish
clientThread.join();
return client;
}
TEST_F(AsyncSocketIntegrationTest, PingPong) {
startServer();
auto pingClient = performPingPongTest(server->address(), folly::none);
......@@ -453,6 +566,15 @@ TEST_F(AsyncSocketIntegrationTest, PingPongNotify) {
ASSERT_TRUE(pingClient->notifyInvoked);
}
TEST_F(AsyncSocketIntegrationTest, PingPongNotifyMmsg) {
startServer();
auto pingClient =
performPingPongNotifyMmsgTest(server->address(), 10, folly::none);
// This should succeed.
ASSERT_EQ(pingClient->pongRecvd(), 10);
ASSERT_TRUE(pingClient->notifyInvoked);
}
TEST_F(AsyncSocketIntegrationTest, ConnectedPingPong) {
server->setChangePortForWrites(false);
startServer();
......
......@@ -316,6 +316,32 @@ ssize_t recvmsg(NetworkSocket s, msghdr* message, int flags) {
#endif
}
int recvmmsg(
NetworkSocket s,
mmsghdr* msgvec,
unsigned int vlen,
unsigned int flags,
timespec* timeout) {
#if FOLLY_HAVE_RECVMMSG
return wrapSocketFunction<int>(::recvmmsg, s, msgvec, vlen, flags, timeout);
#else
// implement via recvmsg
for (unsigned int i = 0; i < vlen; i++) {
ssize_t ret = recvmsg(s, &msgvec[i].msg_hdr, flags);
// in case of an error
// we return the number of msgs received if > 0
// or an error if no msg was sent
if (ret < 0) {
if (i) {
return static_cast<int>(i);
}
return static_cast<int>(ret);
}
}
return static_cast<int>(vlen);
#endif
}
ssize_t send(NetworkSocket s, const void* buf, size_t len, int flags) {
#ifdef _WIN32
return wrapSocketFunction<ssize_t>(
......
......@@ -22,6 +22,7 @@
#include <folly/net/NetworkSocket.h>
#include <folly/portability/IOVec.h>
#include <folly/portability/SysTypes.h>
#include <folly/portability/Time.h>
#include <folly/portability/Windows.h>
#ifndef _WIN32
......@@ -173,6 +174,12 @@ ssize_t recvfrom(
sockaddr* from,
socklen_t* fromlen);
ssize_t recvmsg(NetworkSocket s, msghdr* message, int flags);
int recvmmsg(
NetworkSocket s,
mmsghdr* msgvec,
unsigned int vlen,
unsigned int flags,
timespec* timeout);
ssize_t send(NetworkSocket s, const void* buf, size_t len, int flags);
ssize_t sendto(
NetworkSocket s,
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment