From da5ff1fdabad5922a405f89c9a4da05edd1902ed Mon Sep 17 00:00:00 2001
From: Hiroshi Shirosaki <h.shirosaki@gmail.com>
Date: Tue, 28 Feb 2017 19:39:28 +0900
Subject: [PATCH] Fix open two connections with changed cache key

`args.to_s` is used for cache key of connection pooling.
`ssh_options` in `args` is changed while creating connections in
`connection_factory.call(*args)`.
Replacing cache key with changed `args` prevents cache miss.
---
 CHANGELOG.md                               |  1 +
 lib/sshkit/backends/connection_pool.rb     | 16 ++++++++++++++++
 test/unit/backends/test_connection_pool.rb | 11 +++++++++++
 3 files changed, 28 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2035ba82..b7020bcd 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ appear at the top.
 
   * Your contribution here!
   * [#390](https://github.com/capistrano/sshkit/pull/390): Properly wrap Ruby StandardError w/ add'l context - [@mattbrictson](https://github.com/mattbrictson)
+  * [#392](https://github.com/capistrano/sshkit/pull/392): Fix open two connections with changed cache key - [@shirosaki](https://github.com/shirosaki)
 
 ## [1.12.0][] (2017-02-10)
 
diff --git a/lib/sshkit/backends/connection_pool.rb b/lib/sshkit/backends/connection_pool.rb
index 3b093103..4d172549 100644
--- a/lib/sshkit/backends/connection_pool.rb
+++ b/lib/sshkit/backends/connection_pool.rb
@@ -54,6 +54,9 @@ def initialize(idle_timeout=30)
   # invoking the `connection_factory` proc with the given `args`. The arguments
   # are used to construct a key used for caching.
   def with(connection_factory, *args)
+    ssh_options = args.last
+    ssh_options_keys = ssh_options.keys if ssh_options.is_a? Hash
+    cache_key = args.to_s
     cache = find_cache(args)
     conn = cache.pop || begin
       connection_factory.call(*args)
@@ -61,6 +64,19 @@ def with(connection_factory, *args)
     yield(conn)
   ensure
     cache.push(conn) unless conn.nil?
+    update_cache_key(ssh_options, ssh_options_keys, cache_key, args)
+  end
+
+  # Update cache key with changed args to prevent cache miss
+  def update_cache_key(ssh_options, ssh_options_keys, cache_key, args)
+    return unless ssh_options_keys
+    args[args.size - 1] = ssh_options.reject { |k, _| !ssh_options_keys.include?(k) }
+    new_key = args.to_s
+    if cache_key != new_key
+      caches.synchronize do
+        caches[new_key] = caches.delete(cache_key)
+      end
+    end
   end
 
   # Immediately remove all cached connections, without closing them. This only
diff --git a/test/unit/backends/test_connection_pool.rb b/test/unit/backends/test_connection_pool.rb
index 5ebf8bc3..e6689738 100644
--- a/test/unit/backends/test_connection_pool.rb
+++ b/test/unit/backends/test_connection_pool.rb
@@ -134,6 +134,17 @@ def test_close_connections
           pool.close_connections
         end
       end
+
+      def test_connections_with_changed_args_is_reused
+        connect_change_args = ->(*args) { args[2][:a] << "a"; args[2][:b] = []; Object.new }
+        ssh_options = { :a => [] }
+        option1 = ["conn", "user", ssh_options.dup]
+        conn1 = pool.with(connect_change_args, *option1) { |c| c }
+        option2 = ["conn", "user", ssh_options.dup]
+        conn2 = pool.with(connect_change_args, *option2) { |c| c }
+
+        assert_equal conn1, conn2
+      end
     end
   end
 end