Add setsockopt to TcpServer (#266)

* use std::shared_ptr<std::vector> instead of std::vector for dns cache

* Add setsockopt to TcpServer

Author: fantasy-peak

trantor/net/TcpServer.cc

@@ -47,6 +47,16 @@ TcpServer::~TcpServer()
     LOG_TRACE << "TcpServer::~TcpServer [" << serverName_ << "] destructing";
 }
 
+void TcpServer::setBeforeListenSockOptCallback(SockOptCallback cb)
+{
+    acceptorPtr_->setBeforeListenSockOptCallback(std::move(cb));
+}
+
+void TcpServer::setAfterAcceptSockOptCallback(SockOptCallback cb)
+{
+    acceptorPtr_->setAfterAcceptSockOptCallback(std::move(cb));
+}
+
 void TcpServer::newConnection(int sockfd, const InetAddress &peer)
 {
     LOG_TRACE << "new connection:fd=" << sockfd

trantor/net/TcpServer.h

@@ -160,6 +160,19 @@ class TRANTOR_EXPORT TcpServer : NonCopyable
         writeCompleteCallback_ = std::move(cb);
     }
 
+    /**
+     * @brief Set the before listen setsockopt callback.
+     *
+     * @param cb This callback will be called before the listen
+     */
+    void setBeforeListenSockOptCallback(SockOptCallback cb);
+    /**
+     * @brief Set the after accept setsockopt callback.
+     *
+     * @param cb This callback will be called after accept
+     */
+    void setAfterAcceptSockOptCallback(SockOptCallback cb);
+
     /**
      * @brief Get the name of the server.
      *

trantor/net/inner/Acceptor.cc

@@ -54,6 +54,8 @@ Acceptor::~Acceptor()
 void Acceptor::listen()
 {
     loop_->assertInLoopThread();
+    if (beforeListenSetSockOptCallback_)
+        beforeListenSetSockOptCallback_(sock_.fd());
     sock_.listen();
     acceptChannel_.enableReading();
 }
@@ -64,6 +66,8 @@ void Acceptor::readCallback()
     int newsock = sock_.accept(&peer);
     if (newsock >= 0)
     {
+        if (afterAcceptSetSockOptCallback_)
+            afterAcceptSetSockOptCallback_(newsock);
         if (newConnectionCallback_)
         {
             newConnectionCallback_(newsock, peer);

trantor/net/inner/Acceptor.h

@@ -24,6 +24,7 @@
 namespace trantor
 {
 using NewConnectionCallback = std::function<void(int fd, const InetAddress &)>;
+using AcceptorSockOptCallback = std::function<void(int)>;
 class Acceptor : NonCopyable
 {
   public:
@@ -42,6 +43,16 @@ class Acceptor : NonCopyable
     };
     void listen();
 
+    void setBeforeListenSockOptCallback(AcceptorSockOptCallback cb)
+    {
+        beforeListenSetSockOptCallback_ = std::move(cb);
+    }
+
+    void setAfterAcceptSockOptCallback(AcceptorSockOptCallback cb)
+    {
+        afterAcceptSetSockOptCallback_ = std::move(cb);
+    }
+
   protected:
 #ifndef _WIN32
     int idleFd_;
@@ -52,5 +63,7 @@ class Acceptor : NonCopyable
     NewConnectionCallback newConnectionCallback_;
     Channel acceptChannel_;
     void readCallback();
+    AcceptorSockOptCallback beforeListenSetSockOptCallback_;
+    AcceptorSockOptCallback afterAcceptSetSockOptCallback_;
 };
 }  // namespace trantor

trantor/net/inner/AresResolver.cc

@@ -171,7 +171,7 @@ void AresResolver::onQueryResult(int status,
                                  const ResolverResultsCallback& callback)
 {
     LOG_TRACE << "onQueryResult " << status;
-    std::vector<trantor::InetAddress> inets;
+    auto inets_ptr = std::make_shared<std::vector<trantor::InetAddress>>();
     if (result)
     {
         auto pptr = (struct in_addr**)result->h_addr_list;
@@ -182,25 +182,25 @@ void AresResolver::onQueryResult(int status,
             addr.sin_family = AF_INET;
             addr.sin_port = 0;
             addr.sin_addr = *reinterpret_cast<in_addr*>(*pptr);
-            inets.emplace_back(trantor::InetAddress{addr});
+            inets_ptr->emplace_back(trantor::InetAddress{addr});
         }
     }
-    if (inets.empty())
+    if (inets_ptr->empty())
     {
         struct sockaddr_in addr;
         memset(&addr, 0, sizeof addr);
         addr.sin_family = AF_INET;
         addr.sin_port = 0;
         InetAddress inet(addr);
-        inets.emplace_back(std::move(inet));
+        inets_ptr->emplace_back(std::move(inet));
     }
     {
         std::lock_guard<std::mutex> lock(globalMutex());
         auto& addrItem = globalCache()[hostname];
-        addrItem.first = inets;
+        addrItem.first = inets_ptr;
         addrItem.second = trantor::Date::date();
     }
-    callback(inets);
+    callback(*inets_ptr);
 }
 
 void AresResolver::onSockCreate(int sockfd, int type)

trantor/net/inner/AresResolver.h

@@ -43,7 +43,7 @@ class AresResolver : public Resolver,
                 if (timeout_ == 0 ||
                     cachedAddr.second.after(timeout_) > trantor::Date::date())
                 {
-                    inet = (cachedAddr.first)[0];
+                    inet = (*cachedAddr.first)[0];
                     cached = true;
                 }
             }
@@ -75,8 +75,7 @@ class AresResolver : public Resolver,
     virtual void resolve(const std::string& hostname,
                          const ResolverResultsCallback& cb) override
     {
-        bool cached = false;
-        std::vector<trantor::InetAddress> inets;
+        std::shared_ptr<std::vector<trantor::InetAddress>> inets_ptr{nullptr};
         {
             std::lock_guard<std::mutex> lock(globalMutex());
             auto iter = globalCache().find(hostname);
@@ -86,14 +85,13 @@ class AresResolver : public Resolver,
                 if (timeout_ == 0 ||
                     cachedAddr.second.after(timeout_) > trantor::Date::date())
                 {
-                    inets = cachedAddr.first;
-                    cached = true;
+                    inets_ptr = cachedAddr.first;
                 }
             }
         }
-        if (cached)
+        if (inets_ptr)
         {
-            cb(inets);
+            cb(*inets_ptr);
             return;
         }
         if (loop_->isInLoopThread())
@@ -132,12 +130,14 @@ class AresResolver : public Resolver,
     ChannelList channels_;
     static std::unordered_map<
         std::string,
-        std::pair<std::vector<InetAddress>, trantor::Date>>&
+        std::pair<std::shared_ptr<std::vector<trantor::InetAddress>>,
+                  trantor::Date>>&
     globalCache()
     {
         static std::unordered_map<
             std::string,
-            std::pair<std::vector<InetAddress>, trantor::Date>>
+            std::pair<std::shared_ptr<std::vector<trantor::InetAddress>>,
+                      trantor::Date>>
             dnsCache;
         return dnsCache;
     }

trantor/tests/TcpServerTest.cc

@@ -17,6 +17,12 @@ int main()
     InetAddress addr(8888);
 #endif
     TcpServer server(loopThread.getLoop(), addr, "test");
+    server.setBeforeListenSockOptCallback([](int fd) {
+        std::cout << "setBeforeListenSockOptCallback:" << fd << std::endl;
+    });
+    server.setAfterAcceptSockOptCallback([](int fd) {
+        std::cout << "afterAcceptSockOptCallback:" << fd << std::endl;
+    });
     server.setRecvMessageCallback(
         [](const TcpConnectionPtr &connectionPtr, MsgBuffer *buffer) {
             // LOG_DEBUG<<"recv callback!";