2 files changed, 4 insertions, 4 deletions

diff --git a/kernel/sched.c b/kernel/sched.c
index 1701eaebb79c..adb5e923cc61 100644
--- a/kernel/sched.c
+++ b/kernel/sched.c
@@ -5520,20 +5520,16 @@ need_resched_nonpreemptible:
                 spin_unlock_irq(&rq->lock);
         }
 
-        TS_SCHED2_START(current);
         sched_trace_task_switch_to(current);
 
         post_schedule(rq);
 
         if (unlikely(reacquire_kernel_lock(current) < 0)) {
-                TS_SCHED2_END(current);
                 goto need_resched_nonpreemptible;
         }
 
         preempt_enable_no_resched();
 
-        TS_SCHED2_END(current);
-
         if (need_resched())
                 goto need_resched;
 
diff --git a/litmus/rt_domain.c b/litmus/rt_domain.c
index 0ed6d5cbbfc5..609ff0f82abb 100644
--- a/litmus/rt_domain.c
+++ b/litmus/rt_domain.c
@@ -301,6 +301,10 @@ void __add_release(rt_domain_t* rt, struct task_struct *task)
         task->rt_param.domain = rt;
 
         /* start release timer */
+        TS_SCHED2_START(task);
+
         arm_release_timer(rt);
+
+        TS_SCHED2_END(task);
 }