Skip to content

Commit

Permalink
Fix open two connections with changed cache key
Browse files Browse the repository at this point in the history
`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.
  • Loading branch information
shirosaki committed Mar 1, 2017
1 parent 018a3cc commit da5ff1f
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 0 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
16 changes: 16 additions & 0 deletions lib/sshkit/backends/connection_pool.rb
Original file line number Diff line number Diff line change
Expand Up @@ -54,13 +54,29 @@ 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)
end
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
Expand Down
11 changes: 11 additions & 0 deletions test/unit/backends/test_connection_pool.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit da5ff1f

Please sign in to comment.