folly/io/async/test/AsyncPipeTest.cpp

   1 /*
   2  * Copyright 2014-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 #include <folly/Memory.h>
  18 #include <folly/io/async/AsyncPipe.h>
  19 #include <folly/io/async/EventBase.h>
  20 #include <folly/portability/GTest.h>
  21
  22 #include <fcntl.h>
  23
  24 using namespace testing;
  25
  26 namespace {
  27
  28 class TestReadCallback : public folly::AsyncReader::ReadCallback {
  29  public:
  30   bool isBufferMovable() noexcept override {
  31     return movable_;
  32   }
  33   void setMovable(bool movable) {
  34     movable_ = movable;
  35   }
  36
  37   void readBufferAvailable(
  38       std::unique_ptr<folly::IOBuf> readBuf) noexcept override {
  39     readBuffer_.append(std::move(readBuf));
  40   }
  41
  42   void readDataAvailable(size_t len) noexcept override {
  43     readBuffer_.postallocate(len);
  44   }
  45
  46   void getReadBuffer(void** bufReturn, size_t* lenReturn) noexcept override {
  47     auto res = readBuffer_.preallocate(4000, 65000);
  48     *bufReturn = res.first;
  49     *lenReturn = res.second;
  50   }
  51
  52   void readEOF() noexcept override {}
  53
  54   void readErr(const folly::AsyncSocketException&) noexcept override {
  55     error_ = true;
  56   }
  57
  58   std::string getData() {
  59     auto buf = readBuffer_.move();
  60     buf->coalesce();
  61     return std::string((char *)buf->data(), buf->length());
  62   }
  63
  64   void reset() {
  65     movable_ = false;
  66     error_ = false;
  67     readBuffer_.clear();
  68   }
  69
  70   folly::IOBufQueue readBuffer_{folly::IOBufQueue::cacheChainLength()};
  71   bool error_{false};
  72   bool movable_{false};
  73 };
  74
  75 class TestWriteCallback : public folly::AsyncWriter::WriteCallback {
  76  public:
  77   void writeSuccess() noexcept override { writes_++; }
  78
  79   void writeErr(size_t, const folly::AsyncSocketException&) noexcept override {
  80     error_ = true;
  81   }
  82
  83   void reset() {
  84     writes_ = 0;
  85     error_ = false;
  86   }
  87
  88   uint32_t writes_{0};
  89   bool error_{false};
  90 };
  91
  92 class AsyncPipeTest: public Test {
  93  public:
  94   void reset(bool movable) {
  95     reader_.reset();
  96     readCallback_.reset();
  97     writer_.reset();
  98     writeCallback_.reset();
  99
 100     int rc = pipe(pipeFds_);
 101     EXPECT_EQ(rc, 0);
 102
 103     EXPECT_EQ(::fcntl(pipeFds_[0], F_SETFL, O_NONBLOCK), 0);
 104     EXPECT_EQ(::fcntl(pipeFds_[1], F_SETFL, O_NONBLOCK), 0);
 105     reader_ = folly::AsyncPipeReader::newReader(
 106       &eventBase_, pipeFds_[0]);
 107     writer_ = folly::AsyncPipeWriter::newWriter(
 108       &eventBase_, pipeFds_[1]);
 109
 110     readCallback_.setMovable(movable);
 111   }
 112
 113  protected:
 114   folly::EventBase eventBase_;
 115   int pipeFds_[2];
 116   folly::AsyncPipeReader::UniquePtr reader_;
 117   folly::AsyncPipeWriter::UniquePtr writer_;
 118   TestReadCallback readCallback_;
 119   TestWriteCallback writeCallback_;
 120 };
 121
 122 std::unique_ptr<folly::IOBuf> getBuf(const std::string& data) {
 123   auto buf = folly::IOBuf::copyBuffer(data.c_str(), data.length());
 124   return buf;
 125 }
 126
 127 } // namespace
 128
 129 TEST_F(AsyncPipeTest, simple) {
 130   for (int pass = 0; pass < 2; ++pass) {
 131     reset(pass % 2 != 0);
 132     reader_->setReadCB(&readCallback_);
 133     writer_->write(getBuf("hello"), &writeCallback_);
 134     writer_->closeOnEmpty();
 135     eventBase_.loop();
 136     EXPECT_EQ(readCallback_.getData(), "hello");
 137     EXPECT_FALSE(readCallback_.error_);
 138     EXPECT_EQ(writeCallback_.writes_, 1);
 139     EXPECT_FALSE(writeCallback_.error_);
 140   }
 141 }
 142
 143 TEST_F(AsyncPipeTest, blocked_writes) {
 144   for (int pass = 0; pass < 2; ++pass) {
 145     reset(pass % 2 != 0);
 146     uint32_t writeAttempts = 0;
 147     do {
 148       ++writeAttempts;
 149       writer_->write(getBuf("hello"), &writeCallback_);
 150     } while (writeCallback_.writes_ == writeAttempts);
 151     // there is one blocked write
 152     writer_->closeOnEmpty();
 153
 154     reader_->setReadCB(&readCallback_);
 155
 156     eventBase_.loop();
 157     std::string expected;
 158     for (uint32_t i = 0; i < writeAttempts; i++) {
 159       expected += "hello";
 160     }
 161     EXPECT_EQ(readCallback_.getData(), expected);
 162     EXPECT_FALSE(readCallback_.error_);
 163     EXPECT_EQ(writeCallback_.writes_, writeAttempts);
 164     EXPECT_FALSE(writeCallback_.error_);
 165   }
 166 }
 167
 168 TEST_F(AsyncPipeTest, writeOnClose) {
 169   for (int pass = 0; pass < 2; ++pass) {
 170     reset(pass % 2 != 0);
 171     reader_->setReadCB(&readCallback_);
 172     writer_->write(getBuf("hello"), &writeCallback_);
 173     writer_->closeOnEmpty();
 174     writer_->write(getBuf("hello"), &writeCallback_);
 175     eventBase_.loop();
 176     EXPECT_EQ(readCallback_.getData(), "hello");
 177     EXPECT_FALSE(readCallback_.error_);
 178     EXPECT_EQ(writeCallback_.writes_, 1);
 179     EXPECT_TRUE(writeCallback_.error_);
 180   }
 181 }