Fix single-server set_multi key encoding

CHANGELOG.md

@@ -4,6 +4,10 @@ Dalli Changelog
 Unreleased
 ==========
 
+- Fix single-server `set_multi` storing keys that need base64 (whitespace or non-ASCII) under their base64 text,
+  where reads never find them, and raising `NoMethodError` for non-String keys. It now validates keys (`to_s`,
+  namespace, truncation) before encoding, like every other operation. (mrattle)
+
 - Read all non-pipelined responses through `Dalli::Protocol::BufferedReader`: reads wait with explicit deadlines and
   poll the socket once more before timing out, so replies delayed only by a busy fiber scheduler are not discarded.
   The reader holds one buffer per connection that grows from 8 KiB to 64 KiB once a reply outgrows it, reads up to

lib/dalli/client.rb

@@ -572,7 +572,7 @@ def pipelined_getter
     end
 
     def pipelined_setter
-      PipelinedSetter.new(ring)
+      PipelinedSetter.new(ring, @key_manager)
     end
 
     def pipelined_deleter

lib/dalli/pipelined_setter.rb

@@ -5,8 +5,9 @@ module Dalli
   # Contains logic for the pipelined sets implemented by the client.
   ##
   class PipelinedSetter
-    def initialize(ring)
+    def initialize(ring, key_manager)
       @ring = ring
+      @key_manager = key_manager
     end
 
     ##
@@ -27,7 +28,8 @@ def process(pairs, ttl, req_options = nil)
       # Note: groups_for_keys(pairs.keys) is slow, so we avoid it.
       raise 'not yet implemented' unless @ring.servers.length == 1
 
-      @ring.servers.first.request(:write_multi_storage_req, pairs, ttl, req_options)
+      validated_pairs = pairs.transform_keys { |key| @key_manager.validate_key(key.to_s) }
+      @ring.servers.first.request(:write_multi_storage_req, validated_pairs, ttl, req_options)
     rescue RetryableNetworkError => e
       Dalli.logger.debug { e.inspect }
       Dalli.logger.debug { 'retrying pipelined set because of timeout' }

lib/dalli/protocol/base.rb

@@ -24,7 +24,6 @@ def initialize(attribs, client_options = {})
         hostname, port, socket_type, @weight, user_creds = ServerConfigParser.parse(attribs)
         @options = client_options.merge(user_creds)
         @value_marshaller = ValueMarshaller.new(@options)
-        @key_manager = KeyManager.new(@options)
         @middlewares_stack = Dalli::Middlewares
         @connection_manager = ConnectionManager.new(hostname, port, socket_type, @options)
 

lib/dalli/protocol/meta.rb

@@ -81,8 +81,7 @@ def write_multi_storage_req(pairs, ttl = nil, options = {})
           total_value_bytesize = 0
           pairs.each do |key, raw_value|
             (value, bitflags) = @value_marshaller.store(key, raw_value, options)
-            encoded_key, _base64 = KeyRegularizer.encode(key)
-            encoded_key = @key_manager.validate_key(encoded_key)
+            encoded_key, base64 = KeyRegularizer.encode(key)
 
             value_bytesize = value.bytesize
             total_value_bytesize += value_bytesize
@@ -94,7 +93,7 @@ def write_multi_storage_req(pairs, ttl = nil, options = {})
             # * first chunk uses interpolated values to avoid extra allocation, but << for larger 'value' strings
             # * avoids using the request formatter pattern for single inline builder
             @connection_manager.write(
-              "ms #{encoded_key} #{value_bytesize} c F#{bitflags} T#{ttl} MS q#{routing_suffix}\r\n"
+              "ms #{encoded_key} #{value_bytesize} c#{' b' if base64} F#{bitflags} T#{ttl} MS q#{routing_suffix}\r\n"
             )
             @connection_manager.write(value)
             @connection_manager.write(TERMINATOR)

test/integration/test_concurrency.rb

@@ -51,8 +51,7 @@
 
   it 'supports multithreaded access for single server optimized' do
     memcached_persistent do |_cache, port|
-      # NOTE: we have a bug with set multi and a namespace  namespace: 'some:namspace' fails
-      cache = Dalli::Client.new("localhost:#{port}", raw: true)
+      cache = Dalli::Client.new("localhost:#{port}", raw: true, namespace: 'some:namspace')
       cache.close
       cache.flush
       workers = []

test/integration/test_namespace_and_key.rb

@@ -89,4 +89,41 @@
       assert_equal(1, dc.get(('x' * 251).to_s))
     end
   end
+
+  # Single-server rings take the pipelined set_multi path; larger rings fall back to quiet sets.
+  def with_ring(server_count, namespace)
+    memcached_persistent do |_, first_port|
+      memcached_persistent(21_355) do |_, second_port|
+        servers = ["localhost:#{first_port}", "localhost:#{second_port}"].first(server_count)
+        dc = Dalli::Client.new(servers, namespace: namespace, raw: true)
+        dc.flush
+        yield dc
+      ensure
+        dc&.close
+      end
+    end
+  end
+
+  [1, 2].product([nil, 'some:namspace']).each do |server_count, namespace|
+    it "set_multi keys match get and get_multi (#{server_count} servers, namespace #{namespace.inspect})" do
+      with_ring(server_count, namespace) do |dc|
+        pairs = { 'plain' => 'a', 'has space' => 'b', "tab\tkey" => 'c', 'naïve' => 'd', 'ключ' => 'e' }
+        dc.set_multi(pairs, 60)
+
+        assert_equal pairs, dc.get_multi(pairs.keys)
+        pairs.each { |key, value| assert_equal value, dc.get(key), key }
+      end
+    end
+  end
+
+  [1, 2].each do |server_count|
+    it "set_multi accepts the non-String keys set accepts (#{server_count} servers)" do
+      with_ring(server_count, nil) do |dc|
+        dc.set_multi({ sym: 'a', 42 => 'b' }, 60)
+
+        assert_equal 'a', dc.get(:sym)
+        assert_equal 'b', dc.get(42)
+      end
+    end
+  end
 end

test/integration/test_operations.rb

@@ -35,28 +35,6 @@
       end
     end
 
-    it 'set_multi respects namespaces when using raw single server' do
-      memcached_persistent do |_, port|
-        dc = Dalli::Client.new("localhost:#{port}", namespace: 'some:namspace', raw: true)
-        dc.close
-        dc.flush
-
-        pairs = { 'paira' => 'vala', 'pairb' => 'valb', 'pairc' => 'valc' }
-        dc.set_multi(pairs, 5)
-
-        assert_equal pairs, dc.get_multi(pairs.keys)
-      end
-    end
-
-    it 'set_multi respects namespaces on multiple servers' do
-      memcached_persistent do |dc, _port|
-        pairs = { 'paira' => 'vala', 'pairb' => 'valb', 'pairc' => 'valc' }
-        dc.set_multi(pairs, 5)
-
-        assert_equal pairs, dc.get_multi(pairs.keys)
-      end
-    end
-
     it 'return the value that include TERMINATOR on a hit' do
       memcached_persistent do |dc|
         dc.flush