From 9dc87fa7f2946cfc9f55fea93ae37709c12919c2 Mon Sep 17 00:00:00 2001
From: Vicent Marti <tanoku@gmail.com>
Date: Sun, 1 May 2016 12:54:25 +0200
Subject: [PATCH] lua: Implement an USDT helper

---
 src/lua/bcc/bpf.lua    | 11 +++++--
 src/lua/bcc/libbcc.lua | 12 ++++++++
 src/lua/bcc/run.lua    |  7 +++--
 src/lua/bcc/usdt.lua   | 81 ++++++++++++++++++++++++++++++++++++++++++++++++++
 src/lua/squishy        |  1 +
 5 files changed, 108 insertions(+), 4 deletions(-)
 create mode 100644 src/lua/bcc/usdt.lua

diff --git a/src/lua/bcc/bpf.lua b/src/lua/bcc/bpf.lua
index e76e608..8cda252 100644
--- a/src/lua/bcc/bpf.lua
+++ b/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)
diff --git a/src/lua/bcc/libbcc.lua b/src/lua/bcc/libbcc.lua
index 005809d..1ff91bc 100644
--- a/src/lua/bcc/libbcc.lua
+++ b/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
diff --git a/src/lua/bcc/run.lua b/src/lua/bcc/run.lua
index 587941d..7f09f43 100644
--- a/src/lua/bcc/run.lua
+++ b/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
diff --git a/src/lua/bcc/usdt.lua b/src/lua/bcc/usdt.lua
new file mode 100644
index 0000000..fefe7c7
--- /dev/null
+++ b/src/lua/bcc/usdt.lua
@@ -0,0 +1,81 @@
+--[[
+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
diff --git a/src/lua/squishy b/src/lua/squishy
index 8622b86..bf83d79 100644
--- a/src/lua/squishy
+++ b/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"
-- 
2.7.4