Merge pull request #11561 from Jokser/s3-http-client-memory-optimization

Avoid copying whole response stream into memory in S3 HTTP client.

contrib/aws

@@ -1 +1 @@
-Subproject commit f7d9ce39f41323300044567be007c233338bb94a
+Subproject commit 17e10c0fc77f22afe890fa6d1b283760e5edaa56

src/IO/S3/PocoHTTPClient.cpp

@@ -2,6 +2,8 @@
 
 #include <utility>
 #include <IO/HTTPCommon.h>
+#include <IO/S3/PocoHTTPResponseStream.h>
+#include <IO/S3/PocoHTTPResponseStream.cpp>
 #include <aws/core/http/HttpRequest.h>
 #include <aws/core/http/HttpResponse.h>
 #include <aws/core/http/standard/StandardHttpResponse.h>
@@ -149,8 +151,7 @@ void PocoHTTPClient::MakeRequestInternal(
                 response->SetClientErrorMessage(error_message);
             }
             else
-                /// TODO: Do not copy whole stream.
-                Poco::StreamCopier::copyStream(response_body_stream, response->GetResponseBody());
+                response->GetResponseStream().SetUnderlyingStream(std::make_shared<PocoHTTPResponseStream>(session, response_body_stream));
 
             break;
         }

src/IO/S3/PocoHTTPClientFactory.cpp

@@ -21,10 +21,12 @@ std::shared_ptr<Aws::Http::HttpRequest> PocoHTTPClientFactory::CreateHttpRequest
 }
 
 std::shared_ptr<Aws::Http::HttpRequest> PocoHTTPClientFactory::CreateHttpRequest(
-    const Aws::Http::URI & uri, Aws::Http::HttpMethod method, const Aws::IOStreamFactory & streamFactory) const
+    const Aws::Http::URI & uri, Aws::Http::HttpMethod method, const Aws::IOStreamFactory &) const
 {
     auto request = Aws::MakeShared<Aws::Http::Standard::StandardHttpRequest>("PocoHTTPClientFactory", uri, method);
-    request->SetResponseStreamFactory(streamFactory);
+
+    /// Don't create default response stream. Actual response stream will be set later in PocoHTTPClient.
+    request->SetResponseStreamFactory(null_factory);
 
     return request;
 }

src/IO/S3/PocoHTTPClientFactory.h

@@ -4,22 +4,25 @@
 
 namespace Aws::Http
 {
-    class HttpClient;
-    class HttpRequest;
+class HttpClient;
+class HttpRequest;
 }
 
 namespace DB::S3
 {
-
 class PocoHTTPClientFactory : public Aws::Http::HttpClientFactory
 {
 public:
     ~PocoHTTPClientFactory() override = default;
-    [[nodiscard]] std::shared_ptr<Aws::Http::HttpClient> CreateHttpClient(const Aws::Client::ClientConfiguration & clientConfiguration) const override;
+    [[nodiscard]] std::shared_ptr<Aws::Http::HttpClient>
+    CreateHttpClient(const Aws::Client::ClientConfiguration & clientConfiguration) const override;
     [[nodiscard]] std::shared_ptr<Aws::Http::HttpRequest>
     CreateHttpRequest(const Aws::String & uri, Aws::Http::HttpMethod method, const Aws::IOStreamFactory & streamFactory) const override;
     [[nodiscard]] std::shared_ptr<Aws::Http::HttpRequest>
     CreateHttpRequest(const Aws::Http::URI & uri, Aws::Http::HttpMethod method, const Aws::IOStreamFactory & streamFactory) const override;
+
+private:
+    const Aws::IOStreamFactory null_factory = []() { return nullptr; };
 };
 
 }

src/IO/S3/PocoHTTPResponseStream.cpp

@@ -0,0 +1,12 @@
+#include "PocoHTTPResponseStream.h"
+
+#include <utility>
+
+namespace DB::S3
+{
+PocoHTTPResponseStream::PocoHTTPResponseStream(std::shared_ptr<Poco::Net::HTTPClientSession> session_, std::istream & response_stream_)
+    : Aws::IStream(response_stream_.rdbuf()), session(std::move(session_))
+{
+}
+
+}

src/IO/S3/PocoHTTPResponseStream.h

@@ -0,0 +1,21 @@
+#pragma once
+
+#include <aws/core/utils/stream/ResponseStream.h>
+#include <Poco/Net/HTTPClientSession.h>
+
+namespace DB::S3
+{
+/**
+ * Wrapper of IStream to store response stream and corresponding HTTP session.
+ */
+class PocoHTTPResponseStream : public Aws::IStream
+{
+public:
+    PocoHTTPResponseStream(std::shared_ptr<Poco::Net::HTTPClientSession> session_, std::istream & response_stream_);
+
+private:
+    /// Poco HTTP session is holder of response stream.
+    std::shared_ptr<Poco::Net::HTTPClientSession> session;
+};
+
+}