# Batch EXecutor -- Queues
# (c) 2011-2015 Martin Mares <mj@ucw.cz>

use strict;
use warnings;
use feature 'switch';
use experimental 'smartmatch';

package BEX::Queue;

use IO::File;
use File::Path;
use Fcntl qw(:flock);
use POSIX ();

sub new($;$) {
	my ($class, $name) = @_;
	$name //= 'queue';
	my $path = $BEX::Config::home . '/' . $name;
	-d $path or die "Queue directory $path does not exist (use 'bex qman init' to create it)\n";
	-d "$path/hosts" && -d "$path/jobs" or die "Queue directory $path is misconfigured\n";
	my $queue = {
		'Name' => $name,
		'Path' => $path,
		'MetaCache' => {},
	};
	return bless $queue;
}

sub log_file($$) {
	my ($queue, $machine, $jid) = @_;
	return $queue->host_dir($machine) . '/' . $jid . '.log';
}

# Most actions have to be logged by the caller
sub log($$$$;$) {
	my ($queue, $mach, $jid, $stat, $msg) = @_;
	my $t = POSIX::strftime("%Y-%m-%d %H:%M:%S", localtime);
	my $m = join(" ", $t, $mach, $jid, $stat);
	$m .= " $msg" if defined $msg;

	my $fh = $queue->{'LogFH'} //= new IO::File $queue->{'Path'} . '/log', '>>' or die "Cannot open log: $!";
	print $fh "$m\n";

	# Append to the per-job log file
	if (open L, '>>', $queue->log_file($mach, $jid)) {
		print L "### $m\n";
		close L;
	}
}

sub host_dir($$) {
	my ($queue, $machine) = @_;
	return $queue->{'Path'} . '/hosts/' . $machine;
}

sub queue_file($$) {
	my ($queue, $machine, $jid) = @_;
	return $queue->host_dir($machine) . '/' . $jid . '.job';
}

sub status_file($$) {
	my ($queue, $machine, $jid) = @_;
	return $queue->host_dir($machine) . '/' . $jid . '.stat';
}

sub temp_file($$) {
	my ($queue, $machine, $jid) = @_;
	return $queue->host_dir($machine) . '/' . $jid . '.tmp';
}

sub job_file($$) {
	my ($queue, $jid) = @_;
	return $queue->{'Path'} . '/jobs/' . $jid. '.job';
}

sub attachment_dir($$) {
	my ($queue, $jid) = @_;
	return $queue->{'Path'} . '/jobs/' . $jid. '.attach';
}

sub save_job($$) {
	my ($queue, $job) = @_;
	# If the job already exists, it is shamelessly rewritten by new contents
	$job->save($queue->job_file($job->id));
}

sub all_job_ids($) {
	my ($queue) = @_;
	opendir my $dir, $queue->{'Path'} . '/jobs' or die "Cannot open job directory: $!\n";
	my @jobs = ();
	while (readdir $dir) {
		s{\.job$}{} or next;
		push @jobs, $_;
	}
	closedir $dir;
	return @jobs;
}

sub match_job_id($$) {
	my ($id, $pattern) = @_;
	return $id =~ m{\b$pattern};
}

# Resolve a (possibly partial) job ID given by the user
sub resolve_job_id($$) {
	my ($queue, $name) = @_;
	BEX::Job::check_id($name) or die "Invalid job ID $name\n";
	if (-f $queue->job_file($name)) {
		return $name;
	}

	my @candidates = grep { match_job_id($_, $name) } $queue->all_job_ids();
	@candidates or die "No job called $name exists\n";
	@candidates == 1 or die "Partial job ID $name is not unique\n";
	return $candidates[0];
}

sub enqueue($$$) {
	my ($queue, $machine, $job) = @_;
	# The job must be already saved to the current queue
	my $qf = $queue->queue_file($machine, $job->id);
	if (-f $qf) {
		return 0;
	}
	my $dir = $queue->host_dir($machine);
	-d $dir or mkdir $dir or die "Cannot create directory $dir: $!";
	symlink '../../jobs/' . $job->id . '.job', $qf or die "Cannot create $qf: $!";
	return 1;
}

sub scan($$) {
	my ($queue, $machine) = @_;
	my @list = ();
	if (opendir D, $queue->host_dir($machine)) {
		while ($_ = readdir D) {
			/^\./ and next;
			s{\.job}{} or next;
			push @list, $_;
		}
		closedir D;
	}
	return sort @list;
}

sub remove($$;$) {
	my ($queue, $machine, $jid, $force_remove) = @_;
	if ($BEX::Config::keep_history && !$force_remove) {
		my $s = $queue->{'Path'} . '/hosts/' . $machine;
		my $d = $queue->{'Path'} . '/history/' . $machine;
		File::Path::mkpath($d);
		for my $suff ('job', 'stat', 'log') {
			my $src = "$s/$jid.$suff";
			my $dst = "$d/$jid.$suff";
			if (-f $src) {
				rename $src, $dst or die "Cannot rename $src to $dst: $!";
			} else {
				# Might be present from the previous incarnation of the same job
				unlink $dst;
			}
		}
	} else {
		unlink $queue->queue_file($machine, $jid);
		unlink $queue->status_file($machine, $jid);
		unlink $queue->log_file($machine, $jid);
	}
	unlink $queue->temp_file($machine, $jid);
}

sub job_metadata($$) {
	my ($queue, $jid) = @_;
	my $cache = $queue->{'MetaCache'};
	if (!defined $cache->{$jid}) {
		$cache->{$jid} = BEX::Job->new_from_file($queue->job_file($jid), 1);
	}
	return $cache->{$jid};
}

sub job_name($$) {
	my ($queue, $jid) = @_;
	return $queue->job_metadata($jid)->name;
}

sub read_job_status($$$) {
	my ($queue, $machine, $jid) = @_;
	my %s = ();
	my $sf = $queue->status_file($machine, $jid);
	if (open S, '<', $sf) {
		while (<S>) {
			chomp;
			/^(\w+):\s*(.*)/ or die "Parse error in $sf";
			$s{$1} = $2;
		}
		close S;
	}
	return \%s;
}

sub write_job_status($$$$) {
	my ($queue, $machine, $jid, $stat) = @_;
	my $sf = $queue->status_file($machine, $jid);
	open S, '>', "$sf.$$" or die "Cannot create $sf.$$: $!";
	for my $k (sort keys %$stat) {
		print S "$k: ", $stat->{$k}, "\n" if defined $stat->{$k};
	}
	close S;
	rename "$sf.$$", $sf or die "Cannot rename $sf.$$ to $sf: $!";
}

sub update_job_status($$$$;$) {
	my ($queue, $machine, $jid, $stat, $msg) = @_;
	my $s = {
		'Time' => time,
		'Status' => $stat,
		'Message' => $msg,
	};
	$queue->write_job_status($machine, $jid, $s);
	$queue->log($machine, $jid, $stat, $msg);
}

sub lock_name($$$) {
	my ($queue, $machine, $jid) = @_;
	my $lock = $queue->{'Path'};
	if (defined $jid) {
		$lock .= "/hosts/$machine/$jid.lock";
	} elsif (defined $machine) {
		$lock .= "/hosts/$machine/lock";
	} else {
		$lock .= '/lock';
	}
}

# Whenever we want to run a job on a machine, we must obtain a lock;
# at most one lock can be held at a time by a single BEX::Queue object.
# See the description of locking schemes in BEX::Config.
sub lock($$$) {
	my ($queue, $machine, $jid) = @_;
	my $lock;
	given ($BEX::Config::locking_scheme) {
		when ('queue') {
			$lock = lock_name($queue, undef, undef);
		}
		when ('host') {
			defined($machine) or return 1;
			$lock = lock_name($queue, $machine, undef);
		}
		when ('job') {
			defined($machine) && defined($jid) or return 1;
			$lock = lock_name($queue, $machine, $jid);
		}
		when ('none') { return 1; }
		default { die "Invalid BEX::Config::locking_scheme"; }
	}
	if (defined($queue->{'LockName'})) {
		return 1 if ($queue->{'LockName'} eq $lock);
		$queue->unlock;
	}
	open $queue->{'LockHandle'}, '>>', $lock or die "Cannot create $lock: $!";
	if (!flock($queue->{'LockHandle'}, LOCK_EX | LOCK_NB)) {
		close $queue->{'LockHandle'};
		delete $queue->{'LockHandle'};
		return 0;
	}
	$queue->{'LockName'} = $lock;
	return 1;
}

sub unlock($) {
	my ($queue) = @_;
	defined $queue->{'LockName'} or return;
	unlink $queue->{'LockName'};
	flock $queue->{'LockHandle'}, LOCK_UN;
	close $queue->{'LockHandle'};
	delete $queue->{'LockHandle'};
	delete $queue->{'LockName'};
}

# Unsafe (does not check fcntl, only existence of a lock file), but should be enough for `bex ls'
sub is_locked($$$) {
	my ($queue, $machine, $jid) = @_;
	given ($BEX::Config::locking_scheme) {
		# Shortcuts
		when ('host') { return unless defined $machine; }
		when ('jid') { return unless defined $jid; }
		when ('none') { return; }
	}
	my $lock = lock_name($queue, $machine, $jid);
	return -f $lock;
}

42;
