Fix TFO refused case

[folly.git] / folly / io / async / test / SocketClient.cpp

diff --git a/folly/io/async/test/SocketClient.cpp b/folly/io/async/test/SocketClient.cpp
index 7f20d480e20695eeaf7178c1f4d916078baaf07b..2fbb546bd32d5778fc487a850738925a23d43fc6 100644 (file)
--- a/folly/io/async/test/SocketClient.cpp
+++ b/folly/io/async/test/SocketClient.cpp
@@ -15,6 +15,8 @@
  */
 #include <folly/io/async/test/BlockingSocket.h>
 
+#include <iostream>
+
 #include <folly/ExceptionWrapper.h>
 #include <gflags/gflags.h>
 
@@ -24,6 +26,10 @@ DEFINE_string(host, "localhost", "Host");
 DEFINE_int32(port, 0, "port");
 DEFINE_bool(tfo, false, "enable tfo");
 DEFINE_string(msg, "", "Message to send");
+DEFINE_bool(ssl, false, "use ssl");
+DEFINE_int32(timeout_ms, 0, "timeout");
+DEFINE_int32(sendtimeout_ms, 0, "send timeout");
+DEFINE_int32(num_writes, 1, "number of writes");
 
 int main(int argc, char** argv) {
   gflags::ParseCommandLineFlags(&argc, &argv, true);
@@ -35,7 +41,13 @@ int main(int argc, char** argv) {
 
   // Prep the socket
   EventBase evb;
-  AsyncSocket::UniquePtr socket(new AsyncSocket(&evb));
+  AsyncSocket::UniquePtr socket;
+  if (FLAGS_ssl) {
+    auto sslContext = std::make_shared<SSLContext>();
+    socket = AsyncSocket::UniquePtr(new AsyncSSLSocket(sslContext, &evb));
+  } else {
+    socket = AsyncSocket::UniquePtr(new AsyncSocket(&evb));
+  }
   socket->detachEventBase();
 
   if (FLAGS_tfo) {
@@ -44,6 +56,10 @@ int main(int argc, char** argv) {
 #endif
   }
 
+  if (FLAGS_sendtimeout_ms != 0) {
+    socket->setSendTimeout(FLAGS_sendtimeout_ms);
+  }
+
   // Keep this around
   auto sockAddr = socket.get();
 
@@ -51,13 +67,17 @@ int main(int argc, char** argv) {
   SocketAddress addr;
   addr.setFromHostPort(FLAGS_host, FLAGS_port);
   sock.setAddress(addr);
-  sock.open();
+  std::chrono::milliseconds timeout(FLAGS_timeout_ms);
+  sock.open(timeout);
   LOG(INFO) << "connected to " << addr.getAddressStr();
 
-  sock.write((const uint8_t*)FLAGS_msg.data(), FLAGS_msg.size());
+  for (int32_t i = 0; i < FLAGS_num_writes; ++i) {
+    sock.write((const uint8_t*)FLAGS_msg.data(), FLAGS_msg.size());
+  }
 
-  LOG(ERROR) << "TFO attempted: " << sockAddr->getTFOAttempted();
-  LOG(ERROR) << "TFO finished: " << sockAddr->getTFOFinished();
+  LOG(INFO) << "TFO attempted: " << sockAddr->getTFOAttempted();
+  LOG(INFO) << "TFO finished: " << sockAddr->getTFOFinished();
+  LOG(INFO) << "TFO success: " << sockAddr->getTFOSucceded();
 
   std::array<char, 1024> buf;
   int32_t bytesRead = 0;