3 files changed, 67 insertions, 12 deletions

diff --git a/tools/perf/builtin-sched.c b/tools/perf/builtin-sched.c
index e1df7055ab82..25b91e784332 100644
--- a/tools/perf/builtin-sched.c
+++ b/tools/perf/builtin-sched.c
@@ -1034,6 +1034,36 @@ add_sched_in_event(struct work_atoms *atoms, u64 timestamp)
         atoms->nb_atoms++;
 }
 
+static struct thread *
+threads__findnew_from_ctx(u32 pid, struct trace_switch_event *switch_event)
+{
+        struct thread *th;
+
+        th = threads__findnew_nocomm(pid, &threads, &last_match);
+        if (th->comm)
+                return th;
+
+        if (pid == switch_event->prev_pid)
+                thread__set_comm(th, switch_event->prev_comm);
+        else
+                thread__set_comm(th, switch_event->next_comm);
+        return th;
+}
+
+static struct thread *
+threads__findnew_from_wakeup(struct trace_wakeup_event *wakeup_event)
+{
+        struct thread *th;
+
+        th =  threads__findnew_nocomm(wakeup_event->pid, &threads, &last_match);
+        if (th->comm)
+                return th;
+
+        thread__set_comm(th, wakeup_event->comm);
+
+        return th;
+}
+
 static void
 latency_switch_event(struct trace_switch_event *switch_event,
                      struct event *event __used,
@@ -1059,8 +1089,10 @@ latency_switch_event(struct trace_switch_event *switch_event,
                 die("hm, delta: %Ld < 0 ?\n", delta);
 
 
-        sched_out = threads__findnew(switch_event->prev_pid, &threads, &last_match);
-        sched_in = threads__findnew(switch_event->next_pid, &threads, &last_match);
+        sched_out = threads__findnew_from_ctx(switch_event->prev_pid,
+                                              switch_event);
+        sched_in = threads__findnew_from_ctx(switch_event->next_pid,
+                                             switch_event);
 
         out_events = thread_atoms_search(&atom_root, sched_out, &cmp_pid);
         if (!out_events) {
@@ -1126,7 +1158,7 @@ latency_wakeup_event(struct trace_wakeup_event *wakeup_event,
         if (!wakeup_event->success)
                 return;
 
-        wakee = threads__findnew(wakeup_event->pid, &threads, &last_match);
+        wakee = threads__findnew_from_wakeup(wakeup_event);
         atoms = thread_atoms_search(&atom_root, wakee, &cmp_pid);
         if (!atoms) {
                 thread_atoms_insert(wakee);
@@ -1386,8 +1418,10 @@ map_switch_event(struct trace_switch_event *switch_event,
                 die("hm, delta: %Ld < 0 ?\n", delta);
 
 
-        sched_out = threads__findnew(switch_event->prev_pid, &threads, &last_match);
-        sched_in = threads__findnew(switch_event->next_pid, &threads, &last_match);
+        sched_out = threads__findnew_from_ctx(switch_event->prev_pid,
+                                              switch_event);
+        sched_in = threads__findnew_from_ctx(switch_event->next_pid,
+                                             switch_event);
 
         curr_thread[this_cpu] = sched_in;
 
diff --git a/tools/perf/util/thread.c b/tools/perf/util/thread.c
index 3b56aebb1f4b..8bd5ca2d2f28 100644
--- a/tools/perf/util/thread.c
+++ b/tools/perf/util/thread.c
@@ -6,15 +6,17 @@
 #include "util.h"
 #include "debug.h"
 
-static struct thread *thread__new(pid_t pid)
+static struct thread *thread__new(pid_t pid, int set_comm)
 {
         struct thread *self = calloc(1, sizeof(*self));
 
         if (self != NULL) {
                 self->pid = pid;
-                self->comm = malloc(32);
-                if (self->comm)
-                        snprintf(self->comm, 32, ":%d", self->pid);
+                if (set_comm) {
+                        self->comm = malloc(32);
+                        if (self->comm)
+                                snprintf(self->comm, 32, ":%d", self->pid);
+                }
                 self->maps = RB_ROOT;
                 INIT_LIST_HEAD(&self->removed_maps);
         }
@@ -50,8 +52,10 @@ static size_t thread__fprintf(struct thread *self, FILE *fp)
         return ret;
 }
 
-struct thread *
-threads__findnew(pid_t pid, struct rb_root *threads, struct thread **last_match)
+static struct thread *
+__threads__findnew(pid_t pid, struct rb_root *threads,
+                   struct thread **last_match,
+                   int set_comm)
 {
         struct rb_node **p = &threads->rb_node;
         struct rb_node *parent = NULL;
@@ -80,7 +84,8 @@ threads__findnew(pid_t pid, struct rb_root *threads, struct thread **last_match)
                         p = &(*p)->rb_right;
         }
 
-        th = thread__new(pid);
+        th = thread__new(pid, set_comm);
+
         if (th != NULL) {
                 rb_link_node(&th->rb_node, parent, p);
                 rb_insert_color(&th->rb_node, threads);
@@ -91,6 +96,19 @@ threads__findnew(pid_t pid, struct rb_root *threads, struct thread **last_match)
 }
 
 struct thread *
+threads__findnew(pid_t pid, struct rb_root *threads, struct thread **last_match)
+{
+        return __threads__findnew(pid, threads, last_match, 1);
+}
+
+struct thread *
+threads__findnew_nocomm(pid_t pid, struct rb_root *threads,
+                        struct thread **last_match)
+{
+        return __threads__findnew(pid, threads, last_match, 0);
+}
+
+struct thread *
 register_idle_thread(struct rb_root *threads, struct thread **last_match)
 {
         struct thread *thread = threads__findnew(0, threads, last_match);
diff --git a/tools/perf/util/thread.h b/tools/perf/util/thread.h
index 845d9b62f96f..75bc843950c4 100644
--- a/tools/perf/util/thread.h
+++ b/tools/perf/util/thread.h
@@ -18,6 +18,9 @@ int thread__set_comm(struct thread *self, const char *comm);
 struct thread *
 threads__findnew(pid_t pid, struct rb_root *threads, struct thread **last_match);
 struct thread *
+threads__findnew_nocomm(pid_t pid, struct rb_root *threads,
+                        struct thread **last_match);
+struct thread *
 register_idle_thread(struct rb_root *threads, struct thread **last_match);
 void thread__insert_map(struct thread *self, struct map *map);
 int thread__fork(struct thread *self, struct thread *parent);