#! /usr/bin/perl -w
#
# Distribute v1.1. By Christopher Neufeld. GPL 1998.
#
# A script to run multiple invocations of a named program on several
# different machines, so that each invocation is run only once, and
# optimal use is made of the CPUs.
#
# A daemon mode exists as of v1.1, primarily useful for performing
# distributed makes.
#
# The machine running the script must have rsh authority or some
# equivalent on the remote machines.
#
# Note: if you see some strange comments, like:   # ])}
# it is to clear up some confusion which certain regular expressions
# provoke in the EMACS perl-mode indentation and fontification routines.
#

use IO::Socket;
use IO::Select;

if ( $#ARGV == 0 ) {  # Not daemon mode, commands in config file
    $daemonmode = 0;
} elsif ( $#ARGV == 1 ) {
    $switch = shift @ARGV;
    if ($switch =~ /-d/ ||
	$switch =~ /--daemon/ ) {
	$daemonmode = 1;
    } else {
	die "Unrecognized switch: \"$switch\"\n" ;
    }
} else {
    die "Usage: distribute [ -d | --daemon ] configfile\n" ;
}


if ($daemonmode) {

    $home = $ENV{'HOME'} || $ENV{'LOGDIR'} ||
	(getpwuid($<))[7] || die "You're homeless!\n";

    $localsockpathdir = "$home/.distribute";
    $localsockfname = "distribute";
    $localsockpathname = $localsockpathdir . "/" . $localsockfname ;

    if ( ! -d $localsockpathdir ) {
	mkdir $localsockpathdir, 0700 || die "Can't create socket directory\n";
    } else {
	$mode = (stat $localsockpathdir)[2];
	$newmode = $mode & 0700 ;
	if ( $mode != $newmode ) {
	    chmod $newmode, $localsockpathdir;
	}
	$othermode = (stat $localsockpathdir)[2] & 0077;
	die "Can't set secure permissions on socket directory\n" if $othermode != 0 ;
    }
    $localserver = IO::Socket::UNIX->new( Type   => SOCK_STREAM,
					  Local  => "$localsockpathname",
					  Listen => SOMAXCONN
					);
    $SIG{'INT'} = 'cleanup';
    $SIG{'QUIT'} = 'cleanup';
    $SIG{'TERM'} = 'cleanup';
}


$socket = 2233;			# The socket number we'll use to communicate
                                # with remote clients
$numtries = 10;			# If that socket isn't available to us,
                                # count up until we find one that is.


# Make the socket for communicating with remote machine clients.
for ($i = 0; $i < $numtries; $i++, $socket++) {

    $server = IO::Socket::INET->new( Proto     => 'tcp',
				     LocalPort => $socket,
				     Listen    => SOMAXCONN,
				     Reuse     => 1
				   );
    last if $server ;
}

die "Unable to open communications socket.\n" if ! $server ;


$sel_cl = new IO::Select( );     # The select type for the remote clients



# The command which invokes PERL on remote machines, and draws its
# script from STDIN
$perlcommand = "/usr/bin/perl -x" ;


# The remote perl script. When it is ready to run another process it
# sends the string "READY" to the server. The server responds with the
# string "EXECUTE", followed by a newline, then the command to
# execute, or "FINISHED" if all processes are done. Output from each
# command, if not suppressed, is stored in a temporary file and then
# sent back up the socket, each line prefixed with the sequence "# ".
# After the last line of output, the string "DONE" is sent. The script
# ends after it has seen the string "FINISHED" and all running processes
# have returned. When its last child has exited, it sends the string
# "DISCONNECTING" and exits.
#
# It is given certain constants by the calling invocation: The socket
# number, the number of parallel processes desired on the remote
# machine, whether or not to suppress output, the name of the
# server, whether or not to report when a command completes, and
# whether the program is running in daemon mode.
#

# SCRIPT DOWNLOADED TO REMOTE CLIENTS STARTS HERE:
#
$perlprog = '#! /usr/bin/perl
    use IO::Socket;
    use POSIX;
    use Fcntl;

# Note!!! The following pack works under Linux, it assumes that
# an off_t is a signed long, and a pid_t is a signed int
    $excl_wrlock = pack("sslli", F_WRLCK, 0, 0, 0, 0);
    $unlock_excl = pack("sslli", F_UNLCK, 0, 0, 0, 0);

# Parameters passed from the invoker
    $socket = .SOCKET. ;
    $suppressoutput = .SUPPRESS. ;
    $numparallel = .PARALLEL. ;
    $servername = .SERVERNAME. ;
    $statreports = .STATUSREPORTS. ;
    $daemonmode = .DAEMONMODE. ;


    $remote = IO::Socket::INET->new(
                        Proto    => "tcp",
                        PeerAddr => "$servername",
                        PeerPort => $socket
              );
    die "Unable to connect to remote host.\n" if ! $remote ;
    $remote->autoflush(1);

    close(STDIN); close(STDOUT); close(STDERR);    # disconnect from terminal
    exit if fork ;        # detach from terminal, or die if can\'t fork

    $numkids = 0;
    $seenfinish = 0;

    $SIG{\'CHLD\'} = sub { wait ; $numkids-- ; };

    while (! $seenfinish) {
        if ($numkids < $numparallel) {
            fcntl($remote, F_SETLKW, $returnbuf = $excl_wrlock);
            print $remote "READY\n" ;
            fcntl($remote, F_SETLK, $returnbuf = $unlock_excl);

            $response = get_full_line($remote) ;
            if ( $response =~ m/^EXECUTE [0-9]*$/ ) {
                $job_id = (split(\' \', $response))[1];
                if ( $daemonmode ) {
                    $numkids++;
                    $forkstat = fork ;
                    if ( ! $forkstat ) {
                        close $remote;
                        $remote2 = IO::Socket::INET->new(
                                             Proto    => "tcp",
                                             PeerAddr => "$servername",
                                             PeerPort => $socket
                                   );
                        $remote2->autoflush(1);
                        $command = get_full_line($remote2);
                        do_invocation($remote2, $command);
                        exit ;
                    } elsif ( $forkstat == -1 ) {
                      die "Can\'t fork.\n";
                    }
                } else {
                    $command = get_full_line($remote);
                    $numkids++;
                    $forkstat = fork ;
                    if ( ! $forkstat ) {
                        do_invocation($remote, $command);
                        exit ;
                    } elsif ( $forkstat == -1 ) {
                      die "Can\'t fork.\n";
                    }
                }
            } elsif ( $response =~ m/^FINISHED$/ ) {
                $seenfinish = 1;
            } else {
                die "Unexpected response from server.\n";
            }
        } else {
            sleep ;
        }
    }
    while (wait != -1) { }   # Reap the rest of the kids.

    print $remote "DISCONNECTING\n" ;
    close $remote;
    exit;


sub get_full_line {
    my $HANDLE = pop(@_);
    my $line = \'\';
    my $fragment = \'\';
    my $n = 1;

    while ($n && $fragment ne "\n") {
	$n = sysread($HANDLE, $fragment, 1);
	$line .= $fragment ;
    }

    return $line;
}

sub do_invocation {
    my ($HANDLE, $command) = @_;

    chomp $command;

    if ( ! $suppressoutput ) {
	$tempfile = POSIX::tmpnam();
	$retcode = (system "(" . $command . ") > $tempfile 2>&1" ) / 256;
    } else {
	$retcode = ( system $command ) / 256 ;
    }
    fcntl($HANDLE, F_SETLKW, $returnbuf = $excl_wrlock);
    print $HANDLE "OUTPUT $job_id\n";
    if ( ! $suppressoutput ) {
	open(OSTREAM, "$tempfile");
	while (<OSTREAM>) {
	    if ( ! m/\n$/ ) {
		$_ .= "\n" ;
	    }
	    print $HANDLE "# ", $_;
	}
	close(OSTREAM);
	unlink $tempfile;
    }
    if ( $statreports || $retcode ) {
	print $HANDLE "# Finished: ", $command, ". Return code= $retcode.\n" ;
    }
    print $HANDLE "DONE $retcode\n";
    fcntl($HANDLE, F_SETLK, $returnbuf = $unlock_excl);

}

' ;				# '

# END OF SCRIPT DOWNLOADED TO REMOTE CLIENTS




# Some initializations and defaults
#
@machinenames = ();
@parallel = ();
@srpatterns = ();
$command = "echo hi";
$suppressoutput = 1;
$statusreports = 0;
$rshcommand = '/usr/bin/rsh -l $2 $1';


$nummachines = 0;           # number of remote machines to connect
$numclients = 0;            # number of active connections opened

# Who, and where, am I?
#
$myusername = getlogin || (getpwuid($<))[0] || "nobody";
chomp ( $myhostname = `hostname`);


# Prepare to read in the configuration file
#
open(CONFIG, $ARGV[0]);


# We want to notice multiple instances of some directives
#
$beenhere1 = 0;
$beenhere2 = 0;
$beenhere3 = 0;
$beenhere4 = 0;

# Read setup.
while (<CONFIG>) { #
    next if m/^\s*(#.*)*$/ ;	# Flush blank and comment lines

    if ( m/^\s*Machine:[ \t]/i ) { # Machine: specifier
	@junk = split_into_strings($_);
	shift @junk ;
        ($v1, $v2, @v3) = @junk ;

        die "Bad syntax in \"Machine:\" directive\n$_"
            if ! defined ($v1) ||
            ! defined($v2) ||
            $v1 eq "" ||
            $v2 < 1 ;

	$machinenames[$nummachines] = $v1;    # name
	$parallel[$nummachines] = $v2;        # number of parallel runs
        $srpatterns{$v1} = join(' ', @v3) ;   # search/replace patterns
	$nummachines++;
	next;
    }

    if ( m/^\s*Suppress_output:[ \t]/i ) {    # suppress output of remote cmds?
        die "Multiple \"Suppress_output:\" directives\n" if $beenhere1++;
	$v1 = (split)[1];
	die "Bad syntax in \"Suppress_output:\" directive\n$_"
	    if ! defined ($v1);

	if ( $v1 =~ m/TRUE/i ) {
	    $suppressoutput = 1;
	} elsif ( $v1 =~ m/FALSE/i ) {
	    $suppressoutput = 0;
	} else {
	    die "Bad argument to \"Supress_output:\" directive\n$_" ;
	}
	next;
    }

    if ( m/^\s*Status_reports:[ \t]/i ) {    # suppress reporting of completed cmds
        die "Multiple \"Status_reports:\" directives\n" if $beenhere2++;
	$v1 = (split)[1];
	die "Bad \"Status_reports:\" directive\n" if ! defined ($v1);

	if ( $v1 =~ m/TRUE/i ) {
	    $statusreports = 1;
	} elsif ( $v1 =~ m/FALSE/i ) {
	    $statusreports = 0;
	} else {
	    die "Bad argument to \"Status_reports:\" directive\n" ;
	}
	next;
    }

    if ( m/^\s*Run_as_user:[ \t]/i ) {    # Do we use a different username remotely?
	($v1, $v2, $v3) = split;
	die "Bad arguments to \"Run_as_user:\" directive\n"
	    if ! defined ($v2) || ! defined ($v3) || $v2 eq "" || $v3 eq "" ;
	die "\"Run_as_user:\" directive repeated for machine $v2\n"
	    if defined( $remotenames{$v2} );
	$remotenames{$v2} = $v3;
	next;
    }

    if ( m/^\s*Remote_exec_command:[ \t]/i ) {  # A different rsh command
        die "Multiple \"Remote_exec_command:\" directives\n" if $beenhere3++;
        @v3 = split;
        shift @v3 ;                   # Discard the "Remote_exec_command:" field
	$rshcommand = join(' ', @v3);
	next;
    }

    if ( m/^\s*Command:[ \t]/i ) {    # The command to run
	warn "Command: entry will be ignored in daemon mode.\n" if $daemonmode;
        die "Multiple \"Command:\" directives\n" if $beenhere4++;
	@v3 = split;
        shift @v3 ;                   # Discard the "Command:" field
	$command = join(' ', @v3);
	next;
    }

    if ( m/^\s*Runlist:\s*$/i ) {
	$inrunlist = 1;
	warn "Runlist: entry will be ignored in daemon mode.\n" if $daemonmode;
	last ;                        # Move to the execution stage
    }

    die "Unrecognized directive: \"$_\"\n";
}

if ( ! $inrunlist && ! $daemonmode ) {
    die "No runlist found. Nothing to do!\n" ;
}

die "No remote machines specified\n" if $nummachines == 0;

# Open up connections to the remote machines, and prepare to pipe
# commands down the line at them.

$SIG{'CHLD'} = sub { wait };      # Reap zombies when the appear

for ($i = 0; $i < $nummachines; $i++) {

    $macname = $machinenames[$i];
    $remname = $remotenames{$macname} ;
    if ( ! defined ($remname) ) {
        $remname = $myusername;
    }


    # First, the remote command itself

    $v1 = $rshcommand;
    $v1 =~ s/([^\\])\$1$/$1$macname/g ;	# $1 at end of the line
    $v1 =~ s/([^\\])\$1(\D)/$1$macname$2/g ; # $1 in the line
    $v1 =~ s/([^\\])\$2$/$1$remname/g ;	# $2 at end of the line
    $v1 =~ s/([^\\])\$2(\D)/$1$remname$2/g ; # $2 in the line

# Pretty up the PERL program. Insert the appropriate values in the script.
#
    $v2 = $perlprog;
    $v2 =~ s/\.SOCKET\./$socket/g ;
    $v2 =~ s/\.SUPPRESS\./$suppressoutput/g ;
    $v2 =~ s/\.PARALLEL\./$parallel[$i]/g ;
    $v2 =~ s/\.SERVERNAME\./$myhostname/g ;
    $v2 =~ s/\.STATUSREPORTS\./$statusreports/g ;
    $v2 =~ s/\.DAEMONMODE\./$daemonmode/g ;

    $forkstat = fork;
    if ( ! $forkstat ) {     # invoke the client
	open (SREM, "|-" ) || exec split(' ', join(' ', $v1, $perlcommand));
	print SREM $v2;   # send the script down the pipe to the remote perl
	close SREM;       # disconnect
	exit ;
    } elsif ( $forkstat == -1 ) {
	die "Can't fork\n";
    }

# Collect a client
#
    $client[$i] = $server->accept() || die "Client on $macname failed to connect.\n" ;
    $client[$i]->autoflush(1);
    $numclients++;

    $sel_cl->add($client[$i]);
    $clientname[fileno($client[$i])] = $macname;  # Save the machine name
}


if ( $daemonmode ) {
    handle_daemonmode() ;
} else {
    handle_runlist() ;
    exit;
}


# handle_daemonmode opens a UNIX domain socket on the local machine and waits
# for local clients to attach to the socket and submit jobs. When they do, it
# forks off to pass the data back and forth between the clients and the
# remote scripts. This subroutine never returns.
sub handle_daemonmode {

    print "Daemon configured\n";

    # Detach from session, or die if can't fork.
    exit if fork ;

    close CONFIG;

    while (1) {           # forever

	$localclient = $localserver->accept() ;
# Now we have an incoming local client request. Find a remote client
# which can handle it.
	do {
	    @ready = $sel_cl->can_read ;   # Wait for a remote client to come ready
	} while $#ready == -1;
#	$pickone = int (rand ( $#ready + 1 ));
	$pickone = 0;
	$fh = $ready[$pickone];               # This client is ready.

	$message = get_full_line($fh);
	if ( ! $message =~ m/^READY$/ ) {  # This is the only thing which
	                                   # should come in when in daemon mode
	    die "Unexpected message from remote client.\n";
	}

	# Tell the remote machine something's coming, then wait for it
	# to open the communications socket for I/O.
	print $fh "EXECUTE ", fileno($fh), "\n";
	$iosocket = $server->accept() || die "Failed to open I/O socket.\n";
	$iosocket->autoflush(1);
	$forkstat = fork ;
	if ( ! $forkstat ) {   # I'm the child.
	    $syscall = get_full_line($localclient);  # Get the command to run
	    chomp $syscall;

	    # Perform search/replace commands for this machine
	    @v3 = split_into_strings($srpatterns{$clientname[fileno($fh)]});
	    foreach $i (@v3) {
		$i =~ s/^(["'])(.*)\1$/$2/ ;         # "])
		eval "\$syscall =~ $i" ;
	    }

	    print $iosocket $syscall, "\n" ;   # Send the command

	    # loop on output of command
	    while (1) {
		$message = get_full_line($iosocket);
		print $localclient $message;
		last if $message =~ m/DONE [0-9]*/ ;
	    }

	    close $iosocket;
	    close $localclient;
	    exit;
	} elsif ( $forkstat == -1 ) {
	    die "Can't fork off relaying child.\n";
	} else {    # Detach from the things we don't need to know about
	    close $iosocket;
	    close $localclient;
	}
    }
}


# handle_runlist takes the commands to run from the configuration file
# and passes them to the clients.
sub handle_runlist {

# Now, execute the commands
    while ($numclients) {

	# Wait for a client to send a message
	@ready = $sel_cl->can_read ;

	foreach $fh (@ready) {

	    $message = get_full_line($fh) ;

	    if ( $message =~ m/^READY$/ ) {   # Client asked for another command
		if ($parms = <CONFIG>) {      # There are more commands to run
		    chomp $parms ;
		    @runline = split_into_strings($parms);
		    unshift @runline, '' ;  # Make the first element numbered '1'

		    $syscall = $command;
		    # Do $NUM substitution.
		    for ($i = 1; $i <= $#runline; $i++) {
			$syscall =~ s/([^\\])\$$i$/$1$runline[$i]/g ;
			$syscall =~ s/([^\\])\$$i(\D)/$1$runline[$i]$2/g ;
		    }
		    $syscall =~ s/([^\\])\$\*/$1$parms/g; # Do $* substitution
		    $syscall =~ s/\\\$/\$/g ;

		    # Now, execute any search/replace commands on this string
		    @v3 = split_into_strings($srpatterns{$clientname[fileno($fh)]});
		    foreach $i (@v3) {
			$i =~ s/^(["'])(.*)\1$/$2/ ;         # "])
			eval "\$syscall =~ $i" ;
		    }

		    # Tell the client that it is about to receive a command, then send the command
		    print $fh "EXECUTE 0\n" ;
		    print $fh $syscall, "\n" ;
		} else {
		    print $fh "FINISHED\n" ;  # Tell the client all commands have been sent already
		}
	    } elsif ( $message =~ m/^DISCONNECTING$/ ) {  # Client is shutting down
		$sel_cl->remove($fh);
		$fh->close;
		$numclients--;
	    } elsif ( $message =~ m/^OUTPUT [0-9]*$/ ) {   # Text response from client, echo it
		while (1) {
		    $response = get_full_line($fh);   # Get next line of tex
		    last if $response =~ m/^DONE [0-9]*$/ ;  # End of tex
		    $response =~ s/^\# //g ;    # strip the protecting '# ' string
		    print $response ;
		}
	    } elsif ( $message =~ m/^DONE [0-9]*$/ ) {   # Client has finished command
		# do nothing
	    } else {
		die "Bad response from client.\n";
	    }
	}
    }
}


# Non-blocking socket reads can be annoying. You might not get all of
# the data at first. Another nuisance is that <> might buffer input,
# reading in the next two lines and then presenting the first of those
# lines to the variable calling with <>. Unfortunately, the other line
# has already been read, so a select() call will show that there is no
# data pending on the input filehandle, and everything locks up solid.
# Get around this by using the sysread() call.
#
sub get_full_line {
    my $HANDLE = pop(@_);
    my $line = '';
    my $fragment = '';
    my $n = 1;

    while ($n && $fragment ne "\n") {
	$n = sysread($HANDLE, $fragment, 1);
	$line .= $fragment ;
    }

    return $line;
}



# A subroutine, modified from the PERL FAQ and originally attributed
# to Jeff Friedl (jfriedl@omron.co.jp) to split a line of text over
# blanks and spaces, but not to split fields which are single, double,
# or backtick-quoted. Takes one argument, the string to split. Retains
# the quoting characters at the beginning and end of the quoted string
# to protect the string from the shell, when that time comes.
#
sub split_into_strings { 
    my @result = (); 
    my $text = pop(@_);

    push(@result, defined($2) ? $1 . $2 . $1 : $4)
	while $text =~ m/(["'`])([^\1\\]*(\\.[^\1\\]*)*)\1|([^ \t\n]+)/g;   # "])

    @result;
}


# A signal handler to clean up when the daemon is killed.
#
sub cleanup {
    unlink $localsockpathname ;     # delete the socket.
    rmdir $localsockpathdir ;       # delete the directory, if empty.
    exit ;
}
