#!/usr/bin/env perl

##**************************************************************
##
## Copyright (C) 1990-2018, Condor Team, Computer Sciences Department,
## University of Wisconsin-Madison, WI.
##
## Licensed under the Apache License, Version 2.0 (the "License"); you
## may not use this file except in compliance with the License.  You may
## obtain a copy of the License at
##
##    http://www.apache.org/licenses/LICENSE-2.0
##
## Unless required by applicable law or agreed to in writing, software
## distributed under the License is distributed on an "AS IS" BASIS,
## WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
## See the License for the specific language governing permissions and
## limitations under the License.
##
##**************************************************************

use strict;
use warnings;

#
# This file is an almost-exact copy of cmd_drain_scavenging.  Sorry.
# Modifying the original to conditionally execute this copy wasn't hard,
# but the test suite gets /really/ cranky if you do that (since it
# depends on the text of the runfile for certain things).
#

use Cwd;

use CondorTest;
use CondorUtils;
use CustomMachineResource;
use Check::CondorLog;

#testreq: personal
my $config = <<CONDOR_TESTREQ_CONFIG;
use feature : PartitionableSlot
NEGOTIATOR_INTERVAL = 5
MaxJobRetirementTime = 3600
MachineMaxVacateTime = 5
NUM_CPUS = 3
STARTD_DEBUG = D_FULLDEBUG D_SUB_SECOND
CONDOR_TESTREQ_CONFIG
#endtestreq

my $testName = 'cmd_drain_scavenging_vacation';
my $testDescription = 'Test scavenging with an original job that ages out of retirement.';

#
# This test validates the new scavenging functionality of the drain state.
# In particular, that:
#
#	* the new START expression is advertised;
#	* the new START expression is effective, that is, some jobs will start
#	  and others will not while the startd is draining;
#	* the condor_drain -start option is working;
#	* that when the last of the 'original' jobs exit, all the scavengers
#	  are preempted
#	* after the startd is drained, cancelling the draining results
#	  in effective change to the START expression;
#	* AcceptedWhileDraining is correct.
#
# It does NOT, but should, validate the following:
#	* the condor_defrag daemon's knob to set the START expression;
#	* that when the last of the 'original' jobs exit, the startd sets
#	  MaxJobRetirementTime, but not MaxMachineVacateTime to 0;
#	* that when the last of the 'original' jobs exit, the startd invalidates
#	  all claim IDs;
#	* that when the draining is cancelled, MaxJobRetirementTime returns
#	  to original value (both advertised and effective);
#	* that after the startd is drained, cancelling the draining results
#	  in the original START expression being advertised (and used).
#

#
# To do the validation, we submit five jobs (which have durations which should
# never matter): the first three don't mention MaxJobRetirementTime at all,
# the fourth sets it to zero, and the fifth to nonzero.  This is so that
# the fourth and fifth jobs will look start (or not start) only on the basis
# of the (recommended) START expression we use when we begin draining.
#
# After the first three jobs have started, we'll begin draining with a START
# expression of 'TARGET.IsInterruptible'.  Then we'll verify that
# the slots have all changed state & activity, and that the new START
# expression is being advertised.  Then we'll terminate two of the original
# jobs and call condor_reschedule.
#
# One and only one should match; we'll wait for it to actually start.  Then
# we'll verify the new slot's state & activity.  Then we wait for the last
# original job to be vacated for exceeding its max job retirement time.  Note
# that the scavenger job must have a retirement time exceeding that of the
# original job's to test this properly.  In the setup suggested by
# cmd_drain_scavenging, that will always be true (because the job must
# have an MJRT of 0), but in production (e.g., CHTC) that's not so.
#
# If run against an unpatched startd, this test will fail when the scavenger
# job unexpectedly succeeds (before the last original job is evicted).
#

my $jobFile = cwd() . '/condition-sleep.pl';
my $jobContents = '#!/usr/bin/env perl

use strict;
use warnings;

# Ignore SIGTERM and force a hard-kill to test if MaxJobRetirementTime is
# being altered appropriately.
$SIG{"TERM"} = "IGNORE";

my $killFile = $ARGV[0];
my $maxSleep = $ARGV[1];

for( my $i = 0; $i < $maxSleep; ++$i ) {
	if( -e $killFile ) { exit( 0 ); }
	sleep( 1 );
}

exit( 0 );
';
CondorTest::WriteFileOrDie( $jobFile, $jobContents );
chmod( 0755, $jobFile ) || die( "Failed to make ${jobFile} executable, aborting.\n" );

my $killFile = cwd() . "/killfile-cdsv";
my $submitBody = "

executable				= ${jobFile}
transfer_executable		= false
should_transfer_files	= true
universe				= vanilla

arguments				= ${killFile}.\$(CLUSTER).\$(PROCESS) 3600
MaxJobRetirementTime	= 60

log						= cmd_drain_scavenging_vacation.log

Requirements			= (NumJobStarts?:0) < 1
queue 3

arguments				= ${killFile}.\$(CLUSTER).\$(PROCESS) 60
+IsInterruptible		= true
MaxJobRetirementTime	= 3600
# prevent job from restarting after eviction before we can kill it
Requirements			= (NumJobStarts?:0) < 1

queue 1

arguments				= ${killFile}.\$(CLUSTER).\$(PROCESS) 60
+IsInterruptible		= false
MaxJobRetirementTime	= 60

queue 1

";

my $clusterID;
my $setClusterID = sub {
	my( $cID ) = @_;
	$clusterID = $cID;
};

my $hostName = `condor_config_val full_hostname`;
CondorUtils::fullchomp( $hostName );

my $testingPhase = 0;
my $longJobsRunningCount = 0;

my $execute = sub {
	my %info = @_;
	my $ID = $info{ 'cluster' } . "." . $info{ 'job' };

	if( $info{ 'cluster' } != $clusterID ) {
		die( "Found stale test job or log file, aborting.\n" );
	}

	if( $testingPhase == 0 ) {
		if( 0 <= $info{ 'job' } && $info{ 'job' } <= 2 ) {
			++$longJobsRunningCount;
		} else {
			die( "Job ${ID} executed before it should have.\n" );
		}

		if( $longJobsRunningCount > 3 ) {
			die( "Original jobs started more than three times.\n" );
		} elsif( $longJobsRunningCount < 3 ) {
			return;
		}

		print( "All three original jobs have started.  Waiting for startd to notice...\n" );
		for( my $delay = 0; $delay <= 20; ++$delay ) {
			if( $delay == 20 ) {
				die( "Startd failed to notice in twenty seconds, aborting.\n" ); 
			}

			sleep( 1 );

			my $ads = CustomMachineResource::parseMachineAds( qw(Name State Activity Start AcceptedWhileDraining) );
			if( scalar(@{$ads}) == 4 ) {
				last;
			}
		}
		print( "... startd noticed.\n" );

		print( "Beginning draining...\n" );
		runCommandCarefully( undef, qw( condor_drain -start ),
			'TARGET.IsInterruptible == TRUE', $hostName );

		print( "Verifying state changes...\n" );
		for( my $delay = 0; $delay <= 20; ++$delay ) {
			if( $delay == 20 ) {
				die( "State failed to change by twenty seconds after draining, aborting.\n" ); 
			}

			sleep( 1 );

			my $ads = CustomMachineResource::parseMachineAds( qw(Name State Activity Start AcceptedWhileDraining) );
			if( scalar(@{$ads}) != 4 ) {
				print( "Found unexpected number of ads: " . scalar(@{$ads}) . ", will retry.\n" );
				next;
			}

			my $stateChanged = 1;
			foreach my $ad (@{$ads}) {
				if( $ad->{ "Name" } =~ /^slot1@/ ) { next; }
				print( "Considering '" . $ad->{ "Name" } . "'...\n" );
				if(! $ad->{ "State" } =~ /^Claimed$/i) { print( "Found unexpected state '" . $ad->{ "State" } . "', will retry.\n" ); $stateChanged = 0; last; }
				if(! $ad->{ "Activity" } =~ /^Retiring$/i) { print( "Found unexpected activity '" . $ad->{ "Activity" } . "', will retry.\n" ); $stateChanged = 0; last; }
				if(! $ad->{ "Start" } =~ /TARGET\.IsInterruptible == TRUE/i) { print( "Found unexpected START '" . $ad->{ "Start" } . "', will retry.\n" ); $stateChanged = 0; last; }
				if(! $ad->{ "AcceptedWhileDraining" } =~ /^False$/i) { print( "Found unexpected AcceptedWhileDraining '" . $ad->{ "AcceptedWhileDraining" } . "', will retry.\n" ); $stateChanged = 0; last; }
			}
			if( $stateChanged == 1 ) { last; }
		}
		print( "... state changed successfully.\n" );

		print( "Killing two of the long-running jobs...\n" );
		touch( $killFile . ".${clusterID}.0" );
		touch( $killFile . ".${clusterID}.1" );
		$testingPhase = 1;
	} elsif( $testingPhase == 2 ) {
		if( $info{ 'job' } != 3 ) {
			die( "Job ${ID} executed during testing phase 2, aborting.\n" );
		}

		print( "Scavenger job started.  Verifying state changes...\n" );
		for( my $delay = 0; $delay <= 20; ++$delay ) {
			if( $delay == 20 ) {
				die( "State failed to change by twenty seconds after draining, aborting.\n" ); 
			}

			sleep( 1 );

			my $ads = CustomMachineResource::parseMachineAds( qw(Name State Activity Start AcceptedWhileDraining) );
			if( scalar(@{$ads}) != 3 ) {
				print( "Found unexpected number of ads: " . scalar(@{$ads}) . ", will retry.\n" );
				next;
			}

			my $stateChanged = 0;
			foreach my $ad (@{$ads}) {
				if( $ad->{ "Name" } =~ /^slot1@/ ) { next; }
				print( "Considering '" . $ad->{ "Name" } . "'...\n" );

				if(		$ad->{ "State" } =~ /^Claimed$/i &&
						$ad->{ "Activity" } =~ /^Busy$/i &&
						$ad->{ "AcceptedWhileDraining" } =~ /^True$/i ) {
					$stateChanged = 1;
					last;
				}
			}

			if( $stateChanged ) { last; }
		}
		print( "... state changed successfully.\n" );


		print( "Waiting for the last long-running job to exceed its retirement time...\n" );
		sleep( 60 + 5 );
		for( my $i = 0; $i < 55; $i += 5 ) {
			sleep( 5 );
			my @output;
			runCondorToolCarefully( \@output, 2, { emit_output => 0 },
				undef, 'condor_q', "${clusterID}.2", '-af', 'Cluster' );
			if( (! defined($output[0])) || $output[0] eq "" ) { last; }
		}

		print( "Waiting for job completion notice...\n" );
		$testingPhase = 3;
	} elsif( $testingPhase == 4 ) {
		if( $info{ 'job' } != 4 ) {
			die( "Job ${ID} executed during testing phase 4, aborting.\n" );
		}

		print( "Killing last job...\n" );
		touch( $killFile . ".${clusterID}.4" );
		$testingPhase = 5;
	} else {
		die( "execute() called unexpectedly in testing phase ${testingPhase}, aborting.\n" );
	}
};

my $longJobsSucceededCount = 0;
my $success = sub {
	my %info = @_;
	my $ID = $info{ 'cluster' } . "." . $info{ 'job' };

	if( $info{ 'cluster' } != $clusterID ) {
		die( "Found stale test job or log file, aborting.\n" );
	}

	if( $testingPhase == 1 ) {

		if( 0 <= $info{ 'job' } && $info{ 'job' } <= 1 ) {
			++$longJobsSucceededCount;
		} else {
			die( "Job ${ID} succeeded before it should have.\n" );
		}

		if( $longJobsSucceededCount > 2 ) {
			die( "Original jobs started more than two times.\n" );
		} elsif( $longJobsSucceededCount < 2 ) {
			return;
		}

		print( "Two original jobs have succeeded, calling condor_reschedule...\n" );
		runCommandCarefully( undef, 'condor_reschedule' );

		print( "Waiting for scavenger job to start...\n" );
		$testingPhase = 2;
	} elsif( $testingPhase == 4 ) {
		# We /should/ get the succeeded event before the eviction event,
		# but it's OK if they happen in the other order.
		if( $info{ 'job' } != 2 ) {
			die( "Job ${ID} succeeded during testing phase 3 or 4, aborting.\n" );
		}
	} elsif( $testingPhase == 5 ) {
		if( $info{ 'job' } != 4 ) {
			die( "Job ${ID} succeeded during testing phase 5, aborting.\n" );
		}

		RegisterResult( 1, check_name => 'normal/drain/scavenge/drained/normal cycle', test_name => $testName );
	} else {
		die( "success() called unexpectedly in testing phase ${testingPhase}, aborting.\n" );
	}
};

my $evicted = sub {
	my %info = @_;
	my $ID = $info{ 'cluster' } . "." . $info{ 'job' };

	if( $info{ 'cluster' } != $clusterID ) {
		die( "Found stale test job or log file, aborting.\n" );
	}

	if( $testingPhase == 3 ) {
		if( $info{ 'job' } != 2 && $info{ 'job' } != 3 ) {
			die( "Job ${ID} evicted, aborting.\n" );
		}

		print( "Waiting for machine to fully drain...\n" );
		for( my $delay = 0; $delay <= 20; ++$delay ) {
			if( $delay == 20 ) {
				die( "State failed to change by twenty seconds after draining, aborting.\n" ); 
			}

			sleep( 1 );

			my $ads = CustomMachineResource::parseMachineAds( qw(Name State Activity Start AcceptedWhileDraining) );
			if( scalar(@{$ads}) != 1 ) {
				print( "Found unexpected number of ads: " . scalar(@{$ads}) . ", will retry.\n" );
				next;
			}

			my $drained = 0;
			foreach my $ad (@{$ads}) {
				if(		$ad->{ 'State' } =~ /^Drained$/i &&
						$ad->{ 'Activity' } =~ /^Idle$/i &&
						$ad->{ 'AcceptedWhileDraining' } =~ /^False$/i ) {
					$drained = 1;
				}
			}
			if( $drained ) { last; }
		}
		print( "... machine fully drained.\n" );

		# I'm not veryfing MaxJobRetirementTime and MaxMachineVacateTime
		# because the sequence of events when the last 'original' job is
		# evicted is very different.  Basically, if the last 'original'
		# job gets evicted at all, it's good enough for this test.

		print( "Cancelling draining state...\n" );
		runCommandCarefully( undef, "condor_drain -cancel ${hostName}" );

		print( "Removing evicted job ${ID}...\n" );

		runCommandCarefully( undef, 'condor_rm', $ID );

		print( "Calling condor_reschedule...\n" );
		runCommandCarefully( undef, 'condor_reschedule' );

		print( "Waiting for last job to start...\n" );
		$testingPhase = 4;
	} elsif( $testingPhase == 4 ) {
		if( $info{ 'job' } != 2 && $info{ 'job' } != 3 ) {
			die( "Job ${ID} evicted, aborting.\n" );
		}

		print( "Removing evicted job ${ID}...\n" );
		runCommandCarefully( undef, 'condor_rm', $ID );
	} else {
		die( "evicted() called unexpectedly in testing phase ${testingPhase}, aborting.\n" );
	}
};

my $aborted = sub {
	my %info = @_;
	my $ID = $info{ 'cluster' } . "." . $info{ 'job' };

	if( $info{ 'cluster' } != $clusterID ) {
		die( "Found stale test job or log file, aborting.\n" );
	}

	if( $testingPhase == 4 ) {
		if( $info{ 'job' } != 3 && $info{ 'job' } != 2 ) {
			die( "Job ${ID} aborted, aborting.\n" );
		}
	} else {
		die( "aborted() called unexpectedly in testing phase ${testingPhase}, aborting.\n" );
	}
};

CondorTest::RegisterExecute( $testName, $execute );
CondorTest::RegisterExitedSuccess( $testName, $success );
CondorTest::RegisterEvictedWithoutCheckpoint( $testName, $evicted );
CondorTest::RegisterAbort( $testName, $aborted );
CondorTest::RunTest2( name => $testName, want_checkpoint => 0,
	submit_body => $submitBody, callback => $setClusterID );

exit( 0 );

sub touch {
	my( $filename ) = @_;

	my $rv = open( my $fh, ">", $filename );
	close( $fh );
	return $rv;
}
