2 * Copyright 2017-present Facebook, Inc.
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
8 * http://www.apache.org/licenses/LICENSE-2.0
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.
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>
28 class ZeroCopyTestAsyncSocket {
30 explicit ZeroCopyTestAsyncSocket(
31 folly::EventBase* evb,
37 sock_(new folly::AsyncSocket(evb)),
40 setBufferSize(bufferSize);
41 setZeroCopy(zeroCopy);
44 explicit ZeroCopyTestAsyncSocket(
45 folly::EventBase* evb,
52 sock_(new folly::AsyncSocket(evb, fd)),
55 setBufferSize(bufferSize);
56 setZeroCopy(zeroCopy);
59 sock_->setReadCB(&callback_);
63 ~ZeroCopyTestAsyncSocket() {
67 void connect(const folly::SocketAddress& remote) {
69 sock_->connect(&callback_, remote);
73 bool isZeroCopyWriteInProgress() const {
74 return sock_->isZeroCopyWriteInProgress();
78 void setZeroCopy(bool enable) {
81 sock_->setZeroCopy(zeroCopy_);
85 void setBufferSize(size_t bufferSize) {
87 bufferSize_ = bufferSize;
89 readBuffer_ = new char[bufferSize_];
92 class Callback : public folly::AsyncSocket::ReadCallback,
93 public folly::AsyncSocket::ConnectCallback {
95 explicit Callback(ZeroCopyTestAsyncSocket* parent) : parent_(parent) {}
97 void connectSuccess() noexcept override {
98 parent_->sock_->setReadCB(this);
99 parent_->onConnected();
102 void connectErr(const folly::AsyncSocketException& ex) noexcept override {
103 LOG(ERROR) << "Connect error: " << ex.what();
104 parent_->onDataFinish(folly::exception_wrapper(ex));
107 void getReadBuffer(void** bufReturn, size_t* lenReturn) override {
108 parent_->getReadBuffer(bufReturn, lenReturn);
111 void readDataAvailable(size_t len) noexcept override {
112 parent_->readDataAvailable(len);
115 void readEOF() noexcept override {
116 parent_->onDataFinish(folly::exception_wrapper());
119 void readErr(const folly::AsyncSocketException& ex) noexcept override {
120 parent_->onDataFinish(folly::exception_wrapper(ex));
124 ZeroCopyTestAsyncSocket* parent_{nullptr};
127 void clearBuffers() {
129 delete[] readBuffer_;
133 void getReadBuffer(void** bufReturn, size_t* lenReturn) {
134 *bufReturn = readBuffer_ + readOffset_;
135 *lenReturn = bufferSize_ - readOffset_;
138 void readDataAvailable(size_t len) noexcept {
140 if (readOffset_ == bufferSize_) {
147 setZeroCopy(zeroCopy_);
153 if (client_ && currLoop_ >= numLoops_) {
155 [this] { evb_->terminateLoopSoon(); }, false /*thisIteration*/);
161 void onDataFinish(folly::exception_wrapper) {
163 evb_->terminateLoopSoon();
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
173 folly::IOBuf::takeOwnership(::calloc(1, bufferSize_), bufferSize_);
175 if (sock_ && writeBuffer_) {
178 std::move(writeBuffer_),
179 zeroCopy_ ? WriteFlags::WRITE_MSG_ZEROCOPY : WriteFlags::NONE);
185 folly::EventBase* evb_;
188 bool zeroCopy_{false};
190 folly::AsyncSocket::UniquePtr sock_;
193 size_t bufferSize_{0};
194 size_t readOffset_{0};
195 char* readBuffer_{nullptr};
196 std::unique_ptr<folly::IOBuf> writeBuffer_;
201 class ZeroCopyTestServer : public folly::AsyncServerSocket::AcceptCallback {
203 explicit ZeroCopyTestServer(
204 folly::EventBase* evb,
210 bufferSize_(bufferSize),
211 zeroCopy_(zeroCopy) {}
213 void addCallbackToServerSocket(folly::AsyncServerSocket& sock) {
214 sock.addAcceptCallback(this, evb_);
217 void connectionAccepted(
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;
225 void acceptError(const std::exception&) noexcept override {}
228 folly::EventBase* evb_;
232 std::unique_ptr<ZeroCopyTestAsyncSocket> client_;
234 ZeroCopyTestAsyncSocket*,
235 std::shared_ptr<ZeroCopyTestAsyncSocket>>
241 explicit ZeroCopyTest(int numLoops, bool zeroCopy, size_t bufferSize);
246 SocketAddress addr = listenSock_->getAddress();
247 client_->connect(addr);
255 std::unique_ptr<ZeroCopyTestAsyncSocket> client_;
256 folly::AsyncServerSocket::UniquePtr listenSock_;
257 ZeroCopyTestServer server_;