Allow 'store_content' to be used with a streamable object

New way to call 'store_content' with a
MogileFS::Util::StoreContent allows you to roll your own method
of streaming data to mogile on an upload (instead of using a
string or file)

[ew: this still requires a known content length beforehand]
[ew: applied with --whitespace=strip, rewritten subject,
     80-column wrapping]

Signed-off-by: Eric Wong <normalperson@yhbt.net>

3 files changed, 26 insertions, 2 deletions

diff --git a/lib/mogilefs/httpfile.rb b/lib/mogilefs/httpfile.rb
index 593a494..c3509ec 100644
--- a/lib/mogilefs/httpfile.rb
+++ b/lib/mogilefs/httpfile.rb
@@ -44,6 +44,8 @@ class MogileFS::HTTPFile < StringIO
 
   attr_accessor :big_io
 
+  attr_accessor :streaming_io
+
   ##
   # Works like File.open.  Use MogileFS::MogileFS#new_file instead of this
   # method.
@@ -73,6 +75,7 @@ class MogileFS::HTTPFile < StringIO
     @klass = klass
     @key = key
     @big_io = nil
+    @streaming_io = nil
 
     @dests = dests
     @tried = {}
@@ -88,7 +91,14 @@ class MogileFS::HTTPFile < StringIO
     sock = Socket.mogilefs_new(uri.host, uri.port)
     sock.mogilefs_tcp_cork = true
 
-    if @big_io
+    if @streaming_io
+      file_size = @streaming_io.length
+      syswrite_full(sock, "PUT #{uri.request_uri} HTTP/1.0\r\n" \
+                          "Content-Length: #{file_size}\r\n\r\n")
+      @streaming_io.call(Proc.new do |data_to_write|
+        syswrite_full(sock, data_to_write)
+      end)
+    elsif @big_io
       # Don't try to run out of memory
       File.open(@big_io, "rb") do |fp|
         file_size = fp.stat.size
diff --git a/lib/mogilefs/mogilefs.rb b/lib/mogilefs/mogilefs.rb
index f95cb1c..737f468 100644
--- a/lib/mogilefs/mogilefs.rb
+++ b/lib/mogilefs/mogilefs.rb
@@ -161,7 +161,11 @@ class MogileFS::MogileFS < MogileFS::Client
     raise MogileFS::ReadOnlyError if readonly?
 
     new_file key, klass do |mfp|
-      mfp << content
+      if content.is_a?(MogileFS::Util::StoreContent)
+        mfp.streaming_io = content
+      else
+        mfp << content
+      end
     end
 
     content.length
diff --git a/lib/mogilefs/util.rb b/lib/mogilefs/util.rb
index b9d048c..deaf6b0 100644
--- a/lib/mogilefs/util.rb
+++ b/lib/mogilefs/util.rb
@@ -67,6 +67,16 @@ module MogileFS::Util
     # should never get here
   end
 
+  class StoreContent < Proc
+    def initialize(total_size, &writer_proc)
+      @total_size = total_size
+      super(&writer_proc)
+    end
+    def length
+      @total_size
+    end
+  end
+
 end
 
 require 'timeout'