Add master side file content streaming

This patch allows the puppetmaster to serve file chunks by chunks without
ever reading the file content in RAM.
This allows serving large files directly with the master without impacting
the master memory footprint.

Signed-off-by: Brice Figureau <brice-puppet@daysofwonder.com>

3 files changed, 39 insertions, 4 deletions

diff --git a/lib/puppet/network/http/mongrel/rest.rb b/lib/puppet/network/http/mongrel/rest.rb
index 7b28d880b..8668bf802 100644
--- a/lib/puppet/network/http/mongrel/rest.rb
+++ b/lib/puppet/network/http/mongrel/rest.rb
@@ -54,12 +54,19 @@ class Puppet::Network::HTTP::MongrelREST < Mongrel::HttpHandler
         # we have a failure, unless we're on a version of mongrel that doesn't
         # support this.
         if status < 300
-            response.start(status) { |head, body| body.write(result) }
+            unless result.is_a?(File)
+                response.start(status) { |head, body| body.write(result) }
+            else
+                response.start(status) { |head, body| }
+                response.send_status(result.stat.size)
+                response.send_header
+                response.send_file(result.path)
+            end
         else
             begin
                 response.start(status,false,result) { |head, body| body.write(result) }
             rescue ArgumentError
-	        response.start(status)              { |head, body| body.write(result) }
+                response.start(status)              { |head, body| body.write(result) }
             end
         end
     end
diff --git a/lib/puppet/network/http/rack/rest.rb b/lib/puppet/network/http/rack/rest.rb
index 104751271..5245275dd 100644
--- a/lib/puppet/network/http/rack/rest.rb
+++ b/lib/puppet/network/http/rack/rest.rb
@@ -8,6 +8,24 @@ class Puppet::Network::HTTP::RackREST < Puppet::Network::HTTP::RackHttpHandler
     HEADER_ACCEPT = 'HTTP_ACCEPT'.freeze
     ContentType = 'Content-Type'.freeze
 
+    CHUNK_SIZE = 8192
+
+    class RackFile
+        def initialize(file)
+            @file = file
+        end
+
+        def each
+            while chunk = @file.read(CHUNK_SIZE)
+                yield chunk
+            end
+        end
+
+        def close
+            @file.close
+        end
+    end
+
     def initialize(args={})
         super()
         initialize_for_puppet(args)
@@ -20,7 +38,12 @@ class Puppet::Network::HTTP::RackREST < Puppet::Network::HTTP::RackHttpHandler
     # produce the body of the response
     def set_response(response, result, status = 200)
         response.status = status
-        response.write result
+        unless result.is_a?(File)
+            response.write result
+        else
+            response["Content-Length"] = result.stat.size
+            response.body = RackFile.new(result)
+        end
     end
 
     # Retrieve the accept header from the http request.
diff --git a/lib/puppet/network/http/webrick/rest.rb b/lib/puppet/network/http/webrick/rest.rb
index 274665dcd..9bfd0ce6a 100644
--- a/lib/puppet/network/http/webrick/rest.rb
+++ b/lib/puppet/network/http/webrick/rest.rb
@@ -50,7 +50,12 @@ class Puppet::Network::HTTP::WEBrickREST < WEBrick::HTTPServlet::AbstractServlet
 
     def set_response(response, result, status = 200)
         response.status = status
-        response.body          = result if status >= 200 and status != 304
+        if status >= 200 and status != 304
+            response.body = result
+            if result.is_a?(File)
+                response["content-length"] = result.stat.size
+            end
+        end
         response.reason_phrase = result if status < 200 or status >= 300
     end