MCOUNT_INSN_SIZE);
 }
 
-void ftrace_nmi_enter(void)
+void arch_ftrace_nmi_enter(void)
 {
        atomic_inc(&in_nmi);
        /* Must have in_nmi seen before reading write flag */
        }
 }
 
-void ftrace_nmi_exit(void)
+void arch_ftrace_nmi_exit(void)
 {
        /* Finish all executions before clearing in_nmi */
        smp_wmb();
  */
 static atomic_t in_nmi;
 
-void ftrace_nmi_enter(void)
+void arch_ftrace_nmi_enter(void)
 {
        atomic_inc(&in_nmi);
 }
 
-void ftrace_nmi_exit(void)
+void arch_ftrace_nmi_exit(void)
 {
        atomic_dec(&in_nmi);
 }
 
 #define _LINUX_FTRACE_IRQ_H
 
 
-#if defined(CONFIG_DYNAMIC_FTRACE) || defined(CONFIG_FUNCTION_GRAPH_TRACER)
+#ifdef CONFIG_FTRACE_NMI_ENTER
+extern void arch_ftrace_nmi_enter(void);
+extern void arch_ftrace_nmi_exit(void);
+#else
+static inline void arch_ftrace_nmi_enter(void) { }
+static inline void arch_ftrace_nmi_exit(void) { }
+#endif
+
+#ifdef CONFIG_RING_BUFFER
 extern void ftrace_nmi_enter(void);
 extern void ftrace_nmi_exit(void);
 #else
 
  * Copyright (C) 2008 Steven Rostedt <srostedt@redhat.com>
  */
 #include <linux/ring_buffer.h>
+#include <linux/ftrace_irq.h>
 #include <linux/spinlock.h>
 #include <linux/debugfs.h>
 #include <linux/uaccess.h>
 
 #include "trace.h"
 
+/*
+ * Since the write to the buffer is still not fully lockless,
+ * we must be careful with NMIs. The locks in the writers
+ * are taken when a write crosses to a new page. The locks
+ * protect against races with the readers (this will soon
+ * be fixed with a lockless solution).
+ *
+ * Because we can not protect against NMIs, and we want to
+ * keep traces reentrant, we need to manage what happens
+ * when we are in an NMI.
+ */
+static DEFINE_PER_CPU(int, rb_in_nmi);
+
+void ftrace_nmi_enter(void)
+{
+       __get_cpu_var(rb_in_nmi)++;
+       /* call arch specific handler too */
+       arch_ftrace_nmi_enter();
+}
+
+void ftrace_nmi_exit(void)
+{
+       arch_ftrace_nmi_exit();
+       __get_cpu_var(rb_in_nmi)--;
+       /* NMIs are not recursive */
+       WARN_ON_ONCE(__get_cpu_var(rb_in_nmi));
+}
+
+
 /*
  * A fast way to enable or disable all ring buffers is to
  * call tracing_on or tracing_off. Turning off the ring buffers
        struct ring_buffer *buffer = cpu_buffer->buffer;
        struct ring_buffer_event *event;
        unsigned long flags;
+       bool lock_taken = false;
 
        commit_page = cpu_buffer->commit_page;
        /* we just need to protect against interrupts */
                struct buffer_page *next_page = tail_page;
 
                local_irq_save(flags);
-               __raw_spin_lock(&cpu_buffer->lock);
+               /*
+                * NMIs can happen after we take the lock.
+                * If we are in an NMI, only take the lock
+                * if it is not already taken. Otherwise
+                * simply fail.
+                */
+               if (unlikely(__get_cpu_var(rb_in_nmi))) {
+                       if (!__raw_spin_trylock(&cpu_buffer->lock))
+                               goto out_unlock;
+               } else
+                       __raw_spin_lock(&cpu_buffer->lock);
+
+               lock_taken = true;
 
                rb_inc_page(cpu_buffer, &next_page);
 
        if (tail <= BUF_PAGE_SIZE)
                local_set(&tail_page->write, tail);
 
-       __raw_spin_unlock(&cpu_buffer->lock);
+       if (likely(lock_taken))
+               __raw_spin_unlock(&cpu_buffer->lock);
        local_irq_restore(flags);
        return NULL;
 }