- Bug fix: prevent a thread from migrating between CPUs between the

time it is marked for user space callchain capture in the NMI
  handler and the time the callchain capture callback runs.

- Improve code and control flow clarity by invoking hwpmc(4)'s user
  space callchain capture callback directly from low-level code.

Reviewed by:	jhb (kern/subr_trap.c)
Testing (various patch revisions): gnn,
		Fabien Thomas <fabien dot thomas at netasq dot com>,
		Artem Belevich <artemb at gmail dot com>

sys/amd64/amd64/exception.S

@@ -480,16 +480,20 @@ outofnmi:
 	/*
 	 * At this point the processor has exited NMI mode and is running
 	 * with interrupts turned off on the normal kernel stack.
-	 * We turn interrupts back on, and take the usual 'doreti' exit
-	 * path.
 	 *
 	 * If a pending NMI gets recognized at or after this point, it 
-	 * will cause a kernel callchain to be traced.  Since this path
-	 * is only taken for NMI interrupts from user space, our `swapgs'
-	 * state is correct for taking the doreti path.
+	 * will cause a kernel callchain to be traced.
+	 *
+	 * We turn interrupts back on, and call the user callchain capture hook.
 	 */
+	movq	pmc_hook,%rax
+	orq	%rax,%rax
+	jz	nocallchain
+	movq	PCPU(CURTHREAD),%rdi		/* thread */
+	movq	$PMC_FN_USER_CALLCHAIN,%rsi	/* command */
+	movq	%rsp,%rdx			/* frame */
 	sti
-	jmp	doreti
+	call	*%rax
 nocallchain:
 #endif
 	testl	%ebx,%ebx

sys/dev/hwpmc/hwpmc_mod.c

@@ -1863,8 +1863,11 @@ pmc_hook_handler(struct thread *td, int function, void *arg)
 		/*
 		 * Record a call chain.
 		 */
+		KASSERT(td == curthread, ("[pmc,%d] td != curthread",
+		    __LINE__));
 		pmc_capture_user_callchain(PCPU_GET(cpuid),
 		    (struct trapframe *) arg);
+		td->td_pflags &= ~TDP_CALLCHAIN;
 		break;
 
 	default:
@@ -3794,30 +3797,28 @@ pmc_syscall_handler(struct thread *td, void *syscall_args)
  */
 
 static void
-pmc_post_callchain_ast(void)
+pmc_post_callchain_callback(void)
 {
 	struct thread *td;
 
 	td = curthread;
 
+	KASSERT((td->td_pflags & TDP_CALLCHAIN) == 0,
+	    ("[pmc,%d] thread %p already marked for callchain capture",
+		__LINE__, (void *) td));
+
 	/*
-	 * Mark this thread as needing processing in ast().
-	 * td->td_pflags will be safe to touch as the process was in
-	 * user space when it was interrupted.
+	 * Mark this thread as needing callchain capture.
+	 * `td->td_pflags' will be safe to touch because this thread
+	 * was in user space when it was interrupted.
 	 */
 	td->td_pflags |= TDP_CALLCHAIN;
 
 	/*
-	 * Again, since we've entered this function directly from
-	 * userland, `td' is guaranteed to be not locked by this CPU,
-	 * so its safe to try acquire the thread lock even though we
-	 * are executing in an NMI context.  We need to acquire this
-	 * lock before touching `td_flags' because other CPUs may be
-	 * in the process of touching this field.
+	 * Don't let this thread migrate between CPUs until callchain
+	 * capture completes.
 	 */
-	thread_lock(td);
-	td->td_flags |= TDF_ASTPENDING;
-	thread_unlock(td);
+	sched_pin();
 
 	return;
 }
@@ -3869,6 +3870,10 @@ pmc_process_interrupt(int cpu, struct pmc *pm, struct trapframe *tf,
 	    (int) (psb->ps_write - psb->ps_samples),
 	    (int) (psb->ps_read - psb->ps_samples));
 
+	KASSERT(pm->pm_runcount >= 0,
+	    ("[pmc,%d] pm=%p runcount %d", __LINE__, (void *) pm,
+		pm->pm_runcount));
+
 	atomic_add_rel_32(&pm->pm_runcount, 1);	/* hold onto PMC */
 	ps->ps_pmc = pm;
 	if ((td = curthread) && td->td_proc)
@@ -3876,6 +3881,7 @@ pmc_process_interrupt(int cpu, struct pmc *pm, struct trapframe *tf,
 	else
 		ps->ps_pid = -1;
 	ps->ps_cpu = cpu;
+	ps->ps_td = td;
 	ps->ps_flags = inuserspace ? PMC_CC_F_USERSPACE : 0;
 
 	callchaindepth = (pm->pm_flags & PMC_F_CALLCHAIN) ?
@@ -3893,7 +3899,7 @@ pmc_process_interrupt(int cpu, struct pmc *pm, struct trapframe *tf,
 			    pmc_save_kernel_callchain(ps->ps_pc,
 				callchaindepth, tf);
 		else {
-			pmc_post_callchain_ast();
+			pmc_post_callchain_callback();
 			callchaindepth = PMC_SAMPLE_INUSE;
 		}
 	}
@@ -3925,20 +3931,41 @@ pmc_capture_user_callchain(int cpu, struct trapframe *tf)
 {
 	int i;
 	struct pmc *pm;
+	struct thread *td;
 	struct pmc_sample *ps;
 	struct pmc_samplebuffer *psb;
+#ifdef	INVARIANTS
+	int ncallchains;
+#endif
+
+	sched_unpin();	/* Can migrate safely now. */
 
 	psb = pmc_pcpu[cpu]->pc_sb;
+	td = curthread;
+
+	KASSERT(td->td_pflags & TDP_CALLCHAIN,
+	    ("[pmc,%d] Retrieving callchain for thread that doesn't want it",
+		__LINE__));
+
+#ifdef	INVARIANTS
+	ncallchains = 0;
+#endif
 
 	/*
 	 * Iterate through all deferred callchain requests.
 	 */
 
-	for (i = 0; i < pmc_nsamples; i++) {
+	ps = psb->ps_samples;
+	for (i = 0; i < pmc_nsamples; i++, ps++) {
 
-		ps = &psb->ps_samples[i];
 		if (ps->ps_nsamples != PMC_SAMPLE_INUSE)
 			continue;
+		if (ps->ps_td != td)
+			continue;
+
+		KASSERT(ps->ps_cpu == cpu,
+		    ("[pmc,%d] cpu mismatch ps_cpu=%d pcpu=%d", __LINE__,
+			ps->ps_cpu, PCPU_GET(cpuid)));
 
 		pm = ps->ps_pmc;
 
@@ -3946,14 +3973,26 @@ pmc_capture_user_callchain(int cpu, struct trapframe *tf)
 		    ("[pmc,%d] Retrieving callchain for PMC that doesn't "
 			"want it", __LINE__));
 
+		KASSERT(pm->pm_runcount > 0,
+		    ("[pmc,%d] runcount %d", __LINE__, pm->pm_runcount));
+
 		/*
 		 * Retrieve the callchain and mark the sample buffer
 		 * as 'processable' by the timer tick sweep code.
 		 */
 		ps->ps_nsamples = pmc_save_user_callchain(ps->ps_pc,
 		    pmc_callchaindepth, tf);
+
+#ifdef	INVARIANTS
+		ncallchains++;
+#endif
+
 	}
 
+	KASSERT(ncallchains > 0,
+	    ("[pmc,%d] cpu %d didn't find a sample to collect", __LINE__,
+		cpu));
+
 	return;
 }
 
@@ -3991,6 +4030,11 @@ pmc_process_samples(int cpu)
 		}
 
 		pm = ps->ps_pmc;
+
+		KASSERT(pm->pm_runcount > 0,
+		    ("[pmc,%d] pm=%p runcount %d", __LINE__, (void *) pm,
+			pm->pm_runcount));
+
 		po = pm->pm_owner;
 
 		KASSERT(PMC_IS_SAMPLING_MODE(PMC_TO_MODE(pm)),

sys/i386/i386/exception.s

@@ -438,9 +438,18 @@ doreti_nmi:
 	iret
 outofnmi:
 	/*
-	 * Clear interrupts and jump to AST handling code.
+	 * Call the callchain capture hook after turning interrupts back on.
 	 */
+	movl	pmc_hook,%ecx
+	orl	%ecx,%ecx
+	jz	doreti_exit
+	pushl	%esp			/* frame pointer */
+	pushl	$PMC_FN_USER_CALLCHAIN	/* command */
+	movl	PCPU(CURTHREAD),%eax
+	pushl	%eax			/* curthread */
 	sti
+	call	*%ecx
+	addl	$12,%esp
 	jmp	doreti_ast
 	ENTRY(end_exceptions)
 #endif

sys/kern/subr_trap.c

@@ -44,7 +44,6 @@
 #include <sys/cdefs.h>
 __FBSDID("$FreeBSD$");
 
-#include "opt_hwpmc_hooks.h"
 #include "opt_ktrace.h"
 #include "opt_mac.h"
 #ifdef __i386__
@@ -179,13 +178,6 @@ ast(struct trapframe *framep)
 		td->td_profil_ticks = 0;
 		td->td_pflags &= ~TDP_OWEUPC;
 	}
-#if defined(HWPMC_HOOKS)
-	if (td->td_pflags & TDP_CALLCHAIN) {
-		PMC_CALL_HOOK_UNLOCKED(td, PMC_FN_USER_CALLCHAIN,
-		    (void *) framep);
-		td->td_pflags &= ~TDP_CALLCHAIN;
-	}
-#endif
 	if (flags & TDF_ALRMPEND) {
 		PROC_LOCK(p);
 		psignal(p, SIGVTALRM);