Merge pull request #518 from vmg/vmg/lua-usdt-2

[RFC] New USDT API

examples/lua/usdt_ruby.lua

+#!/usr/bin/env bcc-lua
+--[[
+Copyright 2016 GitHub, Inc
+
+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.
+]]
+
+local program = [[
+#include <uapi/linux/ptrace.h>
+int trace_method(struct pt_regs *ctx) {
+  uint64_t addr;
+  bpf_usdt_readarg(2, ctx, &addr);
+
+  char fn_name[128] = {};
+  bpf_probe_read(&fn_name, sizeof(fn_name), (void *)addr);
+
+  bpf_trace_printk("%s(...)\n", fn_name);
+  return 0;
+};
+]]
+
+return function(BPF, util)
+  if not arg[1] then
+    print("usage: rubysyms.lua PID")
+    return
+  end
+
+  local u = util.USDT:new{pid=tonumber(arg[1])}
+  u:enable_probe{probe="method__entry", fn_name="trace_method"}
+
+  local b = BPF:new{text=program, usdt=u}
+  local pipe = b:pipe()
+  while true do
+    print(pipe:trace_fields())
+  end
+end

src/cc/bcc_syms.cc

@@ -24,17 +24,15 @@
 #include "bcc_syms.h"
 
 #include "syms.h"
+#include "vendor/tinyformat.hpp"
 
 ino_t ProcStat::getinode_() {
   struct stat s;
   return (!stat(procfs_.c_str(), &s)) ? s.st_ino : -1;
 }
 
-ProcStat::ProcStat(int pid) : inode_(-1) {
-  char buffer[128];
-  snprintf(buffer, sizeof(buffer), "/proc/%d/exe", pid);
-  procfs_ = buffer;
-}
+ProcStat::ProcStat(int pid)
+    : procfs_(tfm::format("/proc/%d/exe", pid)), inode_(getinode_()) {}
 
 void KSyms::_add_symbol(const char *symname, uint64_t addr, void *p) {
   KSyms *ks = static_cast<KSyms *>(p);
@@ -84,11 +82,15 @@ bool KSyms::resolve_name(const char *_unused, const char *name,
   return true;
 }
 
-ProcSyms::ProcSyms(int pid) : pid_(pid), procstat_(pid) { refresh(); }
+ProcSyms::ProcSyms(int pid) : pid_(pid), procstat_(pid) { load_modules(); }
+
+bool ProcSyms::load_modules() {
+  return bcc_procutils_each_module(pid_, _add_module, this) == 0;
+}
 
 void ProcSyms::refresh() {
   modules_.clear();
-  bcc_procutils_each_module(pid_, _add_module, this);
+  load_modules();
   procstat_.reset();
 }
 
@@ -199,6 +201,32 @@ void bcc_symcache_refresh(void *resolver) {
   cache->refresh();
 }
 
+struct mod_st {
+  const char *name;
+  uint64_t start;
+};
+
+static int _find_module(const char *modname, uint64_t start, uint64_t end,
+                        void *p) {
+  struct mod_st *mod = (struct mod_st *)p;
+  if (!strcmp(modname, mod->name)) {
+    mod->start = start;
+    return -1;
+  }
+  return 0;
+}
+
+int bcc_resolve_global_addr(int pid, const char *module, const uint64_t address,
+                            uint64_t *global) {
+  struct mod_st mod = {module, 0x0};
+  if (bcc_procutils_each_module(pid, _find_module, &mod) < 0 ||
+      mod.start == 0x0)
+    return -1;
+
+  *global = mod.start + address;
+  return 0;
+}
+
 static int _find_sym(const char *symname, uint64_t addr, uint64_t end,
                      int flags, void *payload) {
   struct bcc_symbol *sym = (struct bcc_symbol *)payload;

src/cc/bcc_syms.h

@@ -33,6 +33,8 @@ int bcc_symcache_resolve(void *symcache, uint64_t addr, struct bcc_symbol *sym);
 int bcc_symcache_resolve_name(void *resolver, const char *name, uint64_t *addr);
 void bcc_symcache_refresh(void *resolver);
 
+int bcc_resolve_global_addr(int pid, const char *module, const uint64_t address,
+                            uint64_t *global);
 int bcc_find_symbol_addr(struct bcc_symbol *sym);
 int bcc_resolve_symname(const char *module, const char *symname,
                         const uint64_t addr, struct bcc_symbol *sym);

src/cc/bcc_usdt.h

+/*
+ * Copyright (c) 2016 GitHub, Inc.
+ *
+ * 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.
+ */
+#ifndef LIBBCC_USDT_H
+#define LIBBCC_USDT_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include <stdint.h>
+
+void *bcc_usdt_new_frompid(int pid);
+void *bcc_usdt_new_frompath(const char *path);
+void bcc_usdt_close(void *usdt);
+
+int bcc_usdt_enable_probe(void *, const char *, const char *);
+char *bcc_usdt_genargs(void *);
+
+typedef void (*bcc_usdt_uprobe_cb)(const char *, const char *, uint64_t, int);
+void bcc_usdt_foreach_uprobe(void *usdt, bcc_usdt_uprobe_cb callback);
+
+#ifdef __cplusplus
+}
+#endif
+#endif

src/cc/export/helpers.h

@@ -408,7 +408,9 @@ int incr_cksum_l3(void *off, u64 oldval, u64 newval) asm("llvm.bpf.extra");
 int incr_cksum_l4(void *off, u64 oldval, u64 newval, u64 flags) asm("llvm.bpf.extra");
 int bpf_num_cpus() asm("llvm.bpf.extra");
 
-#define lock_xadd(ptr, val) ((void)__sync_fetch_and_add(ptr, val))
+struct pt_regs;
+int bpf_usdt_readarg(int argc, struct pt_regs *ctx, void *arg) asm("llvm.bpf.extra");
+int bpf_usdt_readarg_p(int argc, struct pt_regs *ctx, void *buf, u64 len) asm("llvm.bpf.extra");
 
 #ifdef __powerpc__
 #define PT_REGS_PARM1(ctx)	((ctx)->gpr[3])
@@ -434,5 +436,7 @@ int bpf_num_cpus() asm("llvm.bpf.extra");
 #error "bcc does not support this platform yet"
 #endif
 
+#define lock_xadd(ptr, val) ((void)__sync_fetch_and_add(ptr, val))
+
 #endif
 )********"

src/cc/frontends/clang/b_frontend_action.cc

@@ -254,6 +254,7 @@ bool BTypeVisitor::VisitFunctionDecl(FunctionDecl *D) {
   // put each non-static non-inline function decl in its own section, to be
   // extracted by the MemoryManager
   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);
     if (D->param_size() > MAX_CALLING_CONV_REGS + 1) {
@@ -470,6 +471,20 @@ bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
             numcpu = 1;
           text = to_string(numcpu);
           rewriter_.ReplaceText(SourceRange(Call->getLocStart(), Call->getLocEnd()), text);
+        } else if (Decl->getName() == "bpf_usdt_readarg_p") {
+          text = "({ u64 __addr = 0x0; ";
+          text += "_bpf_readarg_" + current_fn_ + "_" + args[0] + "(" +
+                  args[1] + ", &__addr, sizeof(__addr));";
+          text += "bpf_probe_read(" + args[2] + ", " + args[3] +
+                  ", (void *)__addr);";
+          text += "})";
+          rewriter_.ReplaceText(
+              SourceRange(Call->getLocStart(), Call->getLocEnd()), text);
+        } else if (Decl->getName() == "bpf_usdt_readarg") {
+          text = "_bpf_readarg_" + current_fn_ + "_" + args[0] + "(" + args[1] +
+                 ", " + args[2] + ", sizeof(*(" + args[2] + ")))";
+          rewriter_.ReplaceText(
+              SourceRange(Call->getLocStart(), Call->getLocEnd()), text);
         }
       }
     }

src/cc/frontends/clang/b_frontend_action.h

@@ -83,6 +83,7 @@ class BTypeVisitor : public clang::RecursiveASTVisitor<BTypeVisitor> {
   std::vector<TableDesc> &tables_;  /// store the open FDs
   std::vector<clang::ParmVarDecl *> fn_args_;
   std::set<clang::Expr *> visited_;
+  std::string current_fn_;
 };
 
 // Do a depth-first search to rewrite all pointers that need to be probed

src/cc/syms.h

@@ -93,6 +93,7 @@ class ProcSyms : SymbolCache {
   ProcStat procstat_;
 
   static int _add_module(const char *, uint64_t, uint64_t, void *);
+  bool load_modules();
 
 public:
   ProcSyms(int pid);

src/cc/usdt.cc

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+#include <cstring>
 #include <sstream>
 
 #include <fcntl.h>
@@ -29,21 +30,20 @@ namespace USDT {
 Probe::Location::Location(uint64_t addr, const char *arg_fmt) : address_(addr) {
   ArgumentParser_x64 parser(arg_fmt);
   while (!parser.done()) {
-    Argument *arg = new Argument();
-    if (!parser.parse(arg)) {
-      delete arg;  // TODO: report error
+    Argument arg;
+    if (!parser.parse(&arg))
       continue;
-    }
-    arguments_.push_back(arg);
+    arguments_.push_back(std::move(arg));
   }
 }
 
 Probe::Probe(const char *bin_path, const char *provider, const char *name,
-             uint64_t semaphore)
+             uint64_t semaphore, const optional<int> &pid)
     : bin_path_(bin_path),
       provider_(provider),
       name_(name),
-      semaphore_(semaphore) {}
+      semaphore_(semaphore),
+      pid_(pid) {}
 
 bool Probe::in_shared_object() {
   if (!in_shared_object_)
@@ -51,37 +51,36 @@ bool Probe::in_shared_object() {
   return in_shared_object_.value();
 }
 
-bool Probe::lookup_semaphore_addr(uint64_t *address, int pid) {
-  auto it = semaphores_.find(pid);
-  if (it != semaphores_.end()) {
-    *address = it->second;
-    return true;
-  }
-
+bool Probe::resolve_global_address(uint64_t *global, const uint64_t addr) {
   if (in_shared_object()) {
-    uint64_t load_address = 0x0;  // TODO
-    *address = load_address + semaphore_;
-  } else {
-    *address = semaphore_;
+    return (pid_ &&
+            !bcc_resolve_global_addr(*pid_, bin_path_.c_str(), addr, global));
   }
 
-  semaphores_[pid] = *address;
+  *global = addr;
   return true;
 }
 
-bool Probe::add_to_semaphore(int pid, int16_t val) {
-  uint64_t address;
-  if (!lookup_semaphore_addr(&address, pid))
-    return false;
+bool Probe::add_to_semaphore(int16_t val) {
+  assert(pid_ && attached_semaphore_);
+
+  if (!attached_semaphore_) {
+    uint64_t addr;
+    if (!resolve_global_address(&addr, semaphore_))
+      return false;
+    attached_semaphore_ = addr;
+  }
+
+  off_t address = static_cast<off_t>(attached_semaphore_.value());
 
-  std::string procmem = tfm::format("/proc/%d/mem", pid);
+  std::string procmem = tfm::format("/proc/%d/mem", pid_.value());
   int memfd = ::open(procmem.c_str(), O_RDWR);
   if (memfd < 0)
     return false;
 
-  int16_t original;  // TODO: should this be unsigned?
+  int16_t original;
 
-  if (::lseek(memfd, static_cast<off_t>(address), SEEK_SET) < 0 ||
+  if (::lseek(memfd, address, SEEK_SET) < 0 ||
       ::read(memfd, &original, 2) != 2) {
     ::close(memfd);
     return false;
@@ -89,7 +88,7 @@ bool Probe::add_to_semaphore(int pid, int16_t val) {
 
   original = original + val;
 
-  if (::lseek(memfd, static_cast<off_t>(address), SEEK_SET) < 0 ||
+  if (::lseek(memfd, address, SEEK_SET) < 0 ||
       ::write(memfd, &original, 2) != 2) {
     ::close(memfd);
     return false;
@@ -99,66 +98,92 @@ bool Probe::add_to_semaphore(int pid, int16_t val) {
   return true;
 }
 
-bool Probe::enable(int pid) {
-  if (!add_to_semaphore(pid, +1))
+bool Probe::enable(const std::string &fn_name) {
+  if (attached_to_)
     return false;
 
-  // TODO: what happens if we enable this twice?
-  enabled_semaphores_.emplace(pid, std::move(ProcStat(pid)));
+  if (need_enable()) {
+    if (!pid_)
+      return false;
+
+    if (!add_to_semaphore(+1))
+      return false;
+  }
+
+  attached_to_ = fn_name;
   return true;
 }
 
-bool Probe::disable(int pid) {
-  auto it = enabled_semaphores_.find(pid);
-  if (it == enabled_semaphores_.end())
+bool Probe::disable() {
+  if (!attached_to_)
     return false;
 
-  bool result = true;
-  if (!it->second.is_stale())
-    result = add_to_semaphore(pid, -1);
+  attached_to_ = nullopt;
 
-  enabled_semaphores_.erase(it);
-  return result;
+  if (need_enable()) {
+    assert(pid_);
+    return add_to_semaphore(-1);
+  }
+  return true;
 }
 
-bool Probe::usdt_thunks(std::ostream &stream, const std::string &prefix) {
-  assert(!locations_.empty());
-  for (size_t i = 0; i < locations_.size(); ++i) {
-    tfm::format(
-        stream,
-        "int %s_thunk_%d(struct pt_regs *ctx) { return %s(ctx, %d); }\n",
-        prefix, i, prefix, i);
+std::string Probe::largest_arg_type(size_t arg_n) {
+  Argument *largest = nullptr;
+  for (Location &location : locations_) {
+    Argument *candidate = &location.arguments_[arg_n];
+    if (!largest ||
+        std::abs(candidate->arg_size()) > std::abs(largest->arg_size()))
+      largest = candidate;
   }
-  return true;
+
+  assert(largest);
+  return largest->ctype();
 }
 
-bool Probe::usdt_cases(std::ostream &stream, const optional<int> &pid) {
-  assert(!locations_.empty());
+bool Probe::usdt_getarg(std::ostream &stream) {
   const size_t arg_count = locations_[0].arguments_.size();
 
-  for (size_t arg_n = 0; arg_n < arg_count; ++arg_n) {
-    Argument *largest = nullptr;
-    for (Location &location : locations_) {
-      Argument *candidate = location.arguments_[arg_n];
-      if (!largest ||
-          std::abs(candidate->arg_size()) > std::abs(largest->arg_size()))
-        largest = candidate;
-    }
-
-    tfm::format(stream, "%s arg%d = 0;\n", largest->ctype(), arg_n + 1);
-  }
+  if (!attached_to_)
+    return false;
 
-  for (size_t loc_n = 0; loc_n < locations_.size(); ++loc_n) {
-    Location &location = locations_[loc_n];
-    tfm::format(stream, "if (__loc_id == %d) {\n", loc_n);
+  if (arg_count == 0)
+    return true;
 
-    for (size_t arg_n = 0; arg_n < location.arguments_.size(); ++arg_n) {
-      Argument *arg = location.arguments_[arg_n];
-      if (!arg->assign_to_local(stream, tfm::format("arg%d", arg_n + 1),
-                                bin_path_, pid))
+  for (size_t arg_n = 0; arg_n < arg_count; ++arg_n) {
+    std::string ctype = largest_arg_type(arg_n);
+    std::string cptr = tfm::format("*((%s *)dest)", ctype);
+
+    tfm::format(stream,
+                "static inline int _bpf_readarg_%s_%d("
+                "struct pt_regs *ctx, void *dest, size_t len) {\n"
+                "  if (len != sizeof(%s)) return -1;\n",
+                attached_to_.value(), arg_n + 1, ctype);
+
+    if (locations_.size() == 1) {
+      Location &location = locations_.front();
+      stream << "  ";
+      if (!location.arguments_[arg_n].assign_to_local(stream, cptr, bin_path_,
+                                                      pid_))
         return false;
+      stream << "\n  return 0;\n}\n";
+    } else {
+      stream << "  switch(ctx->ip) {\n";
+      for (Location &location : locations_) {
+        uint64_t global_address;
+
+        if (!resolve_global_address(&global_address, location.address_))
+          return false;
+
+        tfm::format(stream, "  case 0x%xULL: ", global_address);
+        if (!location.arguments_[arg_n].assign_to_local(stream, cptr, bin_path_,
+                                                        pid_))
+          return false;
+
+        stream << " return 0;\n";
+      }
+      stream << "  }\n";
+      stream << "  return -1;\n}\n";
     }
-    stream << "}\n";
   }
   return true;
 }
@@ -179,22 +204,16 @@ int Context::_each_module(const char *modpath, uint64_t, uint64_t, void *p) {
 }
 
 void Context::add_probe(const char *binpath, const struct bcc_elf_usdt *probe) {
-  Probe *found_probe = nullptr;
-
-  for (Probe *p : probes_) {
+  for (auto &p : probes_) {
     if (p->provider_ == probe->provider && p->name_ == probe->name) {
-      found_probe = p;
-      break;
+      p->add_location(probe->pc, probe->arg_fmt);
+      return;
     }
   }
 
-  if (!found_probe) {
-    found_probe =
-        new Probe(binpath, probe->provider, probe->name, probe->semaphore);
-    probes_.push_back(found_probe);
-  }
-
-  found_probe->add_location(probe->pc, probe->arg_fmt);
+  probes_.emplace_back(
+      new Probe(binpath, probe->provider, probe->name, probe->semaphore, pid_));
+  probes_.back()->add_location(probe->pc, probe->arg_fmt);
 }
 
 std::string Context::resolve_bin_path(const std::string &bin_path) {
@@ -210,14 +229,44 @@ std::string Context::resolve_bin_path(const std::string &bin_path) {
   return result;
 }
 
-Probe *Context::find_probe(const std::string &probe_name) {
-  for (Probe *p : probes_) {
+Probe *Context::get(const std::string &probe_name) {
+  for (auto &p : probes_) {
     if (p->name_ == probe_name)
-      return p;
+      return p.get();
   }
   return nullptr;
 }
 
+bool Context::generate_usdt_args(std::ostream &stream) {
+  stream << "#include <uapi/linux/ptrace.h>\n";
+  for (auto &p : probes_) {
+    if (p->enabled() && !p->usdt_getarg(stream))
+      return false;
+  }
+  return true;
+}
+
+bool Context::enable_probe(const std::string &probe_name,
+                           const std::string &fn_name) {
+  if (pid_stat_ && pid_stat_->is_stale())
+    return false;
+
+  auto p = get(probe_name);
+  return p && p->enable(fn_name);
+}
+
+void Context::each_uprobe(each_uprobe_cb callback) {
+  for (auto &p : probes_) {
+    if (!p->enabled())
+      continue;
+
+    for (Probe::Location &loc : p->locations_) {
+      callback(p->bin_path_.c_str(), p->attached_to_->c_str(), loc.address_,
+               pid_.value_or(-1));
+    }
+  }
+}
+
 Context::Context(const std::string &bin_path) : loaded_(false) {
   std::string full_path = resolve_bin_path(bin_path);
   if (!full_path.empty()) {
@@ -226,8 +275,60 @@ Context::Context(const std::string &bin_path) : loaded_(false) {
   }
 }
 
-Context::Context(int pid) : loaded_(false) {
+Context::Context(int pid) : pid_(pid), pid_stat_(pid), loaded_(false) {
   if (bcc_procutils_each_module(pid, _each_module, this) == 0)
     loaded_ = true;
 }
+
+Context::~Context() {
+  if (pid_stat_ && !pid_stat_->is_stale()) {
+    for (auto &p : probes_) p->disable();
+  }
+}
+}
+
+extern "C" {
+#include "bcc_usdt.h"
+
+void *bcc_usdt_new_frompid(int pid) {
+  USDT::Context *ctx = new USDT::Context(pid);
+  if (!ctx->loaded()) {
+    delete ctx;
+    return nullptr;
+  }
+  return static_cast<void *>(ctx);
+}
+
+void *bcc_usdt_new_frompath(const char *path) {
+  USDT::Context *ctx = new USDT::Context(path);
+  if (!ctx->loaded()) {
+    delete ctx;
+    return nullptr;
+  }
+  return static_cast<void *>(ctx);
+}
+
+void bcc_usdt_close(void *usdt) {
+  USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
+  delete ctx;
+}
+
+int bcc_usdt_enable_probe(void *usdt, const char *probe_name,
+                          const char *fn_name) {
+  USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
+  return ctx->enable_probe(probe_name, fn_name) ? 0 : -1;
+}
+
+char *bcc_usdt_genargs(void *usdt) {
+  USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
+  std::ostringstream stream;
+  if (!ctx->generate_usdt_args(stream))
+    return nullptr;
+  return strdup(stream.str().c_str());
+}
+
+void bcc_usdt_foreach_uprobe(void *usdt, bcc_usdt_uprobe_cb callback) {
+  USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
+  ctx->each_uprobe(callback);
+}
 }

src/cc/usdt.h

@@ -15,6 +15,7 @@
  */
 #pragma once
 
+#include <memory>
 #include <string>
 #include <unordered_map>
 #include <vector>
@@ -122,32 +123,39 @@ class Probe {
 
   struct Location {
     uint64_t address_;
-    std::vector<Argument *> arguments_;
+    std::vector<Argument> arguments_;
     Location(uint64_t addr, const char *arg_fmt);
   };
 
   std::vector<Location> locations_;
-  std::unordered_map<int, uint64_t> semaphores_;
-  std::unordered_map<int, ProcStat> enabled_semaphores_;
+
+  optional<int> pid_;
   optional<bool> in_shared_object_;
 
-  bool add_to_semaphore(int pid, int16_t val);
-  bool lookup_semaphore_addr(uint64_t *address, int pid);
+  optional<std::string> attached_to_;
+  optional<uint64_t> attached_semaphore_;
+
+  std::string largest_arg_type(size_t arg_n);
+
+  bool add_to_semaphore(int16_t val);
+  bool resolve_global_address(uint64_t *global, const uint64_t addr);
+  bool lookup_semaphore_addr(uint64_t *address);
   void add_location(uint64_t addr, const char *fmt);
 
 public:
   Probe(const char *bin_path, const char *provider, const char *name,
-        uint64_t semaphore);
+        uint64_t semaphore, const optional<int> &pid);
 
   size_t num_locations() const { return locations_.size(); }
   size_t num_arguments() const { return locations_.front().arguments_.size(); }
 
-  bool usdt_thunks(std::ostream &stream, const std::string &prefix);
-  bool usdt_cases(std::ostream &stream, const optional<int> &pid = nullopt);
+  uint64_t address(size_t n = 0) const { return locations_[n].address_; }
+  bool usdt_getarg(std::ostream &stream);
 
   bool need_enable() const { return semaphore_ != 0x0; }
-  bool enable(int pid);
-  bool disable(int pid);
+  bool enable(const std::string &fn_name);
+  bool disable();
+  bool enabled() const { return !!attached_to_; }
 
   bool in_shared_object();
   const std::string &name() { return name_; }
@@ -158,7 +166,10 @@ public:
 };
 
 class Context {
-  std::vector<Probe *> probes_;
+  std::vector<std::unique_ptr<Probe>> probes_;
+
+  optional<int> pid_;
+  optional<ProcStat> pid_stat_;
   bool loaded_;
 
   static void _each_probe(const char *binpath, const struct bcc_elf_usdt *probe,
@@ -171,9 +182,19 @@ class Context {
 public:
   Context(const std::string &bin_path);
   Context(int pid);
+  ~Context();
 
+  optional<int> pid() const { return pid_; }
   bool loaded() const { return loaded_; }
   size_t num_probes() const { return probes_.size(); }
-  Probe *find_probe(const std::string &probe_name);
+
+  Probe *get(const std::string &probe_name);
+  Probe *get(int pos) { return probes_[pos].get(); }
+
+  bool enable_probe(const std::string &probe_name, const std::string &fn_name);
+  bool generate_usdt_args(std::ostream &stream);
+
+  typedef void (*each_uprobe_cb)(const char *, const char *, uint64_t, int);
+  void each_uprobe(each_uprobe_cb callback);
 };
 }

src/cc/usdt_args.cc

@@ -55,37 +55,35 @@ bool Argument::assign_to_local(std::ostream &stream,
                                const std::string &binpath,
                                const optional<int> &pid) const {
   if (constant_) {
-    tfm::format(stream, "%s = %d;\n", local_name, *constant_);
+    tfm::format(stream, "%s = %d;", local_name, *constant_);
     return true;
   }
 
   if (!deref_offset_) {
-    tfm::format(stream, "%s = (%s)ctx->%s;\n", local_name, ctype(),
+    tfm::format(stream, "%s = (%s)ctx->%s;", local_name, ctype(),
                 *register_name_);
     return true;
   }
 
   if (deref_offset_ && !deref_ident_) {
     tfm::format(stream,
-                "{\n"
-                "    u64 __temp = ctx->%s + (%d);\n"
-                "    bpf_probe_read(&%s, sizeof(%s), (void *)__temp);\n"
-                "}\n",
-                *register_name_, *deref_offset_, local_name, local_name);
+                "{ u64 __addr = ctx->%s + (%d); %s __res = 0x0; "
+                "bpf_probe_read(&__res, sizeof(__res), (void *)__addr); "
+                "%s = __res; }",
+                *register_name_, *deref_offset_, ctype(), local_name);
     return true;
   }
 
-  if (deref_offset_ && deref_ident_) {
+  if (deref_offset_ && deref_ident_ && *register_name_ == "ip") {
     uint64_t global_address;
     if (!get_global_address(&global_address, binpath, pid))
       return false;
 
     tfm::format(stream,
-                "{\n"
-                "    u64 __temp = 0x%xull + %d;\n"
-                "    bpf_probe_read(&%s, sizeof(%s), (void *)__temp);\n"
-                "}\n",
-                global_address, *deref_offset_, local_name, local_name);
+                "{ u64 __addr = 0x%xull + %d; %s __res = 0x0; "
+                "bpf_probe_read(&__res, sizeof(__res), (void *)__addr); "
+                "%s = __res; }",
+                global_address, *deref_offset_, ctype(), local_name);
     return true;
   }
 

src/lua/bcc/bpf.lua

@@ -37,7 +37,7 @@ function Bpf.static.check_probe_quota(n)
   assert(cur + n <= Bpf.static.KPROBE_LIMIT, "number of open probes would exceed quota")
 end
 
-function Bpf.static.cleanup_probes()
+function Bpf.static.cleanup()
   local function detach_all(probe_type, all_probes)
     for key, probe in pairs(all_probes) do
       libbcc.perf_reader_free(probe)
@@ -102,10 +102,13 @@ local function _find_file(script_root, filename)
 end
 
 function Bpf:initialize(args)
-  self.do_debug = args.debug or false
   self.funcs = {}
   self.tables = {}
 
+  if args.usdt and args.text then
+    args.text = args.usdt:_get_text() .. args.text
+  end
+
   local cflags = table.join(Bpf.DEFAULT_CFLAGS, args.cflags)
   local cflags_ary = ffi.new("const char *[?]", #cflags, cflags)
 
@@ -127,6 +130,10 @@ function Bpf:initialize(args)
   end
 
   assert(self.module ~= nil, "failed to compile BPF module")
+
+  if args.usdt then
+    args.usdt:_attach_uprobes(self)
+  end
 end
 
 function Bpf:load_funcs(prog_type)

src/lua/bcc/libbcc.lua

@@ -112,6 +112,18 @@ int bcc_symcache_resolve(void *symcache, uint64_t addr, struct bcc_symbol *sym);
 void bcc_symcache_refresh(void *resolver);
 ]]
 
+ffi.cdef[[
+void *bcc_usdt_new_frompid(int pid);
+void *bcc_usdt_new_frompath(const char *path);
+void bcc_usdt_close(void *usdt);
+
+int bcc_usdt_enable_probe(void *, const char *, const char *);
+char *bcc_usdt_genargs(void *);
+
+typedef void (*bcc_usdt_uprobe_cb)(const char *, const char *, uint64_t, int);
+void bcc_usdt_foreach_uprobe(void *usdt, bcc_usdt_uprobe_cb callback);
+]]
+
 if rawget(_G, "BCC_STANDALONE") then
   return ffi.C
 else

src/lua/bcc/run.lua

@@ -57,18 +57,21 @@ return function()
   local BPF = require("bcc.bpf")
   BPF.script_root(tracefile)
 
+  local USDT = require("bcc.usdt")
   local utils = {
     argparse = require("bcc.vendor.argparse"),
     posix = require("bcc.vendor.posix"),
+    USDT = USDT,
   }
 
   local command = dofile(tracefile)
   local res, err = xpcall(command, debug.traceback, BPF, utils)
 
-  if not res then
+  if not res and err ~= "interrupted!" then
     io.stderr:write("[ERROR] "..err.."\n")
   end
 
-  BPF.cleanup_probes()
+  BPF.cleanup()
+  USDT.cleanup()
   return res, err
 end

src/lua/bcc/usdt.lua

+--[[
+Copyright 2016 GitHub, Inc
+
+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.
+]]
+local ffi = require("ffi")
+ffi.cdef "void free(void *ptr);"
+
+local libbcc = require("bcc.libbcc")
+local Usdt = class("USDT")
+
+Usdt.static.open_contexts = {}
+
+function Usdt.static.cleanup()
+  for _, context in ipairs(Usdt.static.open_contexts) do
+    context:_cleanup()
+  end
+end
+
+function Usdt:initialize(args)
+  assert(args.pid or args.path)
+
+  if args.pid then
+    self.pid = args.pid
+    self.context = libbcc.bcc_usdt_new_frompid(args.pid)
+  elseif args.path then
+    self.path = args.path
+    self.context = libbcc.bcc_usdt_new_frompath(args.path)
+  end
+
+  assert(self.context ~= nil, "failed to create USDT context")
+  table.insert(Usdt.open_contexts, self)
+end
+
+function Usdt:enable_probe(args)
+  assert(args.probe and args.fn_name)
+  assert(libbcc.bcc_usdt_enable_probe(
+    self.context, args.probe, args.fn_name) == 0)
+end
+
+function Usdt:_cleanup()
+  libbcc.bcc_usdt_close(self.context)
+  self.context = nil
+end
+
+function Usdt:_get_text()
+  local argc = libbcc.bcc_usdt_genargs(self.context)
+  assert(argc ~= nil)
+
+  local text = ffi.string(argc)
+  ffi.C.free(argc)
+  return text
+end
+
+function Usdt:_attach_uprobes(bpf)
+  local uprobes = {}
+  local cb = ffi.cast("bcc_usdt_uprobe_cb",
+    function(binpath, fn_name, addr, pid)
+      table.insert(uprobes, {name=ffi.string(binpath),
+        addr=addr, fn_name=ffi.string(fn_name), pid=pid})
+    end)
+
+  libbcc.bcc_usdt_foreach_uprobe(self.context, cb)
+  cb:free()
+
+  for _, args in ipairs(uprobes) do
+    bpf:attach_uprobe(args)
+  end
+end
+
+return Usdt

src/lua/squishy

@@ -11,6 +11,7 @@ Module "bcc.sym" "bcc/sym.lua"
 Module "bcc.libbcc" "bcc/libbcc.lua"
 Module "bcc.tracerpipe" "bcc/tracerpipe.lua"
 Module "bcc.table" "bcc/table.lua"
+Module "bcc.usdt" "bcc/usdt.lua"
 
 Main "bcc/run.lua"
 Output "bcc.lua"

tests/cc/test_usdt_probes.cc

@@ -39,8 +39,8 @@ TEST_CASE("test finding a probe in our own process", "[usdt]") {
   REQUIRE(ctx.num_probes() >= 1);
 
   SECTION("our test probe") {
-    USDT::Probe *probe = ctx.find_probe("sample_probe_1");
-    REQUIRE(probe != nullptr);
+    auto probe = ctx.get("sample_probe_1");
+    REQUIRE(probe);
 
     REQUIRE(probe->in_shared_object() == false);
     REQUIRE(probe->name() == "sample_probe_1");
@@ -52,13 +52,6 @@ TEST_CASE("test finding a probe in our own process", "[usdt]") {
     REQUIRE(probe->need_enable() == false);
 
     REQUIRE(a_probed_function() != 0);
-
-    std::ostringstream case_stream;
-    REQUIRE(probe->usdt_cases(case_stream));
-
-    std::string cases = case_stream.str();
-    REQUIRE(cases.find("int32_t arg1") != std::string::npos);
-    REQUIRE(cases.find("uint64_t arg2") != std::string::npos);
   }
 }
 #endif  // HAVE_SDT_HEADER
@@ -115,8 +108,8 @@ TEST_CASE("test listing all USDT probes in Ruby/MRI", "[usdt]") {
     mri_probe_count = ctx.num_probes();
 
     SECTION("GC static probe") {
-      USDT::Probe *probe = ctx.find_probe("gc__mark__begin");
-      REQUIRE(probe != nullptr);
+      auto probe = ctx.get("gc__mark__begin");
+      REQUIRE(probe);
 
       REQUIRE(probe->in_shared_object() == true);
       REQUIRE(probe->name() == "gc__mark__begin");
@@ -129,8 +122,8 @@ TEST_CASE("test listing all USDT probes in Ruby/MRI", "[usdt]") {
     }
 
     SECTION("object creation probe") {
-      USDT::Probe *probe = ctx.find_probe("object__create");
-      REQUIRE(probe != nullptr);
+      auto probe = ctx.get("object__create");
+      REQUIRE(probe);
 
       REQUIRE(probe->in_shared_object() == true);
       REQUIRE(probe->name() == "object__create");
@@ -140,54 +133,16 @@ TEST_CASE("test listing all USDT probes in Ruby/MRI", "[usdt]") {
       REQUIRE(probe->num_locations() == 1);
       REQUIRE(probe->num_arguments() == 3);
       REQUIRE(probe->need_enable() == true);
-
-      std::ostringstream thunks_stream;
-      REQUIRE(probe->usdt_thunks(thunks_stream, "ruby_usdt"));
-
-      std::string thunks = thunks_stream.str();
-      REQUIRE(std::count(thunks.begin(), thunks.end(), '\n') == 1);
-      REQUIRE(thunks.find("ruby_usdt_thunk_0") != std::string::npos);
-
-      std::ostringstream case_stream;
-      REQUIRE(probe->usdt_cases(case_stream));
-
-      std::string cases = case_stream.str();
-      REQUIRE(countsubs(cases, "arg1") == 2);
-      REQUIRE(countsubs(cases, "arg2") == 2);
-      REQUIRE(countsubs(cases, "arg3") == 2);
-
-      REQUIRE(countsubs(cases, "uint64_t") == 4);
-      REQUIRE(countsubs(cases, "int32_t") == 2);
     }
 
     SECTION("array creation probe") {
-      USDT::Probe *probe = ctx.find_probe("array__create");
-      REQUIRE(probe != nullptr);
+      auto probe = ctx.get("array__create");
+      REQUIRE(probe);
       REQUIRE(probe->name() == "array__create");
 
       REQUIRE(probe->num_locations() == 7);
       REQUIRE(probe->num_arguments() == 3);
       REQUIRE(probe->need_enable() == true);
-
-      std::ostringstream thunks_stream;
-      REQUIRE(probe->usdt_thunks(thunks_stream, "ruby_usdt"));
-
-      std::string thunks = thunks_stream.str();
-      REQUIRE(std::count(thunks.begin(), thunks.end(), '\n') == 7);
-      REQUIRE(thunks.find("ruby_usdt_thunk_0") != std::string::npos);
-      REQUIRE(thunks.find("ruby_usdt_thunk_6") != std::string::npos);
-      REQUIRE(thunks.find("ruby_usdt_thunk_7") == std::string::npos);
-
-      std::ostringstream case_stream;
-      REQUIRE(probe->usdt_cases(case_stream));
-
-      std::string cases = case_stream.str();
-      REQUIRE(countsubs(cases, "arg1") == 8);
-      REQUIRE(countsubs(cases, "arg2") == 8);
-      REQUIRE(countsubs(cases, "arg3") == 8);
-
-      REQUIRE(countsubs(cases, "__loc_id") == 7);
-      REQUIRE(cases.find("int64_t arg1 =") != std::string::npos);
     }
   }
 
@@ -203,8 +158,8 @@ TEST_CASE("test listing all USDT probes in Ruby/MRI", "[usdt]") {
     REQUIRE(ctx.num_probes() >= mri_probe_count);
 
     SECTION("get probe in running process") {
-      USDT::Probe *probe = ctx.find_probe("gc__mark__begin");
-      REQUIRE(probe != nullptr);
+      auto probe = ctx.get("gc__mark__begin");
+      REQUIRE(probe);
 
       REQUIRE(probe->in_shared_object() == true);
       REQUIRE(probe->name() == "gc__mark__begin");

tests/lua/test_uprobes.lua

@@ -64,7 +64,7 @@ int count(struct pt_regs *ctx) {
 end
 
 function TestUprobes:teardown()
-  BPF.cleanup_probes()
+  BPF.cleanup()
 end
 
 suite("TestUprobes", TestUprobes)