folly/io/async/test/AsyncPipeTest.cpp

   1 /*
   2  * Copyright 2016 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/io/async/AsyncPipe.h>
  18 #include <folly/io/async/EventBase.h>
  19 #include <folly/Memory.h>
  20 #include <gtest/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   void readDataAvailable(size_t len) noexcept override {
  31     readBuffer_.postallocate(len);
  32   }
  33
  34   void getReadBuffer(void** bufReturn, size_t* lenReturn) noexcept override {
  35     auto res = readBuffer_.preallocate(4000, 65000);
  36     *bufReturn = res.first;
  37     *lenReturn = res.second;
  38   }
  39
  40   void readEOF() noexcept override {}
  41
  42   void readErr(const folly::AsyncSocketException&) noexcept override {
  43     error_ = true;
  44   }
  45
  46   std::string getData() {
  47     auto buf = readBuffer_.move();
  48     buf->coalesce();
  49     return std::string((char *)buf->data(), buf->length());
  50   }
  51
  52   folly::IOBufQueue readBuffer_{folly::IOBufQueue::cacheChainLength()};
  53   bool error_{false};
  54 };
  55
  56 class TestWriteCallback : public folly::AsyncWriter::WriteCallback {
  57  public:
  58   void writeSuccess() noexcept override { writes_++; }
  59
  60   void writeErr(size_t, const folly::AsyncSocketException&) noexcept override {
  61     error_ = true;
  62   }
  63
  64   uint32_t writes_{0};
  65   bool error_{false};
  66 };
  67
  68 class AsyncPipeTest: public Test {
  69  public:
  70   void SetUp() override {
  71     int rc = pipe(pipeFds_);
  72     EXPECT_EQ(rc, 0);
  73
  74     EXPECT_EQ(::fcntl(pipeFds_[0], F_SETFL, O_NONBLOCK), 0);
  75     EXPECT_EQ(::fcntl(pipeFds_[1], F_SETFL, O_NONBLOCK), 0);
  76     reader_ = folly::AsyncPipeReader::newReader(
  77       &eventBase_, pipeFds_[0]);
  78     writer_ = folly::AsyncPipeWriter::newWriter(
  79       &eventBase_, pipeFds_[1]);
  80   }
  81
  82  protected:
  83   folly::EventBase eventBase_;
  84   int pipeFds_[2];
  85   folly::AsyncPipeReader::UniquePtr reader_;
  86   folly::AsyncPipeWriter::UniquePtr writer_;
  87   TestReadCallback readCallback_;
  88   TestWriteCallback writeCallback_;
  89 };
  90
  91 std::unique_ptr<folly::IOBuf> getBuf(const std::string& data) {
  92   auto buf = folly::IOBuf::copyBuffer(data.c_str(), data.length());
  93   return buf;
  94 }
  95
  96 } // anonymous namespace
  97
  98
  99 TEST_F(AsyncPipeTest, simple) {
 100   reader_->setReadCB(&readCallback_);
 101   writer_->write(getBuf("hello"), &writeCallback_);
 102   writer_->closeOnEmpty();
 103   eventBase_.loop();
 104   EXPECT_EQ(readCallback_.getData(), "hello");
 105   EXPECT_FALSE(readCallback_.error_);
 106   EXPECT_EQ(writeCallback_.writes_, 1);
 107   EXPECT_FALSE(writeCallback_.error_);
 108 }
 109
 110 TEST_F(AsyncPipeTest, blocked_writes) {
 111   uint32_t writeAttempts = 0;
 112   do {
 113     ++writeAttempts;
 114     writer_->write(getBuf("hello"), &writeCallback_);
 115   } while (writeCallback_.writes_ == writeAttempts);
 116   // there is one blocked write
 117   writer_->closeOnEmpty();
 118
 119   reader_->setReadCB(&readCallback_);
 120
 121   eventBase_.loop();
 122   std::string expected;
 123   for (uint32_t i = 0; i < writeAttempts; i++) {
 124     expected += "hello";
 125   }
 126   EXPECT_EQ(readCallback_.getData(), expected);
 127   EXPECT_FALSE(readCallback_.error_);
 128   EXPECT_EQ(writeCallback_.writes_, writeAttempts);
 129   EXPECT_FALSE(writeCallback_.error_);
 130 }
 131
 132 TEST_F(AsyncPipeTest, writeOnClose) {
 133   reader_->setReadCB(&readCallback_);
 134   writer_->write(getBuf("hello"), &writeCallback_);
 135   writer_->closeOnEmpty();
 136   writer_->write(getBuf("hello"), &writeCallback_);
 137   eventBase_.loop();
 138   EXPECT_EQ(readCallback_.getData(), "hello");
 139   EXPECT_FALSE(readCallback_.error_);
 140   EXPECT_EQ(writeCallback_.writes_, 1);
 141   EXPECT_TRUE(writeCallback_.error_);
 142 }