From e6fd088b4fb2ae61f92c1cd85fb8ee54d3483760 Mon Sep 17 00:00:00 2001
From: Neel Goyal <ngoyal@fb.com>
Date: Wed, 13 Apr 2016 08:04:45 -0700
Subject: [PATCH] Add method to get the connect timeout used for an AsyncSocket

Summary: Have the AsyncSocket keep track of the timeout used for connecting and add a getter to retrieve it.

Reviewed By: hiteshk

Differential Revision: D3170625

fb-gh-sync-id: 61d0ecd8d975c49978a1cf222671aa16a2160499
fbshipit-source-id: 61d0ecd8d975c49978a1cf222671aa16a2160499
---
 folly/io/async/AsyncSocket.cpp           |  1 +
 folly/io/async/AsyncSocket.h             | 28 ++++++++++++++++--------
 folly/io/async/test/AsyncSocketTest2.cpp |  4 ++++
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/folly/io/async/AsyncSocket.cpp b/folly/io/async/AsyncSocket.cpp
index beabc5d8..6fd3c355 100644
--- a/folly/io/async/AsyncSocket.cpp
+++ b/folly/io/async/AsyncSocket.cpp
@@ -321,6 +321,7 @@ void AsyncSocket::connect(ConnectCallback* callback,
     return invalidState(callback);
   }
 
+  connectTimeout_ = std::chrono::milliseconds(timeout);
   connectStartTime_ = std::chrono::steady_clock::now();
   // Make connect end time at least >= connectStartTime.
   connectEndTime_ = connectStartTime_;
diff --git a/folly/io/async/AsyncSocket.h b/folly/io/async/AsyncSocket.h
index 28f1c65d..ba706747 100644
--- a/folly/io/async/AsyncSocket.h
+++ b/folly/io/async/AsyncSocket.h
@@ -257,15 +257,19 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
    *                  does not succeed within this period,
    *                  callback->connectError() will be invoked.
    */
-  virtual void connect(ConnectCallback* callback,
-               const folly::SocketAddress& address,
-               int timeout = 0,
-               const OptionMap &options = emptyOptionMap,
-               const folly::SocketAddress& bindAddr = anyAddress()
-               ) noexcept;
-  void connect(ConnectCallback* callback, const std::string& ip, uint16_t port,
-               int timeout = 00,
-               const OptionMap &options = emptyOptionMap) noexcept;
+  virtual void connect(
+      ConnectCallback* callback,
+      const folly::SocketAddress& address,
+      int timeout = 0,
+      const OptionMap& options = emptyOptionMap,
+      const folly::SocketAddress& bindAddr = anyAddress()) noexcept;
+
+  void connect(
+      ConnectCallback* callback,
+      const std::string& ip,
+      uint16_t port,
+      int timeout = 0,
+      const OptionMap& options = emptyOptionMap) noexcept;
 
   /**
    * If a connect request is in-flight, cancels it and closes the socket
@@ -406,6 +410,10 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
     return connectEndTime_ - connectStartTime_;
   }
 
+  std::chrono::milliseconds getConnectTimeout() const {
+    return connectTimeout_;
+  }
+
   // Methods controlling socket options
 
   /**
@@ -802,6 +810,8 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
   std::chrono::steady_clock::time_point connectStartTime_;
   std::chrono::steady_clock::time_point connectEndTime_;
 
+  std::chrono::milliseconds connectTimeout_{0};
+
   BufferCallback* bufferCallback_{nullptr};
 };
 #ifdef _MSC_VER
diff --git a/folly/io/async/test/AsyncSocketTest2.cpp b/folly/io/async/test/AsyncSocketTest2.cpp
index 45adfa9e..549dc294 100644
--- a/folly/io/async/test/AsyncSocketTest2.cpp
+++ b/folly/io/async/test/AsyncSocketTest2.cpp
@@ -98,6 +98,7 @@ TEST(AsyncSocketTest, Connect) {
 
   CHECK_EQ(cb.state, STATE_SUCCEEDED);
   EXPECT_LE(0, socket->getConnectTime().count());
+  EXPECT_EQ(socket->getConnectTimeout(), std::chrono::milliseconds(30));
 }
 
 /**
@@ -118,6 +119,7 @@ TEST(AsyncSocketTest, ConnectRefused) {
   CHECK_EQ(cb.state, STATE_FAILED);
   CHECK_EQ(cb.exception.getType(), AsyncSocketException::NOT_OPEN);
   EXPECT_LE(0, socket->getConnectTime().count());
+  EXPECT_EQ(socket->getConnectTimeout(), std::chrono::milliseconds(30));
 }
 
 /**
@@ -156,6 +158,7 @@ TEST(AsyncSocketTest, ConnectTimeout) {
   socket->getPeerAddress(&peer);
   CHECK_EQ(peer, addr);
   EXPECT_LE(0, socket->getConnectTime().count());
+  EXPECT_EQ(socket->getConnectTimeout(), std::chrono::milliseconds(1));
 }
 
 /**
@@ -190,6 +193,7 @@ TEST(AsyncSocketTest, ConnectAndWrite) {
 
   ASSERT_TRUE(socket->isClosedBySelf());
   ASSERT_FALSE(socket->isClosedByPeer());
+  EXPECT_EQ(socket->getConnectTimeout(), std::chrono::milliseconds(30));
 }
 
 /**
-- 
2.34.1