Net::SSHで一つのシェルセッションで複数のコマンドを実行する

Net::SSHの基本的な使い方

Rubyの"net-ssh" gemを使うと、SSH接続でリモートホストに繋いでコマンドを実行できる。
例えば、

require 'net/ssh'

Net::SSH.start('localhost') do |ssh|
  pp ssh.exec!("echo hello")         # "hello"
  pp ssh.exec!("echo world")         # "world"
end

単純なコマンドを実行したいだけであればこれで十分である。

リモートホストでrbenvで管理しているrubyを使いたい場合

リモートホストでrbenvやpyenvなどのツールを使っていて、そこで管理されているrubyを使いたい場合、この方法ではうまくいかない。

require 'net/ssh'

Net::SSH.start('localhost') do |ssh|
  pp ssh.exec!("which ruby")         # "/usr/bin/ruby"
end

なぜか？sshで起動されるのは非インタラクティブシェルかつ非ログインシェルなので、"~/.bashrc"や"~/.bash_porfile"が読まれないからである。
（rbenv以外にも一般的にいうと、.bashrcや.bash_profileで設定したPATHや環境変数を利用するツールはうまく動作しない）

この問題を回避するためにはbashのログインシェル経由でコマンドが実行すればよい。
-lオプションをつけるとログインシェルとして起動するため、.bash_profileが読まれる。（もしrbenvの設定を.bashrcに書いている場合は-lの代わりに-iを使う。）

require 'net/ssh'

Net::SSH.start('localhost') do |ssh|
  pp ssh.exec!("bash -l -c 'ruby --version'")     # => rbenvで設定したrubyのバージョンになっているはず。
end

（ちなみに、シェルを経由するためのgemとして"net-ssh-shell"という名前のものもある。しかし、数年以上更新されておらず最新のNet::SSHと互換性があるのかどうかもわからないので、ここでは導入をやめた。）

シェルを経由する場合の問題点

シェルを経由することで適切のPATHやその他の環境変数が設定されるようになったが、もしリモートホストで複数のコマンドを実行したい場合は、実行に時間がかかりすぎる場合がある。
各コマンドごとにshellを起動すると、その度にrbenv initが実行される。これが結構時間がかかる。典型的には2~5秒くらいかかる。
できれば一つのシェルを起動したら、そのシェルの環境を使いまわして複数のコマンドを実行したい。

コマンドが決まっていて、それぞれのコマンドの標準出力やエラー出力、リターンコードを個別に取らなくてもいい場合は簡単で、単にコマンドを繋げればよい。

require 'net/ssh'

Net::SSH.start('localhost') do |ssh|
  pp ssh.exec!("bash -l -c 'ruby 1.rb; ruby 2.rb; ruby 3.rb'")
end

しかし、この方法は標準エラー出力やリターンコードが得られないし、個別の処理の結果に応じて処理を分けることができない。
またコマンドの数が少数じゃないと、bashの引数が異常に長くなり、スマートではない。
そこで、以下のような方法を試した。

シェル上で複数のコマンドを実行する方法

StackOverflowで以下のような方法が紹介されていた。
https://stackoverflow.com/questions/5051782/ruby-net-ssh-login-shell

コードを抜粋すると

require 'net/ssh'

def execute_in_shell!(session, commands, shell="bash")
  channel = session.open_channel do |ch|
    ch.exec("#{shell} -l") do |ch2, success|
      # Set the terminal type
      ch2.send_data "export TERM=vt100\n"

      # Output each command as if they were entered on the command line
      [commands].flatten.each do |command|
        ch2.send_data "#{command}\n"
      end

      # Remember to exit or we'll hang!
      ch2.send_data "exit\n"

      # Configure to listen to ch2 data so you can grab stdout
      ch2.on_data do |c,data|
        $stdout.puts "o: #{data}"
      end

      ch2.on_extended_data do |c,type,data|
        $stderr.puts "e: #{data}"
      end
    end
  end

  # Wait for everything to complete
  channel.wait
end

Net::SSH.start('localhost') do |ssh|
  execute_in_shell!(ssh, ["echo hello", "sleep 1", "echo world", "which ruby", "echo $SHELL"])
end

SSHのセッションからchannelを作る。このchannelはリモートホストとのI/Oのやりとりをするためのオブジェクト。
Net::SSHの処理はすべて非同期処理が前提で作られている。

ch.exec("bash -l")でbashをログインシェルとして起動している。
そのシェルに対してsend_dataで入力を送り、コマンドを実行している。コマンドは文字列として与えているので、改行文字を末尾につけることが必要。またすべてのコマンドのあとで"exit"を実行してシェルを完了している。

on_dataはコマンドから標準出力が得られた時に呼ばれるコールバック関数。on_extended_dataは標準エラー出力が出力されたときに呼ばれる関数。（extended_dataは定義上は、標準エラー出力以外のものも出せるが、SSHの仕様で明示的に定義されているのはstderrのみ。typeは"1"になり、dataにstderrの出力が出てくる）
詳細についてはAPIのリファレンスを参照 http://net-ssh.github.io/net-ssh/

これでだいぶスマートに複数のコマンドを渡すことができるようにはなった。
この方法の問題点は、各コマンドの個別の実行結果を得られないことである。
複数のコマンドを改行で区切ってshellに渡していて、結果もon_dataでまとめて受けているので、どの出力がどの結果かわからない。また各コマンドのリターンコードも取れない。
以前行なったコマンドの結果に応じて次に行う処理を決めることもできない。

解決策

解決策として以下のようなコードを実装した。

require 'net/ssh'

class ShellSession

  TOKEN = "XXXDONEXXX"
  PATTERN = /XXXDONEXXX (\d+)$/   # 場合によっては行頭にTOKENがこない場合が稀に発生する。/^XXXDONEXXX (\d+)$/ではダメ。

  def initialize(channel)
    @ch = channel
  end

  def exec(command)
    @ch.send_data("#{command}\necho '#{TOKEN}' $?\n")
    Fiber.yield
  end

  def self.start!(session, shell="bash -l", debug=false)
    channel = session.open_channel do |ch|
      ch.exec(shell) do |ch2, success|
        raise "failed to open shell" unless success
        # Set the terminal type
        ch2.send_data "export TERM=vt100\necho '#{TOKEN}' $?\n"

        sh = ShellSession.new(ch2)
        f = Fiber.new do
          yield sh
          ch2.send_data("exit\n")
        end

        output = {stdout: "", stderr: "", rc: nil}

        ch2.on_data do |c,data|
          $stderr.puts "o: #{data}" if debug
          if data =~ PATTERN
            rc = $1.to_i
            $stderr.puts "rc: #{rc}" if debug
            output[:rc] = rc
            o = output
            output = {stdout: "", stderr: "", rc: nil}
            f.resume o
          else
            output[:stdout] += data
          end
        end

        ch2.on_extended_data do |c,type,data|
          $stderr.puts "e: #{data}" if debug
          output[:stderr] += data
        end
      end
    end
    channel.wait
  end
end

Net::SSH.start('localhost') do |ssh|
  ShellSession.start!(ssh) do |sh|
    pp sh.exec("echo hello")                                #=> {:stdout=>"hello\n", :stderr=>"", :rc=>0}
    pp sh.exec("echo world")                               #=> {:stdout=>"world\n", :stderr=>"", :rc=>0}
    pp sh.exec("cd /usr/local")                             #=> {:stdout=>"", :stderr=>"", :rc=>0}
    pp sh.exec("pwd")                                           #=> {:stdout=>"/usr/local\n", :stderr=>"", :rc=>0}
    pp sh.exec("ruby --version; which ruby")    #=> {:stdout=> "ruby 2.5.1p57 (2018-03-29 revision 63029) [x86_64-darwin16]\n/Users/user/.rbenv/shims/ruby\n", :stderr=>"", :rc=>0}
    pp sh.exec("foo bar")                                     #=> {:stdout=>"", :stderr=>"bash: line 13: foo: command not found\n", :rc=>127}
  end
end

なんだかずいぶん大げさなコードになってしまったが、これでうまくいく。仕組みは以下の通り。

• 各コマンドのあとに、echo 'XXXDONEXXX' $? というコマンドを入れる。
  • on_dataのところでこのトークンにマッチしたら、そのコマンドが完了したことがわかる。またリターンコードも取得できる。
• 各コマンドの結果はoutputという名前のHashに格納される。一つのコマンドに対して複数回on_dataや_on_extended_dataが呼ばれることがあるので+=で文字列を追加している。
• 各コマンドの終わりで、outputを新しいHashにして、次のコマンドの結果を格納できるようにする。
• 全体のコードを同期的にかけるようにFiberを使って書いている。execコマンドは実行完了時に処理が返ってくる。
  • shellのプロセスは維持したまま同期的に書くのは結構難しいが、Fiberをうまく使うと同期的であるかのように書ける。Fiberを使わないとネストしたcallbackを書くはめになる。