Index: trap.c
===================================================================
RCS file: /cvsroot/src/sys/arch/amd64/amd64/trap.c,v
retrieving revision 1.60
diff -u -p -r1.60 trap.c
--- trap.c	21 Nov 2009 03:11:01 -0000	1.60
+++ trap.c	23 Feb 2010 00:33:57 -0000
@@ -73,6 +73,7 @@ __KERNEL_RCSID(0, "$NetBSD: trap.c,v 1.6
 #include "opt_ddb.h"
 #include "opt_kgdb.h"
 #include "opt_xen.h"
+#include "opt_dtrace.h"
 
 #include <sys/param.h>
 #include <sys/systm.h>
@@ -111,6 +112,19 @@ __KERNEL_RCSID(0, "$NetBSD: trap.c,v 1.6
 #include <sys/kgdb.h>
 #endif
 
+#ifdef KDTRACE_HOOKS
+#include <sys/dtrace_bsd.h>
+
+/*
+ * This is a hook which is initialized by the dtrace module
+ * to handle traps which might occur during DTrace probe
+ * execution.
+ */
+dtrace_trap_func_t	dtrace_trap_func = NULL;
+
+dtrace_doubletrap_func_t	dtrace_doubletrap_func = NULL;
+#endif
+
 void trap(struct trapframe *);
 
 const char * const trap_type[] = {
@@ -232,6 +246,27 @@ trap(struct trapframe *frame)
 		LWP_CACHE_CREDS(l, p);
 	}
 
+#ifdef KDTRACE_HOOKS
+	/*
+	 * A trap can occur while DTrace executes a probe. Before
+	 * executing the probe, DTrace blocks re-scheduling and sets
+	 * a flag in it's per-cpu flags to indicate that it doesn't
+	 * want to fault. On returning from the the probe, the no-fault
+	 * flag is cleared and finally re-scheduling is enabled.
+	 *
+	 * If the DTrace kernel module has registered a trap handler,
+	 * call it and if it returns non-zero, assume that it has
+	 * handled the trap and modified the trap frame so that this
+	 * function can return normally.
+	 */
+	if ((type == T_PROTFLT || type == T_PAGEFLT) &&
+	    dtrace_trap_func != NULL) {
+		if ((*dtrace_trap_func)(frame, type)) {
+			return;
+		}
+	}
+#endif
+
 	switch (type) {
 
 	default:
Index: vector.S
===================================================================
RCS file: /cvsroot/src/sys/arch/amd64/amd64/vector.S,v
retrieving revision 1.32
diff -u -p -r1.32 vector.S
--- vector.S	23 Feb 2010 00:23:36 -0000	1.32
+++ vector.S	23 Feb 2010 00:33:58 -0000
@@ -69,6 +69,7 @@
 #include "opt_ddb.h"
 #include "opt_multiprocessor.h"
 #include "opt_xen.h"
+#include "opt_dtrace.h"
 
 #define ALIGN_TEXT	.align 16,0x90
 
@@ -112,6 +113,21 @@
 
 #define	BPTTRAP(a)	ZTRAP(a)
 
+#ifdef KDTRACE_HOOKS
+	.bss
+	.globl	dtrace_invop_jump_addr
+	.align	8
+	.type	dtrace_invop_jump_addr, @object
+	.size	dtrace_invop_jump_addr, 8
+dtrace_invop_jump_addr:
+	.zero	8
+	.globl	dtrace_invop_calltrap_addr
+	.align	8
+	.type	dtrace_invop_calltrap_addr, @object
+	.size	dtrace_invop_calltrap_addr, 8
+dtrace_invop_calltrap_addr:
+	.zero	8
+#endif
 	.text
 
 IDTVEC(trap00)
@@ -283,6 +299,30 @@ NENTRY(resume_pop_es)
 NENTRY(alltraps)
 	INTRENTRY
   	STI(si)
+#ifdef KDTRACE_HOOKS
+	/*
+	 * DTrace Function Boundary Trace (fbt) probes are triggered
+	 * by int3 (0xcc) which causes the #BP (T_BPTFLT) breakpoint
+	 * interrupt. For all other trap types, just handle them in
+	 * the usual way.
+	 */
+	cmpl    $T_BPTFLT,TF_TRAPNO(%rsp)
+	jne     calltrap
+
+	/* Check if there is no DTrace hook registered. */
+	cmpq	$0,dtrace_invop_jump_addr
+	je	calltrap
+
+	/*
+	 * Set our jump address for the jump back in the event that
+	 * the exception wasn't caused by DTrace at all.
+	 */
+	movq	$calltrap, dtrace_invop_calltrap_addr(%rip)
+
+	/* Jump to the code hooked in by DTrace. */
+	movq	dtrace_invop_jump_addr, %rax
+	jmpq	*dtrace_invop_jump_addr
+#endif
 calltrap:
 #ifdef DIAGNOSTIC
 	movl	CPUVAR(ILEVEL),%ebx