From 231d1a91292f39d42df2317a2fe4eaa96fd68f41 Mon Sep 17 00:00:00 2001
From: Eric Wong <normalperson@yhbt.net>
Date: Mon, 22 Oct 2012 20:18:06 +0000
Subject: [PATCH] new_file allows optional :info hash to be populated

This allows clients to avoid calling #file_info or #get_uris
immediate after uploading a file to MogileFS.  This can speed
things up for cache-using clients with write-through caching.
---
 lib/mogilefs/mogilefs.rb        |  8 ++++++
 lib/mogilefs/new_file/common.rb | 33 ++++++++++++++---------
 test/test_fresh.rb              | 58 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 87 insertions(+), 12 deletions(-)

diff --git a/lib/mogilefs/mogilefs.rb b/lib/mogilefs/mogilefs.rb
index f5cb096..76325cb 100644
--- a/lib/mogilefs/mogilefs.rb
+++ b/lib/mogilefs/mogilefs.rb
@@ -244,6 +244,13 @@ class MogileFS::MogileFS < MogileFS::Client
   #   Keep in mind most HTTP servers do not support HTTP trailers, so
   #   passing a String is usually the safest way to use this.
   #
+  # [:info => Hash]
+  #
+  #   This is an empty hash that will be filled the same information
+  #   MogileFS::MogileFS#file_info.
+  #
+  #   Additionally, it contains one additional key: :uris,
+  #   an array of URI::HTTP objects to the stored destinations
   def new_file(key, args = nil, bytes = nil) # :yields: file
     raise MogileFS::ReadOnlyError if readonly?
     opts = { :key => key, :multi_dest => 1 }
@@ -275,6 +282,7 @@ class MogileFS::MogileFS < MogileFS::Client
     opts[:content_length] ||= bytes if bytes
     opts[:new_file_max_time] ||= @new_file_max_time
     opts[:start_time] = Time.now
+    info = opts[:info] and info["class"] = klass || "default"
 
     case (dests[0][1] rescue nil)
     when %r{\Ahttp://}
diff --git a/lib/mogilefs/new_file/common.rb b/lib/mogilefs/new_file/common.rb
index e959a21..9b6511d 100644
--- a/lib/mogilefs/new_file/common.rb
+++ b/lib/mogilefs/new_file/common.rb
@@ -34,21 +34,21 @@ module MogileFS::NewFile::Common
   end
 
   def create_close(devid, uri, bytes_uploaded)
-    args = {
-      :fid => @opts[:fid],
-      :devid => devid,
-      :key => @opts[:key],
-      :domain => @opts[:domain],
-      :size => bytes_uploaded,
-      :path => uri.to_s,
-    }
+    dest_info = @opts[:info] ||= {}
+    dest_info["fid"] = @opts[:fid].to_i
+    dest_info["key"] = @opts[:key]
+    dest_info["domain"] = @opts[:domain]
+    dest_info[:devid] = devid
+    dest_info[:path] = uri.to_s
+    dest_info[:size] = bytes_uploaded
     if @md5
-      args[:checksum] = "MD5:#{@md5.hexdigest}"
+      dest_info["checksum"] = "MD5:#{@md5.hexdigest}"
     elsif String === @opts[:content_md5]
       hex = @opts[:content_md5].unpack('m')[0].unpack('H*')[0]
-      args[:checksum] = "MD5:#{hex}"
+      dest_info["checksum"] = "MD5:#{hex}"
     end
-    args[:checksumverify] = 1 if @opts[:checksumverify]
+
+    dest_info[:checksumverify] = 1 if @opts[:checksumverify]
     backend = @opts[:backend]
 
     # upload could've taken a long time, ping and try to ensure socket
@@ -61,7 +61,16 @@ module MogileFS::NewFile::Common
     # twice will hurt us...
     backend.noop
 
-    backend.create_close(args)
+    backend.create_close(dest_info)
+
+    # make this look like file_info + get_uris
+    dest_info.delete(:checksumverify)
+    dest_info.delete(:path)
+    dest_info[:uris] = [ uri ]
+    dest_info["devcount"] = 1
+    dest_info["devids"] = [ dest_info.delete(:devid).to_i ]
+    dest_info["length"] = dest_info.delete(:size)
+
     bytes_uploaded
   end
 
diff --git a/test/test_fresh.rb b/test/test_fresh.rb
index f20aacc..369e363 100644
--- a/test/test_fresh.rb
+++ b/test/test_fresh.rb
@@ -32,4 +32,62 @@ class TestMogFresh < Test::Unit::TestCase
     client = MogileFS::MogileFS.new :hosts => @hosts, :domain => @domain
     assert_equal false, client.exist?("non-existent")
   end
+
+  def test_new_file_info(checksum = nil)
+    add_host_device_domain unless checksum
+    @client = MogileFS::MogileFS.new :hosts => @hosts, :domain => @domain
+    info = {}
+    opts = { :info => info }
+    key = "new_file_info"
+    content = "ZZZZ"
+    if checksum
+      opts[:content_md5] = [ Digest::MD5.digest(content) ].pack('m').rstrip
+      opts[:class] = "check"
+    end
+    rv = @client.new_file(key, opts) do |http_file|
+      http_file << content
+    end
+
+    uris = info.delete(:uris)
+    assert_kind_of Array, uris
+    assert_equal(uris, (@client.get_uris(key) & uris))
+    expect_info = @client.file_info(key, :devices => true)
+    match_keys = %w(class fid key domain length)
+    match_keys << "checksum" if checksum
+    match_keys.each do |field|
+      assert_equal expect_info.delete(field), info.delete(field)
+    end
+    assert_operator expect_info.delete("devcount"), :>=, info.delete("devcount")
+    devids = info.delete("devids")
+    assert_equal(devids, (expect_info.delete("devids") & devids))
+
+    assert info.empty?, info.inspect
+    assert expect_info.empty?, expect_info.inspect
+  ensure
+    @client.delete(key)
+  end
+
+  def test_new_file_info_checksum
+    add_host_device_domain
+    opts = @admin.get_domains[@domain]["default"]
+    opts["hashtype"] = "MD5"
+    @admin.create_class(@domain, "check", opts)
+    yield_for_monitor_update do
+      tmp = @admin.get_domains[@domain]["check"]
+      if tmp
+        case tmp["hashtype"]
+        when "MD5"
+          break
+        when nil
+          warn "skipping checksum test, MogileFS server too old"
+          return
+        else
+          raise "Unhandled hashtype: #{tmp['hashtype']}"
+        end
+      end
+    end
+    test_new_file_info(:md5)
+  ensure
+    @admin.delete_class @domain, "check"
+  end
 end
-- 
2.11.4.GIT