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

Switch getReadBuffers to be IOBufVecQueue::IoVecVec based

Summary:
Switch getReadBuffers to be IOBufVecQueue::IoVecVec based

(Note: this ignores all push blocking failures!)

Reviewed By: simpkins

Differential Revision: D28467957

fbshipit-source-id: 637bf07d516120cb2bab89c3cbc28f58b700748a
parent 2b7d6e1d
...@@ -76,7 +76,7 @@ class IOBufIovecBuilder { ...@@ -76,7 +76,7 @@ class IOBufIovecBuilder {
// This is safe, because it means this is the last reference // This is safe, because it means this is the last reference
// Anything trying to copy it is already undefined behavior. // Anything trying to copy it is already undefined behavior.
if (refcount_.load(std::memory_order_acquire) > 1) { if (refcount_.load(std::memory_order_acquire) > 1) {
uint32_t newcnt = refcount_.fetch_sub(1, std::memory_order_acq_rel) - 1; size_t newcnt = refcount_.fetch_sub(1, std::memory_order_acq_rel) - 1;
if (newcnt > 0) { if (newcnt > 0) {
return; return;
} }
......
...@@ -2396,10 +2396,10 @@ void AsyncSocket::prepareReadBuffer(void** buf, size_t* buflen) { ...@@ -2396,10 +2396,10 @@ void AsyncSocket::prepareReadBuffer(void** buf, size_t* buflen) {
readCallback_->getReadBuffer(buf, buflen); readCallback_->getReadBuffer(buf, buflen);
} }
size_t AsyncSocket::prepareReadBuffers(struct iovec* iovs, size_t num) { void AsyncSocket::prepareReadBuffers(IOBufIovecBuilder::IoVecVec& iovs) {
// no matter what, buffers should be prepared for non-ssl socket // no matter what, buffers should be prepared for non-ssl socket
CHECK(readCallback_); CHECK(readCallback_);
return readCallback_->getReadBuffers(iovs, num); readCallback_->getReadBuffers(iovs);
} }
size_t AsyncSocket::handleErrMessages() noexcept { size_t AsyncSocket::handleErrMessages() noexcept {
...@@ -2608,12 +2608,12 @@ void AsyncSocket::handleRead() noexcept { ...@@ -2608,12 +2608,12 @@ void AsyncSocket::handleRead() noexcept {
// Get the buffer(s) to read into. // Get the buffer(s) to read into.
void* buf = nullptr; void* buf = nullptr;
size_t buflen = 0, offset = 0, num = 0; size_t buflen = 0, offset = 0, num = 0;
static constexpr size_t kNumIov = 16; IOBufIovecBuilder::IoVecVec iovs; // this can be an Asyncsocket member too
std::array<struct iovec, kNumIov> iovs;
try { try {
if (readMode == AsyncReader::ReadCallback::ReadMode::ReadVec) { if (readMode == AsyncReader::ReadCallback::ReadMode::ReadVec) {
num = prepareReadBuffers(iovs.data(), iovs.size()); prepareReadBuffers(iovs);
num = iovs.size();
VLOG(5) << "prepareReadBuffers() bufs=" << iovs.data() VLOG(5) << "prepareReadBuffers() bufs=" << iovs.data()
<< ", num=" << num; << ", num=" << num;
} else { } else {
......
...@@ -27,6 +27,7 @@ ...@@ -27,6 +27,7 @@
#include <folly/SocketAddress.h> #include <folly/SocketAddress.h>
#include <folly/detail/SocketFastOpen.h> #include <folly/detail/SocketFastOpen.h>
#include <folly/io/IOBuf.h> #include <folly/io/IOBuf.h>
#include <folly/io/IOBufIovecBuilder.h>
#include <folly/io/ShutdownSocketSet.h> #include <folly/io/ShutdownSocketSet.h>
#include <folly/io/SocketOptionMap.h> #include <folly/io/SocketOptionMap.h>
#include <folly/io/async/AsyncSocketException.h> #include <folly/io/async/AsyncSocketException.h>
...@@ -1292,7 +1293,7 @@ class AsyncSocket : public AsyncTransport { ...@@ -1292,7 +1293,7 @@ class AsyncSocket : public AsyncTransport {
virtual void checkForImmediateRead() noexcept; virtual void checkForImmediateRead() noexcept;
virtual void handleInitialReadWrite() noexcept; virtual void handleInitialReadWrite() noexcept;
virtual void prepareReadBuffer(void** buf, size_t* buflen); virtual void prepareReadBuffer(void** buf, size_t* buflen);
virtual size_t prepareReadBuffers(struct iovec* iovs, size_t num); virtual void prepareReadBuffers(IOBufIovecBuilder::IoVecVec& iovs);
virtual size_t handleErrMessages() noexcept; virtual size_t handleErrMessages() noexcept;
virtual void handleRead() noexcept; virtual void handleRead() noexcept;
virtual void handleWrite() noexcept; virtual void handleWrite() noexcept;
......
...@@ -21,6 +21,7 @@ ...@@ -21,6 +21,7 @@
#include <folly/Optional.h> #include <folly/Optional.h>
#include <folly/io/IOBuf.h> #include <folly/io/IOBuf.h>
#include <folly/io/IOBufIovecBuilder.h>
#include <folly/io/async/AsyncSocketBase.h> #include <folly/io/async/AsyncSocketBase.h>
#include <folly/io/async/AsyncTransportCertificate.h> #include <folly/io/async/AsyncTransportCertificate.h>
#include <folly/io/async/DelayedDestruction.h> #include <folly/io/async/DelayedDestruction.h>
...@@ -206,7 +207,7 @@ class AsyncReader { ...@@ -206,7 +207,7 @@ class AsyncReader {
* returned different buffers, the ReadCallback is responsible for ensuring * returned different buffers, the ReadCallback is responsible for ensuring
* that they are not leaked. * that they are not leaked.
* *
* If getReadBuffera() throws an exception or returns a zero length array * If getReadBuffers() throws an exception or returns a zero length array
* the ReadCallback will be uninstalled and its readError() method will be * the ReadCallback will be uninstalled and its readError() method will be
* invoked. * invoked.
* *
...@@ -215,14 +216,10 @@ class AsyncReader { ...@@ -215,14 +216,10 @@ class AsyncReader {
* set a different read callback.) * set a different read callback.)
* *
* @param iovs getReadBuffers() will copy up to num iovec entries into * @param iovs getReadBuffers() will copy up to num iovec entries into
* iovs. iovs cannot be nullptr unless num is 0 * iovs
* @param num number of iovec entries in the iovs array
* @return number of entried copied to the iovs array
* this is less than or equal to num
*/ */
virtual size_t getReadBuffers( virtual void getReadBuffers(IOBufIovecBuilder::IoVecVec& iovs) {
FOLLY_MAYBE_UNUSED struct iovec* iovs, FOLLY_MAYBE_UNUSED size_t num) { iovs.clear();
return 0;
} }
/** /**
......
...@@ -207,127 +207,11 @@ class ReadCallback : public folly::AsyncTransport::ReadCallback { ...@@ -207,127 +207,11 @@ class ReadCallback : public folly::AsyncTransport::ReadCallback {
}; };
class ReadvCallback : public folly::AsyncTransport::ReadCallback { class ReadvCallback : public folly::AsyncTransport::ReadCallback {
private:
class IOBufVecQueue {
private:
struct RefCountMem {
explicit RefCountMem(size_t size) {
mem_ = ::malloc(size);
len_ = size;
}
~RefCountMem() { ::free(mem_); }
void* usableMem() const {
return reinterpret_cast<uint8_t*>(mem_) + used_;
}
size_t usableSize() const { return len_ - used_; }
void incUsedMem(size_t len) { used_ += len; }
static void freeMem(void* buf, void* userData) {
std::ignore = buf;
reinterpret_cast<RefCountMem*>(userData)->decRef();
}
void addRef() { ++count_; }
void decRef() {
if (--count_ == 0) {
delete this;
}
}
private:
std::atomic<size_t> count_{1};
void* mem_{nullptr};
size_t len_{0};
size_t used_{0};
};
public:
struct Options {
static constexpr size_t kBlockSize = 16 * 1024;
size_t blockSize_{kBlockSize};
};
IOBufVecQueue() = default;
explicit IOBufVecQueue(const Options& options) : options_(options) {}
~IOBufVecQueue() {
for (auto& buf : buffers_) {
buf->decRef();
}
}
static Options getBlockSizeOptions(size_t blockSize) {
Options options;
options.blockSize_ = blockSize;
return options;
}
size_t preallocate(size_t len, struct iovec* iovs, size_t num) {
size_t total = 0;
size_t i = 0;
for (; (i < num) && (total < len); ++i) {
if (i >= buffers_.size()) {
buffers_.push_back(new RefCountMem(options_.blockSize_));
}
iovs[i].iov_base = buffers_[i]->usableMem();
iovs[i].iov_len = buffers_[i]->usableSize();
total += buffers_[i]->usableSize();
}
return i;
}
std::unique_ptr<folly::IOBuf> postallocate(size_t len) {
std::unique_ptr<folly::IOBuf> ret, tmp;
while (len > 0) {
CHECK(!buffers_.empty());
auto* buf = buffers_.front();
auto size = buf->usableSize();
if (len >= size) {
// no need to inc the ref count since we're transferring ownership
tmp = folly::IOBuf::takeOwnership(
buf->usableMem(), size, RefCountMem::freeMem, buf);
buffers_.pop_front();
len -= size;
} else {
buf->addRef();
tmp = folly::IOBuf::takeOwnership(
buf->usableMem(), len, RefCountMem::freeMem, buf);
buf->incUsedMem(len);
len = 0;
}
CHECK(!tmp->isShared());
if (ret) {
ret->prependChain(std::move(tmp));
} else {
ret = std::move(tmp);
}
}
return ret;
}
private:
Options options_;
std::deque<RefCountMem*> buffers_;
};
public: public:
ReadvCallback(size_t bufferSize, size_t len) ReadvCallback(size_t bufferSize, size_t len)
: state_(STATE_WAITING), : state_(STATE_WAITING),
exception_(folly::AsyncSocketException::UNKNOWN, "none"), exception_(folly::AsyncSocketException::UNKNOWN, "none"),
queue_(IOBufVecQueue::getBlockSizeOptions(bufferSize)), queue_(folly::IOBufIovecBuilder::Options().setBlockSize(bufferSize)),
len_(len) { len_(len) {
setReadMode(folly::AsyncTransport::ReadCallback::ReadMode::ReadVec); setReadMode(folly::AsyncTransport::ReadCallback::ReadMode::ReadVec);
} }
...@@ -341,12 +225,12 @@ class ReadvCallback : public folly::AsyncTransport::ReadCallback { ...@@ -341,12 +225,12 @@ class ReadvCallback : public folly::AsyncTransport::ReadCallback {
CHECK(false); // this should not be called CHECK(false); // this should not be called
} }
size_t getReadBuffers(struct iovec* iovs, size_t num) override { void getReadBuffers(folly::IOBufIovecBuilder::IoVecVec& iovs) override {
return queue_.preallocate(len_, iovs, num); queue_.allocateBuffers(iovs, len_);
} }
void readDataAvailable(size_t len) noexcept override { void readDataAvailable(size_t len) noexcept override {
auto tmp = queue_.postallocate(len); auto tmp = queue_.extractIOBufChain(len);
if (!buf_) { if (!buf_) {
buf_ = std::move(tmp); buf_ = std::move(tmp);
} else { } else {
...@@ -374,7 +258,7 @@ class ReadvCallback : public folly::AsyncTransport::ReadCallback { ...@@ -374,7 +258,7 @@ class ReadvCallback : public folly::AsyncTransport::ReadCallback {
private: private:
StateEnum state_; StateEnum state_;
folly::AsyncSocketException exception_; folly::AsyncSocketException exception_;
IOBufVecQueue queue_; folly::IOBufIovecBuilder queue_;
std::unique_ptr<folly::IOBuf> buf_; std::unique_ptr<folly::IOBuf> buf_;
const size_t len_; const size_t len_;
}; };
......
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