6 files changed, 59 insertions, 33 deletions

diff --git a/include/linux/context_tracking.h b/include/linux/context_tracking.h
index 2821838256b4..b96bd299966f 100644
--- a/include/linux/context_tracking.h
+++ b/include/linux/context_tracking.h
@@ -14,8 +14,6 @@ extern void context_tracking_enter(enum ctx_state state);
 extern void context_tracking_exit(enum ctx_state state);
 extern void context_tracking_user_enter(void);
 extern void context_tracking_user_exit(void);
-extern void __context_tracking_task_switch(struct task_struct *prev,
-                                           struct task_struct *next);
 
 static inline void user_enter(void)
 {
@@ -51,19 +49,11 @@ static inline void exception_exit(enum ctx_state prev_ctx)
         }
 }
 
-static inline void context_tracking_task_switch(struct task_struct *prev,
-                                                struct task_struct *next)
-{
-        if (context_tracking_is_enabled())
-                __context_tracking_task_switch(prev, next);
-}
 #else
 static inline void user_enter(void) { }
 static inline void user_exit(void) { }
 static inline enum ctx_state exception_enter(void) { return 0; }
 static inline void exception_exit(enum ctx_state prev_ctx) { }
-static inline void context_tracking_task_switch(struct task_struct *prev,
-                                                struct task_struct *next) { }
 #endif /* !CONFIG_CONTEXT_TRACKING */
 
 
diff --git a/include/linux/context_tracking_state.h b/include/linux/context_tracking_state.h
index 6b7b96a32b75..678ecdf90cf6 100644
--- a/include/linux/context_tracking_state.h
+++ b/include/linux/context_tracking_state.h
@@ -12,6 +12,7 @@ struct context_tracking {
          * may be further optimized using static keys.
          */
         bool active;
+        int recursion;
         enum ctx_state {
                 CONTEXT_KERNEL = 0,
                 CONTEXT_USER,
diff --git a/include/linux/sched.h b/include/linux/sched.h
index 30364cb58b1f..6633e83e608a 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -2599,6 +2599,9 @@ static inline unsigned long wait_task_inactive(struct task_struct *p,
 }
 #endif
 
+#define tasklist_empty() \
+        list_empty(&init_task.tasks)
+
 #define next_task(p) \
         list_entry_rcu((p)->tasks.next, struct task_struct, tasks)
 
diff --git a/include/linux/tick.h b/include/linux/tick.h
index f8492da57ad3..4191b5623a28 100644
--- a/include/linux/tick.h
+++ b/include/linux/tick.h
@@ -134,6 +134,12 @@ static inline bool tick_nohz_full_cpu(int cpu)
         return cpumask_test_cpu(cpu, tick_nohz_full_mask);
 }
 
+static inline void tick_nohz_full_add_cpus_to(struct cpumask *mask)
+{
+        if (tick_nohz_full_enabled())
+                cpumask_or(mask, mask, tick_nohz_full_mask);
+}
+
 extern void __tick_nohz_full_check(void);
 extern void tick_nohz_full_kick(void);
 extern void tick_nohz_full_kick_cpu(int cpu);
@@ -142,6 +148,7 @@ extern void __tick_nohz_task_switch(struct task_struct *tsk);
 #else
 static inline bool tick_nohz_full_enabled(void) { return false; }
 static inline bool tick_nohz_full_cpu(int cpu) { return false; }
+static inline void tick_nohz_full_add_cpus_to(struct cpumask *mask) { }
 static inline void __tick_nohz_full_check(void) { }
 static inline void tick_nohz_full_kick_cpu(int cpu) { }
 static inline void tick_nohz_full_kick(void) { }
diff --git a/kernel/context_tracking.c b/kernel/context_tracking.c
index 72d59a1a6eb6..0a495ab35bc7 100644
--- a/kernel/context_tracking.c
+++ b/kernel/context_tracking.c
@@ -30,12 +30,23 @@ EXPORT_SYMBOL_GPL(context_tracking_enabled);
 DEFINE_PER_CPU(struct context_tracking, context_tracking);
 EXPORT_SYMBOL_GPL(context_tracking);
 
-void context_tracking_cpu_set(int cpu)
+static bool context_tracking_recursion_enter(void)
 {
-        if (!per_cpu(context_tracking.active, cpu)) {
-                per_cpu(context_tracking.active, cpu) = true;
-                static_key_slow_inc(&context_tracking_enabled);
-        }
+        int recursion;
+
+        recursion = __this_cpu_inc_return(context_tracking.recursion);
+        if (recursion == 1)
+                return true;
+
+        WARN_ONCE((recursion < 1), "Invalid context tracking recursion value %d\n", recursion);
+        __this_cpu_dec(context_tracking.recursion);
+
+        return false;
+}
+
+static void context_tracking_recursion_exit(void)
+{
+        __this_cpu_dec(context_tracking.recursion);
 }
 
 /**
@@ -75,6 +86,9 @@ void context_tracking_enter(enum ctx_state state)
         WARN_ON_ONCE(!current->mm);
 
         local_irq_save(flags);
+        if (!context_tracking_recursion_enter())
+                goto out_irq_restore;
+
         if ( __this_cpu_read(context_tracking.state) != state) {
                 if (__this_cpu_read(context_tracking.active)) {
                         /*
@@ -105,6 +119,8 @@ void context_tracking_enter(enum ctx_state state)
                  */
                 __this_cpu_write(context_tracking.state, state);
         }
+        context_tracking_recursion_exit();
+out_irq_restore:
         local_irq_restore(flags);
 }
 NOKPROBE_SYMBOL(context_tracking_enter);
@@ -139,6 +155,9 @@ void context_tracking_exit(enum ctx_state state)
                 return;
 
         local_irq_save(flags);
+        if (!context_tracking_recursion_enter())
+                goto out_irq_restore;
+
         if (__this_cpu_read(context_tracking.state) == state) {
                 if (__this_cpu_read(context_tracking.active)) {
                         /*
@@ -153,6 +172,8 @@ void context_tracking_exit(enum ctx_state state)
                 }
                 __this_cpu_write(context_tracking.state, CONTEXT_KERNEL);
         }
+        context_tracking_recursion_exit();
+out_irq_restore:
         local_irq_restore(flags);
 }
 NOKPROBE_SYMBOL(context_tracking_exit);
@@ -164,24 +185,26 @@ void context_tracking_user_exit(void)
 }
 NOKPROBE_SYMBOL(context_tracking_user_exit);
 
-/**
- * __context_tracking_task_switch - context switch the syscall callbacks
- * @prev: the task that is being switched out
- * @next: the task that is being switched in
- *
- * The context tracking uses the syscall slow path to implement its user-kernel
- * boundaries probes on syscalls. This way it doesn't impact the syscall fast
- * path on CPUs that don't do context tracking.
- *
- * But we need to clear the flag on the previous task because it may later
- * migrate to some CPU that doesn't do the context tracking. As such the TIF
- * flag may not be desired there.
- */
-void __context_tracking_task_switch(struct task_struct *prev,
-                                    struct task_struct *next)
+void __init context_tracking_cpu_set(int cpu)
 {
-        clear_tsk_thread_flag(prev, TIF_NOHZ);
-        set_tsk_thread_flag(next, TIF_NOHZ);
+        static __initdata bool initialized = false;
+
+        if (!per_cpu(context_tracking.active, cpu)) {
+                per_cpu(context_tracking.active, cpu) = true;
+                static_key_slow_inc(&context_tracking_enabled);
+        }
+
+        if (initialized)
+                return;
+
+        /*
+         * Set TIF_NOHZ to init/0 and let it propagate to all tasks through fork
+         * This assumes that init is the only task at this early boot stage.
+         */
+        set_tsk_thread_flag(&init_task, TIF_NOHZ);
+        WARN_ON_ONCE(!tasklist_empty());
+
+        initialized = true;
 }
 
 #ifdef CONFIG_CONTEXT_TRACKING_FORCE
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index c9a707b59331..c86935a7f1f8 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -2374,7 +2374,6 @@ context_switch(struct rq *rq, struct task_struct *prev,
          */
         spin_release(&rq->lock.dep_map, 1, _THIS_IP_);
 
-        context_tracking_task_switch(prev, next);
         /* Here we just switch the register state and the stack. */
         switch_to(prev, next, prev);
         barrier();
@@ -7068,6 +7067,9 @@ void __init sched_init_smp(void)
         alloc_cpumask_var(&non_isolated_cpus, GFP_KERNEL);
         alloc_cpumask_var(&fallback_doms, GFP_KERNEL);
 
+        /* nohz_full won't take effect without isolating the cpus. */
+        tick_nohz_full_add_cpus_to(cpu_isolated_map);
+
         sched_init_numa();
 
         /*