funccount: Verify probe max limit (iovisor#771)

Because `funccount` doesn't use the direct regex attach infrastructure
in the BPF module, it needs its own checking for a maximum probe
limit that would make sense. We use 1000 because that's what the
BPF module uses as well. When trying to attach to more than 1000
probes, we bail out early.

tools/funccount.py

@@ -27,6 +27,12 @@
 
 debug = False
 
+def verify_limit(num):
+ probe_limit = 1000
+ if num > probe_limit:
+ raise Exception("maximum of %d probes allowed, attempted %d" %
+ (probe_limit, num))
+
 class Probe(object):
  def __init__(self, pattern, use_regex=False, pid=None):
  """Init a new probe.
@@ -117,7 +123,9 @@ def _generate_functions(self, template):
  self.usdt = None
  text = ""
  if self.type == "p" and not self.library:
- for function in BPF.get_kprobe_functions(self.pattern):
+ functions = BPF.get_kprobe_functions(self.pattern)
+ verify_limit(len(functions))
+ for function in functions:
  text += self._add_function(template, function)
  elif self.type == "p" and self.library:
  # uprobes are tricky because the same function may have multiple
@@ -127,25 +135,33 @@ def _generate_functions(self, template):
  # map to an address that we've already seen. Also ignore functions
  # that may repeat multiple times with different addresses.
  addresses, functions = (set(), set())
- for function, address in BPF.get_user_functions_and_addresses(
- self.library, self.pattern):
+ functions_and_addresses = BPF.get_user_functions_and_addresses(
+ self.library, self.pattern)
+ verify_limit(len(functions_and_addresses))
+ for function, address in functions_and_addresses:
  if address in addresses or function in functions:
  continue
  addresses.add(address)
  functions.add(function)
  text += self._add_function(template, function)
  elif self.type == "t":
- for tracepoint in BPF.get_tracepoints(self.pattern):
+ tracepoints = BPF.get_tracepoints(self.pattern)
+ verify_limit(len(tracepoints))
+ for tracepoint in tracepoints:
  text += self._add_function(template, tracepoint)
  elif self.type == "u":
  self.usdt = USDT(path=self.library, pid=self.pid)
+ matches = []
  for probe in self.usdt.enumerate_probes():
  if not self.pid and (probe.bin_path != self.library):
  continue
  if re.match(self.pattern, probe.name):
- new_func = "trace_count_%d" % self.matched
- text += self._add_function(template, probe.name)
- self.usdt.enable_probe(probe.name, new_func)
+ matches.append(probe.name)
+ verify_limit(len(matches))
+ for match in matches:
+ new_func = "trace_count_%d" % self.matched
+ text += self._add_function(template, match)
+ self.usdt.enable_probe(match, new_func)
  if debug:
  print(self.usdt.get_text())
  return text