From 58c77a9d530350a60ae5ec0a0160066364742c6e Mon Sep 17 00:00:00 2001
From: Edward Tomasz Napierala <trasz@FreeBSD.org>
Date: Thu, 31 Mar 2011 19:22:11 +0000
Subject: [PATCH] Enable accounting for RACCT_NPROC and RACCT_NTHR.

Sponsored by:	The FreeBSD Foundation
Reviewed by:	kib (earlier version)
---
 sys/kern/init_main.c |  3 +++
 sys/kern/kern_exit.c |  8 ++++++++
 sys/kern/kern_fork.c | 20 ++++++++++++++++++++
 sys/kern/kern_thr.c  | 29 ++++++++++++++++++++++++-----
 4 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/sys/kern/init_main.c b/sys/kern/init_main.c
index 1977b9679691..be5c26fec029 100644
--- a/sys/kern/init_main.c
+++ b/sys/kern/init_main.c
@@ -557,6 +557,9 @@ proc0_init(void *dummy __unused)
 	 * Charge root for one process.
 	 */
 	(void)chgproccnt(p->p_ucred->cr_ruidinfo, 1, 0);
+	PROC_LOCK(p);
+	racct_add_force(p, RACCT_NPROC, 1);
+	PROC_UNLOCK(p);
 }
 SYSINIT(p0init, SI_SUB_INTRINSIC, SI_ORDER_FIRST, proc0_init, NULL);
 
diff --git a/sys/kern/kern_exit.c b/sys/kern/kern_exit.c
index 01d6b7572c4c..5fec31b060b2 100644
--- a/sys/kern/kern_exit.c
+++ b/sys/kern/kern_exit.c
@@ -177,6 +177,7 @@ exit1(struct thread *td, int rv)
 	}
 	KASSERT(p->p_numthreads == 1,
 	    ("exit1: proc %p exiting with %d threads", p, p->p_numthreads));
+	racct_sub(p, RACCT_NTHR, 1);
 	/*
 	 * Wakeup anyone in procfs' PIOCWAIT.  They should have a hold
 	 * on our vmspace, so we should block below until they have
@@ -745,6 +746,9 @@ proc_reap(struct thread *td, struct proc *p, int *status, int options,
 	 * Destroy resource accounting information associated with the process.
 	 */
 	racct_proc_exit(p);
+	PROC_LOCK(p->p_pptr);
+	racct_sub(p->p_pptr, RACCT_NPROC, 1);
+	PROC_UNLOCK(p->p_pptr);
 
 	/*
 	 * Free credentials, arguments, and sigacts.
@@ -905,7 +909,11 @@ proc_reparent(struct proc *child, struct proc *parent)
 	if (child->p_pptr == parent)
 		return;
 
+	PROC_LOCK(parent);
+	racct_add_force(parent, RACCT_NPROC, 1);
+	PROC_UNLOCK(parent);
 	PROC_LOCK(child->p_pptr);
+	racct_sub(child->p_pptr, RACCT_NPROC, 1);
 	sigqueue_take(child->p_ksi);
 	PROC_UNLOCK(child->p_pptr);
 	LIST_REMOVE(child, p_sibling);
diff --git a/sys/kern/kern_fork.c b/sys/kern/kern_fork.c
index 1dcc5bb1e574..04e635acdaa7 100644
--- a/sys/kern/kern_fork.c
+++ b/sys/kern/kern_fork.c
@@ -734,6 +734,12 @@ fork1(struct thread *td, int flags, int pages, struct proc **procp)
 		return (fork_norfproc(td, flags));
 	}
 
+	PROC_LOCK(p1);
+	error = racct_add(p1, RACCT_NPROC, 1);
+	PROC_UNLOCK(p1);
+	if (error != 0)
+		return (EAGAIN);
+
 	mem_charged = 0;
 	vm2 = NULL;
 	if (pages == 0)
@@ -816,6 +822,17 @@ fork1(struct thread *td, int flags, int pages, struct proc **procp)
 		goto fail;
 	}
 
+	/*
+	 * After fork, there is exactly one thread running.
+	 */
+	PROC_LOCK(newproc);
+	error = racct_set(newproc, RACCT_NTHR, 1);
+	PROC_UNLOCK(newproc);
+	if (error != 0) {
+		error = EAGAIN;
+		goto fail;
+	}
+
 	/*
 	 * Increment the count of procs running with this uid. Don't allow
 	 * a nonprivileged user to exceed their current limit.
@@ -857,6 +874,9 @@ fork1(struct thread *td, int flags, int pages, struct proc **procp)
 		vmspace_free(vm2);
 	uma_zfree(proc_zone, newproc);
 	pause("fork", hz / 2);
+	PROC_LOCK(p1);
+	racct_sub(p1, RACCT_NPROC, 1);
+	PROC_UNLOCK(p1);
 	return (error);
 }
 
diff --git a/sys/kern/kern_thr.c b/sys/kern/kern_thr.c
index 58c38062a6a5..7011a53b7cf4 100644
--- a/sys/kern/kern_thr.c
+++ b/sys/kern/kern_thr.c
@@ -36,6 +36,7 @@ __FBSDID("$FreeBSD$");
 #include <sys/priv.h>
 #include <sys/proc.h>
 #include <sys/posix4.h>
+#include <sys/racct.h>
 #include <sys/resourcevar.h>
 #include <sys/rwlock.h>
 #include <sys/sched.h>
@@ -184,10 +185,18 @@ create_thread(struct thread *td, mcontext_t *ctx,
 		}
 	}
 
+	PROC_LOCK(td->td_proc);
+	error = racct_add(p, RACCT_NTHR, 1);
+	PROC_UNLOCK(td->td_proc);
+	if (error != 0)
+		return (EPROCLIM);
+
 	/* Initialize our td */
 	newtd = thread_alloc(0);
-	if (newtd == NULL)
-		return (ENOMEM);
+	if (newtd == NULL) {
+		error = ENOMEM;
+		goto fail;
+	}
 
 	/*
 	 * Try the copyout as soon as we allocate the td so we don't
@@ -203,7 +212,8 @@ create_thread(struct thread *td, mcontext_t *ctx,
 	    (parent_tid != NULL &&
 	    suword_lwpid(parent_tid, newtd->td_tid))) {
 		thread_free(newtd);
-		return (EFAULT);
+		error = EFAULT;
+		goto fail;
 	}
 
 	bzero(&newtd->td_startzero,
@@ -220,7 +230,7 @@ create_thread(struct thread *td, mcontext_t *ctx,
 		if (error != 0) {
 			thread_free(newtd);
 			crfree(td->td_ucred);
-			return (error);
+			goto fail;
 		}
 	} else {
 		/* Set up our machine context. */
@@ -233,7 +243,7 @@ create_thread(struct thread *td, mcontext_t *ctx,
 		if (error != 0) {
 			thread_free(newtd);
 			crfree(td->td_ucred);
-			return (error);
+			goto fail;
 		}
 	}
 
@@ -265,6 +275,12 @@ create_thread(struct thread *td, mcontext_t *ctx,
 	thread_unlock(newtd);
 
 	return (0);
+
+fail:
+	PROC_LOCK(p);
+	racct_sub(p, RACCT_NTHR, 1);
+	PROC_UNLOCK(p);
+	return (error);
 }
 
 int
@@ -294,7 +310,10 @@ thr_exit(struct thread *td, struct thr_exit_args *uap)
 	}
 
 	rw_wlock(&tidhash_lock);
+
 	PROC_LOCK(p);
+	racct_sub(p, RACCT_NTHR, 1);
+
 	/*
 	 * Shutting down last thread in the proc.  This will actually
 	 * call exit() in the trampoline when it returns.