Add setsockopt to HttpServer (drogonframework#2086)

hwc0919 · Jul 2, 2024 · dfacd1b · dfacd1b
1 parent 85b918f
commit dfacd1b
Show file tree

Hide file tree

Showing 8 changed files with 102 additions and 0 deletions.
diff --git a/examples/helloworld/main.cc b/examples/helloworld/main.cc
@@ -1,3 +1,10 @@
+#ifdef _WIN32
+#include <ws2tcpip.h>
+#else
+#include <netinet/tcp.h>
+#include <sys/socket.h>
+#endif
+
 #include <drogon/drogon.h>
 using namespace drogon;
 
@@ -61,6 +68,23 @@ int main()
         },
         {Get});
 
+    app()
+        .setBeforeListenSockOptCallback([](int fd) {
+            LOG_INFO << "setBeforeListenSockOptCallback:" << fd;
+#ifdef _WIN32
+#elif __linux__
+            int enable = 1;
+            if (setsockopt(
+                    fd, IPPROTO_TCP, TCP_FASTOPEN, &enable, sizeof(enable)) ==
+                -1)
+            {
+                LOG_INFO << "setsockopt TCP_FASTOPEN failed";
+            }
+#else
+#endif
+        })
+        .setAfterAcceptSockOptCallback([](int) {});
+
     // Ask Drogon to listen on 127.0.0.1 port 8848. Drogon supports listening
     // on multiple IP addresses by adding multiple listeners. For example, if
     // you want the server also listen on 127.0.0.1 port 5555. Just add another

diff --git a/lib/inc/drogon/HttpAppFramework.h b/lib/inc/drogon/HttpAppFramework.h
@@ -1590,6 +1590,22 @@ class DROGON_EXPORT HttpAppFramework : public trantor::NonCopyable
      */
     virtual int64_t getConnectionCount() const = 0;
 
+    /**
+     * @brief Set the before listen setsockopt callback.
+     *
+     * @param cb This callback will be called before the listen
+     */
+    virtual HttpAppFramework &setBeforeListenSockOptCallback(
+        std::function<void(int)> cb) = 0;
+
+    /**
+     * @brief Set the after accept setsockopt callback.
+     *
+     * @param cb This callback will be called after accept
+     */
+    virtual HttpAppFramework &setAfterAcceptSockOptCallback(
+        std::function<void(int)> cb) = 0;
+
   private:
     virtual void registerHttpController(
         const std::string &pathPattern,

diff --git a/lib/src/HttpAppFrameworkImpl.cc b/lib/src/HttpAppFrameworkImpl.cc
@@ -1336,3 +1336,17 @@ HttpAppFramework &HttpAppFrameworkImpl::registerPreSendingAdvice(
     AopAdvice::instance().registerPreSendingAdvice(advice);
     return *this;
 }
+
+HttpAppFramework &HttpAppFrameworkImpl::setBeforeListenSockOptCallback(
+    std::function<void(int)> cb)
+{
+    listenerManagerPtr_->setBeforeListenSockOptCallback(std::move(cb));
+    return *this;
+}
+
+HttpAppFramework &HttpAppFrameworkImpl::setAfterAcceptSockOptCallback(
+    std::function<void(int)> cb)
+{
+    listenerManagerPtr_->setAfterAcceptSockOptCallback(std::move(cb));
+    return *this;
+}
diff --git a/lib/src/HttpAppFrameworkImpl.h b/lib/src/HttpAppFrameworkImpl.h
@@ -658,6 +658,11 @@ class HttpAppFrameworkImpl final : public HttpAppFramework
     HttpResponsePtr handleSessionForResponse(const HttpRequestImplPtr &req,
                                              const HttpResponsePtr &resp);
 
+    HttpAppFramework &setBeforeListenSockOptCallback(
+        std::function<void(int)> cb) override;
+    HttpAppFramework &setAfterAcceptSockOptCallback(
+        std::function<void(int)> cb) override;
+
   private:
     void registerHttpController(const std::string &pathPattern,
                                 const internal::HttpBinderBasePtr &binder,

diff --git a/lib/src/HttpServer.cc b/lib/src/HttpServer.cc
@@ -85,6 +85,14 @@ HttpServer::~HttpServer() = default;
 
 void HttpServer::start()
 {
+    if (beforeListenSetSockOptCallback_)
+    {
+        server_.setBeforeListenSockOptCallback(beforeListenSetSockOptCallback_);
+    }
+    if (afterAcceptSetSockOptCallback_)
+    {
+        server_.setAfterAcceptSockOptCallback(afterAcceptSetSockOptCallback_);
+    }
     LOG_TRACE << "HttpServer[" << server_.name() << "] starts listening on "
               << server_.ipPort();
     server_.start();

diff --git a/lib/src/HttpServer.h b/lib/src/HttpServer.h
@@ -54,6 +54,16 @@ class HttpServer : trantor::NonCopyable
         return server_.address();
     }
 
+    void setBeforeListenSockOptCallback(std::function<void(int)> cb)
+    {
+        beforeListenSetSockOptCallback_ = std::move(cb);
+    }
+
+    void setAfterAcceptSockOptCallback(std::function<void(int)> cb)
+    {
+        afterAcceptSetSockOptCallback_ = std::move(cb);
+    }
+
   private:
     friend class HttpInternalForwardHelper;
 
@@ -126,6 +136,9 @@ class HttpServer : trantor::NonCopyable
         trantor::MsgBuffer &buffer);
 
     trantor::TcpServer server_;
+
+    std::function<void(int)> beforeListenSetSockOptCallback_;
+    std::function<void(int)> afterAcceptSetSockOptCallback_;
 };
 
 class HttpInternalForwardHelper

diff --git a/lib/src/ListenerManager.cc b/lib/src/ListenerManager.cc
@@ -111,6 +111,16 @@ void ListenerManager::createListeners(
                 std::make_shared<HttpServer>(ioLoops[i],
                                              listenAddress,
                                              "drogon");
+            if (beforeListenSetSockOptCallback_)
+            {
+                serverPtr->setBeforeListenSockOptCallback(
+                    beforeListenSetSockOptCallback_);
+            }
+            if (afterAcceptSetSockOptCallback_)
+            {
+                serverPtr->setAfterAcceptSockOptCallback(
+                    afterAcceptSetSockOptCallback_);
+            }
 
             if (listener.useSSL_ && utils::supportsTls())
             {

diff --git a/lib/src/ListenerManager.h b/lib/src/ListenerManager.h
@@ -51,6 +51,16 @@ class ListenerManager : public trantor::NonCopyable
     void startListening();
     void stopListening();
 
+    void setBeforeListenSockOptCallback(std::function<void(int)> cb)
+    {
+        beforeListenSetSockOptCallback_ = std::move(cb);
+    }
+
+    void setAfterAcceptSockOptCallback(std::function<void(int)> cb)
+    {
+        afterAcceptSetSockOptCallback_ = std::move(cb);
+    }
+
   private:
     struct ListenerInfo
     {
@@ -87,6 +97,8 @@ class ListenerManager : public trantor::NonCopyable
     // should have value when and only when on OS that one port can only be
     // listened by one thread
     std::unique_ptr<trantor::EventLoopThread> listeningThread_;
+    std::function<void(int)> beforeListenSetSockOptCallback_;
+    std::function<void(int)> afterAcceptSetSockOptCallback_;
 };
 
 }  // namespace drogon