Add setting chunk callback to http request

drogonframework · Feb 9, 2025 · bffc90f · bffc90f
1 parent e155df9
commit bffc90f
Show file tree

Hide file tree

Showing 6 changed files with 87 additions and 3 deletions.
diff --git a/examples/client_example/main.cc b/examples/client_example/main.cc
@@ -14,6 +14,7 @@ int nth_resp = 0;
 
 int main()
 {
+#if 1
     trantor::Logger::setLogLevel(trantor::Logger::kTrace);
     {
         auto client = HttpClient::newHttpClient("http://www.baidu.com");
@@ -75,6 +76,28 @@ int main()
         std::cout << "requestsBufferSize:" << client->requestsBufferSize()
                   << std::endl;
     }
+#else
+    {
+        auto client = HttpClient::newHttpClient("http://127.0.0.1:8848/stream");
 
+        auto req = HttpRequest::newHttpRequest();
+        req->setMethod(drogon::Get);
+        req->setPath("/stream");
+        req->setChunkCallback([](ResponseChunk ch) {
+            std::cout << "recv chunk:" << ch.data() << std::endl;
+        });
+        client->sendRequest(
+            req, [](ReqResult result, const HttpResponsePtr &response) {
+                if (result != ReqResult::Ok)
+                {
+                    std::cout
+                        << "error while sending request to server! result: "
+                        << result << std::endl;
+                    return;
+                }
+                std::cout << "recv chunk done" << std::endl;
+            });
+    }
+#endif
     app().run();
 }
diff --git a/lib/inc/drogon/HttpRequest.h b/lib/inc/drogon/HttpRequest.h
@@ -25,6 +25,8 @@
 #include <trantor/net/InetAddress.h>
 #include <trantor/net/Certificate.h>
 #include <trantor/utils/Date.h>
+#include <cstddef>
+#include <functional>
 #include <memory>
 #include <string>
 #include <unordered_map>
@@ -37,6 +39,32 @@ namespace drogon
 class HttpRequest;
 using HttpRequestPtr = std::shared_ptr<HttpRequest>;
 
+class DROGON_EXPORT ResponseChunk
+{
+  public:
+    ResponseChunk() = default;
+
+    ResponseChunk(const char *s, std::size_t count)
+    {
+        data_.append(s, count);
+    }
+
+    ~ResponseChunk() = default;
+
+    ResponseChunk(const ResponseChunk &other) = default;
+    ResponseChunk(ResponseChunk &&other) noexcept = default;
+    ResponseChunk &operator=(const ResponseChunk &other) noexcept = default;
+    ResponseChunk &operator=(ResponseChunk &&other) noexcept = default;
+
+    auto &data()
+    {
+        return data_;
+    }
+
+  private:
+    std::string data_;
+};
+
 /**
  * @brief This template is used to convert a request object to a custom
  * type object. Users must specialize the template for a particular type.
@@ -510,6 +538,12 @@ class DROGON_EXPORT HttpRequest
     virtual const std::weak_ptr<trantor::TcpConnection> &getConnectionPtr()
         const noexcept = 0;
 
+    virtual void setChunkCallback(
+        std::function<void(ResponseChunk)>) noexcept = 0;
+
+    virtual const std::function<void(ResponseChunk)> &getChunkCallback()
+        const noexcept = 0;
+
     virtual ~HttpRequest()
     {
     }

diff --git a/lib/src/HttpClientImpl.cc b/lib/src/HttpClientImpl.cc
@@ -618,7 +618,8 @@ void HttpClientImpl::onRecvMessage(const trantor::TcpConnectionPtr &connPtr,
         {
             responseParser->setForHeadMethod();
         }
-        if (!responseParser->parseResponse(msg))
+        if (!responseParser->parseResponse(msg,
+                                           firstReq.first->getChunkCallback()))
         {
             onError(ReqResult::BadResponse);
             bytesReceived_ += (msgSize - msg->readableBytes());

diff --git a/lib/src/HttpRequestImpl.h b/lib/src/HttpRequestImpl.h
@@ -579,6 +579,18 @@ class HttpRequestImpl : public HttpRequest
         return connPtr_;
     }
 
+    void setChunkCallback(
+        std::function<void(ResponseChunk)> cb) noexcept override
+    {
+        chunkCb_ = std::move(cb);
+    }
+
+    const std::function<void(ResponseChunk)> &getChunkCallback()
+        const noexcept override
+    {
+        return chunkCb_;
+    }
+
     bool isOnSecureConnection() const noexcept override
     {
         return isOnSecureConnection_;
@@ -731,6 +743,7 @@ class HttpRequestImpl : public HttpRequest
     std::exception_ptr streamExceptionPtr_;
     bool startProcessing_{false};
     std::weak_ptr<trantor::TcpConnection> connPtr_;
+    std::function<void(ResponseChunk)> chunkCb_;
 
   protected:
     std::string content_;

diff --git a/lib/src/HttpResponseParser.cc b/lib/src/HttpResponseParser.cc
@@ -17,6 +17,7 @@
 #include <trantor/utils/Logger.h>
 #include <trantor/utils/MsgBuffer.h>
 #include <algorithm>
+#include <string>
 
 using namespace trantor;
 using namespace drogon;
@@ -84,7 +85,9 @@ bool HttpResponseParser::parseResponseOnClose()
 }
 
 // return false if any error
-bool HttpResponseParser::parseResponse(MsgBuffer *buf)
+bool HttpResponseParser::parseResponse(
+    MsgBuffer *buf,
+    const std::function<void(ResponseChunk)> &chunkCallback)
 {
     bool ok = true;
     bool hasMore = true;
@@ -277,6 +280,11 @@ bool HttpResponseParser::parseResponse(MsgBuffer *buf)
                         responsePtr_->bodyPtr_ =
                             std::make_shared<HttpMessageStringBody>();
                     }
+                    if (chunkCallback)
+                    {
+                        chunkCallback(
+                            ResponseChunk(buf->peek(), currentChunkLength_));
+                    }
                     responsePtr_->bodyPtr_->append(buf->peek(),
                                                    currentChunkLength_);
                     buf->retrieve(currentChunkLength_ + 2);

diff --git a/lib/src/HttpResponseParser.h b/lib/src/HttpResponseParser.h
@@ -15,9 +15,11 @@
 #pragma once
 
 #include "impl_forwards.h"
+#include <drogon/HttpRequest.h>
 #include <trantor/utils/NonCopyable.h>
 #include <trantor/net/TcpConnection.h>
 #include <trantor/utils/MsgBuffer.h>
+#include <functional>
 #include <list>
 #include <mutex>
 
@@ -43,7 +45,10 @@ class HttpResponseParser : public trantor::NonCopyable
     // default copy-ctor, dtor and assignment are fine
 
     // return false if any error
-    bool parseResponse(trantor::MsgBuffer *buf);
+    bool parseResponse(
+        trantor::MsgBuffer *buf,
+        const std::function<void(ResponseChunk)> &chunkCallback = nullptr);
+
     bool parseResponseOnClose();
 
     bool gotAll() const