folly/io/async/test/ZeroCopy.h

   1 /*
   2  * Copyright 2017-present Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #pragma once
  18
  19 #include <folly/ExceptionWrapper.h>
  20 #include <folly/SocketAddress.h>
  21 #include <folly/io/IOBufQueue.h>
  22 #include <folly/io/async/AsyncServerSocket.h>
  23 #include <folly/io/async/AsyncSocket.h>
  24 #include <folly/io/async/EventBase.h>
  25
  26 namespace folly {
  27
  28 class ZeroCopyTestAsyncSocket {
  29  public:
  30   explicit ZeroCopyTestAsyncSocket(
  31       folly::EventBase* evb,
  32       int numLoops,
  33       size_t bufferSize,
  34       bool zeroCopy)
  35       : evb_(evb),
  36         numLoops_(numLoops),
  37         sock_(new folly::AsyncSocket(evb)),
  38         callback_(this),
  39         client_(true) {
  40     setBufferSize(bufferSize);
  41     setZeroCopy(zeroCopy);
  42   }
  43
  44   explicit ZeroCopyTestAsyncSocket(
  45       folly::EventBase* evb,
  46       int fd,
  47       int numLoops,
  48       size_t bufferSize,
  49       bool zeroCopy)
  50       : evb_(evb),
  51         numLoops_(numLoops),
  52         sock_(new folly::AsyncSocket(evb, fd)),
  53         callback_(this),
  54         client_(false) {
  55     setBufferSize(bufferSize);
  56     setZeroCopy(zeroCopy);
  57     // enable reads
  58     if (sock_) {
  59       sock_->setReadCB(&callback_);
  60     }
  61   }
  62
  63   ~ZeroCopyTestAsyncSocket() {
  64     clearBuffers();
  65   }
  66
  67   void connect(const folly::SocketAddress& remote) {
  68     if (sock_) {
  69       sock_->connect(&callback_, remote);
  70     }
  71   }
  72
  73   bool isZeroCopyWriteInProgress() const {
  74     return sock_->isZeroCopyWriteInProgress();
  75   }
  76
  77  private:
  78   void setZeroCopy(bool enable) {
  79     zeroCopy_ = enable;
  80     if (sock_) {
  81       sock_->setZeroCopy(zeroCopy_);
  82     }
  83   }
  84
  85   void setBufferSize(size_t bufferSize) {
  86     clearBuffers();
  87     bufferSize_ = bufferSize;
  88
  89     readBuffer_ = new char[bufferSize_];
  90   }
  91
  92   class Callback : public folly::AsyncSocket::ReadCallback,
  93                    public folly::AsyncSocket::ConnectCallback {
  94    public:
  95     explicit Callback(ZeroCopyTestAsyncSocket* parent) : parent_(parent) {}
  96
  97     void connectSuccess() noexcept override {
  98       parent_->sock_->setReadCB(this);
  99       parent_->onConnected();
 100     }
 101
 102     void connectErr(const folly::AsyncSocketException& ex) noexcept override {
 103       LOG(ERROR) << "Connect error: " << ex.what();
 104       parent_->onDataFinish(folly::exception_wrapper(ex));
 105     }
 106
 107     void getReadBuffer(void** bufReturn, size_t* lenReturn) override {
 108       parent_->getReadBuffer(bufReturn, lenReturn);
 109     }
 110
 111     void readDataAvailable(size_t len) noexcept override {
 112       parent_->readDataAvailable(len);
 113     }
 114
 115     void readEOF() noexcept override {
 116       parent_->onDataFinish(folly::exception_wrapper());
 117     }
 118
 119     void readErr(const folly::AsyncSocketException& ex) noexcept override {
 120       parent_->onDataFinish(folly::exception_wrapper(ex));
 121     }
 122
 123    private:
 124     ZeroCopyTestAsyncSocket* parent_{nullptr};
 125   };
 126
 127   void clearBuffers() {
 128     if (readBuffer_) {
 129       delete[] readBuffer_;
 130     }
 131   }
 132
 133   void getReadBuffer(void** bufReturn, size_t* lenReturn) {
 134     *bufReturn = readBuffer_ + readOffset_;
 135     *lenReturn = bufferSize_ - readOffset_;
 136   }
 137
 138   void readDataAvailable(size_t len) noexcept {
 139     readOffset_ += len;
 140     if (readOffset_ == bufferSize_) {
 141       readOffset_ = 0;
 142       onDataReady();
 143     }
 144   }
 145
 146   void onConnected() {
 147     setZeroCopy(zeroCopy_);
 148     writeBuffer();
 149   }
 150
 151   void onDataReady() {
 152     currLoop_++;
 153     if (client_ && currLoop_ >= numLoops_) {
 154       evb_->runInLoop(
 155           [this] { evb_->terminateLoopSoon(); }, false /*thisIteration*/);
 156       return;
 157     }
 158     writeBuffer();
 159   }
 160
 161   void onDataFinish(folly::exception_wrapper) {
 162     if (client_) {
 163       evb_->terminateLoopSoon();
 164     }
 165   }
 166
 167   bool writeBuffer() {
 168     // use calloc to make sure the memory is touched
 169     // if the memory is just malloc'd, running the zeroCopyOn
 170     // and the zeroCopyOff back to back on a system that does not support
 171     // zerocopy leads to the second test being much slower
 172     writeBuffer_ =
 173         folly::IOBuf::takeOwnership(::calloc(1, bufferSize_), bufferSize_);
 174
 175     if (sock_ && writeBuffer_) {
 176       sock_->writeChain(
 177           nullptr,
 178           std::move(writeBuffer_),
 179           zeroCopy_ ? WriteFlags::WRITE_MSG_ZEROCOPY : WriteFlags::NONE);
 180     }
 181
 182     return true;
 183   }
 184
 185   folly::EventBase* evb_;
 186   int numLoops_{0};
 187   int currLoop_{0};
 188   bool zeroCopy_{false};
 189
 190   folly::AsyncSocket::UniquePtr sock_;
 191   Callback callback_;
 192
 193   size_t bufferSize_{0};
 194   size_t readOffset_{0};
 195   char* readBuffer_{nullptr};
 196   std::unique_ptr<folly::IOBuf> writeBuffer_;
 197
 198   bool client_;
 199 };
 200
 201 class ZeroCopyTestServer : public folly::AsyncServerSocket::AcceptCallback {
 202  public:
 203   explicit ZeroCopyTestServer(
 204       folly::EventBase* evb,
 205       int numLoops,
 206       size_t bufferSize,
 207       bool zeroCopy)
 208       : evb_(evb),
 209         numLoops_(numLoops),
 210         bufferSize_(bufferSize),
 211         zeroCopy_(zeroCopy) {}
 212
 213   void addCallbackToServerSocket(folly::AsyncServerSocket& sock) {
 214     sock.addAcceptCallback(this, evb_);
 215   }
 216
 217   void connectionAccepted(
 218       int fd,
 219       const folly::SocketAddress& /* unused */) noexcept override {
 220     auto client = std::make_shared<ZeroCopyTestAsyncSocket>(
 221         evb_, fd, numLoops_, bufferSize_, zeroCopy_);
 222     clients_[client.get()] = client;
 223   }
 224
 225   void acceptError(const std::exception&) noexcept override {}
 226
 227  private:
 228   folly::EventBase* evb_;
 229   int numLoops_;
 230   size_t bufferSize_;
 231   bool zeroCopy_;
 232   std::unique_ptr<ZeroCopyTestAsyncSocket> client_;
 233   std::unordered_map<
 234       ZeroCopyTestAsyncSocket*,
 235       std::shared_ptr<ZeroCopyTestAsyncSocket>>
 236       clients_;
 237 };
 238
 239 class ZeroCopyTest {
 240  public:
 241   explicit ZeroCopyTest(int numLoops, bool zeroCopy, size_t bufferSize);
 242   bool run();
 243
 244  private:
 245   void connectOne() {
 246     SocketAddress addr = listenSock_->getAddress();
 247     client_->connect(addr);
 248   }
 249
 250   int numLoops_;
 251   bool zeroCopy_;
 252   size_t bufferSize_;
 253
 254   EventBase evb_;
 255   std::unique_ptr<ZeroCopyTestAsyncSocket> client_;
 256   folly::AsyncServerSocket::UniquePtr listenSock_;
 257   ZeroCopyTestServer server_;
 258 };
 259
 260 } // namespace folly