Add setsockopt to TcpServer (#266)

fantasy-peak · web-flow · commit 31cb406fd96b · 2023-07-10T21:24:11.000+08:00
* use std::shared_ptr&lt;std::vector&gt; instead of std::vector for dns cache

* Add setsockopt to TcpServer
diff --git a/trantor/net/TcpServer.cc b/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
diff --git a/trantor/net/TcpServer.h b/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.
      *
diff --git a/trantor/net/inner/Acceptor.cc b/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);
diff --git a/trantor/net/inner/Acceptor.h b/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
diff --git a/trantor/net/inner/AresResolver.cc b/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)
diff --git a/trantor/net/inner/AresResolver.h b/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;
     }
diff --git a/trantor/tests/TcpServerTest.cc b/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!";

Original file line number	Diff line number	Diff line change
`@@ -54,6 +54,8 @@ Acceptor::~Acceptor()`
`54`	`54`	`void Acceptor::listen()`
`55`	`55`	`{`
`56`	`56`	`loop_->assertInLoopThread();`
	`57`	`+ if (beforeListenSetSockOptCallback_)`
	`58`	`+ beforeListenSetSockOptCallback_(sock_.fd());`
`57`	`59`	`sock_.listen();`
`58`	`60`	`acceptChannel_.enableReading();`
`59`	`61`	`}`
`@@ -64,6 +66,8 @@ void Acceptor::readCallback()`
`64`	`66`	`int newsock = sock_.accept(&peer);`
`65`	`67`	`if (newsock >= 0)`
`66`	`68`	`{`
	`69`	`+ if (afterAcceptSetSockOptCallback_)`
	`70`	`+ afterAcceptSetSockOptCallback_(newsock);`
`67`	`71`	`if (newConnectionCallback_)`
`68`	`72`	`{`
`69`	`73`	`newConnectionCallback_(newsock, peer);`
Original file line number	Diff line number	Diff line change
`@@ -171,7 +171,7 @@ void AresResolver::onQueryResult(int status,`
`171`	`171`	`const ResolverResultsCallback& callback)`
`172`	`172`	`{`
`173`	`173`	`LOG_TRACE << "onQueryResult " << status;`
`174`		`- std::vector<trantor::InetAddress> inets;`
	`174`	`+ auto inets_ptr = std::make_shared<std::vector<trantor::InetAddress>>();`
`175`	`175`	`if (result)`
`176`	`176`	`{`
`177`	`177`	`auto pptr = (struct in_addr**)result->h_addr_list;`
`@@ -182,25 +182,25 @@ void AresResolver::onQueryResult(int status,`
`182`	`182`	`addr.sin_family = AF_INET;`
`183`	`183`	`addr.sin_port = 0;`
`184`	`184`	`addr.sin_addr = reinterpret_cast<in_addr>(*pptr);`
`185`		`- inets.emplace_back(trantor::InetAddress{addr});`
	`185`	`+ inets_ptr->emplace_back(trantor::InetAddress{addr});`
`186`	`186`	`}`
`187`	`187`	`}`
`188`		`- if (inets.empty())`
	`188`	`+ if (inets_ptr->empty())`
`189`	`189`	`{`
`190`	`190`	`struct sockaddr_in addr;`
`191`	`191`	`memset(&addr, 0, sizeof addr);`
`192`	`192`	`addr.sin_family = AF_INET;`
`193`	`193`	`addr.sin_port = 0;`
`194`	`194`	`InetAddress inet(addr);`
`195`		`- inets.emplace_back(std::move(inet));`
	`195`	`+ inets_ptr->emplace_back(std::move(inet));`
`196`	`196`	`}`
`197`	`197`	`{`
`198`	`198`	`std::lock_guard<std::mutex> lock(globalMutex());`
`199`	`199`	`auto& addrItem = globalCache()[hostname];`
`200`		`- addrItem.first = inets;`
	`200`	`+ addrItem.first = inets_ptr;`
`201`	`201`	`addrItem.second = trantor::Date::date();`
`202`	`202`	`}`
`203`		`- callback(inets);`
	`203`	`+ callback(*inets_ptr);`
`204`	`204`	`}`
`205`	`205`
`206`	`206`	`void AresResolver::onSockCreate(int sockfd, int type)`
Original file line number	Diff line number	Diff line change
`@@ -43,7 +43,7 @@ class AresResolver : public Resolver,`
`43`	`43`	`if (timeout_ == 0 \|\|`
`44`	`44`	`cachedAddr.second.after(timeout_) > trantor::Date::date())`
`45`	`45`	`{`
`46`		`- inet = (cachedAddr.first)[0];`
	`46`	`+ inet = (*cachedAddr.first)[0];`
`47`	`47`	`cached = true;`
`48`	`48`	`}`
`49`	`49`	`}`
`@@ -75,8 +75,7 @@ class AresResolver : public Resolver,`
`75`	`75`	`virtual void resolve(const std::string& hostname,`
`76`	`76`	`const ResolverResultsCallback& cb) override`
`77`	`77`	`{`
`78`		`- bool cached = false;`
`79`		`- std::vector<trantor::InetAddress> inets;`
	`78`	`+ std::shared_ptr<std::vector<trantor::InetAddress>> inets_ptr{nullptr};`
`80`	`79`	`{`
`81`	`80`	`std::lock_guard<std::mutex> lock(globalMutex());`
`82`	`81`	`auto iter = globalCache().find(hostname);`
`@@ -86,14 +85,13 @@ class AresResolver : public Resolver,`
`86`	`85`	`if (timeout_ == 0 \|\|`
`87`	`86`	`cachedAddr.second.after(timeout_) > trantor::Date::date())`
`88`	`87`	`{`
`89`		`- inets = cachedAddr.first;`
`90`		`- cached = true;`
	`88`	`+ inets_ptr = cachedAddr.first;`
`91`	`89`	`}`
`92`	`90`	`}`
`93`	`91`	`}`
`94`		`- if (cached)`
	`92`	`+ if (inets_ptr)`
`95`	`93`	`{`
`96`		`- cb(inets);`
	`94`	`+ cb(*inets_ptr);`
`97`	`95`	`return;`
`98`	`96`	`}`
`99`	`97`	`if (loop_->isInLoopThread())`
`@@ -132,12 +130,14 @@ class AresResolver : public Resolver,`
`132`	`130`	`ChannelList channels_;`
`133`	`131`	`static std::unordered_map<`
`134`	`132`	`std::string,`
`135`		`- std::pair<std::vector<InetAddress>, trantor::Date>>&`
	`133`	`+ std::pair<std::shared_ptr<std::vector<trantor::InetAddress>>,`
	`134`	`+ trantor::Date>>&`
`136`	`135`	`globalCache()`
`137`	`136`	`{`
`138`	`137`	`static std::unordered_map<`
`139`	`138`	`std::string,`
`140`		`- std::pair<std::vector<InetAddress>, trantor::Date>>`
	`139`	`+ std::pair<std::shared_ptr<std::vector<trantor::InetAddress>>,`
	`140`	`+ trantor::Date>>`
`141`	`141`	`dnsCache;`
`142`	`142`	`return dnsCache;`
`143`	`143`	`}`