list_keys/each_key: better handling of verbose listings pipeline

This allows us to implement "mog ls -l" much more efficiently

2 files changed, 59 insertions, 65 deletions

diff --git a/bin/mog b/bin/mog
index 93434bc..aa20f08 100755
--- a/bin/mog
+++ b/bin/mog
@@ -115,6 +115,18 @@ def store_file_retry(mg, key, storage_class, filepath)
   end
 end
 
+def human_size(size)
+  suff = ''
+  %w(K M G).each do |s|
+    size /= 1024.0
+    if size <= 1024
+      suff = s
+      break
+    end
+  end
+  sprintf("%.1f%s", size, suff)
+end
+
 begin
   case cmd
   when 'cp'
@@ -133,29 +145,16 @@ begin
     end
   when 'ls'
     prefixes = ARGV.empty? ? [ nil ] : ARGV
-    prefixes.each do |prefix|
-      mg.each_key(prefix) do |key|
-        if ls_l
-          path_nr = "% 2d" % mg.get_paths(key).size
-          size = mg.size(key)
-          if ls_h && size > 1024
-            suff = ''
-            %w(K M G).each do |s|
-              size /= 1024.0
-              suff = s
-              break if size <= 1024
-            end
-            size = sprintf("%.1f%s", size, suff)
-          else
-            size = size.to_s
-          end
-          size = (' ' * (12 - size.length)) << size # right justify
-          puts [ path_nr, size, key ].pack("A4 A16 A*")
-        else
-          puts key
-        end
+    if ls_l
+      each_key = lambda do |key, size, devcount|
+        size = ls_h && size > 1024 ? human_size(size) : size.to_s
+        size = (' ' * (12 - size.length)) << size # right justify
+        puts [ sprintf("% 2d", devcount), size, key ].pack("A4 A16 A*")
       end
+    else
+      each_key = lambda { |key| puts key }
     end
+    prefixes.each { |prefix| mg.each_key(prefix, &each_key) }
   when 'rm'
     ARGV.empty? and raise ArgumentError, '<key1> [<key2>]'
     ARGV.each { |key| mg.delete(key) }
diff --git a/lib/mogilefs/mogilefs.rb b/lib/mogilefs/mogilefs.rb
index a40c519..55b5681 100644
--- a/lib/mogilefs/mogilefs.rb
+++ b/lib/mogilefs/mogilefs.rb
@@ -56,16 +56,11 @@ class MogileFS::MogileFS < MogileFS::Client
   end
 
   # Enumerates keys, limited by optional +prefix+
-  def each_key(prefix = "")
+  def each_key(prefix = "", &block)
     after = nil
-
-    keys, after = list_keys prefix
-
-    until keys.nil? or keys.empty? do
-      keys.each { |k| yield k }
-      keys, after = list_keys prefix, after
-    end
-
+    begin
+      keys, after = list_keys(prefix, after, 1000, &block)
+    end while keys && keys[0]
     nil
   end
 
@@ -230,53 +225,53 @@ class MogileFS::MogileFS < MogileFS::Client
 
   # Lists keys starting with +prefix+ following +after+ up to +limit+.  If
   # +after+ is nil the list starts at the beginning.
-  def list_keys(prefix = "", after = nil, limit = 1000)
-    if @backend.respond_to?(:_list_keys)
-      block_given? or return @backend._list_keys(domain, prefix, after, limit)
-      return @backend._list_keys(domain, prefix, after, limit) do |*a|
-        yield(*a)
-      end
-    end
+  def list_keys(prefix = "", after = nil, limit = 1000, &block)
+    @backend.respond_to?(:_list_keys) and
+      return @backend._list_keys(domain, prefix, after, limit, &block)
 
-    res = begin
-      @backend.list_keys(:domain => domain, :prefix => prefix,
-                         :after => after, :limit => limit)
+    begin
+      res = @backend.list_keys(:domain => domain, :prefix => prefix,
+                               :after => after, :limit => limit)
     rescue MogileFS::Backend::NoneMatchError
       return
     end
 
     keys = (1..res['key_count'].to_i).map { |i| res["key_#{i}"] }
-    if block_given?
-      # emulate the MogileFS::Mysql interface, slowly...
-      begin
-        opts = { :domain => @domain }
-        keys.each do |key|
-          opts[:key] = key
-          @backend.pipeline_dispatch(:file_info, opts) do |info|
-            if Hash === info
-              file_info_cleanup(info)
-              yield key, info["length"], info["devcount"]
-            else
-              raise info
-            end
-          end
-        end
-        @backend.pipeline_wait
-      rescue MogileFS::Backend::UnknownCommandError # MogileFS < 2.45
-        @backend.shutdown # reset the socket
-        keys.each do |key|
-          paths = get_paths(key)
-          yield key, paths_size(paths), paths.size
-        end
-      rescue
-        @backend.shutdown
-        raise
+    if block
+      if 1 == block.arity
+        keys.each { |key| block.call(key) }
+      else
+        list_keys_verbose(keys, block)
       end
     end
 
     [ keys, res['next_after'] ]
   end
 
+  def list_keys_verbose(keys, block) # :nodoc:
+    # emulate the MogileFS::Mysql interface, slowly...
+    on_file_info = lambda do |info|
+      Hash === info or raise info
+      file_info_cleanup(info)
+      block.call(info["key"], info["length"], info["devcount"])
+    end
+    opts = { :domain => @domain }
+    keys.each do |key|
+      opts[:key] = key
+      @backend.pipeline_dispatch(:file_info, opts, &on_file_info)
+    end
+    @backend.pipeline_wait
+  rescue MogileFS::Backend::UnknownCommandError # MogileFS < 2.45
+    @backend.shutdown # reset the socket
+    keys.each do |key|
+      paths = get_paths(key)
+      block.call(key, paths_size(paths), paths.size)
+    end
+  rescue
+    @backend.shutdown
+    raise
+  end
+
   # Return metadata about a file as a hash.
   # Returns the domain, class, expected length, devcount, etc.
   # Optionally device ids (not paths) can be returned as