/*
 * locking.c
 * This program demonstrates how to use the file locking
 * features.
 *
 * to compile:  cc -O -pipe   locking.c  -o locking
 * 
 */

#include <stdio.h>
#include <errno.h>
#include <setjmp.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/time.h>
#include <sys/file.h>
#include <signal.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <sys/resource.h>
#include <unistd.h>
#include <fcntl.h>

#define MODE          (S_IRWXU | S_IRWXG | S_IROTH | S_IWOTH)
#define MAX_TIME      (60 * 60 * 2)
#define MONKEY_BOOK   ("monkey_book.out")
#define PAGE_END      ("\n\n----------------------END-OF-PAGE-----------------------\n\n")
#define END_LEN       (60)
#define PAGE_LEN      (500 * 4)
#define LOCK_LEN      (PAGE_LEN + END_LEN);
#define LETTERS       (" abcdefghijklmnopqrstuvwxyz\n ABCDEFGHIJKLMNOPQRSTUVWXYZ.! \n ")

/*  
 * To use FLOCK just uncomment the line below. Otherwise
 * comment it out to use the fcntl function instead.
 */
/* #define USE_FLOCK   */

/* our functions */
void million_monkies(int fd);
static void clean_up(int sig);
static void parent_term(int sig);

/* 
 * Save the enviroment for a signal return to
 * clean up.
 */
jmp_buf  env; 

int main(int argc, char **argv)
{
	int children, fd, i;
	struct rlimit limit;

	/*
	 * Get our args 
	 */
	if( argc > 1 )
	{
		children = atoi(argv[1]);

		/* get our limits */
		if( getrlimit(RLIMIT_NPROC, &limit) == 0 )
		{
			/* Check to see if we are under our limits  */
			if( children < limit.rlim_max )
			{
				/* open our file */
				fd = open(MONKEY_BOOK, (O_RDWR | O_CREAT), MODE );
				if( fd > 0 )
				{
					/* Lets get those monkies typing ! */
					for( i = 0; i < children; i++ )
					{
						if( fork() == 0 )
						{
							signal(SIGHUP, clean_up);
							million_monkies(fd);
							close(fd);

						} /* if(fork) */

					} /* for */

					/* set the parent's signals */
					signal(SIGHUP, SIG_IGN);
					signal(SIGINT, parent_term);
					signal(SIGTERM, parent_term);
					signal(SIGKILL, parent_term);
					if( setjmp(env) == 0 )
					{
						sleep(MAX_TIME);
						/* Exciting: Here we have a use for the raise call ! */
						raise(SIGTERM); 
					}
					else
					{
						printf("\nSignal recieved, killing off children \n");
						kill(0, SIGHUP);

						/*
						 * now we have to wait for our children
						 */
						while( waitpid(-1, (int *)NULL, 0) > 0 )
						{
							;
						}
					}

				} /* if( fd ) */
				else
				{
					printf("ERROR ! cannot open %s\n", MONKEY_BOOK );
				}

			} /* if( children ) */
			else
			{
				printf("ERROR ! rlimits nproc is %u \n", limit.rlim_max );
			}

		} /* if( getrlimit ) */
		else
		{
			printf("ERROR ! the call to getrlimit failed \n");
		}

	} /* if( argc ) */
	else
	{
		printf("ERROR ! you need to supply and interger \n");
	}

	/* FIN */
	return(0);
}

#ifndef USE_FLOCK

/*
 * They will each lock the file, and write a
 * page. 
 */
void million_monkies(int fd)
{
	int i, len;
	char buff[PAGE_LEN], *letters = LETTERS;  
	struct flock lock;

	/* seed our random generator */
	srand( (time(NULL) - getpid()) / getpid() );
	len = strlen(letters);
	while( 1 )
	{
		for( i = 0; i < PAGE_LEN; i++ )
		{
			buff[i] = letters[ (rand() % len) ];
		}
		/*
		 * Fill out our flock structure. Note
		 * that we are just going to the end of the
		 * file for our lock.
		 */
		lock.l_len = LOCK_LEN;
		lock.l_pid = getpid();
		lock.l_start = 0;
		lock.l_type = F_WRLCK;
		lock.l_whence = SEEK_END;

		/* we will wait untill the lock is granted */
		if( fcntl(fd, F_SETLK, &lock) == 0 )
		{
			/* we have the lock now write to it */
			printf("FCNTL lock granted, Process [ %d ] writing \n", getpid() );
			write(fd, buff, PAGE_LEN);
			write(fd, PAGE_END, END_LEN);
			/* now unlock it */
			lock.l_type = F_UNLCK;
			if( fcntl(fd, F_SETLK, &lock) != 0 )
			{
				perror("Error unlocking the file\n");
			}
		} /* if( fcntl) */

		usleep( (rand() % getpid()) * 100 );
	} /* while(1) */

}
#endif


#ifdef USE_FLOCK
/*
 * Same as above except we use flock. As you can
 * see the function is much cleaner
 */
void million_monkies(int fd)
{
	int i, len;
	char buff[PAGE_LEN], *letters = LETTERS;  

	/* seed our random generator */
	srand( (time(NULL) - getpid()) / getpid() );
	len = strlen(letters);
	while( 1 )
	{
		for( i = 0; i < PAGE_LEN; i++ )
		{
			buff[i] = letters[ (rand() % len) ];
		}

		/* now lock the file and write our data */
		if( flock(fd, LOCK_EX) == 0 )
		{
			/* we have the lock now write to the file */
			printf("FLOCK granted, Process [ %d ] writing \n", getpid() );
			write(fd, buff, PAGE_LEN);
			write(fd, PAGE_END, END_LEN);
			/* now unlock it */
			if( flock(fd, LOCK_UN) != 0 )
			{
				perror("Error unlocking the file\n");
			}

		} /* if(flock) */

		usleep( (rand() % getpid()) * 100 );

	} /* while */

}
#endif


/* 
 * Child's exit function. 
 * Actually an empty function, used to demonstrate signal 
 * usage.
 */
static void clean_up(int sig)
{
	/* 
	 * Do nothing, just used to interupt the process.
	 * Once we return, we can clean up nicely.
	 */  
}

/*
 * Parents exit function
 */
static void parent_term(int sig)
{
	/* return and clean up */
	longjmp(env, sig);
}