Francois Gouget : testbot/Engine: Add --shutdown, --kill-tasks and --kill-vms options.

[Alexandre Julliard]

Module: tools
Branch: master
Commit: 9641db82e569662e9cb53188f241678aa77719e3
URL:    http://source.winehq.org/git/tools.git/?a=commit;h=9641db82e569662e9cb53188f...

Author: Francois Gouget <fgouget(a)codeweavers.com>
Date:   Tue Mar 26 16:20:57 2013 +0100

testbot/Engine: Add --shutdown, --kill-tasks and --kill-vms options.

This makes it possible to cleanly stop the Engine and all its related tasks and VMs.

---

 testbot/bin/Engine.pl                    |  102 ++++++++++++++++++++++++++++--
 testbot/lib/WineTestBot/Engine/Notify.pm |   23 ++++++-
 2 files changed, 118 insertions(+), 7 deletions(-)

diff --git a/testbot/bin/Engine.pl b/testbot/bin/Engine.pl
index f5b0661..fec0c67 100755
--- a/testbot/bin/Engine.pl
+++ b/testbot/bin/Engine.pl
@@ -43,6 +43,7 @@ use Socket;
 use ObjectModel::BackEnd;
 use WineTestBot::Config;
 use WineTestBot::Engine::Events;
+use WineTestBot::Engine::Notify;
 use WineTestBot::Jobs;
 use WineTestBot::Log;
 use WineTestBot::Patches;
@@ -50,6 +51,8 @@ use WineTestBot::PendingPatchSets;
 use WineTestBot::Utils;
 use WineTestBot::VMs;
 
+my $RunEngine = 1;
+
 sub FatalError
 {
   LogMsg @_;
@@ -65,7 +68,8 @@ sub FatalError
 =item C<Cleanup()>
 
 The Cleanup() function gets the tasks and VMs in a consistent state on the
-Engine startup. It has to contend with two main scenarios:
+Engine startup or cleanly stops the tasks and VMs on shutdown.
+It has to contend with three main scenarios:
 - The Engine being restarted. Any task started just before that is still
   running should still have its process and powered on VM and should be left
   alone so it can complete normally. If a task died unexpectedly while the
@@ -75,6 +79,8 @@ Engine startup. It has to contend with two main scenarios:
   it is quite possible that they will still be running. Hopefully any running
   process matching a task's ChildPid will belong to another user so we don't
   mistake that case for the previous one.
+- A shutdown of the Engine and its tasks / VMs. In this case it's used to
+  kill the running tasks and requeue them, and/or shut down the VMs.
 
 In all cases we only trust that a VM status field is still valid:
 - It is 'running' and used by a task that is still running.
@@ -86,8 +92,10 @@ In all other cases the VM will be powered off and marked as dirty.
 =back
 =cut
 
-sub Cleanup()
+sub Cleanup(;$$)
 {
+  my ($KillTasks, $KillVMs) = @_;
+
   # Verify that the running tasks are still alive and requeue them if not.
   # Ignore the Job and Step status fields because they may be a bit out of date.
   my %BusyVMs;
@@ -115,6 +123,14 @@ sub Cleanup()
           # Kill the task process if it's still there.
           kill("TERM", $Task->ChildPid);
         }
+        elsif ($KillTasks)
+        {
+          # Kill the task and requeue. Note that since the VM is still running
+          # we're not in the computer reboot case so ChildPid is probably
+          # still valid.
+          kill("TERM", $Task->ChildPid);
+          $Requeue = 1;
+        }
         elsif (!kill(0, $Task->ChildPid))
         {
           # The event that caused the WineTestBot server to restart probably
@@ -165,8 +181,12 @@ sub Cleanup()
     if ($VM->IsPoweredOn())
     {
       next if ($VM->Status eq "idle");
-      if (($VM->Status eq "reverting" or $VM->Status eq "sleeping") and
-          defined $VM->ChildPid and kill(0, $VM->ChildPid))
+      if ($KillVMs)
+      {
+        kill("TERM", $VM->ChildPid) if (defined $VM->ChildPid);
+      }
+      elsif (($VM->Status eq "reverting" or $VM->Status eq "sleeping") and
+             defined $VM->ChildPid and kill(0, $VM->ChildPid))
       {
         # This VM is still being reverted. Let that process run its course.
         LogMsg "$VMKey is being reverted\n";
@@ -188,6 +208,43 @@ sub Cleanup()
 }
 
 
+sub HandleShutdown
+{
+  my ($KillTasks, $KillVMs) = @_;
+
+  if (!defined $KillTasks or !defined $KillVMs)
+  {
+    LogMsg "Missing parameters in shutdown message\n";
+    return "0Missing shutdown parameters";
+  }
+
+  # Untaint parameters
+  if ($KillTasks =~ /^([01])$/)
+  {
+    $KillTasks = $1;
+  }
+  else
+  {
+    LogMsg "Invalid KillTasks $KillTasks in shutdown message\n";
+    return "0Invalid KillTasks shutdown parameter\n";
+  }
+  if ($KillVMs =~ /^([01])$/)
+  {
+    $KillVMs = $1;
+  }
+  else
+  {
+    LogMsg "Invalid KillVMs $KillVMs in shutdown message\n";
+    return "0Invalid KillVMs shutdown parameter\n";
+  }
+
+  Cleanup($KillTasks, $KillVMs);
+  $RunEngine = 0;
+
+  LogMsg "Waiting for the last clients to disconnect...\n";
+  return "1OK\n";
+}
+
 sub HandlePing
 {
   return "1pong\n";
@@ -492,6 +549,7 @@ my %Handlers=(
     "jobstatuschange"          => \&HandleJobStatusChange,
     "jobsubmit"                => \&HandleJobSubmit,
     "ping"                     => \&HandlePing,
+    "shutdown"                 => \&HandleShutdown,
     "taskcomplete"             => \&HandleTaskComplete,
     "vmstatuschange"           => \&HandleVMStatusChange,
     "winepatchmlsubmission"    => \&HandleWinePatchMLSubmission,
@@ -621,6 +679,40 @@ sub REAPER
 
 sub main 
 {
+  my ($Shutdown, $KillTasks, $KillVMs);
+  while (@ARGV)
+  {
+    my $Arg = shift @ARGV;
+    if ($Arg eq "--shutdown")
+    {
+      $Shutdown = 1;
+    }
+    elsif ($Arg eq "--kill-tasks")
+    {
+      $KillTasks = 1;
+      $Shutdown = 1;
+    }
+    elsif ($Arg eq "--kill-vms")
+    {
+      $KillVMs = 1;
+      $Shutdown = 1;
+    }
+    else
+    {
+      die "Usage: Engine.pl [--shutdown] [--kill-tasks] [--kill-vms]";
+    }
+  }
+  if ($Shutdown)
+  {
+    my $ErrMessage = Shutdown($KillTasks, $KillVMs);
+    if (defined $ErrMessage)
+    {
+      print STDERR "$ErrMessage\n";
+      exit 1;
+    }
+    exit 0;
+  }
+
   $ENV{PATH} = "/usr/bin:/bin";
   delete $ENV{ENV};
   $SIG{CHLD} = \&REAPER;
@@ -657,7 +749,7 @@ sub main
   AddEvent("SafetyNet", 600, 1, \&SafetyNet);
 
   my @Clients;
-  while (1)
+  while ($RunEngine or @Clients)
   {
     my $ReadyRead = "";
     my $ReadyWrite = "";
diff --git a/testbot/lib/WineTestBot/Engine/Notify.pm b/testbot/lib/WineTestBot/Engine/Notify.pm
index 8452615..1a1ab6d 100644
--- a/testbot/lib/WineTestBot/Engine/Notify.pm
+++ b/testbot/lib/WineTestBot/Engine/Notify.pm
@@ -33,8 +33,8 @@ use vars qw (@ISA @EXPORT @EXPORT_OK $RunningInEngine);
 
 require Exporter;
 @ISA = qw(Exporter);
-(a)EXPORT = qw(&PingEngine &JobSubmit &JobStatusChange &JobCancel &JobRestart
-             &TaskComplete &VMStatusChange &FoundWinetestUpdate
+(a)EXPORT = qw(&Shutdown &PingEngine &JobSubmit &JobStatusChange &JobCancel
+             &JobRestart &TaskComplete &VMStatusChange &FoundWinetestUpdate
              &WinePatchMLSubmission &WinePatchWebSubmission &GetScreenshot);
 @EXPORT_OK = qw($RunningInEngine);
 
@@ -74,6 +74,25 @@ sub SendCmdReceiveReply
   return $Reply;
 }
 
+sub Shutdown
+{
+  my ($KillTasks, $KillVMs) = @_;
+
+  $KillTasks ||= 0;
+  $KillVMs ||= 0;
+  my $Reply = SendCmdReceiveReply("shutdown $KillTasks $KillVMs\n");
+  if (length($Reply) < 1)
+  {
+    return "Unrecognized reply received from engine";
+  }
+  if (substr($Reply, 0, 1) eq "1")
+  {
+    return undef;
+  }
+
+  return substr($Reply, 1);
+}
+
 sub PingEngine
 {
   my $Reply = SendCmdReceiveReply("ping\n");