folly/wangle/channel/test/FileRegionTest.cpp

   1 /*
   2  * Copyright 2015 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 #include <folly/wangle/channel/FileRegion.h>
  17 #include <folly/io/async/test/AsyncSocketTest.h>
  18 #include <gtest/gtest.h>
  19
  20 using namespace folly;
  21 using namespace folly::wangle;
  22 using namespace testing;
  23
  24 struct FileRegionTest : public Test {
  25   FileRegionTest() {
  26     // Connect
  27     socket = AsyncSocket::newSocket(&evb);
  28     socket->connect(&ccb, server.getAddress(), 30);
  29
  30     // Accept the connection
  31     acceptedSocket = server.acceptAsync(&evb);
  32     acceptedSocket->setReadCB(&rcb);
  33
  34     // Create temp file
  35     char path[] = "/tmp/AsyncSocketTest.WriteFile.XXXXXX";
  36     fd = mkostemp(path, O_RDWR);
  37     EXPECT_TRUE(fd > 0);
  38     EXPECT_EQ(0, unlink(path));
  39   }
  40
  41   ~FileRegionTest() override {
  42     // Close up shop
  43     close(fd);
  44     acceptedSocket->close();
  45     socket->close();
  46   }
  47
  48   TestServer server;
  49   EventBase evb;
  50   std::shared_ptr<AsyncSocket> socket;
  51   std::shared_ptr<AsyncSocket> acceptedSocket;
  52   ConnCallback ccb;
  53   ReadCallback rcb;
  54   int fd;
  55 };
  56
  57 TEST_F(FileRegionTest, Basic) {
  58   size_t count = 1000000000; // 1 GB
  59   void* zeroBuf = calloc(1, count);
  60   write(fd, zeroBuf, count);
  61
  62   FileRegion fileRegion(fd, 0, count);
  63   auto f = fileRegion.transferTo(socket);
  64   try {
  65     f.getVia(&evb);
  66   } catch (std::exception& e) {
  67     LOG(FATAL) << exceptionStr(e);
  68   }
  69
  70   // Let the reads run to completion
  71   socket->shutdownWrite();
  72   evb.loop();
  73
  74   ASSERT_EQ(rcb.state, STATE_SUCCEEDED);
  75
  76   size_t receivedBytes = 0;
  77   for (auto& buf : rcb.buffers) {
  78     receivedBytes += buf.length;
  79     ASSERT_EQ(memcmp(buf.buffer, zeroBuf, buf.length), 0);
  80   }
  81   ASSERT_EQ(receivedBytes, count);
  82 }
  83
  84 TEST_F(FileRegionTest, Repeated) {
  85   size_t count = 1000000;
  86   void* zeroBuf = calloc(1, count);
  87   write(fd, zeroBuf, count);
  88
  89   int sendCount = 1000;
  90
  91   FileRegion fileRegion(fd, 0, count);
  92   std::vector<Future<void>> fs;
  93   for (int i = 0; i < sendCount; i++) {
  94     fs.push_back(fileRegion.transferTo(socket));
  95   }
  96   auto f = collect(fs);
  97   ASSERT_NO_THROW(f.getVia(&evb));
  98
  99   // Let the reads run to completion
 100   socket->shutdownWrite();
 101   evb.loop();
 102
 103   ASSERT_EQ(rcb.state, STATE_SUCCEEDED);
 104
 105   size_t receivedBytes = 0;
 106   for (auto& buf : rcb.buffers) {
 107     receivedBytes += buf.length;
 108   }
 109   ASSERT_EQ(receivedBytes, sendCount*count);
 110 }