list_keys/each_key: better handling of verbose listings pipeline

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

1 files changed, 39 insertions, 44 deletions

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