Merge pull request #602 from goldshtn/auto-tp

Full tracepoint support in Clang front-end

src/cc/export/helpers.h

@@ -451,5 +451,8 @@ int bpf_usdt_readarg_p(int argc, struct pt_regs *ctx, void *buf, u64 len) asm("l
 
 #define lock_xadd(ptr, val) ((void)__sync_fetch_and_add(ptr, val))
 
+#define TRACEPOINT_PROBE(category, event) \
+int tracepoint__##category##__##event(struct tracepoint__##category##__##event *args)
+
 #endif
 )********"

src/cc/frontends/clang/CMakeLists.txt

@@ -4,4 +4,4 @@
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DKERNEL_MODULES_DIR='\"${BCC_KERNEL_MODULES_DIR}\"'")
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DKERNEL_MODULES_SUFFIX='\"${BCC_KERNEL_MODULES_SUFFIX}\"'")
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DKERNEL_HAS_SOURCE_DIR=${BCC_KERNEL_HAS_SOURCE_DIR}")
-add_library(clang_frontend loader.cc b_frontend_action.cc kbuild_helper.cc)
+add_library(clang_frontend loader.cc b_frontend_action.cc tp_frontend_action.cc kbuild_helper.cc)

src/cc/frontends/clang/b_frontend_action.cc

@@ -253,10 +253,11 @@ BTypeVisitor::BTypeVisitor(ASTContext &C, Rewriter &rewriter, vector<TableDesc>
 bool BTypeVisitor::VisitFunctionDecl(FunctionDecl *D) {
   // put each non-static non-inline function decl in its own section, to be
   // extracted by the MemoryManager
+  auto real_start_loc = rewriter_.getSourceMgr().getFileLoc(D->getLocStart());
   if (D->isExternallyVisible() && D->hasBody()) {
     current_fn_ = D->getName();
     string attr = string("__attribute__((section(\"") + BPF_FN_PREFIX + D->getName().str() + "\")))\n";
-    rewriter_.InsertText(D->getLocStart(), attr);
+    rewriter_.InsertText(real_start_loc, attr);
     if (D->param_size() > MAX_CALLING_CONV_REGS + 1) {
       error(D->getParamDecl(MAX_CALLING_CONV_REGS + 1)->getLocStart(),
             "too many arguments, bcc only supports in-register parameters");
@@ -295,10 +296,10 @@ bool BTypeVisitor::VisitFunctionDecl(FunctionDecl *D) {
     if (CompoundStmt *S = dyn_cast<CompoundStmt>(D->getBody()))
       rewriter_.ReplaceText(S->getLBracLoc(), 1, preamble);
   } else if (D->hasBody() &&
-             rewriter_.getSourceMgr().getFileID(D->getLocStart())
+             rewriter_.getSourceMgr().getFileID(real_start_loc)
                == rewriter_.getSourceMgr().getMainFileID()) {
     // rewritable functions that are static should be always treated as helper
-    rewriter_.InsertText(D->getLocStart(), "__attribute__((always_inline))\n");
+    rewriter_.InsertText(real_start_loc, "__attribute__((always_inline))\n");
   }
   return true;
 }

src/cc/frontends/clang/loader.cc

@@ -50,6 +50,7 @@
 #include "exported_files.h"
 #include "kbuild_helper.h"
 #include "b_frontend_action.h"
+#include "tp_frontend_action.h"
 #include "loader.h"
 
 using std::map;
@@ -166,6 +167,34 @@ int ClangLoader::parse(unique_ptr<llvm::Module> *mod, unique_ptr<vector<TableDes
     llvm::errs() << "\n";
   }
 
+  // pre-compilation pass for generating tracepoint structures
+  auto invocation0 = make_unique<CompilerInvocation>();
+  if (!CompilerInvocation::CreateFromArgs(*invocation0, const_cast<const char **>(ccargs.data()),
+                                          const_cast<const char **>(ccargs.data()) + ccargs.size(), diags))
+    return -1;
+
+  invocation0->getPreprocessorOpts().RetainRemappedFileBuffers = true;
+  for (const auto &f : remapped_files_)
+    invocation0->getPreprocessorOpts().addRemappedFile(f.first, &*f.second);
+
+  if (in_memory) {
+    invocation0->getPreprocessorOpts().addRemappedFile(main_path, &*main_buf);
+    invocation0->getFrontendOpts().Inputs.clear();
+    invocation0->getFrontendOpts().Inputs.push_back(FrontendInputFile(main_path, IK_C));
+  }
+  invocation0->getFrontendOpts().DisableFree = false;
+
+  CompilerInstance compiler0;
+  compiler0.setInvocation(invocation0.release());
+  compiler0.createDiagnostics(new IgnoringDiagConsumer());
+
+  // capture the rewritten c file
+  string out_str;
+  llvm::raw_string_ostream os(out_str);
+  TracepointFrontendAction tpact(os);
+  compiler0.ExecuteAction(tpact); // ignore errors, they will be reported later
+  unique_ptr<llvm::MemoryBuffer> out_buf = llvm::MemoryBuffer::getMemBuffer(out_str);
+
   // first pass
   auto invocation1 = make_unique<CompilerInvocation>();
   if (!CompilerInvocation::CreateFromArgs(*invocation1, const_cast<const char **>(ccargs.data()),
@@ -178,12 +207,9 @@ int ClangLoader::parse(unique_ptr<llvm::Module> *mod, unique_ptr<vector<TableDes
   invocation1->getPreprocessorOpts().RetainRemappedFileBuffers = true;
   for (const auto &f : remapped_files_)
     invocation1->getPreprocessorOpts().addRemappedFile(f.first, &*f.second);
-
-  if (in_memory) {
-    invocation1->getPreprocessorOpts().addRemappedFile(main_path, &*main_buf);
-    invocation1->getFrontendOpts().Inputs.clear();
-    invocation1->getFrontendOpts().Inputs.push_back(FrontendInputFile(main_path, IK_C));
-  }
+  invocation1->getPreprocessorOpts().addRemappedFile(main_path, &*out_buf);
+  invocation1->getFrontendOpts().Inputs.clear();
+  invocation1->getFrontendOpts().Inputs.push_back(FrontendInputFile(main_path, IK_C));
   invocation1->getFrontendOpts().DisableFree = false;
 
   CompilerInstance compiler1;
@@ -191,12 +217,12 @@ int ClangLoader::parse(unique_ptr<llvm::Module> *mod, unique_ptr<vector<TableDes
   compiler1.createDiagnostics();
 
   // capture the rewritten c file
-  string out_str;
-  llvm::raw_string_ostream os(out_str);
-  BFrontendAction bact(os, flags_);
+  string out_str1;
+  llvm::raw_string_ostream os1(out_str1);
+  BFrontendAction bact(os1, flags_);
   if (!compiler1.ExecuteAction(bact))
     return -1;
-  unique_ptr<llvm::MemoryBuffer> out_buf = llvm::MemoryBuffer::getMemBuffer(out_str);
+  unique_ptr<llvm::MemoryBuffer> out_buf1 = llvm::MemoryBuffer::getMemBuffer(out_str1);
   // this contains the open FDs
   *tables = bact.take_tables();
 
@@ -209,7 +235,7 @@ int ClangLoader::parse(unique_ptr<llvm::Module> *mod, unique_ptr<vector<TableDes
   invocation2->getPreprocessorOpts().RetainRemappedFileBuffers = true;
   for (const auto &f : remapped_files_)
     invocation2->getPreprocessorOpts().addRemappedFile(f.first, &*f.second);
-  invocation2->getPreprocessorOpts().addRemappedFile(main_path, &*out_buf);
+  invocation2->getPreprocessorOpts().addRemappedFile(main_path, &*out_buf1);
   invocation2->getFrontendOpts().Inputs.clear();
   invocation2->getFrontendOpts().Inputs.push_back(FrontendInputFile(main_path, IK_C));
   invocation2->getFrontendOpts().DisableFree = false;

src/cc/frontends/clang/tp_frontend_action.cc

+/*
+ * Copyright (c) 2016 Sasha Goldshtein
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <linux/bpf.h>
+#include <linux/version.h>
+#include <sys/utsname.h>
+#include <unistd.h>
+
+#include <fstream>
+#include <regex>
+
+#include <clang/AST/ASTConsumer.h>
+#include <clang/AST/ASTContext.h>
+#include <clang/AST/RecordLayout.h>
+#include <clang/Frontend/CompilerInstance.h>
+#include <clang/Frontend/MultiplexConsumer.h>
+#include <clang/Rewrite/Core/Rewriter.h>
+
+#include "tp_frontend_action.h"
+
+namespace ebpf {
+
+using std::map;
+using std::set;
+using std::string;
+using std::to_string;
+using std::unique_ptr;
+using std::vector;
+using std::regex;
+using std::smatch;
+using std::regex_search;
+using std::ifstream;
+using namespace clang;
+
+TracepointTypeVisitor::TracepointTypeVisitor(ASTContext &C, Rewriter &rewriter)
+    : C(C), diag_(C.getDiagnostics()), rewriter_(rewriter), out_(llvm::errs()) {
+}
+
+static inline bool _is_valid_field(string const& line,
+                                   string& field_type,
+                                   string& field_name) {
+  auto field_pos = line.find("field:");
+  if (field_pos == string::npos)
+    return false;
+
+  auto semi_pos = line.find(';', field_pos);
+  if (semi_pos == string::npos)
+    return false;
+
+  auto size_pos = line.find("size:", semi_pos);
+  if (size_pos == string::npos)
+    return false;
+
+  auto field = line.substr(field_pos + 6/*"field:"*/,
+                           semi_pos - field_pos - 6);
+  auto pos = field.find_last_of("\t ");
+  if (pos == string::npos)
+    return false;
+
+  field_type = field.substr(0, pos);
+  field_name = field.substr(pos + 1);
+  if (field_type.find("__data_loc") != string::npos)
+    return false;
+  if (field_name.find("common_") == 0)
+    return false;
+
+  return true;
+}
+
+string TracepointTypeVisitor::GenerateTracepointStruct(
+    SourceLocation loc, string const& category, string const& event) {
+  string format_file = "/sys/kernel/debug/tracing/events/" +
+    category + "/" + event + "/format";
+  ifstream input(format_file.c_str());
+  if (!input)
+    return "";
+
+  string tp_struct = "struct tracepoint__" + category + "__" + event + " {\n";
+  tp_struct += "\tu64 __do_not_use__;\n";
+  for (string line; getline(input, line); ) {
+    string field_type, field_name;
+    if (!_is_valid_field(line, field_type, field_name))
+      continue;
+    tp_struct += "\t" + field_type + " " + field_name + ";\n";
+  }
+
+  tp_struct += "};\n";
+  return tp_struct;
+}
+
+static inline bool _is_tracepoint_struct_type(string const& type_name,
+                                              string& tp_category,
+                                              string& tp_event) {
+  // We are looking to roughly match the regex:
+  //    (?:struct|class)\s+tracepoint__(\S+)__(\S+)
+  // Not using std::regex because older versions of GCC don't support it yet.
+  // E.g., the libstdc++ that ships with Ubuntu 14.04.
+  
+  auto first_space_pos = type_name.find_first_of("\t ");
+  if (first_space_pos == string::npos)
+    return false;
+  auto first_tok = type_name.substr(0, first_space_pos);
+  if (first_tok != "struct" && first_tok != "class")
+    return false;
+
+  auto non_space_pos = type_name.find_first_not_of("\t ", first_space_pos);
+  auto second_space_pos = type_name.find_first_of("\t ", non_space_pos);
+  auto second_tok = type_name.substr(non_space_pos,
+                                     second_space_pos - non_space_pos);
+  if (second_tok.find("tracepoint__") != 0)
+    return false;
+
+  auto tp_event_pos = second_tok.rfind("__");
+  if (tp_event_pos == string::npos)
+    return false;
+  tp_event = second_tok.substr(tp_event_pos + 2);
+
+  auto tp_category_pos = second_tok.find("__");
+  if (tp_category_pos == tp_event_pos)
+    return false;
+  tp_category = second_tok.substr(tp_category_pos + 2,
+                                  tp_event_pos - tp_category_pos - 2);
+  return true;
+}
+
+
+bool TracepointTypeVisitor::VisitFunctionDecl(FunctionDecl *D) {
+  if (D->isExternallyVisible() && D->hasBody()) {
+    // If this function has a tracepoint structure as an argument,
+    // add that structure declaration based on the structure name.
+    for (auto it = D->param_begin(); it != D->param_end(); ++it) {
+      auto arg = *it;
+      auto type = arg->getType();
+      if (type->isPointerType() &&
+          type->getPointeeType()->isStructureOrClassType()) {
+        auto type_name = QualType::getAsString(type.split());
+        string tp_cat, tp_evt;
+        if (_is_tracepoint_struct_type(type_name, tp_cat, tp_evt)) {
+          string tp_struct = GenerateTracepointStruct(
+              D->getLocStart(), tp_cat, tp_evt);
+          // Get the actual function declaration point (the macro instantiation
+          // point if using the TRACEPOINT_PROBE macro instead of the macro
+          // declaration point in bpf_helpers.h).
+          auto insert_loc = D->getLocStart();
+          insert_loc = rewriter_.getSourceMgr().getFileLoc(insert_loc);
+          rewriter_.InsertText(insert_loc, tp_struct);
+        }
+      }
+    }
+  }
+  return true;
+}
+
+TracepointTypeConsumer::TracepointTypeConsumer(ASTContext &C, Rewriter &rewriter)
+    : visitor_(C, rewriter) {
+}
+
+bool TracepointTypeConsumer::HandleTopLevelDecl(DeclGroupRef Group) {
+  for (auto D : Group)
+    visitor_.TraverseDecl(D);
+  return true;
+}
+
+TracepointFrontendAction::TracepointFrontendAction(llvm::raw_ostream &os)
+    : os_(os), rewriter_(new Rewriter) {
+}
+
+void TracepointFrontendAction::EndSourceFileAction() {
+  rewriter_->getEditBuffer(rewriter_->getSourceMgr().getMainFileID()).write(os_);
+  os_.flush();
+}
+
+unique_ptr<ASTConsumer> TracepointFrontendAction::CreateASTConsumer(
+        CompilerInstance &Compiler, llvm::StringRef InFile) {
+  rewriter_->setSourceMgr(Compiler.getSourceManager(), Compiler.getLangOpts());
+  return unique_ptr<ASTConsumer>(new TracepointTypeConsumer(
+              Compiler.getASTContext(), *rewriter_));
+}
+
+}

src/cc/frontends/clang/tp_frontend_action.h

+/*
+ * Copyright (c) 2016 Sasha Goldshtein
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <map>
+#include <memory>
+#include <set>
+#include <string>
+#include <vector>
+
+#include <clang/AST/RecursiveASTVisitor.h>
+#include <clang/Frontend/FrontendAction.h>
+#include <clang/Rewrite/Core/Rewriter.h>
+
+namespace clang {
+class ASTConsumer;
+class ASTContext;
+class CompilerInstance;
+}
+
+namespace llvm {
+class raw_ostream;
+class StringRef;
+}
+
+namespace ebpf {
+
+// Visit functions that have a tracepoint argument structure in their signature
+// and automatically generate the structure on-the-fly.
+class TracepointTypeVisitor :
+  public clang::RecursiveASTVisitor<TracepointTypeVisitor> {
+ public:
+  explicit TracepointTypeVisitor(clang::ASTContext &C,
+                                 clang::Rewriter &rewriter);
+  bool VisitFunctionDecl(clang::FunctionDecl *D);
+
+ private:
+  std::string GenerateTracepointStruct(clang::SourceLocation loc,
+          std::string const& category, std::string const& event);
+
+  clang::ASTContext &C;
+  clang::DiagnosticsEngine &diag_;
+  clang::Rewriter &rewriter_;
+  llvm::raw_ostream &out_; 
+};
+
+class TracepointTypeConsumer : public clang::ASTConsumer {
+ public:
+  explicit TracepointTypeConsumer(clang::ASTContext &C,
+                                  clang::Rewriter &rewriter);
+  bool HandleTopLevelDecl(clang::DeclGroupRef Group) override;
+ private:
+  TracepointTypeVisitor visitor_;
+};
+
+class TracepointFrontendAction : public clang::ASTFrontendAction {
+ public:
+  TracepointFrontendAction(llvm::raw_ostream &os);
+
+  void EndSourceFileAction() override;
+
+  std::unique_ptr<clang::ASTConsumer>
+      CreateASTConsumer(clang::CompilerInstance &Compiler, llvm::StringRef InFile) override;
+
+ private:
+  llvm::raw_ostream &os_;
+  std::unique_ptr<clang::Rewriter> rewriter_;
+};
+
+}  // namespace visitor

src/python/bcc/__init__.py

@@ -182,8 +182,8 @@ class BPF(object):
         if not self.module:
             raise Exception("Failed to compile BPF module %s" % src_file)
 
-        # If any "kprobe__" prefixed functions were defined, they will be
-        # loaded and attached here.
+        # If any "kprobe__" or "tracepoint__" prefixed functions were defined,
+        # they will be loaded and attached here.
         self._trace_autoload()
 
     def load_funcs(self, prog_type=KPROBE):
@@ -608,17 +608,18 @@ class BPF(object):
         del open_uprobes[ev_name]
 
     def _trace_autoload(self):
-        # Cater to one-liner case where attach_kprobe is omitted and C function
-        # name matches that of the kprobe.
-        if len(open_kprobes) == 0:
-            for i in range(0, lib.bpf_num_functions(self.module)):
-                func_name = lib.bpf_function_name(self.module, i).decode()
-                if func_name.startswith("kprobe__"):
-                    fn = self.load_func(func_name, BPF.KPROBE)
-                    self.attach_kprobe(event=fn.name[8:], fn_name=fn.name)
-                elif func_name.startswith("kretprobe__"):
-                    fn = self.load_func(func_name, BPF.KPROBE)
-                    self.attach_kretprobe(event=fn.name[11:], fn_name=fn.name)
+        for i in range(0, lib.bpf_num_functions(self.module)):
+            func_name = lib.bpf_function_name(self.module, i).decode()
+            if len(open_kprobes) == 0 and func_name.startswith("kprobe__"):
+                fn = self.load_func(func_name, BPF.KPROBE)
+                self.attach_kprobe(event=fn.name[8:], fn_name=fn.name)
+            elif len(open_kprobes) == 0 and func_name.startswith("kretprobe__"):
+                fn = self.load_func(func_name, BPF.KPROBE)
+                self.attach_kretprobe(event=fn.name[11:], fn_name=fn.name)
+            elif func_name.startswith("tracepoint__"):
+                fn = self.load_func(func_name, BPF.TRACEPOINT)
+                tp = fn.name[len("tracepoint__"):].replace("__", ":")
+                self.attach_tracepoint(tp=tp, fn_name=fn.name)
 
     def trace_open(self, nonblocking=False):
         """trace_open(nonblocking=False)

tests/python/test_tracepoint.py

@@ -22,21 +22,9 @@ def kernel_version_ge(major, minor):
 @unittest.skipUnless(kernel_version_ge(4,7), "requires kernel >= 4.7")
 class TestTracepoint(unittest.TestCase):
     def test_tracepoint(self):
-        text = """#include <linux/ptrace.h>
-        struct tp_args {
-            unsigned long long __unused__;
-            char prev_comm[16];
-            pid_t prev_pid;
-            int prev_prio;
-            long prev_state;
-            char next_comm[16];
-            pid_t next_pid;
-            int next_prio;
-        };
+        text = """
         BPF_HASH(switches, u32, u64);
-        int probe_switch(struct tp_args *args) {
-            if (args == 0)
-                return 0;
+        TRACEPOINT_PROBE(sched, sched_switch) {
             u64 val = 0;
             u32 pid = args->next_pid;
             u64 *existing = switches.lookup_or_init(&pid, &val);
@@ -45,13 +33,11 @@ class TestTracepoint(unittest.TestCase):
         }
         """
         b = bcc.BPF(text=text)
-        b.attach_tracepoint("sched:sched_switch", "probe_switch")
         sleep(1)
         total_switches = 0
         for k, v in b["switches"].items():
             total_switches += v.value
         self.assertNotEqual(0, total_switches)
-        b.detach_tracepoint("sched:sched_switch")
 
 if __name__ == "__main__":
     unittest.main()