pty-keeper

#!/usr/bin/env ruby
require "pty"
require "optparse"
require "tmpdir"

module PtyKeeper
  module Core
    def parse!(argv)
      options = {}

      opt = ::OptionParser.new
      opt.banner = "Usage: pty-keeper [options] command [-- command_args]"
      opt.on("-d", "--debug") { options[:debug] = true }
      opt.on("-D", "--no-daemon") { options[:no_daemon] = true }
      opt.on("-pp PPIDFILE", "--ppidfile=PPIDFILE") {|val| options[:ppidfile] = File.expand_path(val) }
      opt.on("-p PIDFILE", "--pidfile=PIDFILE") {|val| options[:pidfile] = File.expand_path(val) }
      opt.on("-cp CPIDFILE", "--cpidfile=CPIDFILE") {|val| options[:cpidfile] = File.expand_path(val) }
      opt.on("-i INTERVAL", "--interval=INTERVAL") {|val| options[:interval] = val.to_f }
      opt.on("-s", "--sigstop") { options[:sigstop] = true }
      opt.parse! argv

      if argv.size < 1
        $stderr.puts "command is required"
        $stderr.puts opt.help
        exit 1
      end
      if options[:debug]
        tmpdir = ENV['TMPDIR'] || ENV['TMP'] || ENV['TEMP'] || (Dir.exist?("/tmp") ? "/tmp" : Dir.mktmpdir)
        options[:ppidfile] ||= tmpdir + "/" + File.basename(argv[0]) + ".ppid"
        options[:pidfile] ||= tmpdir + "/" + File.basename(argv[0]) + ".pid"
        options[:cpidfile] ||= tmpdir + "/" + File.basename(argv[0]) + ".cpid"
      end
      options[:interval] ||= 1
      options[:command] = argv

      options
    end

    def dp(*strs)
      p strs if @options[:debug]
    end

    def ep(e)
      if @options[:debug]
        $stderr.puts e.inspect
        $stderr.puts caller(2)
      end
      stop 1
    end

    def start(argv)
      @options = parse!(argv.dup)
      dp @options
      spawn
      set_trap
      daemonize
      create_pidfiles
      Process.kill(:STOP, Process.ppid) if @options[:sigstop]
      loop do
        check
        sleep @options[:interval]
      end
    end

    def spawn
      @ppid = Process.ppid
      @r, @w, @cpid = PTY.spawn(*@options[:command])
      @alive = true
    end

    def daemonize
      unless @options[:no_daemon]
        if @options[:debug]
          Process.daemon(true, true)
        else
          Process.daemon(true)
        end
      end
      @pid = Process.pid
    end

    def create_pidfile file, pid
      dp "create file:", file
      File.open(file, "w"){|f| f.write pid.to_s } if file
    rescue => e
      ep e
    end

    def create_pidfiles
      create_pidfile(@options[:ppidfile], @ppid)
      create_pidfile(@options[:pidfile], @pid)
      create_pidfile(@options[:cpidfile], @cpid)
    end

    def set_trap
      Signal.trap("TERM") do
        dp "catch signal:", "TERM"
        stop
      end
      %w|CHLD CLD|.each do |sig|
        Signal.trap(sig) do
          dp "catch signal:", sig
          status = PTY.check(@cpid)
          if status.is_a? Process::Status
            dp status.inspect
            if status.exited?
              @alive = false
              stop 1
            end
          end
        end
      end
      %w|INT QUIT|.each do |sig|
        Signal.trap(sig) do
          dp "catch signal:", sig
          stop 1
        end
      end
    rescue => e
      ep e
    end

    def check
      dp :checking,
      Process.kill(0, @cpid)
    rescue Errno::ESRCH => e
      dp e.inspect
      @alive = false
      stop 1
    rescue => e
      ep e
    end

    def close
      @r.close
      @w.close
    rescue => e
    end

    def wait
      if @alive
        Process.kill(:TERM, @cpid)
        Process.waitpid(@cpid)
      end
    rescue => e
    end

    def delete_pidfile file
      File.delete(file) if file
    rescue => e
    end

    def delete_pidfiles
      delete_pidfile(@options[:ppidfile])
      delete_pidfile(@options[:pidfile])
      delete_pidfile(@options[:cpidfile])
    end

    def stop code=0
      unless @stopping
        @stopping = true
	close
        wait
        delete_pidfiles
        exit code
      end
    end
  end
  extend Core
end

PtyKeeper.start(ARGV)