PT-167 Still need to add tests

Description here

bin/pt-kill

@@ -3297,6 +3297,17 @@ sub new {
    foreach my $arg ( qw(MasterSlave) ) {
       die "I need a $arg argument" unless $args{$arg};
    }
+   my $kill_busy_commands = {};
+   if ($args{kill_busy_commands}) {
+       for my $command (split /,/,$args{kill_busy_commands}) {
+           $command =~ s/^\s+|\s+$//g;
+           $kill_busy_commands->{$command} = 1;
+       }
+   } else {
+       $kill_busy_commands->{Query} = 1;
+   }
+   $args{kill_busy_commands} = $kill_busy_commands;
+
    my $self = {
       %args,
       polls       => 0,
@@ -3523,7 +3534,7 @@ sub find {
          next QUERY;
       }
 
-      if ( $find_spec{busy_time} && ($query->{Command} || '') eq 'Query' ) {
+      if ( $find_spec{busy_time} && exists($self->{kill_busy_commands}->{$query->{Command} || ''}) ) {
          next QUERY unless defined($query->{Time});
          if ( $query->{Time} < $find_spec{busy_time} ) {
             PTDEBUG && _d("Query isn't running long enough");
@@ -6788,7 +6799,7 @@ sub main {
       DSNParser    => $dp,
       Quoter       => "Quoter",
    );
-   my $pl = new Processlist(MasterSlave => $ms);
+   my $pl = new Processlist(MasterSlave => $ms, kill_busy_commands => $o->get('kill-busy-commands'));
    my $qr = new QueryRewriter();
 
    my $cxn;
@@ -8198,6 +8209,21 @@ that pt-kill matched and killed a query.
 
 See also L<"--wait-before-kill"> and L<"--wait-after-kill">.
 
+=item --kill-busy-commands
+
+default: Query
+
+group: Actions
+
+Comma sepatated list of commands that will be watched/killed if they ran for
+more than L<"--busy-time"> seconds. Default: C<Query>
+
+By default, L<"--busy-time"> kills only C<Query> commands but in some cases, it
+is needed to make L<"--busy-time"> to watch and kill other commands. For example,
+a prepared statement execution command is C<Execute> instead of C<Query>. In this
+case, specifying C<--kill-busy-commands=Query,Execute> will also kill the prepared
+stamente execution.
+
 =item --kill-query
 
 group: Actions

bin/pt-query-digest

@@ -3251,6 +3251,17 @@ sub new {
    foreach my $arg ( qw(MasterSlave) ) {
       die "I need a $arg argument" unless $args{$arg};
    }
+   my $kill_busy_commands = {};
+   if ($args{kill_busy_commands}) {
+       for my $command (split /,/,$args{kill_busy_commands}) {
+           $command =~ s/^\s+|\s+$//g;
+           $kill_busy_commands->{$command} = 1;
+       }
+   } else {
+       $kill_busy_commands->{Query} = 1;
+   }
+   $args{kill_busy_commands} = $kill_busy_commands;
+
    my $self = {
       %args,
       polls       => 0,
@@ -3465,6 +3476,7 @@ sub find {
    my $ms  = $self->{MasterSlave};
 
    my @matches;
+   $self->{_reasons_for_matching} = undef;
    QUERY:
    foreach my $query ( @$proclist ) {
       PTDEBUG && _d('Checking query', Dumper($query));
@@ -3476,7 +3488,7 @@ sub find {
          next QUERY;
       }
 
-      if ( $find_spec{busy_time} && ($query->{Command} || '') eq 'Query' ) {
+      if ( $find_spec{busy_time} && exists($self->{kill_busy_commands}->{$query->{Command} || ''}) ) {
          next QUERY unless defined($query->{Time});
          if ( $query->{Time} < $find_spec{busy_time} ) {
             PTDEBUG && _d("Query isn't running long enough");

lib/Processlist.pm

@@ -69,6 +69,19 @@ sub new {
    foreach my $arg ( qw(MasterSlave) ) {
       die "I need a $arg argument" unless $args{$arg};
    }
+   # Convert the list of kill commands (Query, Execute, etc) to a hashref for
+   # faster check later
+   my $kill_busy_commands = {};
+   if ($args{kill_busy_commands}) {
+       for my $command (split /,/,$args{kill_busy_commands}) {
+           $command =~ s/^\s+|\s+$//g;
+           $kill_busy_commands->{$command} = 1;
+       }
+   } else {
+       $kill_busy_commands->{Query} = 1;
+   }
+   $args{kill_busy_commands} = $kill_busy_commands;
+
    my $self = {
       %args,
       polls       => 0,
@@ -471,6 +484,7 @@ sub find {
    my $ms  = $self->{MasterSlave};
 
    my @matches;
+   $self->{_reasons_for_matching} = undef;
    QUERY:
    foreach my $query ( @$proclist ) {
       PTDEBUG && _d('Checking query', Dumper($query));
@@ -484,7 +498,8 @@ sub find {
       }
 
       # Match special busy_time.
-      if ( $find_spec{busy_time} && ($query->{Command} || '') eq 'Query' ) {
+      #if ( $find_spec{busy_time} && ($query->{Command} || '') eq 'Query' ) {
+      if ( $find_spec{busy_time} && exists($self->{kill_busy_commands}->{$query->{Command} || ''}) ) {
          next QUERY unless defined($query->{Time});
          if ( $query->{Time} < $find_spec{busy_time} ) {
             PTDEBUG && _d("Query isn't running long enough");

t/pt-kill/pt_167.t

@@ -0,0 +1,93 @@
+#!/usr/bin/env perl
+
+BEGIN {
+   die "The PERCONA_TOOLKIT_BRANCH environment variable is not set.\n"
+      unless $ENV{PERCONA_TOOLKIT_BRANCH} && -d $ENV{PERCONA_TOOLKIT_BRANCH};
+   unshift @INC, "$ENV{PERCONA_TOOLKIT_BRANCH}/lib";
+};
+
+use strict;
+use warnings FATAL => 'all';
+use English qw(-no_match_vars);
+use Time::HiRes qw(sleep);
+use Test::More;
+
+use PerconaTest;
+use Sandbox;
+require "$trunk/bin/pt-kill";
+
+use Data::Dumper;
+$Data::Dumper::Indent    = 1;
+$Data::Dumper::Sortkeys  = 1;
+$Data::Dumper::Quotekeys = 0;
+
+my $dp  = new DSNParser(opts=>$dsn_opts);
+my $sb  = new Sandbox(basedir => '/tmp', DSNParser => $dp);
+my $dbh = $sb->get_dbh_for('master');
+
+if ( !$dbh ) {
+   plan skip_all => 'Cannot connect to sandbox master';
+}
+
+my $output;
+my $dsn = $sb->dsn_for('master');
+my $cnf = '/tmp/12345/my.sandbox.cnf';
+
+# #############################################################################
+# Test that --kill-query only kills the query, not the connection.
+# #############################################################################
+
+# Here's how this works.  This cmd is going to try 2 queries on the same
+# connection: sleep5 and sleep3.  --kill-query will kill sleep5 causing
+# sleep3 to start using the same connection id (pid).
+system("/tmp/12345/use -e 'select sleep(5); select sleep(2)' >/dev/null&");
+sleep 0.5;
+
+SKIP: {
+                                         
+    skip "TODO";
+    my $iterations=99;
+    my $query = 'select benchmark(?, md5("when will it end?"))';
+    my $ps = $dbh->prepare($query);
+    $ps->execute($iterations);
+    my $rows = $dbh->selectall_hashref('show processlist', 'id');
+    my $pid = 0;  # reuse, reset
+    map  { $pid = $_->{id} }
+    grep { $_->{info} && $_->{info} =~ m/select sleep\(15\)/ }
+    values %$rows;
+    
+    ok(
+       $pid,
+       'Got proc id of sleeping query'
+    );
+    
+    $output = output(
+       sub { pt_kill::main('-F', $cnf, qw(--busy-time 4s --print --match-info "^(select|SELECT)")), },
+    );
+    like(
+       $output,
+       qr/KILL QUERY $pid /,
+       '--kill-query'
+    );
+    
+    sleep 1;
+    $rows = $dbh->selectall_hashref('show processlist', 'id');
+    my $con_alive = grep { $_->{id} eq $pid } values %$rows;
+    ok(
+       $con_alive,
+       'Killed query, not connection'
+    );
+    
+    is(
+       ($rows->{$pid}->{info} || ''),
+       'select sleep(3)',
+       'Connection is still alive'
+    );
+}
+
+# #############################################################################
+# Done.
+# #############################################################################
+$sb->wipe_clean($dbh);
+ok($sb->ok(), "Sandbox servers") or BAIL_OUT(__FILE__ . " broke the sandbox");
+done_testing;