Merge pull request #100 from iovisor/multiple_kprobes

attach to multiple identical probes

src/ast/ast.cpp

@@ -124,6 +124,15 @@ std::string AttachPoint::name(const std::string &attach_point) const
   return n;
 }
 
+int AttachPoint::index(std::string name) {
+  if (index_.count(name) == 0) return 0;
+  return index_[name];
+}
+
+void AttachPoint::set_index(std::string name, int index) {
+  index_[name] = index;
+}
+
 std::string Probe::name() const
 {
   std::string n = "";
@@ -141,5 +150,13 @@ std::string Probe::name() const
   return n.substr(0, n.size()-1);
 }
 
+int Probe::index() {
+  return index_;
+}
+
+void Probe::set_index(int index) {
+  index_ = index;
+}
+
 } // namespace ast
 } // namespace bpftrace

src/ast/ast.h

@@ -2,6 +2,7 @@
 
 #include <string>
 #include <vector>
+#include <map>
 
 #include "types.h"
 
@@ -193,6 +194,11 @@ public:
 
   void accept(Visitor &v) override;
   std::string name(const std::string &attach_point) const;
+
+  int index(std::string name);
+  void set_index(std::string name, int index);
+private:
+  std::map<std::string, int> index_;
 };
 using AttachPointList = std::vector<AttachPoint *>;
 
@@ -208,6 +214,11 @@ public:
   void accept(Visitor &v) override;
   std::string name() const;
   bool need_expansion = false;	// must build a BPF program per wildcard match
+
+  int index();
+  void set_index(int index);
+private:
+  int index_ = 0;
 };
 using ProbeList = std::vector<Probe *>;
 

src/ast/codegen_llvm.cpp

@@ -885,7 +885,8 @@ void CodegenLLVM::visit(Probe &probe)
   if (probe.need_expansion == false) {
     // build a single BPF program pre-wildcards
     Function *func = Function::Create(func_type, Function::ExternalLinkage, probe.name(), module_.get());
-    func->setSection("s_" + probe.name());
+    probe.set_index(getNextIndexForProbe(probe.name()));
+    func->setSection(getSectionNameForProbe(probe.name(), probe.index()));
     BasicBlock *entry = BasicBlock::Create(module_->getContext(), "entry", func);
     b_.SetInsertPoint(entry);
 
@@ -930,8 +931,10 @@ void CodegenLLVM::visit(Probe &probe)
         printf_id_ = starting_printf_id_;
         time_id_ = starting_time_id_;
         probefull_ = attach_point->name(match);
-        Function *func = Function::Create(func_type, Function::ExternalLinkage, attach_point->name(match), module_.get());
-        func->setSection("s_" + attach_point->name(match));
+        int index = getNextIndexForProbe(probe.name());
+        attach_point->set_index(match, index);
+        Function *func = Function::Create(func_type, Function::ExternalLinkage, probefull_, module_.get());
+        func->setSection(getSectionNameForProbe(probefull_, index));
         BasicBlock *entry = BasicBlock::Create(module_->getContext(), "entry", func);
         b_.SetInsertPoint(entry);
 
@@ -947,6 +950,7 @@ void CodegenLLVM::visit(Probe &probe)
       }
     }
   }
+  bpftrace_.add_probe(probe);
 }
 
 void CodegenLLVM::visit(Program &program)
@@ -955,6 +959,18 @@ void CodegenLLVM::visit(Program &program)
     probe->accept(*this);
 }
 
+int CodegenLLVM::getNextIndexForProbe(std::string probe_name) {
+  if (next_probe_index_.count(probe_name) == 0)
+    next_probe_index_[probe_name] = 1;
+  int index = next_probe_index_[probe_name];
+  next_probe_index_[probe_name] += 1;
+  return index;
+}
+
+std::string CodegenLLVM::getSectionNameForProbe(std::string probe_name, int index) {
+  return "s_" + probe_name + "_" + std::to_string(index);
+}
+
 AllocaInst *CodegenLLVM::getMapKey(Map &map)
 {
   AllocaInst *key;

src/ast/codegen_llvm.h

@@ -47,6 +47,8 @@ public:
   void visit(Program &program) override;
   AllocaInst *getMapKey(Map &map);
   AllocaInst *getHistMapKey(Map &map, Value *log2);
+  int         getNextIndexForProbe(std::string probe_name);
+  std::string getSectionNameForProbe(std::string probe_name, int index);
   Value      *createLogicalAnd(Binop &binop);
   Value      *createLogicalOr(Binop &binop);
 
@@ -67,6 +69,7 @@ private:
   BPFtrace &bpftrace_;
   std::string probefull_;
   std::string path_;
+  std::map<std::string, int> next_probe_index_;
 
   std::map<std::string, Value *> variables_;
   int printf_id_ = 0;

src/ast/semantic_analyser.cpp

@@ -731,9 +731,6 @@ void SemanticAnalyser::visit(Probe &probe)
     stmt->accept(*this);
   }
 
-  if (is_final_pass()) {
-    bpftrace_.add_probe(probe);
-  }
 }
 
 void SemanticAnalyser::visit(Program &program)

src/attached_probe.cpp

@@ -156,18 +156,19 @@ std::string AttachedProbe::eventprefix() const
 std::string AttachedProbe::eventname() const
 {
   std::ostringstream offset_str;
+  std::string index_str = "_" + std::to_string(probe_.index);
   switch (probe_.type)
   {
     case ProbeType::kprobe:
     case ProbeType::kretprobe:
-      return eventprefix() + probe_.attach_point;
+      return eventprefix() + probe_.attach_point + index_str;
     case ProbeType::uprobe:
     case ProbeType::uretprobe:
     case ProbeType::usdt:
       offset_str << std::hex << offset();
-      return eventprefix() + sanitise(probe_.path) + "_" + offset_str.str();
+      return eventprefix() + sanitise(probe_.path) + "_" + offset_str.str() + index_str;
     case ProbeType::tracepoint:
-      return probe_.attach_point;
+      return probe_.attach_point + index_str;
     default:
       abort();
   }

src/bpftrace.cpp

@@ -89,6 +89,8 @@ int BPFtrace::add_probe(ast::Probe &p)
       probe.name = attach_point->name(func);
       probe.freq = attach_point->freq;
       probe.loc = 0;
+      probe.index = attach_point->index(func) > 0 ? 
+          attach_point->index(func) : p.index();
       probes_.push_back(probe);
     }
   }
@@ -307,9 +309,10 @@ std::unique_ptr<AttachedProbe> BPFtrace::attach_probe(Probe &probe, const BpfOrc
   // and the name builtin, which must be expanded into separate programs per
   // probe), else try to find a the program based on the original probe name
   // that includes wildcards.
-  auto func = bpforc.sections_.find("s_" + probe.name);
+  std::string index_str = "_" + std::to_string(probe.index);
+  auto func = bpforc.sections_.find("s_" + probe.name + index_str);
   if (func == bpforc.sections_.end())
-    func = bpforc.sections_.find("s_" + probe.orig_name);
+    func = bpforc.sections_.find("s_" + probe.orig_name + index_str);
   if (func == bpforc.sections_.end())
   {
     if (probe.name != probe.orig_name)
@@ -334,9 +337,10 @@ std::unique_ptr<AttachedProbe> BPFtrace::attach_probe(Probe &probe, const BpfOrc
 
 int BPFtrace::run(std::unique_ptr<BpfOrc> bpforc)
 {
-  for (Probe &probe : special_probes_)
+  auto r_special_probes = special_probes_.rbegin();
+  for (; r_special_probes != special_probes_.rend(); ++r_special_probes)
   {
-    auto attached_probe = attach_probe(probe, *bpforc.get());
+    auto attached_probe = attach_probe(*r_special_probes, *bpforc.get());
     if (attached_probe == nullptr)
       return -1;
     special_attached_probes_.push_back(std::move(attached_probe));
@@ -348,9 +352,13 @@ int BPFtrace::run(std::unique_ptr<BpfOrc> bpforc)
 
   BEGIN_trigger();
 
-  for (Probe &probe : probes_)
+  // NOTE (mmarchini): Apparently the kernel fires kprobe_events in the reverse
+  // order they were attached, so we insert them backwards to make sure blocks
+  // are executed in the same order they were declared.
+  auto r_probes = probes_.rbegin();
+  for (; r_probes != probes_.rend(); ++r_probes)
   {
-    auto attached_probe = attach_probe(probe, *bpforc.get());
+    auto attached_probe = attach_probe(*r_probes, *bpforc.get());
     if (attached_probe == nullptr)
       return -1;
     attached_probes_.push_back(std::move(attached_probe));

src/types.h

@@ -83,6 +83,7 @@ public:
 				// before wildcard expansion
   std::string name;		// full probe name
   uint64_t loc;			// for USDT probes
+  int index = 0;
   int freq;
 };
 

tests/semantic_analyser.cpp

@@ -9,11 +9,6 @@ namespace bpftrace {
 namespace test {
 namespace semantic_analyser {
 
-class MockBPFtrace : public BPFtrace {
-public:
-  MOCK_METHOD1(add_probe, int(ast::Probe &p));
-};
-
 using ::testing::_;
 
 void test(BPFtrace &bpftrace, Driver &driver, const std::string &input, int expected_result=0)
@@ -97,14 +92,6 @@ TEST(semantic_analyser, builtin_functions)
   test("kprobe:f { fake() }", 1);
 }
 
-TEST(semantic_analyser, probe_count)
-{
-  MockBPFtrace bpftrace;
-  EXPECT_CALL(bpftrace, add_probe(_)).Times(2);
-
-  test(bpftrace, "kprobe:f { 1; } kprobe:d { 1; }");
-}
-
 TEST(semantic_analyser, undefined_map)
 {
   test("kprobe:f / @mymap == 123 / { @mymap = 0 }", 0);