#!/bin/bash
#Copyright 1999, Bascom Global Internet Services 
#http://www.bascom.com and William Stearns wstearns@pobox.com
#Released under the GNU General Public License (GPL)
#
#This set of functions is an attempt to enforce exclusive access in
#bash.  In a setting where only one copy of a script can run, that
#script tests the return code from "getlock scriptname".  If two 
#copies of the script are run simultaneously, only one will 
#acquire the lock and do its work; the other will exit.

#Approach: A script that wants the lock appends its PID to a lock file.
#The lock is acquired by the script that writes first; all others
#get a false return value from the getlock() function.  The losing
#scripts may exit or wait and try again, as appropriate.

#Possible issues:
#- If the script that gets the lock dies or fails to release it, no more
#scripts will be able to run until the lock is manually deleted.  We trap
#SIGHUP, SIGINT and exit once we own the lock and clean up on the way out
#to minimize this possibility.
#- The locking code can check to see if the written PID is ever
#corrupted, just in case bash's ">>" append is not atomic.
#- The delay between the write and the test should not be an issue; we're
#using append, so no value gets overwritten.
#- If copy2 writes to the file, and copy1 (which already has the lock) 
#releases the lock by deleting the file, then copy2 will be catting an 
#empty file, thus not getting the lock.  No problem.
#- Anyone truly familiar with formal locking may wish to have a barf 
#bag handy while reading this.  There are more race conditions here than
#at the Indianapolis Speedway.  That said, it seems to work generally as
#expected, with the exception of it's inability to safely clean up a
#stale lock (which really shouldn't happen unless you "kill -9" a script
#that owns the lock).  If you need real locking, use C, perl, or assembly.
#- This is _not_ a first come, first serve, queue.  Spawn off 10 of these
#in the background and you're as likely to have the 4th copy of the script
#get the lock as the first copy.

warn () {
#Notify someone of a problem
	echo $* >/dev/stderr
	echo $* >>/tmp/debuglog
	#mail -s "$*" root@localhost	#Hope you like a lot of mail!
}

stalelock () {
#Return True if the lock owner (first line in the lock file) has _no_ corresponding running PID, False otherwise.
	if [ -n "$1" ]; then
		CURROWNER="`cat ${LOCKDIR}${1}.lock 2>/dev/null | head -1`"
		#If: there _is_ a first line, that line has no matching running task and that lockfile still exists, then
		if [ -n "$CURROWNER" ] && ! ps -p "$CURROWNER" >/dev/null 2>/dev/null && [ -f ${LOCKDIR}${1}.lock ]; then
			#Last double check, race here still
			if [ "`cat ${LOCKDIR}${1}.lock 2>/dev/null | head -1`" = "$CURROWNER" ]; then
				#rm -f ${LOCKDIR}${1}.lock	#not safe to do		#Clean up dead lock if an old lockholder died before releasing.
				return 0	#True, lock is stale
			fi
		fi
	else
		warn Missing lock name in stalelock! ; exit
	fi
	return 1	#False, lock is not stale
}

getlock () {
#Attempt to acquire the lock by writing our PID to the lock file.  Return True if our PID is the first line, false otherwise.
	if [ -n "$1" ]; then
		if [ -f ${LOCKDIR}${1}.lock ] && [ ! -w "${LOCKDIR}${1}.lock" ]; then warn Lock \"$LOCKDIR${1}.lock\" not writable.  Exiting. ; exit ; fi

		MYPID="$$"
		#Hopefully the following is atomic.
		echo $MYPID >>${LOCKDIR}${1}.lock
		if [ "`cat ${LOCKDIR}${1}.lock 2>/dev/null | head -1`" = "$MYPID" ]; then
			return 0	#True, we acquired the lock
		#elif [ `cat ${LOCKDIR}${1}.lock 2>/dev/null | grep "^$MYPID\$" | wc -l` -eq 1 ]; then
		#	return 1	#False, we failed to acquire the lock, but at least our uncorrupted pid is in the file.
		#elif [ ! -f ${LOCKDIR}${1}.lock ]; then
		#	return 1	#False, we failed to acquire the lock, a previous owner deleted the file after we appended.
		else
			#If we've done the above tests, our pid probably got corrupted in the file.  Whine.
			#cat ${LOCKDIR}${1}.lock 2>&1 | mail -s "$MYPID corrupted" root@localhost
			return 1	#False, we failed to acquire the lock
		fi
	else
		warn Missing lock name in getlock! ; exit
	fi
}

releaselock () {
#For use by the lockholder only!
#Release the lock by erasing the lockfile
	if [ -n "$1" ]; then
		if [ -f ${LOCKDIR}${1}.lock ]; then
			rm -f ${LOCKDIR}${1}.lock
		else
			warn ${LOCKDIR}${1}.lock has been stolen from under us!
		fi
	else
		warn Missing lock name in releaselock! ; exit
	fi
}

cleanup () {
#If we catch SIGHUP, SIGINT, or exit from the script, release the lock.
#Do not trap to this unless we own the lock.
#This needs to be generalized to handle arbitrary locknames.
	warn Cleaning up on the way out
	releaselock locktest
}

#Set up locking
LOCKDIR="/var/run/"		#Needs trailing slash
if [ ! -d "${LOCKDIR}" ]; then warn Lock directory \"$LOCKDIR\" nonexistant.  Exiting. ; exit ; fi
if [ ! -w "${LOCKDIR}" ]; then warn Lock directory \"$LOCKDIR\" not writable.  Exiting. ; exit ; fi

#Beginning of main code
if stalelock locktest ; then
	warn Lock stale
	cat ${LOCKDIR}locktest.lock | mail -s "dead pid locking locktest" root@localhost	#Someone has to clean up the file manually.
elif getlock locktest; then
	trap cleanup 0,SIGHUP,SIGINT		#If we die, release the lock on the way out.
	warn Acquired lock
	#Do exclusive access stuff here.
	echo Yup, doing stuff
	sleep 10							#I can't _imagine_ the system corruption if two copies of sleep ran at the same time! *smile*
	echo Yup, did stuff
	releaselock locktest
	trap - 0,SIGHUP,SIGINT				#No need to clean up anymore.
else
	warn Failed to acquire lock, exiting
	#Do _not_ call releaselock here - someone else owns it.
	#If this copy _must_ run sometime (but we still need to avoid multiple
	#simultaneous copies), the program could wait some $RANDOM number of 
	#seconds and try again instead of exiting.
fi