new_file: add :largefile => :tempfile support

This is the most-compatible way to support largefiles
with the new_file interface.  This will unfortunately
generate disk I/O, though...

Also moving the "mog" util to use this with the "tee"
subcommand since it already included its own private
implementation of this before.

3 files changed, 42 insertions, 35 deletions

diff --git a/bin/mog b/bin/mog
index b4acd18..4121459 100755
--- a/bin/mog
+++ b/bin/mog
@@ -43,6 +43,7 @@ config_file = nil
 ls_l = false
 ls_h = false
 chunk = false
+range = false
 test = {}
 cat = { :raw => false }
 
@@ -70,6 +71,7 @@ ARGV.options do |x|
   x.on('-h', '--human-readable',
        "print sizes in human-readable format (`ls' command)") { ls_h = true }
   x.on('--chunk', "chunk uploads (`tee' command)") { chunk = true }
+  x.on('--range', "stream partial uploads (`tee' command)") { range = true }
   x.separator ''
   x.on('--help', 'Show this help message.') { puts x; exit }
   x.on('--version', 'Show --version') { puts "#$0 #{MogileFS::VERSION}"; exit }
@@ -201,49 +203,27 @@ begin
     end
     exit(ok)
   when 'tee'
-    require 'tempfile'
+    abort "--range and --chunk are incompatible" if range && chunk
     dkey = ARGV.shift or raise ArgumentError, '<key>'
     ARGV.shift and raise ArgumentError, '<key>'
     cfg[:noclobber] && mg.exist?(dkey) and
       abort "`#{dkey}' already exists and -n/--no-clobber was specified"
     skip_tee = File.stat('/dev/null') == $stdout.stat
+    largefile = :tempfile
+    largefile = :content_range if range
+    largefile = :chunked if chunk
 
-    if chunk
-      if skip_tee
-        tee_obj = $stdin
-      else
-        tee_obj = lambda do |*args|
-          buf = $stdin.readpartial(*args)
-          $stdout.write(buf)
-          buf
-        end
-        class << tee_obj
-          alias readpartial call
-        end
-      end
-      mg.store_file(dkey, cfg[:class], tee_obj)
-    else # buffer input, first
-      tmp = Tempfile.new('mog-tee')
-      tmp.sync = true
-
-      # if stdout is pointing to /dev/null, don't bother installing the filter.
-      tee_obj = tmp
-      unless skip_tee
-        tee_obj = lambda do |buf|
-          $stdout.write(buf)
-          tmp.write(buf)
-        end
-        class << tee_obj
-          alias write call
-        end
-      end
+    io = mg.new_file(dkey, :class => cfg[:class], :largefile => largefile)
+    begin
+      buf = $stdin.readpartial(16384)
       begin
-        MogileFS.io.copy_stream($stdin, tee_obj)
-        store_file_retry(mg, dkey, cfg[:class], tmp.path)
-      ensure
-        tmp.close!
-      end
+        io.write(buf)
+        $stdout.write(buf) unless skip_tee
+        $stdin.readpartial(16384, buf)
+      end while true
+    rescue EOFError
     end
+    io.close
   when 'test'
     truth, ok = true, nil
     raise ArgumentError, "-e must be specified" unless (test.size == 1)
diff --git a/lib/mogilefs/http_tempfile.rb b/lib/mogilefs/http_tempfile.rb
new file mode 100644
index 0000000..5019522
--- /dev/null
+++ b/lib/mogilefs/http_tempfile.rb
@@ -0,0 +1,24 @@
+# -*- encoding: binary -*-
+# here are internal implementation details, do not rely on them in your code
+require 'tempfile'
+require 'mogilefs/http_file'
+
+class MogileFS::HTTPTempfile < Tempfile
+  def initialize(*args)
+    @mogilefs_httpfile_args = args
+    super("mogilefs-client")
+    unlink
+  end
+
+  def commit
+    rewind
+    tmp = MogileFS::HTTPFile.new(*@mogilefs_httpfile_args)
+    tmp.big_io = to_io
+    tmp.commit
+  end
+
+  def close
+    commit
+    super
+  end
+end
diff --git a/lib/mogilefs/mogilefs.rb b/lib/mogilefs/mogilefs.rb
index f536a57..40927d4 100644
--- a/lib/mogilefs/mogilefs.rb
+++ b/lib/mogilefs/mogilefs.rb
@@ -164,6 +164,9 @@ class MogileFS::MogileFS < MogileFS::Client
       http_file = case opts[:largefile]
                   when :chunked
                     MogileFS::HTTPStream
+                  when :tempfile
+                    require 'mogilefs/http_tempfile'
+                    MogileFS::HTTPTempfile
                   when :content_range
                     require 'mogilefs/http_range_put'
                     MogileFS::HTTPRangePut