1 files changed, 3 insertions, 7 deletions

diff --git a/litmus/sched_task_trace.c b/litmus/sched_task_trace.c
index 8c1ca188bce1..67b01c1dd51b 100644
--- a/litmus/sched_task_trace.c
+++ b/litmus/sched_task_trace.c
@@ -7,10 +7,10 @@
 #include <linux/module.h>
 #include <linux/sched.h>
 #include <linux/percpu.h>
+#include <linux/math64.h>
 
 #include <litmus/ftdev.h>
 #include <litmus/litmus.h>
-
 #include <litmus/sched_trace.h>
 #include <litmus/feather_trace.h>
 #include <litmus/ftdev.h>
@@ -235,13 +235,9 @@ feather_callback void do_sched_trace_task_exit(unsigned long id,
                                                unsigned long _task)
 {
         struct task_struct *t = (struct task_struct*) _task;
-#ifdef CONFIG_PLUGIN_COLOR
         const lt_t max_exec_time = tsk_rt(t)->max_exec_time;
-        const lt_t avg_exec_time = tsk_rt(t)->tot_exec_time / (get_rt_job(t) - 1);
-#else
-        const lt_t max_exec_time = 0;
-        const lt_t avg_exec_time = 0;
-#endif
+        const lt_t avg_exec_time = div64_u64(tsk_rt(t)->tot_exec_time, (get_job_no(t) - 1));
+
         struct st_event_record *rec = get_record(ST_TASK_EXIT, t);
         if (rec) {
                 rec->data.task_exit.avg_exec_time = avg_exec_time;