[PATCH] Kprobes: preempt_disable/enable() simplification

Reorganize the preempt_disable/enable calls to eliminate the extra preempt
depth.  Changes based on Paul McKenney's review suggestions for the kprobes
RCU changeset.

Signed-off-by: Ananth N Mavinakayanahalli <ananth@in.ibm.com>
Signed-off-by: Anil S Keshavamurthy <anil.s.keshavamurthy@intel.com>
Signed-off-by: Andrew Morton <akpm@osdl.org>
Signed-off-by: Linus Torvalds <torvalds@osdl.org>

7 files changed, 80 insertions, 61 deletions

diff --git a/arch/i386/kernel/kprobes.c b/arch/i386/kernel/kprobes.c
index ad469299267a..32b0c24ab9a6 100644
--- a/arch/i386/kernel/kprobes.c
+++ b/arch/i386/kernel/kprobes.c
@@ -153,7 +153,14 @@ static int __kprobes kprobe_handler(struct pt_regs *regs)
         int ret = 0;
         kprobe_opcode_t *addr = NULL;
         unsigned long *lp;
-        struct kprobe_ctlblk *kcb = get_kprobe_ctlblk();
+        struct kprobe_ctlblk *kcb;
+
+        /*
+         * We don't want to be preempted for the entire
+         * duration of kprobe processing
+         */
+        preempt_disable();
+        kcb = get_kprobe_ctlblk();
 
         /* Check if the application is using LDT entry for its code segment and
          * calculate the address by reading the base address from the LDT entry.
@@ -221,11 +228,6 @@ static int __kprobes kprobe_handler(struct pt_regs *regs)
                 goto no_kprobe;
         }
 
-        /*
-         * This preempt_disable() matches the preempt_enable_no_resched()
-         * in post_kprobe_handler()
-         */
-        preempt_disable();
         set_current_kprobe(p, regs, kcb);
         kcb->kprobe_status = KPROBE_HIT_ACTIVE;
 
@@ -239,6 +241,7 @@ ss_probe:
         return 1;
 
 no_kprobe:
+        preempt_enable_no_resched();
         return ret;
 }
 
@@ -310,8 +313,8 @@ int __kprobes trampoline_probe_handler(struct kprobe *p, struct pt_regs *regs)
 
         /*
          * By returning a non-zero value, we are telling
-         * kprobe_handler() that we have handled unlocking
-         * and re-enabling preemption
+         * kprobe_handler() that we don't want the post_handler
+         * to run (and have re-enabled preemption)
          */
         return 1;
 }
@@ -455,7 +458,6 @@ int __kprobes kprobe_exceptions_notify(struct notifier_block *self,
         struct die_args *args = (struct die_args *)data;
         int ret = NOTIFY_DONE;
 
-        rcu_read_lock();
         switch (val) {
         case DIE_INT3:
                 if (kprobe_handler(args->regs))
@@ -467,14 +469,16 @@ int __kprobes kprobe_exceptions_notify(struct notifier_block *self,
                 break;
         case DIE_GPF:
         case DIE_PAGE_FAULT:
+                /* kprobe_running() needs smp_processor_id() */
+                preempt_disable();
                 if (kprobe_running() &&
                     kprobe_fault_handler(args->regs, args->trapnr))
                         ret = NOTIFY_STOP;
+                preempt_enable();
                 break;
         default:
                 break;
         }
-        rcu_read_unlock();
         return ret;
 }
 
@@ -537,6 +541,7 @@ int __kprobes longjmp_break_handler(struct kprobe *p, struct pt_regs *regs)
                 *regs = kcb->jprobe_saved_regs;
                 memcpy((kprobe_opcode_t *) stack_addr, kcb->jprobes_stack,
                        MIN_STACK_SIZE(stack_addr));
+                preempt_enable_no_resched();
                 return 1;
         }
         return 0;
diff --git a/arch/ia64/kernel/kprobes.c b/arch/ia64/kernel/kprobes.c
index fddbac32d44a..96736a119c91 100644
--- a/arch/ia64/kernel/kprobes.c
+++ b/arch/ia64/kernel/kprobes.c
@@ -389,11 +389,11 @@ int __kprobes trampoline_probe_handler(struct kprobe *p, struct pt_regs *regs)
         spin_unlock_irqrestore(&kretprobe_lock, flags);
         preempt_enable_no_resched();
 
-        /*
-         * By returning a non-zero value, we are telling
-         * kprobe_handler() that we have handled unlocking
-         * and re-enabling preemption
-         */
+        /*
+         * By returning a non-zero value, we are telling
+         * kprobe_handler() that we don't want the post_handler
+         * to run (and have re-enabled preemption)
+         */
         return 1;
 }
 
@@ -604,7 +604,14 @@ static int __kprobes pre_kprobes_handler(struct die_args *args)
         int ret = 0;
         struct pt_regs *regs = args->regs;
         kprobe_opcode_t *addr = (kprobe_opcode_t *)instruction_pointer(regs);
-        struct kprobe_ctlblk *kcb = get_kprobe_ctlblk();
+        struct kprobe_ctlblk *kcb;
+
+        /*
+         * We don't want to be preempted for the entire
+         * duration of kprobe processing
+         */
+        preempt_disable();
+        kcb = get_kprobe_ctlblk();
 
         /* Handle recursion cases */
         if (kprobe_running()) {
@@ -659,11 +666,6 @@ static int __kprobes pre_kprobes_handler(struct die_args *args)
                 goto no_kprobe;
         }
 
-        /*
-         * This preempt_disable() matches the preempt_enable_no_resched()
-         * in post_kprobes_handler()
-         */
-        preempt_disable();
         set_current_kprobe(p, kcb);
         kcb->kprobe_status = KPROBE_HIT_ACTIVE;
 
@@ -681,6 +683,7 @@ ss_probe:
         return 1;
 
 no_kprobe:
+        preempt_enable_no_resched();
         return ret;
 }
 
@@ -716,9 +719,6 @@ static int __kprobes kprobes_fault_handler(struct pt_regs *regs, int trapnr)
         struct kprobe *cur = kprobe_running();
         struct kprobe_ctlblk *kcb = get_kprobe_ctlblk();
 
-        if (!cur)
-                return 0;
-
         if (cur->fault_handler && cur->fault_handler(cur, regs, trapnr))
                 return 1;
 
@@ -737,7 +737,6 @@ int __kprobes kprobe_exceptions_notify(struct notifier_block *self,
         struct die_args *args = (struct die_args *)data;
         int ret = NOTIFY_DONE;
 
-        rcu_read_lock();
         switch(val) {
         case DIE_BREAK:
                 if (pre_kprobes_handler(args))
@@ -748,12 +747,15 @@ int __kprobes kprobe_exceptions_notify(struct notifier_block *self,
                         ret = NOTIFY_STOP;
                 break;
         case DIE_PAGE_FAULT:
-                if (kprobes_fault_handler(args->regs, args->trapnr))
+                /* kprobe_running() needs smp_processor_id() */
+                preempt_disable();
+                if (kprobe_running() &&
+                        kprobes_fault_handler(args->regs, args->trapnr))
                         ret = NOTIFY_STOP;
+                preempt_enable();
         default:
                 break;
         }
-        rcu_read_unlock();
         return ret;
 }
 
@@ -785,6 +787,7 @@ int __kprobes longjmp_break_handler(struct kprobe *p, struct pt_regs *regs)
         struct kprobe_ctlblk *kcb = get_kprobe_ctlblk();
 
         *regs = kcb->jprobe_saved_regs;
+        preempt_enable_no_resched();
         return 1;
 }
 
diff --git a/arch/ppc64/kernel/kprobes.c b/arch/ppc64/kernel/kprobes.c
index e0a25b35437f..511af54e6230 100644
--- a/arch/ppc64/kernel/kprobes.c
+++ b/arch/ppc64/kernel/kprobes.c
@@ -148,7 +148,14 @@ static inline int kprobe_handler(struct pt_regs *regs)
         struct kprobe *p;
         int ret = 0;
         unsigned int *addr = (unsigned int *)regs->nip;
-        struct kprobe_ctlblk *kcb = get_kprobe_ctlblk();
+        struct kprobe_ctlblk *kcb;
+
+        /*
+         * We don't want to be preempted for the entire
+         * duration of kprobe processing
+         */
+        preempt_disable();
+        kcb = get_kprobe_ctlblk();
 
         /* Check we're not actually recursing */
         if (kprobe_running()) {
@@ -207,11 +214,6 @@ static inline int kprobe_handler(struct pt_regs *regs)
                 goto no_kprobe;
         }
 
-        /*
-         * This preempt_disable() matches the preempt_enable_no_resched()
-         * in post_kprobe_handler().
-         */
-        preempt_disable();
         kcb->kprobe_status = KPROBE_HIT_ACTIVE;
         set_current_kprobe(p, regs, kcb);
         if (p->pre_handler && p->pre_handler(p, regs))
@@ -224,6 +226,7 @@ ss_probe:
         return 1;
 
 no_kprobe:
+        preempt_enable_no_resched();
         return ret;
 }
 
@@ -296,8 +299,8 @@ int __kprobes trampoline_probe_handler(struct kprobe *p, struct pt_regs *regs)
 
         /*
          * By returning a non-zero value, we are telling
-         * kprobe_handler() that we have handled unlocking
-         * and re-enabling preemption.
+         * kprobe_handler() that we don't want the post_handler
+         * to run (and have re-enabled preemption)
          */
         return 1;
 }
@@ -385,7 +388,6 @@ int __kprobes kprobe_exceptions_notify(struct notifier_block *self,
         struct die_args *args = (struct die_args *)data;
         int ret = NOTIFY_DONE;
 
-        rcu_read_lock();
         switch (val) {
         case DIE_BPT:
                 if (kprobe_handler(args->regs))
@@ -396,14 +398,16 @@ int __kprobes kprobe_exceptions_notify(struct notifier_block *self,
                         ret = NOTIFY_STOP;
                 break;
         case DIE_PAGE_FAULT:
+                /* kprobe_running() needs smp_processor_id() */
+                preempt_disable();
                 if (kprobe_running() &&
                     kprobe_fault_handler(args->regs, args->trapnr))
                         ret = NOTIFY_STOP;
+                preempt_enable();
                 break;
         default:
                 break;
         }
-        rcu_read_unlock();
         return ret;
 }
 
@@ -440,6 +444,7 @@ int __kprobes longjmp_break_handler(struct kprobe *p, struct pt_regs *regs)
          * saved regs...
          */
         memcpy(regs, &kcb->jprobe_saved_regs, sizeof(struct pt_regs));
+        preempt_enable_no_resched();
         return 1;
 }
 
diff --git a/arch/sparc64/kernel/kprobes.c b/arch/sparc64/kernel/kprobes.c
index 58a815e90373..96bd09b098f4 100644
--- a/arch/sparc64/kernel/kprobes.c
+++ b/arch/sparc64/kernel/kprobes.c
@@ -113,7 +113,14 @@ static int __kprobes kprobe_handler(struct pt_regs *regs)
         struct kprobe *p;
         void *addr = (void *) regs->tpc;
         int ret = 0;
-        struct kprobe_ctlblk *kcb = get_kprobe_ctlblk();
+        struct kprobe_ctlblk *kcb;
+
+        /*
+         * We don't want to be preempted for the entire
+         * duration of kprobe processing
+         */
+        preempt_disable();
+        kcb = get_kprobe_ctlblk();
 
         if (kprobe_running()) {
                 p = get_kprobe(addr);
@@ -159,11 +166,6 @@ static int __kprobes kprobe_handler(struct pt_regs *regs)
                 goto no_kprobe;
         }
 
-        /*
-         * This preempt_disable() matches the preempt_enable_no_resched()
-         * in post_kprobes_handler()
-         */
-        preempt_disable();
         set_current_kprobe(p, regs, kcb);
         kcb->kprobe_status = KPROBE_HIT_ACTIVE;
         if (p->pre_handler && p->pre_handler(p, regs))
@@ -175,6 +177,7 @@ ss_probe:
         return 1;
 
 no_kprobe:
+        preempt_enable_no_resched();
         return ret;
 }
 
@@ -321,7 +324,6 @@ int __kprobes kprobe_exceptions_notify(struct notifier_block *self,
         struct die_args *args = (struct die_args *)data;
         int ret = NOTIFY_DONE;
 
-        rcu_read_lock();
         switch (val) {
         case DIE_DEBUG:
                 if (kprobe_handler(args->regs))
@@ -333,14 +335,16 @@ int __kprobes kprobe_exceptions_notify(struct notifier_block *self,
                 break;
         case DIE_GPF:
         case DIE_PAGE_FAULT:
+                /* kprobe_running() needs smp_processor_id() */
+                preempt_disable();
                 if (kprobe_running() &&
                     kprobe_fault_handler(args->regs, args->trapnr))
                         ret = NOTIFY_STOP;
+                preempt_enable();
                 break;
         default:
                 break;
         }
-        rcu_read_unlock();
         return ret;
 }
 
@@ -426,6 +430,7 @@ int __kprobes longjmp_break_handler(struct kprobe *p, struct pt_regs *regs)
                        &(kcb->jprobe_saved_stack),
                        sizeof(kcb->jprobe_saved_stack));
 
+                preempt_enable_no_resched();
                 return 1;
         }
         return 0;
diff --git a/arch/x86_64/kernel/kprobes.c b/arch/x86_64/kernel/kprobes.c
index 9bef2c8dc12c..dddeb678b440 100644
--- a/arch/x86_64/kernel/kprobes.c
+++ b/arch/x86_64/kernel/kprobes.c
@@ -286,16 +286,19 @@ void __kprobes arch_prepare_kretprobe(struct kretprobe *rp,
         }
 }
 
-/*
- * Interrupts are disabled on entry as trap3 is an interrupt gate and they
- * remain disabled thorough out this function.
- */
 int __kprobes kprobe_handler(struct pt_regs *regs)
 {
         struct kprobe *p;
         int ret = 0;
         kprobe_opcode_t *addr = (kprobe_opcode_t *)(regs->rip - sizeof(kprobe_opcode_t));
-        struct kprobe_ctlblk *kcb = get_kprobe_ctlblk();
+        struct kprobe_ctlblk *kcb;
+
+        /*
+         * We don't want to be preempted for the entire
+         * duration of kprobe processing
+         */
+        preempt_disable();
+        kcb = get_kprobe_ctlblk();
 
         /* Check we're not actually recursing */
         if (kprobe_running()) {
@@ -359,11 +362,6 @@ int __kprobes kprobe_handler(struct pt_regs *regs)
                 goto no_kprobe;
         }
 
-        /*
-         * This preempt_disable() matches the preempt_enable_no_resched()
-         * in post_kprobe_handler()
-         */
-        preempt_disable();
         set_current_kprobe(p, regs, kcb);
         kcb->kprobe_status = KPROBE_HIT_ACTIVE;
 
@@ -377,6 +375,7 @@ ss_probe:
         return 1;
 
 no_kprobe:
+        preempt_enable_no_resched();
         return ret;
 }
 
@@ -448,8 +447,8 @@ int __kprobes trampoline_probe_handler(struct kprobe *p, struct pt_regs *regs)
 
         /*
          * By returning a non-zero value, we are telling
-         * kprobe_handler() that we have handled unlocking
-         * and re-enabling preemption
+         * kprobe_handler() that we don't want the post_handler
+         * to run (and have re-enabled preemption)
          */
         return 1;
 }
@@ -594,7 +593,6 @@ int __kprobes kprobe_exceptions_notify(struct notifier_block *self,
         struct die_args *args = (struct die_args *)data;
         int ret = NOTIFY_DONE;
 
-        rcu_read_lock();
         switch (val) {
         case DIE_INT3:
                 if (kprobe_handler(args->regs))
@@ -606,14 +604,16 @@ int __kprobes kprobe_exceptions_notify(struct notifier_block *self,
                 break;
         case DIE_GPF:
         case DIE_PAGE_FAULT:
+                /* kprobe_running() needs smp_processor_id() */
+                preempt_disable();
                 if (kprobe_running() &&
                     kprobe_fault_handler(args->regs, args->trapnr))
                         ret = NOTIFY_STOP;
+                preempt_enable();
                 break;
         default:
                 break;
         }
-        rcu_read_unlock();
         return ret;
 }
 
@@ -675,6 +675,7 @@ int __kprobes longjmp_break_handler(struct kprobe *p, struct pt_regs *regs)
                 *regs = kcb->jprobe_saved_regs;
                 memcpy((kprobe_opcode_t *) stack_addr, kcb->jprobes_stack,
                        MIN_STACK_SIZE(stack_addr));
+                preempt_enable_no_resched();
                 return 1;
         }
         return 0;
diff --git a/include/linux/kprobes.h b/include/linux/kprobes.h
index cff281cf70cf..e373c4a9de53 100644
--- a/include/linux/kprobes.h
+++ b/include/linux/kprobes.h
@@ -159,7 +159,7 @@ extern void show_registers(struct pt_regs *regs);
 extern kprobe_opcode_t *get_insn_slot(void);
 extern void free_insn_slot(kprobe_opcode_t *slot);
 
-/* Get the kprobe at this addr (if any) - called under a rcu_read_lock() */
+/* Get the kprobe at this addr (if any) - called with preemption disabled */
 struct kprobe *get_kprobe(void *addr);
 struct hlist_head * kretprobe_inst_table_head(struct task_struct *tsk);
 
diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index cfef426e4cdc..5beda378cc75 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -167,7 +167,7 @@ static inline void reset_kprobe_instance(void)
  * This routine is called either:
  *      - under the kprobe_lock spinlock - during kprobe_[un]register()
  *                              OR
- *      - under an rcu_read_lock() - from arch/xxx/kernel/kprobes.c
+ *      - with preemption disabled - from arch/xxx/kernel/kprobes.c
  */
 struct kprobe __kprobes *get_kprobe(void *addr)
 {