Merge pull request #399 from mrsked/manage-ssh-connection-starts

Manage SSH connection starts

Gemfile.lock

@@ -4,6 +4,7 @@ PATH
     mrsk (0.15.1)
       activesupport (>= 7.0)
       bcrypt_pbkdf (~> 1.0)
+      concurrent-ruby (~> 1.2)
       dotenv (~> 2.8)
       ed25519 (~> 1.2)
       net-ssh (~> 7.0)

README.md

@@ -990,6 +990,20 @@ That'll post a line like the following to a preconfigured chatbot in Basecamp:
 
 Set `--skip_hooks` to avoid running the hooks.
 
+## SSH connection management
+
+Creating SSH connections concurrently can be an issue when deploying to many servers. By default MRSK will limit concurrent connection starts to 30 at a time.
+
+It also sets a long idle timeout of 900 seconds on connections to prevent re-connection storms after a long idle period, like building an image or waiting for CI.
+
+You can configure both of these settings:
+
+```yaml
+sshkit:
+  max_concurrent_starts: 10
+  pool_idle_timeout: 300
+```
+
 ## Stage of development
 
 This is beta software. Commands may still move around. But we're live in production at [37signals](https://37signals.com).

lib/mrsk/commander.rb

@@ -143,7 +143,11 @@ def hold_lock_on_error?
   private
     # Lazy setup of SSHKit
     def configure_sshkit_with(config)
-      SSHKit::Backend::Netssh.configure { |ssh| ssh.ssh_options = config.ssh_options }
+      SSHKit::Backend::Netssh.pool.idle_timeout = config.sshkit.pool_idle_timeout
+      SSHKit::Backend::Netssh.configure do |sshkit|
+        sshkit.max_concurrent_starts = config.sshkit.max_concurrent_starts
+        sshkit.ssh_options = config.ssh.options
+      end
       SSHKit.config.command_map[:docker] = "docker" # No need to use /usr/bin/env, just clogs up the logs
       SSHKit.config.output_verbosity = verbosity
     end

lib/mrsk/commands/base.rb

@@ -13,12 +13,12 @@ def initialize(config)
 
     def run_over_ssh(*command, host:)
       "ssh".tap do |cmd|
-        if config.ssh_proxy && config.ssh_proxy.is_a?(Net::SSH::Proxy::Jump)
-          cmd << " -J #{config.ssh_proxy.jump_proxies}"
-        elsif config.ssh_proxy && config.ssh_proxy.is_a?(Net::SSH::Proxy::Command)
-          cmd << " -o ProxyCommand='#{config.ssh_proxy.command_line_template}'"
+        if config.ssh.proxy && config.ssh.proxy.is_a?(Net::SSH::Proxy::Jump)
+          cmd << " -J #{config.ssh.proxy.jump_proxies}"
+        elsif config.ssh.proxy && config.ssh.proxy.is_a?(Net::SSH::Proxy::Command)
+          cmd << " -o ProxyCommand='#{config.ssh.proxy.command_line_template}'"
         end
-        cmd << " -t #{config.ssh_user}@#{host} '#{command.join(" ")}'"
+        cmd << " -t #{config.ssh.user}@#{host} '#{command.join(" ")}'"
       end
     end
 

lib/mrsk/configuration.rb

@@ -135,25 +135,12 @@ def logging_args
   end
 
 
-  def ssh_user
-    if raw_config.ssh.present?
-      raw_config.ssh["user"] || "root"
-    else
-      "root"
-    end
-  end
-
-  def ssh_proxy
-    if raw_config.ssh.present? && raw_config.ssh["proxy"]
-      Net::SSH::Proxy::Jump.new \
-        raw_config.ssh["proxy"].include?("@") ? raw_config.ssh["proxy"] : "root@#{raw_config.ssh["proxy"]}"
-    elsif raw_config.ssh.present? && raw_config.ssh["proxy_command"]
-      Net::SSH::Proxy::Command.new(raw_config.ssh["proxy_command"])
-    end
+  def ssh
+    Mrsk::Configuration::Ssh.new(config: self)
   end
 
-  def ssh_options
-    { user: ssh_user, proxy: ssh_proxy, auth_methods: [ "publickey" ] }.compact
+  def sshkit
+    Mrsk::Configuration::Sshkit.new(config: self)
   end
 
 
@@ -185,7 +172,8 @@ def to_h
       service_with_version: service_with_version,
       env_args: env_args,
       volume_args: volume_args,
-      ssh_options: ssh_options,
+      ssh_options: ssh.options,
+      sshkit: sshkit.to_h,
       builder: builder.to_h,
       accessories: raw_config.accessories,
       logging: logging_args,

lib/mrsk/configuration/ssh.rb

@@ -0,0 +1,24 @@
+class Mrsk::Configuration::Ssh
+  def initialize(config:)
+    @config = config.raw_config.ssh || {}
+  end
+
+  def user
+    config.fetch("user", "root")
+  end
+
+  def proxy
+    if (proxy = config["proxy"])
+      Net::SSH::Proxy::Jump.new(proxy.include?("@") ? proxy : "root@#{proxy}")
+    elsif (proxy_command = config["proxy_command"])
+      Net::SSH::Proxy::Command.new(proxy_command)
+    end
+  end
+
+  def options
+    { user: user, proxy: proxy, auth_methods: [ "publickey" ], keepalive: true, keepalive_interval: 30 }.compact
+  end
+
+  private
+    attr_accessor :config
+end

lib/mrsk/configuration/sshkit.rb

@@ -0,0 +1,20 @@
+class Mrsk::Configuration::Sshkit
+  def initialize(config:)
+    @options = config.raw_config.sshkit || {}
+  end
+
+  def max_concurrent_starts
+    options.fetch("max_concurrent_starts", 30)
+  end
+
+  def pool_idle_timeout
+    options.fetch("pool_idle_timeout", 900)
+  end
+
+  def to_h
+    options
+  end
+
+  private
+    attr_accessor :options
+end

lib/mrsk/sshkit_with_ext.rb

@@ -54,3 +54,51 @@ def env_for(env)
   end
   prepend CommandEnvMerge
 end
+
+class SSHKit::Backend::Netssh::Configuration
+  attr_accessor :max_concurrent_starts
+end
+
+class SSHKit::Backend::Netssh
+  module LimitConcurrentStartsClass
+    attr_reader :start_semaphore
+
+    def configure(&block)
+      super &block
+      # Create this here to avoid lazy creation by multiple threads
+      if config.max_concurrent_starts
+        @start_semaphore = Concurrent::Semaphore.new(config.max_concurrent_starts)
+      end
+    end
+  end
+
+  class << self
+    prepend LimitConcurrentStartsClass
+  end
+
+  module LimitConcurrentStartsInstance
+    private
+      def with_ssh(&block)
+        host.ssh_options = self.class.config.ssh_options.merge(host.ssh_options || {})
+        self.class.pool.with(
+          method(:start_with_concurrency_limit),
+          String(host.hostname),
+          host.username,
+          host.netssh_options,
+          &block
+        )
+      end
+
+      def start_with_concurrency_limit(*args)
+        if self.class.start_semaphore
+          self.class.start_semaphore.acquire do
+            Net::SSH.start(*args)
+          end
+        else
+          Net::SSH.start(*args)
+        end
+      end
+  end
+
+  prepend LimitConcurrentStartsInstance
+end

mrsk.gemspec

@@ -20,6 +20,7 @@ Gem::Specification.new do |spec|
   spec.add_dependency "zeitwerk", "~> 2.5"
   spec.add_dependency "ed25519", "~> 1.2"
   spec.add_dependency "bcrypt_pbkdf", "~> 1.0"
+  spec.add_dependency "concurrent-ruby", "~> 1.2"
 
   spec.add_development_dependency "debug"
   spec.add_development_dependency "mocha"

test/configuration/ssh_test.rb

@@ -0,0 +1,32 @@
+require "test_helper"
+
+class ConfigurationSshTest < ActiveSupport::TestCase
+  setup do
+    @deploy = {
+      service: "app", image: "dhh/app",
+      registry: { "username" => "dhh", "password" => "secret" },
+      env: { "REDIS_URL" => "redis://x/y" },
+      servers: [ "1.1.1.1", "1.1.1.2" ],
+      volumes: ["/local/path:/container/path"]
+    }
+
+    @config = Mrsk::Configuration.new(@deploy)
+  end
+
+  test "ssh options" do
+    assert_equal "root", @config.ssh.options[:user]
+
+    config = Mrsk::Configuration.new(@deploy.tap { |c| c.merge!(ssh: { "user" => "app" }) })
+    assert_equal "app", @config.ssh.options[:user]
+  end
+
+  test "ssh options with proxy host" do
+    config = Mrsk::Configuration.new(@deploy.tap { |c| c.merge!(ssh: { "proxy" => "1.2.3.4" }) })
+    assert_equal "[email protected]", @config.ssh.options[:proxy].jump_proxies
+  end
+
+  test "ssh options with proxy host and user" do
+    config = Mrsk::Configuration.new(@deploy.tap { |c| c.merge!(ssh: { "proxy" => "[email protected]" }) })
+    assert_equal "[email protected]", @config.ssh.options[:proxy].jump_proxies
+  end
+end

test/configuration/sshkit_test.rb

@@ -0,0 +1,27 @@
+require "test_helper"
+
+class ConfigurationSshkitTest < ActiveSupport::TestCase
+  setup do
+    @deploy = {
+      service: "app", image: "dhh/app",
+      registry: { "username" => "dhh", "password" => "secret" },
+      env: { "REDIS_URL" => "redis://x/y" },
+      servers: [ "1.1.1.1", "1.1.1.2" ],
+      volumes: ["/local/path:/container/path"]
+    }
+
+    @config = Mrsk::Configuration.new(@deploy)
+  end
+
+  test "sshkit max concurrent starts" do
+    assert_equal 30, @config.sshkit.max_concurrent_starts
+    @config = Mrsk::Configuration.new(@deploy.tap { |c| c.merge!(sshkit: { "max_concurrent_starts" => 50 }) })
+    assert_equal 50, @config.sshkit.max_concurrent_starts
+  end
+
+  test "sshkit pool idle timeout" do
+    assert_equal 900, @config.sshkit.pool_idle_timeout
+    @config = Mrsk::Configuration.new(@deploy.tap { |c| c.merge!(sshkit: { "pool_idle_timeout" => 600 }) })
+    assert_equal 600, @config.sshkit.pool_idle_timeout
+  end
+end

test/configuration_test.rb

@@ -207,23 +207,6 @@ class ConfigurationTest < ActiveSupport::TestCase
     end
   end
 
-  test "ssh options" do
-    assert_equal "root", @config.ssh_options[:user]
-
-    config = Mrsk::Configuration.new(@deploy.tap { |c| c.merge!(ssh: { "user" => "app" }) })
-    assert_equal "app", @config.ssh_options[:user]
-  end
-
-  test "ssh options with proxy host" do
-    config = Mrsk::Configuration.new(@deploy.tap { |c| c.merge!(ssh: { "proxy" => "1.2.3.4" }) })
-    assert_equal "[email protected]", @config.ssh_options[:proxy].jump_proxies
-  end
-
-  test "ssh options with proxy host and user" do
-    config = Mrsk::Configuration.new(@deploy.tap { |c| c.merge!(ssh: { "proxy" => "[email protected]" }) })
-    assert_equal "[email protected]", @config.ssh_options[:proxy].jump_proxies
-  end
-
   test "volume_args" do
     assert_equal ["--volume", "/local/path:/container/path"], @config.volume_args
   end
@@ -266,7 +249,23 @@ class ConfigurationTest < ActiveSupport::TestCase
   end
 
   test "to_h" do
-    assert_equal({ :roles=>["web"], :hosts=>["1.1.1.1", "1.1.1.2"], :primary_host=>"1.1.1.1", :version=>"missing", :repository=>"dhh/app", :absolute_image=>"dhh/app:missing", :service_with_version=>"app-missing", :env_args=>["-e", "REDIS_URL=\"redis://x/y\""], :ssh_options=>{:user=>"root", :auth_methods=>["publickey"]}, :volume_args=>["--volume", "/local/path:/container/path"], :builder=>{}, :logging=>["--log-opt", "max-size=\"10m\""], :healthcheck=>{"path"=>"/up", "port"=>3000, "max_attempts" => 7 }}, @config.to_h)
+    expected_config = \
+      { :roles=>["web"],
+        :hosts=>["1.1.1.1", "1.1.1.2"],
+        :primary_host=>"1.1.1.1",
+        :version=>"missing",
+        :repository=>"dhh/app",
+        :absolute_image=>"dhh/app:missing",
+        :service_with_version=>"app-missing",
+        :env_args=>["-e", "REDIS_URL=\"redis://x/y\""],
+        :ssh_options=>{ :user=>"root", :auth_methods=>["publickey"], keepalive: true, keepalive_interval: 30 },
+        :sshkit=>{},
+        :volume_args=>["--volume", "/local/path:/container/path"],
+        :builder=>{},
+        :logging=>["--log-opt", "max-size=\"10m\""],
+        :healthcheck=>{ "path"=>"/up", "port"=>3000, "max_attempts" => 7 }}
+
+    assert_equal expected_config, @config.to_h
   end
 
   test "min version is lower" do

test/integration/main_test.rb

@@ -51,7 +51,7 @@ class MainTest < IntegrationTest
     assert_equal "app-#{version}", config[:service_with_version]
     assert_equal [], config[:env_args]
     assert_equal [], config[:volume_args]
-    assert_equal({ user: "root", auth_methods: [ "publickey" ] }, config[:ssh_options])
+    assert_equal({ user: "root", auth_methods: [ "publickey" ], keepalive: true, keepalive_interval: 30 }, config[:ssh_options])
     assert_equal({ "multiarch" => false, "args" => { "COMMIT_SHA" => version } }, config[:builder])
     assert_equal [ "--log-opt", "max-size=\"10m\"" ], config[:logging]
     assert_equal({ "path" => "/up", "port" => 3000, "max_attempts" => 7, "cmd" => "wget -qO- http://localhost > /dev/null" }, config[:healthcheck])