Merge pull request #5 from opsengine/develop

Fix issue #4

6 files changed, 70 insertions, 29 deletions

diff --git a/src/cpulimit.c b/src/cpulimit.c
index 5353b90..e7d5309 100644
--- a/src/cpulimit.c
+++ b/src/cpulimit.c
@@ -27,22 +27,20 @@
  *
  */
 
-#include <getopt.h>
 #include <stdio.h>
-#include <fcntl.h>
 #include <stdlib.h>
-#include <time.h>
 #include <unistd.h>
+#include <getopt.h>
+#include <time.h>
 #include <signal.h>
 #include <string.h>
-#include <dirent.h>
 #include <errno.h>
 #include <string.h>
 #include <sys/stat.h>
 #include <sys/time.h>
-#include <sys/types.h>
 #include <sys/sysctl.h>
 #include <sys/resource.h>
+#include <sys/types.h>
 #include <sys/wait.h>
 
 #include "process_group.h"
@@ -226,13 +224,19 @@ void limit_process(pid_t pid, double limit, int include_children)
 		}
 
 		//resume processes
-		for (node = pgroup.proclist->first; node != NULL; node = node->next) {
+		node = pgroup.proclist->first;
+		while (node != NULL)
+		{
+			struct list_node *next_node = node->next;
 			struct process *proc = (struct process*)(node->data);
-			if (kill(proc->pid,SIGCONT)!=0) {
+			if (kill(proc->pid,SIGCONT) != 0) {
 				//process is dead, remove it from family
-				if (verbose) fprintf(stderr, "Process %d dead!\n", proc->pid);
-				//remove_process_from_family(&pf, proc->pid);
+				if (verbose) fprintf(stderr, "SIGCONT failed. Process %d dead!\n", proc->pid);
+				//remove process from group
+				delete_node(pgroup.proclist, node);
+				remove_process(&pgroup, proc->pid);
 			}
+			node = next_node;
 		}
 
 		//now processes are free to run (same working slice for all)
@@ -248,14 +252,20 @@ void limit_process(pid_t pid, double limit, int include_children)
 		}
 
 		if (tsleep.tv_nsec>0) {
-			//stop only if tsleep>0
-			for (node = pgroup.proclist->first; node != NULL; node = node->next) {
+			//stop processes only if tsleep>0
+			node = pgroup.proclist->first;
+			while (node != NULL)
+			{
+				struct list_node *next_node = node->next;
 				struct process *proc = (struct process*)(node->data);
 				if (kill(proc->pid,SIGSTOP)!=0) {
 					//process is dead, remove it from family
-					if (verbose) fprintf(stderr,"Process %d dead!\n", proc->pid);
-					//remove_process_from_family(&pf, proc->pid);
+					if (verbose) fprintf(stderr, "SIGSTOP failed. Process %d dead!\n", proc->pid);
+					//remove process from group
+					delete_node(pgroup.proclist, node);
+					remove_process(&pgroup, proc->pid);
 				}
+				node = next_node;
 			}
 			//now the processes are sleeping
 			nanosleep(&tsleep,NULL);
@@ -432,7 +442,7 @@ int main(int argc, char **argv) {
 			else {
 				//limiter code
 				if (verbose) printf("Limiting process %d\n",child);
-				limit_process(child, limit, ignore_children);
+				limit_process(child, limit, !ignore_children);
 				exit(0);
 			}
 		}
diff --git a/src/process_group.c b/src/process_group.c
index 671dd73..cb84deb 100644
--- a/src/process_group.c
+++ b/src/process_group.c
@@ -172,7 +172,7 @@ void update_process_group(struct process_group *pgroup)
 //		struct timeval t;
 //		gettimeofday(&t, NULL);
 //		printf("T=%ld.%ld PID=%d PPID=%d START=%d CPUTIME=%d\n", t.tv_sec, t.tv_usec, tmp_process.pid, tmp_process.ppid, tmp_process.starttime, tmp_process.cputime);
-		int hashkey = pid_hashfn(tmp_process.pid + tmp_process.starttime);
+		int hashkey = pid_hashfn(tmp_process.pid);
 		if (pgroup->proctable[hashkey] == NULL)
 		{
 			//empty bucket
@@ -222,3 +222,13 @@ void update_process_group(struct process_group *pgroup)
 	if (dt < MIN_DT) return;
 	pgroup->last_update = now;
 }
+
+int remove_process(struct process_group *pgroup, int pid)
+{
+	int hashkey = pid_hashfn(pid);
+	if (pgroup->proctable[hashkey] == NULL) return 1; //nothing to delete
+	struct list_node *node = (struct list_node*)locate_node(pgroup->proctable[hashkey], &pid);
+	if (node == NULL) return 2;
+	delete_node(pgroup->proctable[hashkey], node);
+	return 0;
+}
diff --git a/src/process_group.h b/src/process_group.h
index ab86816..5a5b581 100644
--- a/src/process_group.h
+++ b/src/process_group.h
@@ -50,4 +50,6 @@ int find_process_by_pid(pid_t pid);
 
 int find_process_by_name(const char *process_name);
 
+int remove_process(struct process_group *pgroup, int pid);
+
 #endif
diff --git a/src/process_iterator_freebsd.c b/src/process_iterator_freebsd.c
index a4bf38f..a638112 100644
--- a/src/process_iterator_freebsd.c
+++ b/src/process_iterator_freebsd.c
@@ -35,24 +35,23 @@ int init_process_iterator(struct process_iterator *it, struct process_filter *fi
 	/* Get the list of processes. */
 	if ((it->procs = kvm_getprocs(it->kd, KERN_PROC_PROC, 0, &it->count)) == NULL) {
 		kvm_close(it->kd);
-		fprintf(stderr, "kvm_getprocs: %s\n", kvm_geterr(it->kd));
+//		fprintf(stderr, "kvm_getprocs: %s\n", kvm_geterr(it->kd));
 		return -1;
 	}
 	it->filter = filter;
 	return 0;
 }
 
-static void kproc2proc(kvm_t *kd, struct kinfo_proc *kproc, struct process *proc)
+static int kproc2proc(kvm_t *kd, struct kinfo_proc *kproc, struct process *proc)
 {
 	proc->pid = kproc->ki_pid;
 	proc->ppid = kproc->ki_ppid;
 	proc->cputime = kproc->ki_runtime / 1000;
 	proc->starttime = kproc->ki_start.tv_sec;
 	char **args = kvm_getargv(kd, kproc, sizeof(proc->command));
-	if (args != NULL)
-	{
-		memcpy(proc->command, args[0], strlen(args[0]) + 1);
-	}
+	if (args == NULL) return -1;
+	memcpy(proc->command, args[0], strlen(args[0]) + 1);
+	return 0;
 }
 
 static int get_single_process(kvm_t *kd, pid_t pid, struct process *process)
@@ -61,7 +60,7 @@ static int get_single_process(kvm_t *kd, pid_t pid, struct process *process)
 	struct kinfo_proc *kproc = kvm_getprocs(kd, KERN_PROC_PID, pid, &count);
 	if (count == 0 || kproc == NULL)
 	{
-		fprintf(stderr, "kvm_getprocs: %s\n", kvm_geterr(kd));
+//		fprintf(stderr, "kvm_getprocs: %s\n", kvm_geterr(kd));
 		return -1;
 	}
 	kproc2proc(kd, kproc, process);
@@ -73,9 +72,13 @@ int get_next_process(struct process_iterator *it, struct process *p) {
 	{
 		return -1;
 	}
-	if (it->filter->pid > 0 && !it->filter->include_children)
+	if (it->filter->pid != 0 && !it->filter->include_children)
 	{
-		get_single_process(it->kd, it->filter->pid, p);
+		if (get_single_process(it->kd, it->filter->pid, p) != 0)
+		{
+			it->i = it->count = 0;
+			return -1;
+		}
 		it->i = it->count = 1;
 		return 0;
 	}
@@ -88,7 +91,7 @@ int get_next_process(struct process_iterator *it, struct process *p) {
 			it->i++;
 			continue;
 		}
-		if (it->filter->pid > 0 && it->filter->include_children)
+		if (it->filter->pid != 0 && it->filter->include_children)
 		{
 			kproc2proc(it->kd, kproc, p);
 			it->i++;
diff --git a/src/process_iterator_linux.c b/src/process_iterator_linux.c
index 9e2f7f2..43c2771 100644
--- a/src/process_iterator_linux.c
+++ b/src/process_iterator_linux.c
@@ -141,21 +141,22 @@ int get_next_process(struct process_iterator *it, struct process *p)
 		//end of processes
 		return -1;
 	}
-	if (it->filter->pid > 0 && !it->filter->include_children)
+	if (it->filter->pid != 0 && !it->filter->include_children)
 	{
-		read_process_info(it->filter->pid, p);
+		int ret = read_process_info(it->filter->pid, p);
 		//p->starttime += it->boot_time;
 		closedir(it->dip);
 		it->dip = NULL;
+		if (ret != 0) return -1;
 		return 0;
 	}
-	struct dirent *dit;
+	struct dirent *dit = NULL;
 	//read in from /proc and seek for process dirs
 	while ((dit = readdir(it->dip)) != NULL) {
 		if(strtok(dit->d_name, "0123456789") != NULL)
 			continue;
 		p->pid = atoi(dit->d_name);
-		if (it->filter->pid > 0 && it->filter->pid != p->pid && !is_child_of(p->pid, it->filter->pid)) continue;
+		if (it->filter->pid != 0 && it->filter->pid != p->pid && !is_child_of(p->pid, it->filter->pid)) continue;
 		read_process_info(p->pid, p);
 		//p->starttime += it->boot_time;
 		break;
diff --git a/tests/process_iterator_test.c b/tests/process_iterator_test.c
index 7b1c6b5..3d5e885 100644
--- a/tests/process_iterator_test.c
+++ b/tests/process_iterator_test.c
@@ -203,6 +203,20 @@ void test_process_name(const char * command)
 	close_process_iterator(&it);
 }
 
+void test_process_group_wrong_pid()
+{
+	struct process_group pgroup;
+	assert(init_process_group(&pgroup, -1, 0) == 0);
+	assert(pgroup.proclist->count == 0);
+	update_process_group(&pgroup);
+	assert(pgroup.proclist->count == 0);
+	assert(init_process_group(&pgroup, 9999999, 0) == 0);
+	assert(pgroup.proclist->count == 0);
+	update_process_group(&pgroup);
+	assert(pgroup.proclist->count == 0);
+	assert(close_process_group(&pgroup) == 0);
+}
+
 int main(int argc, char **argv)
 {
 //	printf("Pid %d\n", getpid());
@@ -212,6 +226,7 @@ int main(int argc, char **argv)
 	test_process_group_all();
 	test_process_group_single(0);
 	test_process_group_single(1);
+	test_process_group_wrong_pid();
 	test_process_name(argv[0]);
 	return 0;
 }