Commit 754409b4 authored by Takshak Chahande's avatar Takshak Chahande Committed by Facebook Github Bot

AsyncSocket's ConnectCallback provide additional preConnect

Summary:
Currently, AsyncSocket's `connect` API creates new socket, set various
options, bind it to passed argument address and used it for connecting to
server.

While most of the things are handled inside that single connect routine; sometimes
if the caller would like to do few more additional operations before actual
connect happens; caller would like to have hold on the underneath used socket.

In order to do that additional operation on the socket, this diff provides the
facility to provide custom pre-connect hook `preConnect(NetworkSocket fd)` in
the `ConnectCallback` which will be invoked just before the actual server
connect.

Reviewed By: yfeldblum

Differential Revision: D15282727

fbshipit-source-id: 5beac55c77a9537ee4f26856dd5f78c5224a1f49
parent 088fab7c
...@@ -92,6 +92,13 @@ class AsyncSSLSocketConnector : public AsyncSocket::ConnectCallback, ...@@ -92,6 +92,13 @@ class AsyncSSLSocketConnector : public AsyncSocket::ConnectCallback,
timeout_(timeout), timeout_(timeout),
startTime_(std::chrono::steady_clock::now()) {} startTime_(std::chrono::steady_clock::now()) {}
void preConnect(folly::NetworkSocket fd) override {
VLOG(7) << "client preConnect hook is invoked";
if (callback_) {
callback_->preConnect(fd);
}
}
void connectSuccess() noexcept override { void connectSuccess() noexcept override {
VLOG(7) << "client socket connected"; VLOG(7) << "client socket connected";
......
...@@ -557,6 +557,11 @@ void AsyncSocket::connect( ...@@ -557,6 +557,11 @@ void AsyncSocket::connect(
} }
} }
// Call preConnect hook if any.
if (connectCallback_) {
connectCallback_->preConnect(fd_);
}
// Perform the connect() // Perform the connect()
address.getAddress(&addrStorage); address.getAddress(&addrStorage);
......
...@@ -98,6 +98,15 @@ class AsyncSocket : virtual public AsyncTransportWrapper { ...@@ -98,6 +98,15 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
* @param ex An exception describing the error that occurred. * @param ex An exception describing the error that occurred.
*/ */
virtual void connectErr(const AsyncSocketException& ex) noexcept = 0; virtual void connectErr(const AsyncSocketException& ex) noexcept = 0;
/**
* preConnect() will be invoked just before the actual connect happens,
* default is no-ops.
*
* @param fd An underneath created socket, use for connection.
*
*/
virtual void preConnect(NetworkSocket /*fd*/) {}
}; };
class EvbChangeCallback { class EvbChangeCallback {
......
...@@ -19,7 +19,6 @@ ...@@ -19,7 +19,6 @@
#include <folly/ExceptionWrapper.h> #include <folly/ExceptionWrapper.h>
#include <folly/Random.h> #include <folly/Random.h>
#include <folly/SocketAddress.h> #include <folly/SocketAddress.h>
#include <folly/io/async/AsyncSocket.h>
#include <folly/io/async/AsyncTimeout.h> #include <folly/io/async/AsyncTimeout.h>
#include <folly/io/async/EventBase.h> #include <folly/io/async/EventBase.h>
#include <folly/io/async/ScopedEventBaseThread.h> #include <folly/io/async/ScopedEventBaseThread.h>
...@@ -3670,6 +3669,30 @@ TEST(AsyncSocketTest, V6TosReflectTest) { ...@@ -3670,6 +3669,30 @@ TEST(AsyncSocketTest, V6TosReflectTest) {
netops::getsockopt(fd, IPPROTO_IPV6, IPV6_TCLASS, &value, &valueLength); netops::getsockopt(fd, IPPROTO_IPV6, IPV6_TCLASS, &value, &valueLength);
ASSERT_EQ(rc, 0); ASSERT_EQ(rc, 0);
ASSERT_EQ(value, 0x2c); ASSERT_EQ(value, 0x2c);
// Additional Test for ConnectCallback without bindAddr
serverSocket->addAcceptCallback(&acceptCallback, &eventBase);
serverSocket->startAccepting();
auto newClientSock = AsyncSocket::newSocket(&eventBase);
TestConnectCallback callback;
// connect call will not set this SO_REUSEADDR if we do not
// pass the bindAddress in its call; so we can safely verify this.
newClientSock->connect(&callback, serverAddress, 30);
// Collect events
eventBase.loop();
auto acceptedFd = acceptCallback.getEvents()->at(1).fd;
ASSERT_NE(acceptedFd, NetworkSocket());
int reuseAddrVal;
socklen_t reuseAddrValLen = sizeof(reuseAddrVal);
// Get the socket created underneath connect call of AsyncSocket
auto usedSockFd = newClientSock->getNetworkSocket();
int getOptRet = netops::getsockopt(
usedSockFd, SOL_SOCKET, SO_REUSEADDR, &reuseAddrVal, &reuseAddrValLen);
ASSERT_EQ(getOptRet, 0);
ASSERT_EQ(reuseAddrVal, 1 /* configured through preConnect*/);
} }
TEST(AsyncSocketTest, V4TosReflectTest) { TEST(AsyncSocketTest, V4TosReflectTest) {
......
...@@ -23,6 +23,7 @@ ...@@ -23,6 +23,7 @@
#include <folly/synchronization/RWSpinLock.h> #include <folly/synchronization/RWSpinLock.h>
#include <folly/io/async/AsyncServerSocket.h> #include <folly/io/async/AsyncServerSocket.h>
#include <folly/io/async/AsyncSocket.h>
namespace folly { namespace folly {
namespace test { namespace test {
...@@ -221,5 +222,16 @@ class TestAcceptCallback : public AsyncServerSocket::AcceptCallback { ...@@ -221,5 +222,16 @@ class TestAcceptCallback : public AsyncServerSocket::AcceptCallback {
std::deque<EventInfo> events_; std::deque<EventInfo> events_;
}; };
class TestConnectCallback : public AsyncSocket::ConnectCallback {
public:
void preConnect(NetworkSocket fd) override {
int one = 1;
netops::setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
}
void connectSuccess() noexcept override {}
void connectErr(const AsyncSocketException& /*ex*/) noexcept override {}
};
} // namespace test } // namespace test
} // namespace folly } // namespace folly
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment