Merge pull request #302 from iovisor/params

basic positional paramater support

README.md

@@ -255,8 +255,9 @@ Variables:
 - `retval` - Return value from function being traced
 - `func` - Name of the function currently being traced
 - `probe` - Full name of the probe
-- `curtask` - Current task_struct as a u64.
-- `rand` - Random number of type u32.
+- `curtask` - Current task_struct as a u64
+- `rand` - Random number of type u32
+- `$1`, `$2`, ... etc. - Positional parameters to the bpftrace program
 
 Functions:
 - `hist(int n)` - Produce a log2 histogram of values of `n`

docs/reference_guide.md

@@ -23,7 +23,6 @@ This is a work in progress. If something is missing, check the bpftrace source t
     - [5. `? :`: ternary operators](#5---ternary-operators)
     - [6. `if () {...} else {...}`: if-else statements](#6-if---else--if-else-statements)
     - [7. `unroll () {...}`: unroll](#7-unroll---unroll)
-
 - [Probes](#probes)
     - [1. `kprobe`/`kretprobe`: Dynamic Tracing, Kernel-Level](#1-kprobekretprobe-dynamic-tracing-kernel-level)
     - [2. `kprobe`/`kretprobe`: Dynamic Tracing, Kernel-Level Arguments](#2-kprobekretprobe-dynamic-tracing-kernel-level-arguments)
@@ -46,6 +45,7 @@ This is a work in progress. If something is missing, check the bpftrace source t
     - [6. `nsecs`: Timestamps and Time Deltas](#6-nsecs-timestamps-and-time-deltas)
     - [7. `stack`: Stack Traces, Kernel](#7-stack-stack-traces-kernel)
     - [8. `ustack`: Stack Traces, User](#8-ustack-stack-traces-user)
+    - [9. `$1`, ..., `$N`: Positional Parameters](#9-1--n-positional-parameters)
 - [Functions](#functions)
     - [1. Builtins](#1-builtins-1)
     - [2. `printf()`: Print Formatted](#2-printf-Printing)
@@ -922,6 +922,7 @@ That would fire once for every 1000000 cache misses. This usually indicates the
 - `curtask` - Current task struct as a u64
 - `rand` - Random number as a u32
 - `cgroup` - Cgroup ID of the current process
+- `$1`, `$2`, ..., `$N`. - Positional parameters for the bpftrace program
 
 Many of these are discussed in other sections (use search).
 
@@ -1168,6 +1169,73 @@ __libc_start_main+231
 
 Note that for this example to work, bash had to be recompiled with frame pointers.
 
+## 9. `$1`, ..., `$N`: Positional Parameters
+
+Syntax: `$1`, `$2`, ..., `$N`
+
+These are the positional parameters to the bpftrace program, also referred to as command line arguments. If the parameter is numeric (entirerly digits), it can be used as a number. If it is non-numeric, it must be used as a string in the `str()` call. If a parameter is used that was not provided, it will default to zero for numeric context, and "" for string context.
+
+This allows scripts to be written that use basic arguments to change their behavior. If you develop a script that requires more complex argument processing, it may be better suited for bcc instead, which supports Python's argparse and completely custom argument processing.
+
+One-liner example:
+
+```
+# bpftrace -e 'BEGIN { printf("I got %d, %s\n", $1, str($2)); }' 42 "hello"
+Attaching 1 probe...
+I got 42, hello
+```
+
+Script example, bsize.d:
+
+```
+#!/usr/local/bin/bpftrace
+
+BEGIN
+{
+	printf("Tracing block I/O sizes > %d bytes\n", $1);
+}
+
+tracepoint:block:block_rq_issue
+/args->bytes > $1/
+{
+	@ = hist(args->bytes);
+}
+```
+
+When run with a 65536 argument:
+
+```
+# ./bsize.bt 65536
+Attaching 2 probes...
+Tracing block I/O sizes > 65536 bytes
+^C
+
+@:
+[512K, 1M)             1 |@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@|
+
+```
+
+It has passed the argument in as $1, and used it as a filter.
+
+With no arguments, $1 defaults to zero:
+
+```
+# ./bsize.bt
+Attaching 2 probes...
+Tracing block I/O sizes > 0 bytes
+^C
+
+@:
+[4K, 8K)             115 |@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@|
+[8K, 16K)             35 |@@@@@@@@@@@@@@@                                     |
+[16K, 32K)             5 |@@                                                  |
+[32K, 64K)             3 |@                                                   |
+[64K, 128K)            1 |                                                    |
+[128K, 256K)           0 |                                                    |
+[256K, 512K)           0 |                                                    |
+[512K, 1M)             1 |                                                    |
+```
+
 # Functions
 
 ## 1. Builtins

src/ast/ast.cpp

@@ -16,6 +16,10 @@ void Builtin::accept(Visitor &v) {
   v.visit(*this);
 }
 
+void PositionalParameter::accept(Visitor &v) {
+  v.visit(*this);
+}
+
 void Call::accept(Visitor &v) {
   v.visit(*this);
 }

src/ast/ast.h

@@ -38,6 +38,14 @@ public:
   void accept(Visitor &v) override;
 };
 
+class PositionalParameter : public Expression {
+public:
+  explicit PositionalParameter(long n) : n(n) { is_literal = true; }
+  long n;
+
+  void accept(Visitor &v) override;
+};
+
 class String : public Expression {
 public:
   explicit String(std::string str) : str(str) { is_literal = true; }
@@ -260,6 +268,7 @@ class Visitor {
 public:
   virtual ~Visitor() { }
   virtual void visit(Integer &integer) = 0;
+  virtual void visit(PositionalParameter &integer) = 0;
   virtual void visit(String &string) = 0;
   virtual void visit(Builtin &builtin) = 0;
   virtual void visit(Call &call) = 0;

src/ast/codegen_llvm.cpp

@@ -20,6 +20,20 @@ void CodegenLLVM::visit(Integer &integer)
   expr_ = b_.getInt64(integer.n);
 }
 
+void CodegenLLVM::visit(PositionalParameter &param)
+{
+  std::string pstr = bpftrace_.get_param(param.n);
+  if (bpftrace_.is_numeric(pstr)) {
+    expr_ = b_.getInt64(std::stoll(pstr));
+  } else {
+    Constant *const_str = ConstantDataArray::getString(module_->getContext(), pstr, true);
+    AllocaInst *buf = b_.CreateAllocaBPF(ArrayType::get(b_.getInt8Ty(), pstr.length() + 1), "str");
+    b_.CreateMemSet(buf, b_.getInt8(0), pstr.length() + 1, 1);
+    b_.CreateStore(b_.CreateGEP(const_str, b_.getInt64(0)), buf);
+    expr_ = buf;
+  }
+}
+
 void CodegenLLVM::visit(String &string)
 {
   string.str.resize(string.type.size-1);

src/ast/codegen_llvm.h

@@ -28,6 +28,7 @@ public:
     { }
 
   void visit(Integer &integer) override;
+  void visit(PositionalParameter &param) override;
   void visit(String &string) override;
   void visit(Builtin &builtin) override;
   void visit(Call &call) override;

src/ast/printer.cpp

@@ -12,6 +12,12 @@ void Printer::visit(Integer &integer)
   out_ << indent << "int: " << integer.n << std::endl;
 }
 
+void Printer::visit(PositionalParameter &param)
+{
+  std::string indent(depth_, ' ');
+  out_ << indent << "builtin: $" << param.n << std::endl;
+}
+
 void Printer::visit(String &string)
 {
   std::string indent(depth_, ' ');

src/ast/printer.h

@@ -11,6 +11,7 @@ public:
   explicit Printer(std::ostream &out) : out_(out) { }
 
   void visit(Integer &integer) override;
+  void visit(PositionalParameter &param) override;
   void visit(String &string) override;
   void visit(Builtin &builtin) override;
   void visit(Call &call) override;

src/ast/semantic_analyser.cpp

@@ -18,6 +18,24 @@ void SemanticAnalyser::visit(Integer &integer)
   integer.type = SizedType(Type::integer, 8);
 }
 
+void SemanticAnalyser::visit(PositionalParameter &param)
+{
+  param.type = SizedType(Type::integer, 8);
+  std::string pstr = bpftrace_.get_param(param.n);
+  if (is_final_pass()) {
+    if (!bpftrace_.is_numeric(pstr)) {
+      if (!call_ || call_->func != "str")
+        /*
+         * call_ was added just for this test: ensuring a string parameter is
+         * only used inside str(). Without it, string parameters used as
+         * integers would return their buffer address. Maybe that's ok?
+         * If this behavior is changed, codegen needs to support it.
+         */
+        err_ << "$" << param.n << " used numerically, but given \"" << pstr << "\". Try using str($" << param.n << ")." << std::endl;
+    }
+  }
+}
+
 void SemanticAnalyser::visit(String &string)
 {
   if (string.str.size() > STRING_SIZE-1) {
@@ -109,6 +127,9 @@ void SemanticAnalyser::visit(Builtin &builtin)
 
 void SemanticAnalyser::visit(Call &call)
 {
+  // needed for positional parameters context:
+  call_ = &call;
+
   if (call.vargs) {
     for (Expression *expr : *call.vargs) {
       expr->accept(*this);

src/ast/semantic_analyser.h

@@ -19,6 +19,7 @@ public:
       out_(out) { }
 
   void visit(Integer &integer) override;
+  void visit(PositionalParameter &param) override;
   void visit(String &string) override;
   void visit(Builtin &builtin) override;
   void visit(Call &call) override;
@@ -60,6 +61,7 @@ private:
   bool check_alpha_numeric(const Call &call, int arg_num);
 
   Probe *probe_;
+  Call *call_;
   std::map<std::string, SizedType> variable_val_;
   std::map<std::string, SizedType> map_val_;
   std::map<std::string, MapKey> map_key_;

src/bpftrace.cpp

@@ -414,6 +414,31 @@ std::vector<uint64_t> BPFtrace::get_arg_values(std::vector<Field> args, uint8_t*
   return arg_values;
 }
 
+bool BPFtrace::is_numeric(std::string str)
+{
+  int i = 0;
+  while (str[i]) {
+    if (str[i] < '0' || str[i] > '9')
+      return false;
+    i++;
+  }
+  if (i == 0)
+    return false;
+  return true;
+}
+
+void BPFtrace::add_param(const std::string &param)
+{
+  params_.emplace_back(param);
+}
+
+std::string BPFtrace::get_param(int i)
+{
+  if (i > 0 && i < params_.size() + 1)
+      return params_[i - 1];
+  return "0";
+}
+
 void perf_event_lost(void *cb_cookie, uint64_t lost)
 {
   printf("Lost %lu events\n", lost);

src/bpftrace.h

@@ -69,6 +69,9 @@ public:
   std::string resolve_probe(uint64_t probe_id);
   uint64_t resolve_cgroupid(const std::string &path);
   std::vector<uint64_t> get_arg_values(std::vector<Field> args, uint8_t* arg_data);
+  void add_param(const std::string &param);
+  bool is_numeric(std::string str);
+  std::string get_param(int index);
   std::string cmd_;
   int pid_{0};
 
@@ -103,6 +106,7 @@ private:
   int ncpus_;
   int online_cpus_;
   std::vector<int> child_pids_;
+  std::vector<std::string> params_;
 
   std::unique_ptr<AttachedProbe> attach_probe(Probe &probe, const BpfOrc &bpforc);
   int setup_perf_events();

src/lexer.l

@@ -78,6 +78,7 @@ pid|tid|cgroup|uid|gid|nsecs|cpu|comm|stack|ustack|arg[0-9]|retval|func|probe|cu
 "||"                    { return Parser::make_LOR(loc); }
 "+"                     { return Parser::make_PLUS(loc); }
 "-"                     { return Parser::make_MINUS(loc); }
+"$"                     { return Parser::make_DOLLAR(loc); }
 "*"                     { return Parser::make_MUL(loc); }
 "/"                     { return Parser::make_DIV(loc); }
 "%"                     { return Parser::make_MOD(loc); }

src/main.cpp

@@ -109,6 +109,11 @@ int main(int argc, char *argv[])
     }
   }
 
+  if (argc == 1) {
+    usage();
+    return 1;
+  }
+
   if (bt_verbose && (bt_debug != DebugLevel::kNone))
   {
     // TODO: allow both
@@ -142,23 +147,14 @@ int main(int argc, char *argv[])
 
   if (script.empty())
   {
-    // There should only be 1 non-option argument (the script file)
-    if (optind != argc-1)
-    {
-      usage();
-      return 1;
-    }
+    // Script file
     char *file_name = argv[optind];
     err = driver.parse_file(file_name);
+    optind++;
   }
   else
   {
     // Script is provided as a command line argument
-    if (optind != argc)
-    {
-      usage();
-      return 1;
-    }
     err = driver.parse_str(script);
   }
 
@@ -174,6 +170,12 @@ int main(int argc, char *argv[])
 
   BPFtrace bpftrace;
 
+  // positional parameters
+  while (optind < argc) {
+    bpftrace.add_param(argv[optind]);
+    optind++;
+  }
+
   // defaults
   bpftrace.join_argnum_ = 16;
   bpftrace.join_argsize_ = 1024;

src/parser.yy

@@ -60,6 +60,7 @@ void yyerror(bpftrace::Driver &driver, const char *s);
   LOR      "||"
   PLUS     "+"
   MINUS    "-"
+  DOLLAR   "$"
   MUL      "*"
   DIV      "/"
   MOD      "%"
@@ -99,6 +100,7 @@ void yyerror(bpftrace::Driver &driver, const char *s);
 %type <ast::ExpressionList *> vargs
 %type <ast::AttachPointList *> attach_points
 %type <ast::AttachPoint *> attach_point
+%type <ast::PositionalParameter *> param
 %type <std::string> wildcard
 %type <std::string> ident
 
@@ -161,6 +163,8 @@ pred : DIV expr ENDPRED { $$ = new ast::Predicate($2); }
 ternary : expr QUES expr COLON expr { $$ = new ast::Ternary($1, $3, $5); }
      ;
 
+param : DOLLAR INT { $$ = new ast::PositionalParameter($2); }
+
 block : "{" stmts "}"     { $$ = $2; }
       | "{" stmts ";" "}" { $$ = $2; }
       ;
@@ -181,6 +185,7 @@ expr : INT             { $$ = new ast::Integer($1); }
      | STRING          { $$ = new ast::String($1); }
      | BUILTIN         { $$ = new ast::Builtin($1); }
      | ternary         { $$ = $1; }
+     | param           { $$ = $1; }
      | map             { $$ = $1; }
      | var             { $$ = $1; }
      | call            { $$ = $1; }

tests/parser.cpp

@@ -42,6 +42,7 @@ TEST(Parser, builtin_variables)
   test("kprobe:f { func }", "Program\n kprobe:f\n  builtin: func\n");
   test("kprobe:f { probe }", "Program\n kprobe:f\n  builtin: probe\n");
   test("kprobe:f { args }", "Program\n kprobe:f\n  builtin: args\n");
+  test("kprobe:f { $1 }", "Program\n kprobe:f\n  builtin: $1\n");
 }
 
 TEST(Parser, map_assign)

tests/semantic_analyser.cpp

@@ -64,6 +64,7 @@ TEST(semantic_analyser, builtin_variables)
   test("kprobe:f { retval }", 0);
   test("kprobe:f { func }", 0);
   test("kprobe:f { probe }", 0);
+  test("kprobe:f { $1 }", 0);
   test("tracepoint:a:b { args }", 0);
 //  test("kprobe:f { fake }", 1);
 }
@@ -661,6 +662,13 @@ TEST(semantic_analyser, probe_short_name)
   test("i:s:1 { 1 }", 0);
 }
 
+TEST(semantic_analyser, positional_parameters)
+{
+  // $1 won't be defined, will be tested more in runtime.
+  test("kprobe:f { printf(\"%d\", $1); }", 0);
+  test("kprobe:f { printf(\"%s\", str($1)); }", 0);
+}
+
 } // namespace semantic_analyser
 } // namespace test
 } // namespace bpftrace