Enable accounting for RACCT_NPROC and RACCT_NTHR.

Sponsored by:	The FreeBSD Foundation
Reviewed by:	kib (earlier version)

4 files changed, 55 insertions, 5 deletions

diff --git a/sys/kern/init_main.c b/sys/kern/init_main.c
index 1977b96..be5c26f 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 01d6b75..5fec31b 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 1dcc5bb..04e635a 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)
@@ -817,6 +823,17 @@ fork1(struct thread *td, int flags, int pages, struct proc **procp)
 	}
 
 	/*
+	 * 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 @@ fail1:
 		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 58c3806..7011a53 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.