commit f293028cd131a16f8ba2504fe766317509af9d42
Author: Andrey Zonov <zont@FreeBSD.org>
Date:   Wed Oct 31 01:12:24 2012 +0400

    - Fix locked memory accounting for maps with MAP_WIREFUTURE flag.
    
    Reviewed by:	avg, trasz
    Approved by:	kib (mentor)
    MFC after:	1 week

diff --git a/sys/vm/vm_map.c b/sys/vm/vm_map.c
index 5033980..a885367 100644
--- a/sys/vm/vm_map.c
+++ b/sys/vm/vm_map.c
@@ -3247,7 +3247,7 @@ vm_map_stack(vm_map_t map, vm_offset_t addrbos, vm_size_t max_ssize,
 	vm_offset_t bot, top;
 	vm_size_t growsize, init_ssize;
 	int orient, rv;
-	rlim_t vmemlim;
+	rlim_t lmemlim, vmemlim;
 
 	/*
 	 * The stack orientation is piggybacked with the cow argument.
@@ -3267,9 +3267,10 @@ vm_map_stack(vm_map_t map, vm_offset_t addrbos, vm_size_t max_ssize,
 	growsize = sgrowsiz;
 	init_ssize = (max_ssize < growsize) ? max_ssize : growsize;
 
-	PROC_LOCK(curthread->td_proc);
-	vmemlim = lim_cur(curthread->td_proc, RLIMIT_VMEM);
-	PROC_UNLOCK(curthread->td_proc);
+	PROC_LOCK(curproc);
+	lmemlim = lim_cur(curproc, RLIMIT_MEMLOCK);
+	vmemlim = lim_cur(curproc, RLIMIT_VMEM);
+	PROC_UNLOCK(curproc);
 
 	vm_map_lock(map);
 
@@ -3279,6 +3280,14 @@ vm_map_stack(vm_map_t map, vm_offset_t addrbos, vm_size_t max_ssize,
 		return (KERN_NO_SPACE);
 	}
 
+	if (map->flags & MAP_WIREFUTURE) {
+		if (ptoa(vmspace_wired_count(curproc->p_vmspace)) +
+		    init_ssize > lmemlim) {
+			vm_map_unlock(map);
+			return (KERN_NO_SPACE);
+		}
+	}
+
 	/* If we would blow our VMEM resource limit, no go */
 	if (map->size + init_ssize > vmemlim) {
 		vm_map_unlock(map);
@@ -3360,7 +3369,7 @@ vm_map_growstack(struct proc *p, vm_offset_t addr)
 	vm_offset_t end;
 	vm_size_t growsize;
 	size_t grow_amount, max_grow;
-	rlim_t stacklim, vmemlim;
+	rlim_t lmemlim, stacklim, vmemlim;
 	int is_procstack, rv;
 	struct ucred *cred;
 #ifdef notyet
@@ -3372,6 +3381,7 @@ vm_map_growstack(struct proc *p, vm_offset_t addr)
 
 Retry:
 	PROC_LOCK(p);
+	lmemlim = lim_cur(p, RLIMIT_MEMLOCK);
 	stacklim = lim_cur(p, RLIMIT_STACK);
 	vmemlim = lim_cur(p, RLIMIT_VMEM);
 	PROC_UNLOCK(p);
@@ -3494,7 +3504,25 @@ Retry:
 	if (is_procstack && (ctob(vm->vm_ssize) + grow_amount > limit))
 		grow_amount = limit - ctob(vm->vm_ssize);
 #endif
-
+	if (map->flags & MAP_WIREFUTURE) {
+		if (ptoa(vmspace_wired_count(p->p_vmspace)) + grow_amount >
+		    lmemlim) {
+			vm_map_unlock_read(map);
+			rv = KERN_NO_SPACE;
+			goto out;
+		}
+#ifdef RACCT
+		PROC_LOCK(p);
+		if (racct_set(p, RACCT_MEMLOCK,
+		    ptoa(vmspace_wired_count(p->p_vmspace)) + grow_amount)) {
+			PROC_UNLOCK(p);
+			vm_map_unlock_read(map);
+			rv = KERN_NO_SPACE;
+			goto out;
+		}
+		PROC_UNLOCK(p);
+#endif
+	}
 	/* If we would blow our VMEM resource limit, no go */
 	if (map->size + grow_amount > vmemlim) {
 		vm_map_unlock_read(map);
@@ -3615,6 +3643,9 @@ out:
 		PROC_LOCK(p);
 		error = racct_set(p, RACCT_VMEM, map->size);
 		KASSERT(error == 0, ("decreasing RACCT_VMEM failed"));
+		error = racct_set(p, RACCT_MEMLOCK,
+		    ptoa(vmspace_wired_count(p->p_vmspace)));
+		KASSERT(error == 0, ("decreasing RACCT_MEMLOCK failed"));
 	    	error = racct_set(p, RACCT_STACK, ctob(vm->vm_ssize));
 		KASSERT(error == 0, ("decreasing RACCT_STACK failed"));
 		PROC_UNLOCK(p);
diff --git a/sys/vm/vm_mmap.c b/sys/vm/vm_mmap.c
index 0e72577..2b7ba58 100644
--- a/sys/vm/vm_mmap.c
+++ b/sys/vm/vm_mmap.c
@@ -1096,27 +1096,25 @@ sys_mlockall(td, uap)
 	int error;
 
 	map = &td->td_proc->p_vmspace->vm_map;
-	error = 0;
+	error = priv_check(td, PRIV_VM_MLOCK);
+	if (error)
+		return (error);
 
 	if ((uap->how == 0) || ((uap->how & ~(MCL_CURRENT|MCL_FUTURE)) != 0))
 		return (EINVAL);
 
-#if 0
 	/*
 	 * If wiring all pages in the process would cause it to exceed
 	 * a hard resource limit, return ENOMEM.
 	 */
-	PROC_LOCK(td->td_proc);
-	if (map->size > lim_cur(td->td_proc, RLIMIT_MEMLOCK)) {
+	if (uap->how & MCL_CURRENT) {
+		PROC_LOCK(td->td_proc);
+		if (map->size > lim_cur(td->td_proc, RLIMIT_MEMLOCK)) {
+			PROC_UNLOCK(td->td_proc);
+			return (ENOMEM);
+		}
 		PROC_UNLOCK(td->td_proc);
-		return (ENOMEM);
 	}
-	PROC_UNLOCK(td->td_proc);
-#else
-	error = priv_check(td, PRIV_VM_MLOCK);
-	if (error)
-		return (error);
-#endif
 #ifdef RACCT
 	PROC_LOCK(td->td_proc);
 	error = racct_set(td->td_proc, RACCT_MEMLOCK, map->size);
@@ -1480,6 +1478,24 @@ vm_mmap(vm_map_t map, vm_offset_t *addr, vm_size_t size, vm_prot_t prot,
 			PROC_UNLOCK(td->td_proc);
 			return (ENOMEM);
 		}
+		if (map->flags & MAP_WIREFUTURE) {
+			if (ptoa(vmspace_wired_count(td->td_proc->p_vmspace)) +
+			    size > lim_cur(td->td_proc, RLIMIT_MEMLOCK)) {
+				racct_set_force(td->td_proc, RACCT_VMEM,
+				    map->size);
+				PROC_UNLOCK(td->td_proc);
+				return (ENOMEM);
+			}
+			error = racct_set(td->td_proc, RACCT_MEMLOCK,
+			    ptoa(vmspace_wired_count(td->td_proc->p_vmspace)) +
+			    size);
+			if (error != 0) {
+				racct_set_force(td->td_proc, RACCT_VMEM,
+				    map->size);
+				PROC_UNLOCK(td->td_proc);
+				return (error);
+			}
+		}
 		PROC_UNLOCK(td->td_proc);
 	}
 
diff --git a/sys/vm/vm_unix.c b/sys/vm/vm_unix.c
index 0894dc7..6f6c44d 100644
--- a/sys/vm/vm_unix.c
+++ b/sys/vm/vm_unix.c
@@ -77,13 +77,14 @@ sys_obreak(td, uap)
 {
 	struct vmspace *vm = td->td_proc->p_vmspace;
 	vm_offset_t new, old, base;
-	rlim_t datalim, vmemlim;
+	rlim_t datalim, lmemlim, vmemlim;
 	int prot, rv;
 	int error = 0;
 	boolean_t do_map_wirefuture;
 
 	PROC_LOCK(td->td_proc);
 	datalim = lim_cur(td->td_proc, RLIMIT_DATA);
+	lmemlim = lim_cur(td->td_proc, RLIMIT_MEMLOCK);
 	vmemlim = lim_cur(td->td_proc, RLIMIT_VMEM);
 	PROC_UNLOCK(td->td_proc);
 
@@ -116,6 +117,13 @@ sys_obreak(td, uap)
 		goto done;
 	}
 	if (new > old) {
+		if (vm->vm_map.flags & MAP_WIREFUTURE) {
+			if (ptoa(vmspace_wired_count(td->td_proc->p_vmspace)) +
+			    (new - old) > lmemlim) {
+				error = ENOMEM;
+				goto done;
+			}
+		}
 		if (vm->vm_map.size + (new - old) > vmemlim) {
 			error = ENOMEM;
 			goto done;
@@ -136,6 +144,20 @@ sys_obreak(td, uap)
 			error = ENOMEM;
 			goto done;
 		}
+		if (vm->vm_map.flags & MAP_WIREFUTURE) {
+			error = racct_set(td->td_proc, RACCT_MEMLOCK,
+			    ptoa(vmspace_wired_count(td->td_proc->p_vmspace)) +
+			    (new - old));
+			if (error != 0) {
+				racct_set_force(td->td_proc, RACCT_DATA,
+				    old - base);
+				racct_set_force(td->td_proc, RACCT_VMEM,
+				    vm->vm_map.size);
+				PROC_UNLOCK(td->td_proc);
+				error = ENOMEM;
+				goto done;
+			}
+		}
 		PROC_UNLOCK(td->td_proc);
 #endif
 		prot = VM_PROT_RW;
@@ -151,7 +173,13 @@ sys_obreak(td, uap)
 #ifdef RACCT
 			PROC_LOCK(td->td_proc);
 			racct_set_force(td->td_proc, RACCT_DATA, old - base);
-			racct_set_force(td->td_proc, RACCT_VMEM, vm->vm_map.size);
+			racct_set_force(td->td_proc, RACCT_VMEM,
+			    vm->vm_map.size);
+			if (vm->vm_map.flags & MAP_WIREFUTURE) {
+				racct_set_force(td->td_proc, RACCT_MEMLOCK,
+				    ptoa(vmspace_wired_count(
+				    td->td_proc->p_vmspace)));
+			}
 			PROC_UNLOCK(td->td_proc);
 #endif
 			error = ENOMEM;
@@ -183,6 +211,10 @@ sys_obreak(td, uap)
 		PROC_LOCK(td->td_proc);
 		racct_set_force(td->td_proc, RACCT_DATA, new - base);
 		racct_set_force(td->td_proc, RACCT_VMEM, vm->vm_map.size);
+		if (vm->vm_map.flags & MAP_WIREFUTURE) {
+			racct_set_force(td->td_proc, RACCT_MEMLOCK,
+			    ptoa(vmspace_wired_count(td->td_proc->p_vmspace)));
+		}
 		PROC_UNLOCK(td->td_proc);
 #endif
 	}