funccount: Fix on-CPU hang when attaching to SyS_*

src/python/bcc/__init__.py

@@ -422,7 +422,7 @@ def get_kprobe_functions(event_re):
  fn = line.rstrip().split()[0]
  if re.match(event_re, fn) and fn not in blacklist:
  fns.append(fn)
- return fns
+ return set(fns) # Some functions may appear more than once
 
  def _check_probe_quota(self, num_new_probes):
  global _num_open_probes

src/python/bcc/table.py

@@ -200,7 +200,11 @@ def clear(self):
  self.__delitem__(k)
 
  def zero(self):
- for k in self.keys():
+ # Even though this is not very efficient, we grab the entire list of
+ # keys before enumerating it. This helps avoid a potential race where
+ # the leaf assignment changes a hash table bucket that is being
+ # enumerated by the same loop, and may lead to a hang.
+ for k in list(self.keys()):
  self[k] = self.Leaf()
 
  def __iter__(self):

tools/funccount.py

@@ -170,18 +170,18 @@ def load(self):
  trace_count_text = """
 int PROBE_FUNCTION(void *ctx) {
  FILTER
- u64 loc = LOCATION;
- u64 *val = counts.lookup(&loc); // prepopulated on Python side
- if (val) {
- (*val)++;
+ int loc = LOCATION;
+ u64 *val = counts.lookup(&loc);
+ if (!val) {
+ return 0; // Should never happen, # of locations is known
  }
+ (*val)++;
  return 0;
 }
  """
  bpf_text = """#include <uapi/linux/ptrace.h>
 
-BPF_HASH(counts, u64, u64); // map location number to number of calls
-
+BPF_TABLE("array", int, u64, counts, NUMLOCATIONS);
  """
 
  # We really mean the tgid from the kernel's perspective, which is in
@@ -194,15 +194,21 @@ def load(self):
  trace_count_text = trace_count_text.replace('FILTER', '')
 
  bpf_text += self._generate_functions(trace_count_text)
+ bpf_text = bpf_text.replace("NUMLOCATIONS",
+ str(len(self.trace_functions)))
  if debug:
  print(bpf_text)
 
  self.bpf = BPF(text=bpf_text,
  usdt_contexts=[self.usdt] if self.usdt else [])
+ self.clear() # Initialize all array items to zero
+
+ def counts(self):
+ return self.bpf["counts"]
 
-  # Initialize all map entries to zero
+ def clear(self):
  counts = self.bpf["counts"]
- for location, function in self.trace_functions.items():
+ for location, _ in list(self.trace_functions.items()):
  counts[counts.Key(location)] = counts.Leaf()
 
 class Tool(object):
@@ -260,18 +266,19 @@ def run(self):
  print("%-8s\n" % strftime("%H:%M:%S"), end="")
 
  print("%-36s %8s" % ("FUNC", "COUNT"))
- counts = self.probe.bpf["counts"]
+ counts = self.probe.counts()
  for k, v in sorted(counts.items(),
  key=lambda counts: counts[1].value):
  if v.value == 0:
  continue
  print("%-36s %8d" %
  (self.probe.trace_functions[k.value], v.value))
- counts.zero()
 
  if exiting:
  print("Detaching...")
  exit()
+ else:
+ self.probe.clear()
 
 if __name__ == "__main__":
  try: