#!/usr/bin/perl -w
#
# Sun Grid Engine information provider for LCG.
# Based on the existing LCG PBS information provider.
#
# See LICENSE for licensing information.
#
# David McBride <dwm@doc.ic.ac.uk>
# London e-Science Centre, Department of Computing,
# Imperial College, London

# --- Compiler directives ------------------------------------------------

use strict;

# --- Constants ----------------------------------------------------------

$ENV{SGE_ROOT}	= "/opt/sge";
$ENV{SGE_CELL}	= "default";
$ENV{SGE_ARCH} = "lx24-amd64";
$ENV{SGE_EXECD_PORT} = "537";
$ENV{SGE_QMASTER_PORT} = "536";

my $qstat = "$ENV{SGE_ROOT}/bin/$ENV{SGE_ARCH}/qstat";
my $qconf = "$ENV{SGE_ROOT}/bin/$ENV{SGE_ARCH}/qconf";

my @voViewStatic = ("GlueCEStateWorstResponseTime",
                    "GlueCEStateEstimatedResponseTime",
                    "GlueCEStateFreeJobSlots");

# --- Function definitions -----------------------------------------------

# Simply converts an SGE value into one that the Glue schema uses.
# At present, simply:
#  - replaces "INFINITY" with "0", and
#  - replaces AA:BB:CC with ((AA * 60) + BB ) * 60 + CC.
#    (ie converts from hours:mins:seconds to seconds)
sub sgeToGlue($) {
    my ($value) = shift;
    $value =~ s/INFINITY/0/;
    if ($value =~ /^(\d+):(\d+):(\d+)$/) {
        $value = (($1 * 60) + $2) * 60 + $3;
    }
    return $value;
}

# Process our commandline.
# The only valid configurations are with exactly 1 or 2 arguments;
#
# The first argument is always the path to our configuration file.
#
# The second argument (if specified) will be used as the FQDN of the 
# SGE control host.  If this is not specified, the FQDN of the local 
# machine is used.

# TODO: The SGE tools do not trivially allow the query of arbitrary remote
# 	server nodes; at present, any hostname specification is ignored.
#	Determine whether or not this needs to be fixed.

#	One possible way of dealing with this problem is by keeping a separate
#	configuration file which keeps a mapping between master hostnames
#	and local cell identifiers.  Then we could source our environment
#	selectively.

sub processCommandLine() {
	if (scalar(@ARGV) == 0) {
		print STDERR "Usage: $0 <config file>\n";
		exit 1;
	} else {
		my ($configfile) = @ARGV;
		return ($configfile);
	}
}

# Given the path to the local configuration file, extracts all of the 
# configuration information we're interested in.
#
# At the moment, all we want is the path to our static LDIF data.

sub parseConfig($) {
	my ($configpath) = shift;	# Path to the configuration file.
    my $ldif_file;
    my @vos;

	# Open the file.
	open(CONFIG, $configpath) ||
		die "Unable to open configuration file $configpath: $!";
	
	# For every config file entry:
	while(<CONFIG>) {
		# If the line contains the ldif_file definition
		if (/^ldif_file=(.*)$/) {
			# Record it, and stop looking.
			$ldif_file = $1;
			chomp($ldif_file);
		}
        if (/^vos=(.*)$/) {
            # Record it, and stop looking.
            my $vo_list = $1;
            chomp($vo_list);
            @vos = split (/\s*[,\s]\s*/, $vo_list);
        }
	}
    close CONFIG;

	# If we get this far, we never found the ldif_file entry.
	# Treat this as a fatal error.
	die "Failed to find 'ldif_file' entry in configuration file " .
	    "'$configpath'.  Aborting." unless (defined $ldif_file);

    return ($ldif_file, [@vos]);
}

# Given the path to the local static LDIF information file, extract
# a list of DNs as defined in that file.
#
# The each DN will include the 'dn: GlueCEUniqueID=' prefix.

sub parseLDIF($ $ $) {
    my ($qRef) = shift;
    my ($vosRef) = shift;
	my ($ldifpath) = shift;

	open(LDIFDATA, $ldifpath) ||
		die "Unable to open static LDIF information file $ldifpath: $!";
	
    my $voExp = join ("|",@{$vosRef});
    my $queueExp = join ("|",keys %{$qRef});

	while(<LDIFDATA>) {
		# If the line contains a DN definition
		if (/dn:\s*GlueCEUniqueID=\S+($queueExp)\,/) {
            push @{$qRef->{$1}->{dn}}, $_;
            next;
		}
        if (/dn:\s+GlueVOViewLocalID=($voExp)\,\s*GlueCEUniqueID=\S+($queueExp)\,/) {
            $qRef->{$2}->{vos}->{$1} = {GlueCEStateRunningJobs=>0, 
                                        GlueCEStateWaitingJobs=>0, 
                                        GlueCEStateTotalJobs=>0};
            #push @{$qRef->{$2}->{vos}->{$1}->{dn}}, $_;
            $qRef->{$2}->{vos}->{$1}->{dn} = $_;
        }
	}
	close(LDIFDATA);

	1;
}

# This function will lookup the local SGE version number.
# It works by running the command 'qstat -help'; the first line of the output
# of this command will contain the local version number.  eg:
# 	SGEEE 5.3p3
# 	SGE 6.0u1
# However, we just need the version number -- so we strip off the first word
# from this string before returning.

sub lookupSGEVersion() {
	open(QSTAT, "$qstat -help |") || 
		die "Unable to lookup SGE version number; " .
		    "qstat command failed: $!";

	# Read in just the _first_ line of the qstat -help output.
	# This line will contain our version string.
	my ($version_string) = <QSTAT>;
	close QSTAT;
	
	# For the moment, we don't care whether we're running SGEEE or just 
	# SGE.  (Indeed, the Enterprise Edition was discontinued as of SGE 6.)
	# Extract _just_ the version number.
	if ($version_string =~ /^\S+\s(\S+)$/) {
		return $1;
	}
	else {
		# Uh-oh.  The version string (if that is, indeed, what we're 
		# looking at) doesn't appear to follow the format we expect.  
		# This suggests that either something has gone wrong _or_ 
		# we're talking to a new version of SGE that doesn't follow our
		# expected conventions.  Abort.
		die "Version string '$version_string' does not follow the " .
		    "expected format of 'NAME VERSION' -- Aborting!\n" .
		    "This script probably needs to be updated!"
	}
}

# uses qstat -g c to fetch available queues and their status.
#
# this function also creates queues in master queueData hash so it
# has to be invoked prior to any other method dealing with queueData
# Glue attributes:
#  - GlueCEStateFreeJobSlots
#  - GlueCEStateStatus
#  - GlueCEStateFreeCPUs (depricated)
#  - GlueCEInfoTotalCPUs (depricated)
#  - GlueCEPolicyMaxRunningJobs (remove later?)
#  - GlueCEPolicyAssignedJobSlots
# Glue attributes initialized only:
#  - GlueCEStateRunningJobs
#  - GlueCEStateWaitingJobs
#  - GlueCEStateTotalJobs
sub lookupQueueStatus ($) {
    my ($qRef) = shift;

    open(QSTAT, "$qstat -g c |") 
            || die "Unable to run qstat to lookup queue load: $!";

    foreach (<QSTAT>) {
        #CLUSTER QUEUE                   CQLOAD   USED    RES  AVAIL  TOTAL aoACDS  cdsuE
        next if (/^CLUSTER\s+QUEUE\s+CQLOAD\s+USED\s+RES\s+AVAIL\s+TOTAL\s+aoACDS\s+cdsuE\s*$/);
        next if (/^-+$/);
        # dteam                             0.96     23      1     24      0      0
        if (/^(\S+)\s+[\d\.]+\s+(\d+)\s+\d+\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s*$/) {
            my $name = $1;
            my $used = $2;
            my $avail = $3;
            my $total = $4;
            my $error = $5 + $6;

            $qRef->{$name} = {GlueCEStateRunningJobs=>0, 
                                GlueCEStateWaitingJobs=>0, 
                                GlueCEStateTotalJobs=>0};

            $qRef->{$name}->{GlueCEStateFreeJobSlots} = $avail;

            # Others are setting these depricated attributes 
            # Should we disable it? Is it used by some middleware?
            $qRef->{$name}->{GlueCEStateFreeCPUs} = $avail;
            $qRef->{$name}->{TotalSlots} = $total;
            $qRef->{$name}->{GlueCEInfoTotalCPUs} = $total;

            # SGE doesn't have mechanism for limiting number of jobs run by 
            # specific group of user in a single queue. Currently (6.0v4) is 
            # only possible to limit total number of running jobs per single 
            # user. We'll set it to maxium number of slots becuase it cannot
            # run more jobs than number of slots.
            $qRef->{$name}->{GlueCEPolicyMaxRunningJobs} = $total;

            # According to GLUE specification: "maximum number of single-processor
            # jobs that can be running at a given time". This is similar to 
            # GlueCEPolicyMaxRunningJobs. Reasonable value is total amount of CPUs
            $qRef->{$name}->{GlueCEPolicyAssignedJobSlots} = $total;

            # all queue instances are unavailable
            if ($total == $error) {
                # some jobs are still running - cluster is in draining state
                if ($used) {
                    $qRef->{$name}->{GlueCEStateStatus} = "Draining";
                } 
                # cluster is closed
                else {
                    $qRef->{$name}->{GlueCEStateStatus} = "Closed";
                }
            } else {
                # there are some available nodes 
                #if ($avail) {
                    $qRef->{$name}->{GlueCEStateStatus} = "Production";
                #} 
                # all nodes are either used or unavailable - cluster is in Queueing state
                #else {
                #    $qRef->{$name}->{GlueCEStateStatus} = "Queueing";
                #}
            }

        }
    }

    close QSTAT;
}

# function sets various global attributes
# uses qconf -sconf to get global values
# Glue attributes:
#  - GlueCEPolicyMaxTotalJobs
#  - GlueCEInfoLRMSType
#  - GlueCEInfoLRMSVersion
sub lookupGlobalQueueParams ($) {
    my $qRef = shift;
    my %globals;

    open(QCONF, "$qconf -sconf |") ||
        die "Unable to lookup global policy with qconf: $!";
    foreach (<QCONF>) {
        if (/^max_jobs\s+(\d+)\s*$/) {
            $globals{GlueCEPolicyMaxTotalJobs} = $1;
            last;
        }
    }
    close QCONF;

    $globals{GlueCEInfoLRMSType}     = "SGE";
    $globals{GlueCEInfoLRMSVersion}  = lookupSGEVersion();

    foreach my $queue (values %{$qRef}) {
        foreach my $key (keys %globals) {
            $queue->{$key} = $globals{$key};
        }
    }
}

# This function also get ACL and xACL needed for VOView
# Glue attributes:
#  - MaxWallClockTime
#  - MaxCPUTime
sub lookupQueueTimePolicy($) {
	my $qRef = shift;
    my $cqueue;
    my $queues = join (",",keys %{$qRef});

    open(QCONF, "$qconf -sq $queues |") || 
        die "Unable to lookup queue policy with qconf: $!";

    foreach (<QCONF>) {

        $cqueue = $1 if (/^qname\s+(\S+)\s+$/);

        if (/^h_rt\s+([\d\:]+|INFINITY)\s*$/) {
            $qRef->{$cqueue}->{"GlueCEPolicyMaxWallClockTime"} = sgeToGlue($1);
            next;
        }

        if (/^h_cpu\s+([\d\:]+|INFINITY)\s*$/) {
            $qRef->{$cqueue}->{"GlueCEPolicyMaxCPUTime"} = sgeToGlue($1);
            next;
        }

        if (/^(user_lists|projects)\s+(\S.+)\s*$/) {
            exists $qRef->{$cqueue}->{"ACL"} or
                $qRef->{$cqueue}->{"ACL"} = "";

            $qRef->{$cqueue}->{"ACL"} and 
                $qRef->{$cqueue}->{"ACL"} .= " ";

            $qRef->{$cqueue}->{"ACL"} .= $2 if ($2 ne "NONE");
                next;
        }

        if (/^(xuser_lists|xprojects)\s+(\S.+)\s*$/) {
            exists $qRef->{$cqueue}->{"xACL"} or
                $qRef->{$cqueue}->{"xACL"} = "";

            $qRef->{$cqueue}->{"xACL"} and 
                $qRef->{$cqueue}->{"xACL"} .= " ";

            $qRef->{$cqueue}->{"xACL"} .= $2 if ($2 ne "NONE");
            next;
        }
    }

    close QCONF;
}


# This function calculates response times
# Glue attributes:
#  - GlueCEStateWorstResponseTime
#  - GlueCEStateEstimatedResponseTime
sub calculateResponseEstimate($) {
    my $qRef = shift;

    foreach my $dataref (values %{$qRef}) {
       my ($defaultworstcase)   = 2 * 24 * 60 * 60; # 2 days.
       my ($worstjobtime)   = $dataref->{"GlueCEPolicyMaxWallClockTime"};
       my ($jobcount)       = $dataref->{"GlueCEStateTotalJobs"};
       my ($cpucount)       = $dataref->{"GlueCEInfoTotalCPUs"};
       my ($cpufree)        = $dataref->{"GlueCEStateFreeCPUs"};
       
       if ($worstjobtime == 0) {
          # There was no maximum wallclock time set.
          # Assume the $defaultworstcase under these circumstances.
          $worstjobtime = $defaultworstcase;
       }
       
       # The predicted worst-case response time is if every job will run 
       # ahead of any new job, and it will take all the time available to run.
       if ($cpufree > 0)
       {
          # If there are free CPUs job will executed straight away/
          $dataref->{"GlueCEStateWorstResponseTime"} = 0;
       }
       elsif ($cpucount == 0) {
          # No CPUs are available.  Put in a large wait time.
          $dataref->{"GlueCEStateWorstResponseTime"} = 99999999;
       }
       else {
          # The predicted worst-case response time is if every job 
          # will run ahead of any new job, and it will take all the 
          # time available to run.
          # (Algorithm taken from PBS information reporter, minus the
          #  job-time-already-used component as we don't have this 
          #  information.)
          # TODO: Is this sufficient?
          $dataref->{"GlueCEStateWorstResponseTime"} = 
             int(($jobcount * $worstjobtime) / $cpucount);
       }

       # The estimated response time is simply half of the worst-case 
       # response time.
       $dataref->{"GlueCEStateEstimatedResponseTime"} = 
          int($dataref->{"GlueCEStateWorstResponseTime"} / 2);  

    }
}

# This function retrieves information about jobs
# Glue attributes (CE):
#  - GlueCEStateRunningJobs
#  - GlueCEStateTotalJobs (partially, see lookupQueuedJobs)
# Glue attributes (VOView):
#  - GlueCEStateRunningJobs
#  - GlueCEStateTotalJobs
sub lookupJobState($ $) {
	my $qRef = shift;
	my $vosRef = shift;
    my %pend;

	open(QSTAT, "$qstat -u \"*\" -ext |")
		|| die "Unable to lookup jobs' data: $!";

    #my $voExp = join ("|",@{$vosRef});
    my $queueExp = join ("|",keys %{$qRef});
	
	foreach (<QSTAT>) {
		next if (/^job-ID/);
		next if (/^-+/);

		# If the line begins with a number, increment the job count.
		# It's a job entry.  (The number is the job's ID.)
        # 29910 0.25003 0.50000 sge_job_sc voce003      voce
        # Problem is how to determine the groups of user. In this case we assume
        # that each user is assigned to project with name VO.
        if (/^\s*(\d+)\s+\S+\s+\S+\s+\S+\s+\S+\s+(\S+?)\s+\S+\s+(\S+)(.*($queueExp)@)?/) {
            my $jobID = $1;
            my $vo = $2;
            my $state = $3;
            my $queue = $5; 

            $vo = 'crongi' if ($vo eq 'grid');

            if ($state =~ /r/) {
                if (defined $queue) {
                    $qRef->{$queue}->{GlueCEStateRunningJobs}++ ;
                    $qRef->{$queue}->{GlueCEStateTotalJobs}++;
                    $qRef->{$queue}->{vos}->{$vo}->{GlueCEStateRunningJobs}++;
                    $qRef->{$queue}->{vos}->{$vo}->{GlueCEStateTotalJobs}++;
                } 
            } else {
                # at this moment we don't know to which queue does this job belong
                $pend{$jobID} = $vo;
            }

        }
	}
    close QSTAT;

    \%pend;
}

# workaround for qstat non showing correctly queue for pending jobs
# only way to get this information is with qstat -j <jobID>
# Glue attributes (CE):
#  - GlueCEStateWaitingJobs
#  - GlueCEStateTotalJobs (partially, see lookupJobState)
sub lookupQueuedJobs($ $) {
    my $qRef = shift;
    my $pend = shift;
    my $jobId;

    my $queueExp = join ("|",keys %{$qRef});
    my $pending = join (",",keys %{$pend});

    open(QSTAT, "$qstat -j $pending |")
        || die "Unable to lookup waiting jobs' data: $!";

    foreach (<QSTAT>) {
        if (/^job_number:\s+(\d+)/) {
            $jobId = $1;
            next;
        }
        if (/^hard_queue_list:\s+($queueExp)/) {
            $qRef->{$1}->{GlueCEStateWaitingJobs}++;
            $qRef->{$1}->{GlueCEStateTotalJobs}++;
            $qRef->{$1}->{vos}->{$pend->{$jobId}}->{GlueCEStateWaitingJobs}++;
            $qRef->{$1}->{vos}->{$pend->{$jobId}}->{GlueCEStateTotalJobs}++;
            delete $pend->{$jobId};
        }
    }
    close QSTAT;

    # if the queue is not defined add the job to all queues that accept this VO
    foreach my $job (keys %{$pend}) {
        foreach my $queue (values %{$qRef}) {
            if (exists $queue->{vos}->{$pend->{$job}}) {
                $queue->{GlueCEStateWaitingJobs}++;
                $queue->{GlueCEStateTotalJobs}++;
                $queue->{vos}->{$pend->{$job}}->{GlueCEStateWaitingJobs}++;
                $queue->{vos}->{$pend->{$job}}->{GlueCEStateTotalJobs}++;
            }
        }
    }
}

# --- Main program ------------------------------------------------------

# Lookup the path to our configuration file and the hostname we should use
# from the commandline.


my $configpath = processCommandLine();

my @vos;
my $ldifpath;
my $voRef; 

($ldifpath, $voRef) = parseConfig($configpath);
@vos = @{$voRef};

my %queueData;

# get queue data
lookupQueueStatus(\%queueData);
parseLDIF(\%queueData,\@vos,$ldifpath);
lookupGlobalQueueParams(\%queueData);
lookupQueueTimePolicy(\%queueData);

my $pend = lookupJobState(\%queueData, \@vos);
lookupQueuedJobs (\%queueData, $pend) if $pend;

calculateResponseEstimate(\%queueData);

my $output = "";

foreach my $queue (values %queueData) {
    my $cevalues = "";

    # print GlueCE DNs and attributes
    foreach my $key ( sort keys %{$queue} ) {
        $cevalues .= "$key" . ": " . $queue->{$key} . "\n" if ($key =~ /^GlueCE/);
    }
    foreach my $dn (@{$queue->{dn}}) {
        $output .= $dn . $cevalues . "\n";
    }

    # static values for VO
    #  - GlueCEStateWorstResponseTime
    #  - GlueCEStateEstimatedResponseTime
    #  - GlueCEStateFreeJobSlots
    #  - GlueCEStateFreeCPUs
    $cevalues = "";
    foreach my $key (@voViewStatic) {
        $cevalues .= "$key" . ": " . $queue->{$key} . "\n"
    }

    foreach my $vo (values %{$queue->{vos}}) {
	next unless (defined $vo->{dn});
        $output .= $vo->{dn};
        foreach my $key ( sort keys %{$vo} ) {
            $output .= "$key" . ": " . $vo->{$key} . "\n" if ($key =~ /^GlueCE/);
        }
        $output .= $cevalues . "\n";
    }
}

print $output;