From a06be0b93b4fc25ef25e52b374e1d2fef4854973 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Janko=20Marohni=C4=87?= <janko.marohnic@gmail.com>
Date: Fri, 11 May 2018 14:47:17 +0200
Subject: [PATCH] Accept output buffer in Response::Body#readpartial

This makes it possible to use a `HTTP::Response::Body` object as a
source argument in `IO.copy_stream`, allowing easy streaming.

  response = HTTP.get("http://example.com/download")
  IO.copy_stream(response.body, "destination.txt")

Since `IO.copy_stream` uses an output buffer, it makes retrieving the
response body much more memory efficient, because we can now deallocate
retrieved chunks.
---
 .travis.yml                         |  2 +-
 lib/http/connection.rb              |  3 ++-
 lib/http/response/body.rb           | 16 +++++++++++++---
 spec/lib/http/response/body_spec.rb | 24 +++++++++++++++++++++++-
 4 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 92f9549d..6a778e2f 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -17,7 +17,7 @@ env: JRUBY_OPTS="$JRUBY_OPTS --debug"
 
 rvm:
   # Include JRuby first because it takes the longest
-  - jruby-9.1.16.0
+  - jruby-9.2.0.0
   - 2.3
   - 2.4
   - 2.5
diff --git a/lib/http/connection.rb b/lib/http/connection.rb
index 15678cf2..d53aaf1e 100644
--- a/lib/http/connection.rb
+++ b/lib/http/connection.rb
@@ -86,6 +86,7 @@ def send_request(req)
     def readpartial(size = BUFFER_SIZE)
       return unless @pending_response
 
+      # return chunk of body that was retrieved when reading response headers
       chunk = @parser.read(size)
       return chunk if chunk
 
@@ -93,7 +94,7 @@ def readpartial(size = BUFFER_SIZE)
       chunk    = @parser.read(size)
       finish_response if finished
 
-      chunk.to_s
+      chunk
     end
 
     # Reads data from socket up until headers are loaded
diff --git a/lib/http/response/body.rb b/lib/http/response/body.rb
index 8acc368b..33e68bb1 100644
--- a/lib/http/response/body.rb
+++ b/lib/http/response/body.rb
@@ -24,11 +24,21 @@ def initialize(stream, encoding: Encoding::BINARY)
         @encoding   = find_encoding(encoding)
       end
 
-      # (see HTTP::Client#readpartial)
-      def readpartial(*args)
+      # (see HTTP::Connection#readpartial)
+      def readpartial(length = nil, outbuf = nil)
         stream!
-        chunk = @stream.readpartial(*args)
+        chunk = @stream.readpartial(*length)
         chunk.force_encoding(@encoding) if chunk
+
+        if outbuf
+          outbuf.clear.force_encoding(@encoding)
+          raise EOFError if chunk.nil? # IO.copy_stream expects this to be raised
+          outbuf << chunk
+          chunk.clear
+          outbuf
+        else
+          chunk
+        end
       end
 
       # Iterate over the body, allowing it to be enumerable
diff --git a/spec/lib/http/response/body_spec.rb b/spec/lib/http/response/body_spec.rb
index 663eae43..f52b27b8 100644
--- a/spec/lib/http/response/body_spec.rb
+++ b/spec/lib/http/response/body_spec.rb
@@ -6,7 +6,6 @@
 
   before do
     allow(connection).to receive(:readpartial) { chunks.shift }
-    allow(connection).to receive(:body_completed?) { chunks.empty? }
   end
 
   subject(:body) { described_class.new(connection, :encoding => Encoding::UTF_8) }
@@ -42,6 +41,29 @@
       end
     end
 
+    context "with size and output buffer given" do
+      it "fills the output buffer" do
+        expect(connection).to receive(:readpartial).with(3).and_return(String.new("foo"))
+        outbuf = String.new
+        chunk = body.readpartial 3, outbuf
+        expect(outbuf).to eq "foo"
+        expect(chunk).to equal outbuf
+      end
+
+      it "returns nil when there is no more content" do
+        outbuf = String.new
+        expect(body.readpartial(7, outbuf)).to eq "Hello, "
+        expect(body.readpartial(6, outbuf)).to eq "World!"
+        expect { body.readpartial(1, outbuf) }.to raise_error(EOFError)
+      end
+
+      it "supports IO.copy_stream" do
+        stringio = StringIO.new
+        IO.copy_stream(body, stringio)
+        expect(stringio.string).to eq "Hello, World!"
+      end
+    end
+
     it "returns content in specified encoding" do
       body = described_class.new(connection)
       expect(connection).to receive(:readpartial).