From 655e3c302ba39362be04bdab237b28814fe14b48 Mon Sep 17 00:00:00 2001
From: Sasha Goldshtein <goldshtn@gmail.com>
Date: Mon, 30 Jan 2017 11:08:12 +0000
Subject: [PATCH] argdist: Fix -p behavior to filter tgid and not pid

argdist remained one of the last holdouts to use the `-p` switch
inconsistently with other tools, filtering for kernel pid (thread
id from user space perspective) and not kernel tgid (process id
from user space perspective). This is now fixed.

Additionally, minor nits around generating pid filters were fixed,
and a potential collision with user-provided argument names was
fixed too (in general, script-generated arguments/locals should
probably stick to reserved identifiers, such as `__whatever` rather
than `whatever`).
---
 tools/argdist.py | 22 ++++++++++++----------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/tools/argdist.py b/tools/argdist.py
index f5422d90c827..474af49efea9 100755
--- a/tools/argdist.py
+++ b/tools/argdist.py
@@ -20,7 +20,7 @@
 class Probe(object):
         next_probe_index = 0
         streq_index = 0
-        aliases = {"$PID": "bpf_get_current_pid_tgid()"}
+        aliases = {"$PID": "(bpf_get_current_pid_tgid() >> 32)"}
 
         def _substitute_aliases(self, expr):
                 if expr is None:
@@ -47,7 +47,9 @@ def _generate_entry(self):
                 text = """
 int PROBENAME(struct pt_regs *ctx SIGNATURE)
 {
-        u32 pid = bpf_get_current_pid_tgid();
+        u64 __pid_tgid = bpf_get_current_pid_tgid();
+        u32 __pid      = __pid_tgid;        // lower 32 bits
+        u32 __tgid     = __pid_tgid >> 32;  // upper 32 bits
         PID_FILTER
         COLLECT
         return 0;
@@ -56,19 +58,17 @@ def _generate_entry(self):
                 text = text.replace("PROBENAME", self.entry_probe_func)
                 text = text.replace("SIGNATURE",
                      "" if len(self.signature) == 0 else ", " + self.signature)
-                pid_filter = "" if self.is_user or self.pid is None \
-                                else "if (pid != %d) { return 0; }" % self.pid
-                text = text.replace("PID_FILTER", pid_filter)
+                text = text.replace("PID_FILTER", self._generate_pid_filter())
                 collect = ""
                 for pname in self.args_to_probe:
                         param_hash = self.hashname_prefix + pname
                         if pname == "__latency":
                                 collect += """
 u64 __time = bpf_ktime_get_ns();
-%s.update(&pid, &__time);
+%s.update(&__pid, &__time);
                         """ % param_hash
                         else:
-                                collect += "%s.update(&pid, &%s);\n" % \
+                                collect += "%s.update(&__pid, &%s);\n" % \
                                            (param_hash, pname)
                 text = text.replace("COLLECT", collect)
                 return text
@@ -108,7 +108,7 @@ def _generate_retprobe_prefix(self):
                 # argument we needed to probe using $entry(name), and they all
                 # have values (which isn't necessarily the case if we missed
                 # the method entry probe).
-                text = "u32 __pid = bpf_get_current_pid_tgid();\n"
+                text = ""
                 self.param_val_names = {}
                 for pname in self.args_to_probe:
                         val_name = "__%s_val" % pname
@@ -345,8 +345,7 @@ def _generate_pid_filter(self):
                 # Kernel probes need to explicitly filter pid, because the
                 # attach interface doesn't support pid filtering
                 if self.pid is not None and not self.is_user:
-                        return "u32 pid = bpf_get_current_pid_tgid();\n" + \
-                               "if (pid != %d) { return 0; }" % self.pid
+                        return "if (__tgid != %d) { return 0; }" % self.pid
                 else:
                         return ""
 
@@ -360,6 +359,9 @@ def generate_text(self):
                     if self.probe_type == "t"
                     else "int PROBENAME(struct pt_regs *ctx SIGNATURE)") + """
 {
+        u64 __pid_tgid = bpf_get_current_pid_tgid();
+        u32 __pid      = __pid_tgid;        // lower 32 bits
+        u32 __tgid     = __pid_tgid >> 32;  // upper 32 bits
         PID_FILTER
         PREFIX
         if (!(FILTER)) return 0;