C++ TEventBase::loop方法代码示例

TEST(ThriftServer, GetLoadTest) {
  apache::thrift::TestThriftServerFactory<TestInterface> factory;
  auto serv = factory.create();
  ScopedServerThread sst(serv);
  TEventBase base;
  std::shared_ptr<TAsyncSocket> socket(
    TAsyncSocket::newSocket(&base, *sst.getAddress()));

  TestServiceAsyncClient client(
    std::unique_ptr<HeaderClientChannel,
                    apache::thrift::async::TDelayedDestruction::Destructor>(
                      new HeaderClientChannel(socket)));

  auto header_channel = boost::polymorphic_downcast<HeaderClientChannel*>(
    client.getChannel());
  RpcOptions rpcOptions;
  rpcOptions.setWriteHeader("load", "thrift.active_requests");
  auto callback = std::unique_ptr<RequestCallback>(
      new FunctionReplyCallback([&](ClientReceiveState&& state) {
    std::string response;
    auto headers = state.header()->getHeaders();
    auto load = headers.find("load");
    EXPECT_NE(load, headers.end());
    EXPECT_EQ(load->second, "0");
    TestServiceAsyncClient::recv_wrapped_sendResponse(response, state);
    EXPECT_EQ(response, "test64");
  }));
  client.sendResponse(rpcOptions, std::move(callback), 64);
  base.loop();

  serv->setGetLoad([&](std::string counter){
    EXPECT_EQ(counter, "thrift.active_requests");
    return 1;
  });

  rpcOptions.setWriteHeader("load", "thrift.active_requests");
  callback = std::unique_ptr<RequestCallback>(
      new FunctionReplyCallback([&](ClientReceiveState&& state) {
    std::string response;
    auto headers = state.header()->getHeaders();
    auto load = headers.find("load");
    EXPECT_NE(load, headers.end());
    EXPECT_EQ(load->second, "1");
    TestServiceAsyncClient::recv_wrapped_sendResponse(response, state);
    EXPECT_EQ(response, "test64");
  }));
  client.sendResponse(rpcOptions, std::move(callback), 64);
  base.loop();
}

TEST(TAsyncSSLSocketTest, NpnTestUnset) {
  // Identical to above test, except that we want unset NPN before
  // looping.
  TEventBase eventBase;
  std::shared_ptr<SSLContext> clientCtx(new SSLContext);
  std::shared_ptr<SSLContext> serverCtx(new SSLContext);;
  int fds[2];
  getfds(fds);
  getctx(clientCtx, serverCtx);

  clientCtx->setAdvertisedNextProtocols({"blub","baz"});
  serverCtx->setAdvertisedNextProtocols({"foo","bar","baz"});

  TAsyncSSLSocket::UniquePtr clientSock(
    new TAsyncSSLSocket(clientCtx, &eventBase, fds[0], false));
  TAsyncSSLSocket::UniquePtr serverSock(
    new TAsyncSSLSocket(serverCtx, &eventBase, fds[1], true));

  // unsetting NPN for any of [client, server] is enought to make NPN not
  // work
  clientCtx->unsetNextProtocols();

  NpnClient client(std::move(clientSock));
  NpnServer server(std::move(serverSock));

  eventBase.loop();

  EXPECT_TRUE(client.nextProtoLength == 0);
  EXPECT_TRUE(server.nextProtoLength == 0);
  EXPECT_TRUE(client.nextProto == nullptr);
  EXPECT_TRUE(server.nextProto == nullptr);
}

void AsyncCpp2Test(bool enable_security) {
  apache::thrift::TestThriftServerFactory<TestInterface> factory;
  ScopedServerThread sst(factory.create());
  TEventBase base;
  std::shared_ptr<TAsyncSocket> socket(
    TAsyncSocket::newSocket(&base, *sst.getAddress()));

  auto client_channel = HeaderClientChannel::newChannel(socket);
  if (enable_security) {
    client_channel->setSecurityPolicy(THRIFT_SECURITY_PERMITTED);
    client_channel->setSaslClient(std::unique_ptr<SaslClient>(
      new StubSaslClient(socket->getEventBase())
    ));
  }
  TestServiceAsyncClient client(std::move(client_channel));

  boost::polymorphic_downcast<HeaderClientChannel*>(
    client.getChannel())->setTimeout(10000);
  client.sendResponse([&](ClientReceiveState&& state) {
    std::string response;
    try {
      TestServiceAsyncClient::recv_sendResponse(
          response, state);
    } catch(const std::exception& ex) {
    }
    EXPECT_EQ(response, "test64");
  }, 64);
  base.loop();
}

TEST(ThriftServer, OnewayFutureClientTest) {
  using std::chrono::steady_clock;

  ScopedServerThread sst(getServer());
  TEventBase base;
  std::shared_ptr<TAsyncSocket> socket(
    TAsyncSocket::newSocket(&base, *sst.getAddress()));

  FutureServiceAsyncClient client(
    std::unique_ptr<HeaderClientChannel,
                    apache::thrift::async::TDelayedDestruction::Destructor>(
                      new HeaderClientChannel(socket)));

  auto future = client.future_noResponse(1000);
  steady_clock::time_point sent = steady_clock::now();

  // wait for future to finish.
  base.loop();
  steady_clock::time_point waited = steady_clock::now();

  future.value();
  steady_clock::time_point got = steady_clock::now();

  steady_clock::duration waitTime = waited - sent;
  steady_clock::duration gotTime = got - waited;

  int factor = 1;
  EXPECT_GE(waitTime, factor * gotTime);
}

TEST(ThriftServer, useExistingSocketAndConnectionIdleTimeout) {
  // This is ConnectionIdleTimeoutTest, but with an existing socket
  apache::thrift::TestThriftServerFactory<TestInterface> factory;
  auto server = factory.create();
  TAsyncServerSocket::UniquePtr serverSocket(new TAsyncServerSocket);
  serverSocket->bind(0);
  server->useExistingSocket(std::move(serverSocket));

  server->setIdleTimeout(std::chrono::milliseconds(20));
  apache::thrift::util::ScopedServerThread st(server);

  TEventBase base;
  std::shared_ptr<TAsyncSocket> socket(
    TAsyncSocket::newSocket(&base, *st.getAddress()));

  TestServiceAsyncClient client(
    std::unique_ptr<HeaderClientChannel,
                    apache::thrift::async::TDelayedDestruction::Destructor>(
                      new HeaderClientChannel(socket)));

  std::string response;
  client.sync_sendResponse(response, 200);
  EXPECT_EQ(response, "test200");
  base.loop();
}

/**
 * Test SSL server accept timeout with cache path
 */
TEST(TAsyncSSLSocketTest, SSLServerAsyncCacheTimeoutTest) {
  // Start listening on a local port
  WriteCallbackBase writeCallback;
  ReadCallback readCallback(&writeCallback);
  HandshakeCallback handshakeCallback(&readCallback);
  SSLServerAsyncCacheAcceptCallback acceptCallback(&handshakeCallback, 50);
  TestSSLAsyncCacheServer server(&acceptCallback);

  // Set up SSL client
  TEventBase eventBase;
  std::shared_ptr<SSLClient> client(new SSLClient(&eventBase, server.getAddress(),
                                             2));

  client->connect();
  EventBaseAborter eba(&eventBase, 3000);
  eventBase.loop();

  EXPECT_EQ(server.getAsyncCallbacks(), 1);
  EXPECT_EQ(server.getAsyncLookups(), 1);
  EXPECT_EQ(client->getErrors(), 1);
  EXPECT_EQ(client->getMiss(), 1);
  EXPECT_EQ(client->getHit(), 0);

  cerr << "SSLServerAsyncCacheTimeoutTest test completed" << endl;
}

/**
 * Test SSL client socket timeout
 */
TEST(TAsyncSSLSocketTest, SSLClientTimeoutTest) {
  // Start listening on a local port
  EmptyReadCallback readCallback;
  HandshakeCallback handshakeCallback(&readCallback,
                                      HandshakeCallback::EXPECT_ERROR);
  HandshakeTimeoutCallback acceptCallback(&handshakeCallback);
  TestSSLServer server(&acceptCallback);

  // Set up SSL client
  TEventBase eventBase;
  std::shared_ptr<SSLClient> client(new SSLClient(&eventBase, server.getAddress(),
                                             1, 10));
  client->connect(true /* write before connect completes */);
  EventBaseAborter eba(&eventBase, 3000);
  eventBase.loop();

  usleep(100000);
  // This is checking that the connectError callback precedes any queued
  // writeError callbacks.  This matches TAsyncSocket's behavior
  EXPECT_EQ(client->getWriteAfterConnectErrors(), 1);
  EXPECT_EQ(client->getErrors(), 1);
  EXPECT_EQ(client->getMiss(), 0);
  EXPECT_EQ(client->getHit(), 0);

  cerr << "SSLClientTimeoutTest test completed" << endl;
}

/*
 * Test some timeouts that are scheduled on one timeout set, then moved to
 * another timeout set.
 */
TEST(TAsyncTimeoutSetTest, SwitchTimeoutSet) {
  TEventBase eventBase;
  StackTimeoutSet ts10(&eventBase, milliseconds(10));
  StackTimeoutSet ts5(&eventBase, milliseconds(5));

  TestTimeout t1(&ts5, &ts10, &ts5);
  TestTimeout t2(&ts10, &ts10, &ts5);
  TestTimeout t3(&ts5, &ts5, &ts10, &ts5);

  ts5.scheduleTimeout(&t1);

  TimePoint start;
  eventBase.loop();
  TimePoint end;

  ASSERT_EQ(t1.timestamps.size(), 3);
  ASSERT_EQ(t2.timestamps.size(), 3);
  ASSERT_EQ(t3.timestamps.size(), 4);

  T_CHECK_TIMEOUT(start, t1.timestamps[0], milliseconds(5));
  T_CHECK_TIMEOUT(t1.timestamps[0], t1.timestamps[1], milliseconds(10));
  T_CHECK_TIMEOUT(t1.timestamps[1], t1.timestamps[2], milliseconds(5));

  T_CHECK_TIMEOUT(start, t2.timestamps[0], milliseconds(10));
  T_CHECK_TIMEOUT(t2.timestamps[0], t2.timestamps[1], milliseconds(10));
  T_CHECK_TIMEOUT(t2.timestamps[1], t2.timestamps[2], milliseconds(5));

  T_CHECK_TIMEOUT(start, t3.timestamps[0], milliseconds(5));
  T_CHECK_TIMEOUT(t3.timestamps[0], t3.timestamps[1], milliseconds(5));
  T_CHECK_TIMEOUT(t3.timestamps[1], t3.timestamps[2], milliseconds(10));
  T_CHECK_TIMEOUT(t3.timestamps[2], t3.timestamps[3], milliseconds(5));

  // 10ms fudge factor to account for loaded machines
  T_CHECK_TIMEOUT(start, end, milliseconds(25), milliseconds(10));
}

void runTest(TEventBase& evb, LoadTestCobClient& client) {
  // Test sending a few requests to the server
  TestCallback callback;
  client.add(std::bind(&TestCallback::addDone, &callback,
                       std::placeholders::_1),
                       9, 12);
  evb.loop();
  BOOST_CHECK(callback.done);
  BOOST_CHECK_EQUAL(callback.addResult, 21);

  callback.reset();
  string testData;
  randomizeString(&testData, 3*1024);
  client.echo(std::bind(&TestCallback::echoDone, &callback,
                        std::placeholders::_1),
                        testData);
  evb.loop();
  BOOST_CHECK(callback.done);
  BOOST_CHECK(callback.echoResult == testData);

  callback.reset();
  randomizeString(&testData, 1024*1024);
  client.echo(std::bind(&TestCallback::echoDone, &callback,
                        std::placeholders::_1),
                        testData);
  evb.loop();
  BOOST_CHECK(callback.done);
  BOOST_CHECK(callback.echoResult == testData);

  callback.reset();
  client.add(std::bind(&TestCallback::addDone, &callback,
                       std::placeholders::_1),
                       321, 987);
  evb.loop();
  BOOST_CHECK(callback.done);
  BOOST_CHECK_EQUAL(callback.addResult, 1308);
}

TEST(HeaderClientChannelHttpTest, SimpleTest) {
  std::unique_ptr<ScopedServerThread> serverThread = createHttpServer();

  TEventBase eb;
  const TSocketAddress* addr = serverThread->getAddress();
  std::shared_ptr<TAsyncSocket> socket = TAsyncSocket::newSocket(&eb, *addr);
  std::unique_ptr<HeaderClientChannel, TDelayedDestruction::Destructor> channel(
      new HeaderClientChannel(socket));
  channel->getHeader()->useAsHttpClient("127.0.0.1", "meh");
  TestServiceAsyncClient client(std::move(channel));
  client.sendResponse(
    [] (apache::thrift::ClientReceiveState&& state) {
      if (state.exception()) {
        try {
          std::rethrow_exception(state.exception());
        } catch (const std::exception& e) {
          LOG(INFO) << e.what();
        }
      }
      EXPECT_TRUE(state.exception() == nullptr);
      std::string res;
      TestServiceAsyncClient::recv_sendResponse(res, state);
      EXPECT_EQ(res, "test24");
    },
    24);
  eb.loop();

  client.eventBaseAsync(
    [] (apache::thrift::ClientReceiveState&& state) {
      EXPECT_TRUE(state.exception() == nullptr);
      std::string res;
      TestServiceAsyncClient::recv_eventBaseAsync(res, state);
      EXPECT_EQ(res, "hello world");
    });
  eb.loop();
}

int main() {
  TEventBase base;

  const int ports[] = { 9090, 9091, 9092, 9093, 9094 };
  const char *hosts[] = {"127.0.0.1", "127.0.0.1", "127.0.0.1", "127.0.0.1", "127.0.0.1"};
  unsigned char* images[5];
  int cur_pos = 0;

  std::vector<std::shared_ptr<TAsyncSocket>> sockets;
  std::vector<std::shared_ptr<aobench::cpp2::AobenchServiceAsyncClient>> clients;
  for (int i = 0; i < 5; ++i) {
    std::shared_ptr<TAsyncSocket> socket(
      TAsyncSocket::newSocket(&base, hosts[i], ports[i]));
    sockets.push_back(socket);

    auto client_channel = HeaderClientChannel::newChannel(socket);
    auto client = std::make_shared<aobench::cpp2::AobenchServiceAsyncClient>(std::move(client_channel));
    clients.push_back(client);

    client->render(
        [&](ClientReceiveState&& state) {
          std::string result;
          fprintf(stderr, "received\n");
          try {
            aobench::cpp2::AobenchServiceAsyncClient::recv_render(result, state);

            unsigned char* img = new unsigned char [result.size()];
            for (int i = 0; i < static_cast<int>(result.size()); ++i) {
              img[i] = static_cast<unsigned char>(result[i]);
            }
            images[cur_pos] = img;
            ++cur_pos;
            if (cur_pos == 5) {
              saveppm_sum("ao.ppm", 256, 256, images, 5);
              for (int i = 0; i < 5; ++i) {
                delete[] images[i];
              }
              fprintf(stderr, "accumulated\n");
            }
          } catch(const std::exception& ex) {
            fprintf(stderr, "exception thrown %s\n", ex.what());
          }
        }, 256, 256, 2);
  }
  fprintf(stderr, "started\n");
  base.loop();
  fprintf(stderr, "finished\n");
}

void QueueTest::putMessages() {
  TEventBase eventBase;

  QueueConsumer consumer;
  QueueConsumer consumer2;
  consumer.fn = [&](int msg) {
    // Stop consuming after we receive a message with value 0, and start
    // consumer2
    if (msg == 0) {
      consumer.stopConsuming();
      consumer2.startConsuming(&eventBase, &queue);
    }
  };
  consumer2.fn = [&](int msg) {
    // Stop consuming after we receive a message with value 0
    if (msg == 0) {
      consumer2.stopConsuming();
    }
  };
  consumer.startConsuming(&eventBase, &queue);

  list<int> msgList = { 1, 2, 3, 4 };
  vector<int> msgVector = { 5, 0, 9, 8, 7, 6, 7, 7,
                            8, 8, 2, 9, 6, 6, 10, 2, 0 };
  // Call putMessages() several times to add messages to the queue
  queue.putMessages(msgList.begin(), msgList.end());
  queue.putMessages(msgVector.begin() + 2, msgVector.begin() + 4);
  // Test sending 17 messages, the pipe-based queue calls write in 16 byte
  // chunks
  queue.putMessages(msgVector.begin(), msgVector.end());

  // Loop until the consumer has stopped
  eventBase.loop();

  vector<int> expectedMessages = { 1, 2, 3, 4, 9, 8, 7, 5, 0 };
  vector<int> expectedMessages2 = { 9, 8, 7, 6, 7, 7, 8, 8, 2, 9, 6, 10, 2, 0 };
  BOOST_CHECK_EQUAL(expectedMessages.size(), consumer.messages.size());
  for (unsigned int idx = 0; idx < expectedMessages.size(); ++idx) {
    BOOST_CHECK_EQUAL(expectedMessages[idx], consumer.messages.at(idx));
  }
  BOOST_CHECK_EQUAL(expectedMessages2.size(), consumer2.messages.size());
  for (unsigned int idx = 0; idx < expectedMessages2.size(); ++idx) {
    BOOST_CHECK_EQUAL(expectedMessages2[idx], consumer2.messages.at(idx));
  }
}

/*
 * Test destroying a TAsyncTimeoutSet with timeouts outstanding
 */
TEST(TAsyncTimeoutSetTest, DestroyTimeoutSet) {
  TEventBase eventBase;

  TAsyncTimeoutSet::UniquePtr ts5(new TAsyncTimeoutSet(
        &eventBase, milliseconds(5)));
  TAsyncTimeoutSet::UniquePtr ts10(new TAsyncTimeoutSet(
        &eventBase, milliseconds(10)));

  TestTimeout t5_1(ts5.get());
  TestTimeout t5_2(ts5.get());
  TestTimeout t5_3(ts5.get());

  TestTimeout t10_1(ts10.get());
  TestTimeout t10_2(ts10.get());

  // Have t5_1 destroy ts10
  t5_1.fn = [&] { ts10.reset(); };
  // Have t5_2 destroy ts5
  // Note that this will call destroy() on ts5 inside ts5's timeoutExpired()
  // method.
  t5_2.fn = [&] { ts5.reset(); };

  TimePoint start;
  eventBase.loop();
  TimePoint end;

  ASSERT_EQ(t5_1.timestamps.size(), 1);
  T_CHECK_TIMEOUT(start, t5_1.timestamps[0], milliseconds(5));
  ASSERT_EQ(t5_2.timestamps.size(), 1);
  T_CHECK_TIMEOUT(start, t5_2.timestamps[0], milliseconds(5));

  ASSERT_EQ(t5_3.timestamps.size(), 0);
  ASSERT_EQ(t10_1.timestamps.size(), 0);
  ASSERT_EQ(t10_2.timestamps.size(), 0);

  T_CHECK_TIMEOUT(start, end, milliseconds(5));
}

void QueueTest::sendOne() {
  // Create a notification queue and a callback in this thread
  TEventBase eventBase;

  QueueConsumer consumer;
  consumer.fn = [&](int) {
    // Stop consuming after we receive 1 message
    consumer.stopConsuming();
  };
  consumer.startConsuming(&eventBase, &queue);

  // Start a new TEventBase thread to put a message on our queue
  ScopedEventBaseThread t1;
  t1.getEventBase()->runInEventBaseThread([&] {
    queue.putMessage(5);
  });

  // Loop until we receive the message
  eventBase.loop();

  const auto& messages = consumer.messages;
  BOOST_CHECK_EQUAL(messages.size(), 1);
  BOOST_CHECK_EQUAL(messages.at(0), 5);
}

void QueueTest::maxReadAtOnce() {
  // Add 100 messages to the queue
  for (int n = 0; n < 100; ++n) {
    queue.putMessage(n);
  }

  TEventBase eventBase;

  // Record how many messages were processed each loop iteration.
  uint32_t messagesThisLoop = 0;
  std::vector<uint32_t> messagesPerLoop;
  std::function<void()> loopFinished = [&] {
    // Record the current number of messages read this loop
    messagesPerLoop.push_back(messagesThisLoop);
    // Reset messagesThisLoop to 0 for the next loop
    messagesThisLoop = 0;

    // To prevent use-after-free bugs when eventBase destructs,
    // prevent calling runInLoop any more after the test is finished.
    // 55 == number of times loop should run.
    if (messagesPerLoop.size() != 55) {
      // Reschedule ourself to run at the end of the next loop
      eventBase.runInLoop(loopFinished);
    }
  };
  // Schedule the first call to loopFinished
  eventBase.runInLoop(loopFinished);

  QueueConsumer consumer;
  // Read the first 50 messages 10 at a time.
  consumer.setMaxReadAtOnce(10);
  consumer.fn = [&](int value) {
    ++messagesThisLoop;
    // After 50 messages, drop to reading only 1 message at a time.
    if (value == 50) {
      consumer.setMaxReadAtOnce(1);
    }
    // Terminate the loop when we reach the end of the messages.
    if (value == 99) {
      eventBase.terminateLoopSoon();
    }
  };
  consumer.startConsuming(&eventBase, &queue);

  // Run the event loop until the consumer terminates it
  eventBase.loop();

  // The consumer should have read all 100 messages in order
  BOOST_CHECK_EQUAL(consumer.messages.size(), 100);
  for (int n = 0; n < 100; ++n) {
    BOOST_CHECK_EQUAL(consumer.messages.at(n), n);
  }

  // Currently TEventBase happens to still run the loop callbacks even after
  // terminateLoopSoon() is called.  However, we don't really want to depend on
  // this behavior.  In case this ever changes in the future, add
  // messagesThisLoop to messagesPerLoop in loop callback isn't invoked for the
  // last loop iteration.
  if (messagesThisLoop > 0) {
    messagesPerLoop.push_back(messagesThisLoop);
    messagesThisLoop = 0;
  }

  // For the first 5 loops it should have read 10 messages each time.
  // After that it should have read 1 messages per loop for the next 50 loops.
  BOOST_CHECK_EQUAL(messagesPerLoop.size(), 55);
  for (int n = 0; n < 5; ++n) {
    BOOST_CHECK_EQUAL(messagesPerLoop.at(n), 10);
  }
  for (int n = 5; n < 55; ++n) {
    BOOST_CHECK_EQUAL(messagesPerLoop.at(n), 1);
  }
}