size/list_keys improvements

We can use the file_info command to get things faster, now.

5 files changed, 57 insertions, 83 deletions

diff --git a/lib/mogilefs/mogilefs.rb b/lib/mogilefs/mogilefs.rb
index bbf8740..2b02fb7 100644
--- a/lib/mogilefs/mogilefs.rb
+++ b/lib/mogilefs/mogilefs.rb
@@ -38,7 +38,7 @@ class MogileFS::MogileFS < MogileFS::Client
   ##
   # Enumerates keys starting with +key+.
 
-  def each_key(prefix)
+  def each_key(prefix = "")
     after = nil
 
     keys, after = list_keys prefix
@@ -181,8 +181,11 @@ class MogileFS::MogileFS < MogileFS::Client
   # Returns the size of +key+.
   def size(key)
     @backend.respond_to?(:_size) and return @backend._size(domain, key)
-    paths = get_paths(key)
-    paths_size(paths)
+    begin
+      file_info(key)["length"].to_i
+    rescue MogileFS::Backend::UnknownCommandError
+      paths_size(get_paths(key))
+    end
   end
 
   def paths_size(paths)
@@ -194,25 +197,32 @@ class MogileFS::MogileFS < MogileFS::Client
   # Lists keys starting with +prefix+ follwing +after+ up to +limit+.  If
   # +after+ is nil the list starts at the beginning.
 
-  def list_keys(prefix, after = nil, limit = 1000, &block)
+  def list_keys(prefix = "", after = nil, limit = 1000)
     if @backend.respond_to?(:_list_keys)
-      return @backend._list_keys(domain, prefix, after, limit, &block)
+      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
 
     res = begin
       @backend.list_keys(:domain => domain, :prefix => prefix,
                          :after => after, :limit => limit)
     rescue MogileFS::Backend::NoneMatchError
-      return nil
+      return
     end
 
     keys = (1..res['key_count'].to_i).map { |i| res["key_#{i}"] }
     if block_given?
       # emulate the MogileFS::Mysql interface, slowly...
       keys.each do |key|
-        paths = get_paths(key)
-        length = paths_size(paths)
-        yield key, length, paths.size
+        begin
+          res = file_info(key)
+        rescue MogileFS::Backend::UnknownCommandError # MogileFS < 2.45
+          paths = get_paths(key)
+          res = { "length" => paths_size(paths), "devcount" => paths.size }
+        end
+        yield key, res["length"], res["devcount"]
       end
     end
 
@@ -231,7 +241,8 @@ class MogileFS::MogileFS < MogileFS::Client
     args and devices = args[:devices] and opts[:devices] = devices ? 1 : 0
     rv = @backend.file_info(opts)
     %w(fid length devcount).each { |f| rv[f] = rv[f].to_i }
-    devids = rv["devids"] and rv["devids"] = devids.split(/,/).map! { |x| x.to_i }
+    devids = rv["devids"] and
+      rv["devids"] = devids.split(/,/).map! { |x| x.to_i }
     rv
   end
 
diff --git a/lib/mogilefs/mysql.rb b/lib/mogilefs/mysql.rb
index 4d06bc6..4f0dde8 100644
--- a/lib/mogilefs/mysql.rb
+++ b/lib/mogilefs/mysql.rb
@@ -26,7 +26,7 @@ class MogileFS::Mysql
   ##
   # Lists keys starting with +prefix+ follwing +after+ up to +limit+.  If
   # +after+ is nil the list starts at the beginning.
-  def _list_keys(domain, prefix = '', after = '', limit = 1000, &block)
+  def _list_keys(domain, prefix = '', after = '', limit = 1000)
     # this code is based on server/lib/MogileFS/Worker/Query.pm
     dmid = get_dmid(domain)
 
diff --git a/test/test_mogilefs.rb b/test/test_mogilefs.rb
index a8ec7b5..6e9ec7f 100644
--- a/test/test_mogilefs.rb
+++ b/test/test_mogilefs.rb
@@ -208,40 +208,6 @@ class TestMogileFS__MogileFS < TestMogileFS
     assert_equal 'new_key_2', next_after
   end
 
-  def test_list_keys_block
-    @backend.list_keys = { 'key_count' => '2', 'next_after' => 'new_key_2',
-                           'key_1' => 'new_key_1', 'key_2' => 'new_key_2' }
-    http_resp = "HTTP/1.0 200 OK\r\nContent-Length: %u\r\n"
-    srv = Proc.new do |serv, port, size|
-      client, _ = serv.accept
-      client.sync = true
-      readed = client.readpartial(4096)
-      assert %r{\AHEAD } =~ readed
-      client.send(http_resp % size, 0)
-      client.close
-    end
-    t1 = TempServer.new(Proc.new { |serv, port| srv.call(serv, port, 5) })
-    t2 = TempServer.new(Proc.new { |serv, port| srv.call(serv, port, 5) })
-    t3 = TempServer.new(Proc.new { |serv, port| srv.call(serv, port, 10) })
-    @backend.get_paths = { 'paths' => '2',
-                           'path1' => "http://127.0.0.1:#{t1.port}/",
-                           'path2' => "http://127.0.0.1:#{t2.port}/" }
-    @backend.get_paths = { 'paths' => '1',
-                           'path1' => "http://127.0.0.1:#{t3.port}/" }
-
-    res = []
-    keys, next_after = @client.list_keys('new') do |key,length,devcount|
-      res << [ key, length, devcount ]
-    end
-
-    expect_res = [ [ 'new_key_1', 5, 2 ], [ 'new_key_2', 10, 1 ] ]
-    assert_equal expect_res, res
-    assert_equal ['new_key_1', 'new_key_2'], keys.sort
-    assert_equal 'new_key_2', next_after
-    ensure
-      TempServer.destroy_all!
-  end
-
   def test_new_file_http
     @client.readonly = true
     assert_raises MogileFS::ReadOnlyError do
@@ -256,44 +222,6 @@ class TestMogileFS__MogileFS < TestMogileFS
     end
   end
 
-  def test_size_http
-    accept = Tempfile.new('accept')
-    t = TempServer.new(Proc.new do |serv,port|
-      client, _ = serv.accept
-      client.sync = true
-      readed = client.recv(4096, 0) rescue nil
-      accept.syswrite('.')
-      assert_equal "HEAD /path HTTP/1.0\r\n\r\n", readed
-      client.send("HTTP/1.0 200 OK\r\nContent-Length: 5\r\n\r\n", 0)
-      client.close
-    end)
-
-    path = "http://127.0.0.1:#{t.port}/path"
-    @backend.get_paths = { 'paths' => 1, 'path1' => path }
-
-    assert_equal 5, @client.size('key')
-    assert_equal 1, accept.stat.size
-  end
-
-  def test_bad_size_http
-    tmp = Tempfile.new('accept')
-    t = TempServer.new(Proc.new do |serv,port|
-      client, _ = serv.accept
-      client.sync = true
-      readed = client.recv(4096, 0) rescue nil
-      assert_equal "HEAD /path HTTP/1.0\r\n\r\n", readed
-      tmp.syswrite('.')
-      client.send("HTTP/1.0 404 Not Found\r\nContent-Length: 5\r\n\r\n", 0)
-      client.close
-    end)
-
-    path = "http://127.0.0.1:#{t.port}/path"
-    @backend.get_paths = { 'paths' => 1, 'path1' => path }
-
-    assert_raises(MogileFS::Error) { @client.size('key') }
-    assert_equal 1, tmp.stat.size
-  end
-
   def test_store_file_small_http
     received = Tempfile.new('received')
     to_store = Tempfile.new('small')
diff --git a/test/test_mogilefs_integration.rb b/test/test_mogilefs_integration.rb
index f5915bd..47dcb37 100644
--- a/test/test_mogilefs_integration.rb
+++ b/test/test_mogilefs_integration.rb
@@ -9,6 +9,7 @@ class TestMogileFSIntegration < TestMogIntegration
 
   def test_CRUD
     assert_equal 4, @client.store_content("CRUD", "default", "DATA")
+    assert_equal 4, @client.size("CRUD")
     assert_equal "DATA", @client.get_file_data("CRUD")
     sio = StringIO.new("")
     rv = @client.get_file_data("CRUD") do |rd|
diff --git a/test/test_mogilefs_integration_list_keys.rb b/test/test_mogilefs_integration_list_keys.rb
new file mode 100644
index 0000000..bb3f94a
--- /dev/null
+++ b/test/test_mogilefs_integration_list_keys.rb
@@ -0,0 +1,34 @@
+
+# -*- encoding: binary -*-
+require './test/integration'
+
+class TestMogileFSIntegrationListKeys < TestMogIntegration
+  def setup
+    super
+    @client = MogileFS::MogileFS.new(:hosts => @trackers, :domain => @domain)
+  end
+
+  def test_list_keys
+    k = %w(a b c d e f g)
+    k.each { |x| @client.store_content("lk_#{x}", nil, x) }
+    expect = k.map { |x| "lk_#{x}" }
+    rv = @client.list_keys
+    assert_equal([ expect, expect.last ] , rv)
+    nr = 0
+    @client.list_keys do |key, length, devcount|
+      assert_equal 1, length
+      assert_kind_of Integer, devcount
+      assert_equal expect[nr], key
+      nr += 1
+    end
+  end
+
+  def test_each_key
+    9.times { |i| @client.store_content("ek_#{i}", nil, i.to_s) }
+    n = 0
+    @client.each_key do |key|
+      assert_equal "ek_#{n.to_s}", key
+      n += 1
+    end
+  end
+end