Commit a02ba06a authored by Dan Melnic's avatar Dan Melnic Committed by Facebook GitHub Bot

Add support for Linux TCP zerocopy RX

Summary:
Add support for Linux TCP zerocopy RX

(Note: this ignores all push blocking failures!)

Reviewed By: simpkins

Differential Revision: D32741177

fbshipit-source-id: 412981fdd4a27437d6f22a80dc7c5fc864335ddf
parent 886ed94b
This diff is collapsed.
......@@ -651,6 +651,31 @@ class AsyncSocket : public AsyncTransport {
}
}
/**
* Create a memory store to use for zero copy reads.
*
* The memory store contains a fixed number of entries, each with a fixed
* size. When data is read using zero-copy the kernel will place it in one
* of these entries, and it will be returned to the callback with
* readZeroCopyDataAvailable(). The callback must release the IOBuf
* reference to make the entry available again for future zero-copy reads.
* If all entries are exhausted the read code will fall back to non-zero-copy
* reads.
*
* Note: it is the caller's responsibility to ensure that they do not destroy
* the ZeroCopyMemStore while it still has any outstanding entries in use.
* The caller must ensure the ZeroCopyMemStore exists until all callers have
* finished using any data returned via zero-copy reads, and released the
* IOBuf objects containing that data.
*
* @param entries The number of entries to allocate in the memory store.
* @param size The size of each entry, in bytes. This should be a
* multiple of the kernel page size.
*/
static std::unique_ptr<AsyncReader::ReadCallback::ZeroCopyMemStore>
createDefaultZeroCopyMemStore(size_t entries, size_t size);
bool setZeroCopy(bool enable) override;
bool getZeroCopy() const override { return zeroCopyEnabled_; }
......@@ -1467,13 +1492,19 @@ class AsyncSocket : public AsyncTransport {
void doClose();
// error handling methods
enum class ReadCode {
READ_NOT_SUPPORTED = 0,
READ_CONTINUE = 1,
READ_DONE = 2,
};
void startFail();
void finishFail();
void finishFail(const AsyncSocketException& ex);
void invokeAllErrors(const AsyncSocketException& ex);
void fail(const char* fn, const AsyncSocketException& ex);
void failConnect(const char* fn, const AsyncSocketException& ex);
void failRead(const char* fn, const AsyncSocketException& ex);
ReadCode failRead(const char* fn, const AsyncSocketException& ex);
void failErrMessageRead(const char* fn, const AsyncSocketException& ex);
void failWrite(
const char* fn,
......@@ -1517,6 +1548,8 @@ class AsyncSocket : public AsyncTransport {
void releaseIOBuf(
std::unique_ptr<folly::IOBuf> buf, ReleaseIOBufCallback* callback);
ReadCode processZeroCopyRead();
ReadCode processNormalRead();
/**
* Attempt to enable Observer ByteEvents for this socket.
*
......@@ -1630,6 +1663,9 @@ class AsyncSocket : public AsyncTransport {
size_t zeroCopyReenableThreshold_{0};
size_t zeroCopyReenableCounter_{0};
// zerocopy read
bool zerocopyReadSupported_{true};
// subclasses may cache these on first call to get
mutable std::unique_ptr<const AsyncTransportCertificate> peerCertData_{
nullptr};
......
......@@ -238,6 +238,79 @@ class AsyncReader {
virtual void readDataAvailable(size_t len) noexcept = 0;
class ZeroCopyMemStore {
public:
struct Entry {
void* data{nullptr};
size_t len{0}; // in use
size_t capacity{0}; // capacity
ZeroCopyMemStore* store{nullptr};
void put() {
DCHECK(store);
store->put(this);
}
};
struct EntryDeleter {
void operator()(Entry* entry) { entry->put(); }
};
using EntryPtr = std::unique_ptr<Entry, EntryDeleter>;
virtual ~ZeroCopyMemStore() = default;
virtual EntryPtr get() = 0;
virtual void put(Entry*) = 0;
};
/* the next 4 methods can be used if the callback wants to support zerocopy
* RX on Linux as described in https://lwn.net/Articles/754681/ If the
* current kernel version does not support zerocopy RX, the callback will
* revert to regular recv processing
* In case we support zerocopy RX, the callback might be notified of buffer
* chains composed of mmap memory and also memory allocated via the
* getZeroCopyReadBuffer method
*/
/**
* Return a ZeroCopyMemStore to use if the callback would like to enable
* zero-copy reads. Return nullptr to disable zero-copy reads.
*
* The caller must ensure that the ZeroCopyMemStore remains valid for as
* long as this callback is installed and reading data, and until put()
* has been called for every outstanding Entry allocated with get().
*/
virtual ZeroCopyMemStore* readZeroCopyEnabled() noexcept { return nullptr; }
/**
* Get a buffer to read data into when using zero-copy reads if some data
* cannot be read using a zero-copy page.
*
* When data is available, some data may be returned in zero-copy pages,
* followed by some amount of data in this fallback buffer.
*/
virtual void getZeroCopyFallbackBuffer(
void** /*bufReturn*/, size_t* /*lenReturn*/) noexcept {
CHECK(false);
}
/**
* readZeroCopyDataAvailable() will be called when data is available from a
* zero-copy read.
*
* The data returned may be in two separate parts: data that was actually
* read using zero copy pages will be in zeroCopyData. Additionally, some
* number of bytes may have been placed in the fallback buffer returned by
* getZeroCopyFallbackBuffer(). additionalBytes indicates the number of
* bytes placed in getZeroCopyFallbackBuffer().
*/
virtual void readZeroCopyDataAvailable(
std::unique_ptr<IOBuf>&& /*zeroCopyData*/,
size_t /*additionalBytes*/) noexcept {
CHECK(false);
}
/**
* When data becomes available, isBufferMovable() will be invoked to figure
* out which API will be used, readBufferAvailable() or
......@@ -411,6 +484,17 @@ class AsyncWriter {
virtual bool getZeroCopy() const { return false; }
struct RXZerocopyParams {
bool enable{false};
size_t mapSize{0};
};
FOLLY_NODISCARD virtual bool setRXZeroCopy(RXZerocopyParams /*params*/) {
return false;
}
FOLLY_NODISCARD virtual bool getRXZeroCopy() const { return false; }
using ZeroCopyEnableFunc =
std::function<bool(const std::unique_ptr<folly::IOBuf>& buf)>;
......
......@@ -298,6 +298,147 @@ class BufferCallback : public folly::AsyncTransport::BufferCallback {
bool bufferCleared_{false};
};
class ZeroCopyReadCallback : public folly::AsyncTransport::ReadCallback {
public:
explicit ZeroCopyReadCallback(
folly::AsyncTransport::ReadCallback::ZeroCopyMemStore* memStore,
size_t _maxBufferSz = 4096)
: memStore_(memStore),
state(STATE_WAITING),
exception(folly::AsyncSocketException::UNKNOWN, "none"),
maxBufferSz(_maxBufferSz) {}
~ZeroCopyReadCallback() override { currentBuffer.free(); }
// zerocopy
folly::AsyncTransport::ReadCallback::ZeroCopyMemStore*
readZeroCopyEnabled() noexcept override {
return memStore_;
}
void getZeroCopyFallbackBuffer(
void** bufReturn, size_t* lenReturn) noexcept override {
if (!currentZeroCopyBuffer.buffer) {
currentZeroCopyBuffer.allocate(maxBufferSz);
}
*bufReturn = currentZeroCopyBuffer.buffer;
*lenReturn = currentZeroCopyBuffer.length;
}
void readZeroCopyDataAvailable(
std::unique_ptr<folly::IOBuf>&& zeroCopyData,
size_t additionalBytes) noexcept override {
auto ioBuf = std::move(zeroCopyData);
if (additionalBytes) {
auto tmp = folly::IOBuf::takeOwnership(
currentZeroCopyBuffer.buffer,
currentZeroCopyBuffer.length,
0,
additionalBytes);
currentZeroCopyBuffer.reset();
if (ioBuf) {
ioBuf->prependChain(std::move(tmp));
} else {
ioBuf = std::move(tmp);
}
}
if (!data_) {
data_ = std::move(ioBuf);
} else {
data_->prependChain(std::move(ioBuf));
}
}
void getReadBuffer(void** bufReturn, size_t* lenReturn) override {
if (!currentBuffer.buffer) {
currentBuffer.allocate(maxBufferSz);
}
*bufReturn = currentBuffer.buffer;
*lenReturn = currentBuffer.length;
}
void readDataAvailable(size_t len) noexcept override {
auto ioBuf = folly::IOBuf::takeOwnership(
currentBuffer.buffer, currentBuffer.length, 0, len);
currentBuffer.reset();
if (!data_) {
data_ = std::move(ioBuf);
} else {
data_->prependChain(std::move(ioBuf));
}
}
void readEOF() noexcept override { state = STATE_SUCCEEDED; }
void readErr(const folly::AsyncSocketException& ex) noexcept override {
state = STATE_FAILED;
exception = ex;
}
void verifyData(const std::string& expected) const {
verifyData((const unsigned char*)expected.data(), expected.size());
}
void verifyData(const unsigned char* expected, size_t expectedLen) const {
CHECK(!!data_);
auto len = data_->computeChainDataLength();
CHECK_EQ(len, expectedLen);
auto* buf = data_.get();
auto* current = buf;
size_t offset = 0;
do {
size_t cmpLen = std::min(current->length(), expectedLen - offset);
CHECK_EQ(cmpLen, current->length());
CHECK_EQ(memcmp(current->data(), expected + offset, cmpLen), 0);
offset += cmpLen;
current = current->next();
} while (current != buf);
std::ignore = expected;
CHECK_EQ(offset, expectedLen);
}
class Buffer {
public:
Buffer() = default;
Buffer(char* buf, size_t len) : buffer(buf), length(len) {}
~Buffer() {
if (buffer) {
::free(buffer);
}
}
void reset() {
buffer = nullptr;
length = 0;
}
void allocate(size_t len) {
CHECK(buffer == nullptr);
buffer = static_cast<char*>(malloc(len));
length = len;
}
void free() {
::free(buffer);
reset();
}
char* buffer{nullptr};
size_t length{0};
};
folly::AsyncTransport::ReadCallback::ZeroCopyMemStore* memStore_;
StateEnum state;
folly::AsyncSocketException exception;
Buffer currentBuffer, currentZeroCopyBuffer;
VoidCallback dataAvailableCallback;
const size_t maxBufferSz;
std::unique_ptr<folly::IOBuf> data_;
};
class ReadVerifier {};
class TestSendMsgParamsCallback
......
......@@ -631,6 +631,59 @@ TEST_P(AsyncSocketConnectTest, ConnectAndReadv) {
ASSERT_FALSE(socket->isClosedByPeer());
}
TEST_P(AsyncSocketConnectTest, ConnectAndZeroCopyRead) {
TestServer server;
// connect()
EventBase evb;
std::shared_ptr<AsyncSocket> socket = AsyncSocket::newSocket(&evb);
if (GetParam() == TFOState::ENABLED) {
socket->enableTFO();
}
ConnCallback ccb;
socket->connect(&ccb, server.getAddress(), 30);
static constexpr size_t kBuffSize = 4096;
static constexpr size_t kDataSize = 128 * 1024;
static constexpr size_t kNumEntries = 1024;
static constexpr size_t kEntrySize = 128 * 1024;
auto memStore =
AsyncSocket::createDefaultZeroCopyMemStore(kNumEntries, kEntrySize);
ZeroCopyReadCallback rcb(memStore.get(), kBuffSize);
socket->setReadCB(&rcb);
if (GetParam() == TFOState::ENABLED) {
// Trigger a connection
socket->writeChain(nullptr, IOBuf::copyBuffer("hey"));
}
// Even though we haven't looped yet, we should be able to accept
// the connection and send data to it.
std::shared_ptr<BlockingSocket> acceptedSocket = server.accept();
std::string data(kDataSize, ' ');
// generate random data
std::mt19937 rng(folly::randomNumberSeed());
for (size_t i = 0; i < data.size(); ++i) {
data[i] = static_cast<char>(rng());
}
acceptedSocket->write(
reinterpret_cast<unsigned char*>(data.data()), data.size());
acceptedSocket->flush();
acceptedSocket->close();
// Loop
evb.loop();
ASSERT_EQ(ccb.state, STATE_SUCCEEDED);
rcb.verifyData(data);
ASSERT_FALSE(socket->isClosedBySelf());
ASSERT_FALSE(socket->isClosedByPeer());
}
/**
* Test installing a read callback and then closing immediately before the
* connect attempt finishes.
......
......@@ -40,6 +40,7 @@ using sa_family_t = ADDRESS_FAMILY;
#define SOL_UDP 0x0
#define UDP_SEGMENT 0x0
#define IP_BIND_ADDRESS_NO_PORT 0
#define TCP_ZEROCOPY_RECEIVE 0
// We don't actually support either of these flags
// currently.
......@@ -139,6 +140,11 @@ struct mmsghdr {
#define FOLLY_HAVE_MSG_ERRQUEUE 1
#ifndef FOLLY_HAVE_SO_TIMESTAMPING
#define FOLLY_HAVE_SO_TIMESTAMPING 1
#ifndef TCP_ZEROCOPY_RECEIVE
#define TCP_ZEROCOPY_RECEIVE 35
#endif
#else
#define TCP_ZEROCOPY_RECEIVE 0
#endif
/* for struct sock_extended_err*/
#include <linux/errqueue.h>
......@@ -207,6 +213,24 @@ struct sock_txtime {
__kernel_clockid_t clockid; /* reference clockid */
__u32 flags; /* as defined by enum txtime_flags */
};
/* Copied from uapi/linux/tcp.h */
/* setsockopt(fd, IPPROTO_TCP, TCP_ZEROCOPY_RECEIVE, ...) */
struct tcp_zerocopy_receive {
__u64 address; /* in: address of mapping */
__u32 length; /* in/out: number of bytes to map/mapped */
__u32 recv_skip_hint; /* out: amount of bytes to skip */
__u32 inq; /* out: amount of bytes in read queue */
__s32 err; /* out: socket error */
__u64 copybuf_address; /* in: copybuf address (small reads) */
__s32 copybuf_len; /* in/out: copybuf bytes avail/used or error */
__u32 flags; /* in: flags */
__u64 msg_control; /* ancillary data */
__u64 msg_controllen;
__u32 msg_flags;
__u32 reserved; /* set to 0 for now */
};
} // namespace netops
} // namespace folly
#endif
......
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