2018-03-08

capistrano/sshkitコードリーディング

capistrano (3.10.1)のSSHKitを利用する部分をコードリーディングしてみました。

#onや#run_locallyはCapistrano::DSLに定義されており、lib/capistrano/dsl.rb内で当モジュールをextendしています。

module Capistrano
  module DSL
    def on(hosts, options={}, &block)
      subset_copy = Marshal.dump(Configuration.env.filter(hosts))
      SSHKit::Coordinator.new(Marshal.load(subset_copy)).each(options, &block)
    end

    def run_locally(&block)
      SSHKit::Backend::Local.new(&block).run
    end

#onはSSHKit::Coordinator#eachを呼び出し、#run_locallyはSSHKit::Backend::Local#runを呼び出します。まずは前者の#onから追っていきます。

SSHKit::Coordinator#eachではランナーの#executeを呼び出します。

module SSHKit
  class Coordinator
    attr_accessor :hosts

    def initialize(raw_hosts)
      @raw_hosts = Array(raw_hosts)
      @hosts = @raw_hosts.any? ? resolve_hosts : []
    end

    def each(options={}, &block)
      if hosts
        options = default_options.merge(options)
        case options[:in]
        when :parallel then Runner::Parallel
        when :sequence then Runner::Sequential
        when :groups   then Runner::Group
        else
          options[:in]
        end.new(hosts, options, &block).execute
      else
        Runner::Null.new(hosts, options, &block).execute
      end
    end

    private

    def default_options
      SSHKit.config.default_runner_config
    end

    def resolve_hosts
      @raw_hosts.collect { |rh| rh.is_a?(Host) ? rh : Host.new(rh) }.uniq
    end

以下、デフォルトのSSHKit::Runner::Parallel#executeのケースを追っていきます。

SSHKit::Runner::Parallel#executeはThreadをホスト数分呼び出して、スレッド内でbackendのrunメソッドを実行しています。

module SSHKit
  module Runner
    class Parallel < Abstract
      def execute
        threads = hosts.map do |host|
          Thread.new(host) do |h|
            begin
              backend(h, &block).run
            rescue ::StandardError => e
              e2 = ExecuteError.new e
              raise e2, "Exception while executing #{host.user ? "as #{host.user}@" : "on host "}#{host}: #{e.message}"
            end
          end
        end
        threads.each(&:join)
      end

backendメソッドはSSHKit::Runner::Abstract#backendに定義されており、localでなければSSHKit.config.backend=SSHKit::Backend::Netsshのインスタンスを返します。

module SSHKit
  module Runner
    class Abstract
# ...
      private

      def backend(host, &block)
        if host.local?
          SSHKit::Backend::Local.new(&block)
        else
          SSHKit.config.backend.new(host, &block)
        end
      end

SSHKit::Backend::Netssh#runメソッドはまず、SSHKit::Backend::Abstract#runを呼び出します。#runはinstance_execでブロックを評価します。

module SSHKit
  module Backend
    class Abstract
      attr_reader :host

      def run
        Thread.current["sshkit_backend"] = self
        instance_exec(@host, &@block)
      ensure
        Thread.current["sshkit_backend"] = nil
      end

      def execute(*args)
        options = args.extract_options!
        create_command_and_execute(args, options).success?
      end

      private

      def create_command_and_execute(args, options)
        command(args, options).tap { |cmd| execute_command(cmd) }
      end

      def command(args, options)
        SSHKit::Command.new(*args, options.merge({in: pwd_path, env: @env, host: @host, user: @user, group: @group}))
      end

SSHKit::Backend::Netsshのインスタンスのコンテキストで実行されるため、 ブロック内のexecuteはSSHKit::Backend::Netsshのメソッド呼び出しになります。#executeは#create_command_and_executeを呼び出し、さらに#execute_commandを呼び出します。

SSHKit::Backend::Netssh#execute_commandはcmd.to_commandで取得したコマンド文字列をSSHで実行します。

module SSHKit
  module Backend
    class Netssh < Abstract
# ...
      def execute_command(cmd)
        output.log_command_start(cmd)
        cmd.started = true
        exit_status = nil
        with_ssh do |ssh|
          ssh.open_channel do |chan|
            chan.request_pty if Netssh.config.pty
            chan.exec cmd.to_command do |_ch, _success|
              chan.on_data do |ch, data|
                cmd.on_stdout(ch, data)
                output.log_command_data(cmd, :stdout, data)
              end
              chan.on_extended_data do |ch, _type, data|
                cmd.on_stderr(ch, data)
                output.log_command_data(cmd, :stderr, data)
              end
# ...
            end
            chan.wait
          end
          ssh.loop
        end

cmd変数はSSHKit::Backend::Abstract#command経由で生成されたSSHKit::Commandのインスタンスです。

module SSHKit
  class Command
# ...
    def to_command
      return command.to_s unless should_map?
      within do
        umask do
          with do
            user do
              in_background do
                group do
                  to_s
                end
              end
            end
          end
        end
      end
    end

    def to_s
      if should_map?
        [SSHKit.config.command_map[command.to_sym], *Array(args)].join(' ')
      else
        command.to_s
      end
    end

SSHKit::Command#to_commandのネスト部分は以下のような処理になっています

最後のコマンド文字列化では、SSHKit::CommandMapのインスタンスであるSSHKit.config.command_map経由で第一引数を変換し、第二引数以降の文字列をArray#joinで連結しています。第一引数の変換によって具体的なパス指定やbundlerなどのprefixを自動付与することを実現しています。
module SSHKit
  class CommandMap
# ...
    def [](command)
      if prefix[command].any?
        prefixes = prefix[command].map(&TO_VALUE)
        prefixes = prefixes.join(" ")

        "#{prefixes} #{command}"
      else
        TO_VALUE.(@map[command])
      end
    end

    def prefix
      @prefix ||= PrefixProvider.new
    end

    def []=(command, new_command)
      @map[command] = new_command
    end

おまけ

ちなみにwithinやwithなどのメソッドもSSHKit::Backend::Abstractで定義されています。#withinの場合は@pwdに設定し、#executeのinオプションに入れることでカレントディレクトリの操作をしています。#withも同様に@envによって環境変数の操作を行います。
      def within(directory, &_block)
        (@pwd ||= []).push directory.to_s
        execute <<-EOTEST, verbosity: Logger::DEBUG
          if test ! -d #{File.join(@pwd)}
            then echo "Directory does not exist '#{File.join(@pwd)}'" 1>&2
            false
          fi
          EOTEST
          yield
      ensure
        @pwd.pop
      end

      def with(environment, &_block)
        env_old = (@env ||= {})
        @env = env_old.merge environment
        yield
      ensure
        @env = env_old
      end

#makeや#rakeも定義されており、単純にexecuteのショートカットになっています。

      def make(commands=[])
        execute :make, commands
      end

      def rake(commands=[])
        execute :rake, commands
      end

#captureは#executeとほぼ同じですが、SSHKit::Command#full_stdoutによって標準出力を取得して返しています。

      def capture(*args)
        options = { verbosity: Logger::DEBUG, strip: true }.merge(args.extract_options!)
        result = create_command_and_execute(args, options).full_stdout
        options[:strip] ? result.strip : result
      end
このエントリーをはてなブックマークに追加