lhm-kill-queue

#!/usr/bin/env ruby

require 'active_record'
require 'lhm/sql_helper'
require 'optparse'

module Lhm
  class KillQueue

    def initialize
      @port = 3306
      @grace = 10
      @tiny = 0.1
      @marker = "%#{ SqlHelper.annotation }%"

      OptionParser.new do |opts|
        opts.on("-h", "--hostname HOSTNAME") { |v| @hostname = v }
        opts.on("-u", "--username USERNAME") { |v| @username = v }
        opts.on("-p", "--password PASSWORD") { |v| @password = v }
        opts.on("-d", "--database DATABASE") { |v| @database = v }
        opts.on("-m", "--mode MODE") { |v| @mode = v.to_sym }
        opts.on("-y", "--confirm") { |v| @confirm = true }
      end.parse!

      unless(@hostname && @username && @password && @database)
        abort usage
      end

      unless([:kill, :master, :slave].include?(@mode))
        abort "specify -m kill OR -m master OR -m slave"
      end

      connect
    end

    def usage
      <<-desc.gsub(/^      /, '')
        kills queries on the given server after detecting 'lock table#{ @marker }'.
        usage:
          lhm-kill-queue -h hostname -u username -p password -d database \\
          (-m kill | -m master | -m slave) [--confirm]

      desc
    end

    def run
      case @mode
      when :kill then kill
      when :master then master
      when :slave then slave
      end
    end

    def kill
      lock = trip
      kill_process(lock)
    end

    def master
      lock = trip
      puts "starting to kill non lhm processes in #{ @grace } seconds"
      sleep(@grace + @tiny)

      [list_non_lhm].flatten.each do |process|
        kill_process(process)
        sleep(@tiny)
      end
    end

    def slave
      lock = trip
      puts "starting to kill non lhm SELECT processes in #{ @grace } seconds"
      sleep(@grace + @tiny)

      [list_non_lhm].flatten.each do |process|
        if(select?(process))
          kill_process(process)
          sleep(@tiny)
        end
      end
    end

  private

    def connect
      ActiveRecord::Base.establish_connection({
        :adapter => 'mysql',
        :host => @hostname,
        :port => @port,
        :username => @username,
        :password => @password,
        :database => @database
      })
    end

    def connection
      ActiveRecord::Base.connection
    end

    def list_non_lhm
      select_processes %Q(
        info not like '#{ @marker }' and time > #{ @grace } and command = 'Query'
      )
    end

    def trip
      until res = select_processes("info like 'lock table#{ @marker }'").first
        sleep @tiny
        print '.'
      end

      res
    end

    def kill_process(process_id)
      puts "killing #{ select_statement(process_id) }"

      if(@confirm)
        print "confirm ('y' to confirm): "

        if(gets.strip != 'y')
          puts "skipped."
          return
        end
      end

      connection.execute("kill #{ process_id }")
      puts "killed #{ process_id }"
    end

    def select?(process)
      if statement = select_statement(process)
        case statement
        when /delete/i then false
        when /update/i then false
        when /insert/i then false
        else
          !!statement.match(/select/i)
        end
      end
    end

    def select_statement(process)
      if process
        value %Q(
          select info from information_schema.processlist where id = #{ process }
        )
      end
    end

    def select_processes(predicate)
      values %Q(
        select id from information_schema.processlist
         where db = '#{ @database }'
           and user = '#{ @username }'
           and #{ predicate }
      )
    end

    def value(statement)
      connection.select_value(statement)
    end

    def values(statement)
      connection.select_values(statement)
    end
  end
end

killer = Lhm::KillQueue.new
killer.run