Add setsockopt to TcpClient (#255)

Author: fantasy-peak

trantor/net/TcpClient.cc

@@ -133,6 +133,16 @@ void TcpClient::stop()
     connector_->stop();
 }
 
+void TcpClient::setSockOptCallback(SockOptCallback &&cb)
+{
+    connector_->setSockOptCallback(std::move(cb));
+}
+
+void TcpClient::setSockOptCallback(const SockOptCallback &cb)
+{
+    connector_->setSockOptCallback(cb);
+}
+
 void TcpClient::newConnection(int sockfd)
 {
     loop_->assertInLoopThread();

trantor/net/TcpClient.h

@@ -190,6 +190,13 @@ class TRANTOR_EXPORT TcpClient : NonCopyable,
         sslErrorCallback_ = std::move(cb);
     }
 
+    /**
+     * @brief Set the callback for set socket option
+     * @param cb The callback is called, before connect
+     */
+    void setSockOptCallback(const SockOptCallback &cb);
+    void setSockOptCallback(SockOptCallback &&cb);
+
     /**
      * @brief Enable SSL encryption.
      * @param useOldTLS If true, the TLS 1.0 and 1.1 are supported by the

trantor/net/callbacks.h

@@ -40,5 +40,6 @@ using WriteCompleteCallback = std::function<void(const TcpConnectionPtr &)>;
 using HighWaterMarkCallback =
     std::function<void(const TcpConnectionPtr &, const size_t)>;
 using SSLErrorCallback = std::function<void(SSLError)>;
+using SockOptCallback = std::function<void(int)>;
 
 }  // namespace trantor

trantor/net/inner/Connector.cc

@@ -79,6 +79,8 @@ void Connector::connect()
 {
     socketHanded_ = false;
     fd_ = Socket::createNonblockingSocketOrDie(serverAddr_.family());
+    if (sockOptCallback_)
+        sockOptCallback_(fd_);
     errno = 0;
     int ret = Socket::connect(fd_, serverAddr_);
     int savedErrno = (ret == 0) ? 0 : errno;

trantor/net/inner/Connector.h

@@ -28,6 +28,7 @@ class Connector : public NonCopyable,
   public:
     using NewConnectionCallback = std::function<void(int sockfd)>;
     using ConnectionErrorCallback = std::function<void()>;
+    using SockOptCallback = std::function<void(int sockfd)>;
     Connector(EventLoop *loop, const InetAddress &addr, bool retry = true);
     Connector(EventLoop *loop, InetAddress &&addr, bool retry = true);
     ~Connector();
@@ -47,6 +48,14 @@ class Connector : public NonCopyable,
     {
         errorCallback_ = std::move(cb);
     }
+    void setSockOptCallback(const SockOptCallback &cb)
+    {
+        sockOptCallback_ = cb;
+    }
+    void setSockOptCallback(SockOptCallback &&cb)
+    {
+        sockOptCallback_ = std::move(cb);
+    }
     const InetAddress &serverAddress() const
     {
         return serverAddr_;
@@ -58,6 +67,7 @@ class Connector : public NonCopyable,
   private:
     NewConnectionCallback newConnectionCallback_;
     ConnectionErrorCallback errorCallback_;
+    SockOptCallback sockOptCallback_;
     enum class Status
     {
         Disconnected,

trantor/tests/TcpClientTest.cc

@@ -4,6 +4,13 @@
 #include <string>
 #include <iostream>
 #include <atomic>
+#ifdef _WIN32
+#include <ws2tcpip.h>
+#else
+#include <sys/socket.h>
+#include <netinet/tcp.h>
+#endif
+
 using namespace trantor;
 #define USE_IPV6 0
 int main()
@@ -24,6 +31,29 @@ int main()
         client[i] = std::make_shared<trantor::TcpClient>(&loop,
                                                          serverAddr,
                                                          "tcpclienttest");
+        client[i]->setSockOptCallback([](int fd) {
+            LOG_DEBUG << "setSockOptCallback!";
+#ifdef _WIN32
+#elif __linux__
+            int optval = 10;
+            ::setsockopt(fd,
+                         SOL_TCP,
+                         TCP_KEEPCNT,
+                         &optval,
+                         static_cast<socklen_t>(sizeof optval));
+            ::setsockopt(fd,
+                         SOL_TCP,
+                         TCP_KEEPIDLE,
+                         &optval,
+                         static_cast<socklen_t>(sizeof optval));
+            ::setsockopt(fd,
+                         SOL_TCP,
+                         TCP_KEEPINTVL,
+                         &optval,
+                         static_cast<socklen_t>(sizeof optval));
+#else
+#endif
+        });
         client[i]->setConnectionCallback(
             [i, &loop, &connCount](const TcpConnectionPtr &conn) {
                 if (conn->connected())