add map functions: sum(), min(), max(), avg(), stats()

README.md

@@ -222,6 +222,11 @@ Variables:
 Functions:
 - `quantize(int n)` - Produce a log2 histogram of values of `n`
 - `count()` - Count the number of times this function is called
+- `sum(int n)` - Sum this value
+- `min(int n)` - Record the minimum value seen
+- `max(int n)` - Record the maximum value seen
+- `avg(int n)` - Average this value
+- `stats(int n)` - Return the count, average, and total for this value
 - `delete(@x)` - Delete the map element passed in as an argument
 - `str(char *s)` - Returns the string pointed to by `s`
 - `printf(char *fmt, ...)` - Print formatted to stdout

src/ast/codegen_llvm.cpp

@@ -119,6 +119,96 @@ void CodegenLLVM::visit(Call &call)
     b_.CreateLifetimeEnd(newval);
     expr_ = nullptr;
   }
+  else if (call.func == "sum")
+  {
+    Map &map = *call.map;
+    AllocaInst *key = getMapKey(map);
+    Value *oldval = b_.CreateMapLookupElem(map, key);
+    AllocaInst *newval = b_.CreateAllocaBPF(map.type, map.ident + "_val");
+
+    call.vargs->front()->accept(*this);
+    b_.CreateStore(b_.CreateAdd(expr_, oldval), newval);
+    b_.CreateMapUpdateElem(map, key, newval);
+
+    // oldval can only be an integer so won't be in memory and doesn't need lifetime end
+    b_.CreateLifetimeEnd(key);
+    b_.CreateLifetimeEnd(newval);
+    expr_ = nullptr;
+  }
+  else if (call.func == "min")
+  {
+    Map &map = *call.map;
+    AllocaInst *key = getMapKey(map);
+    Value *oldval = b_.CreateMapLookupElem(map, key);
+    AllocaInst *newval = b_.CreateAllocaBPF(map.type, map.ident + "_val");
+
+    // Store the max of (0xffffffff - val), so that our SGE comparison with uninitialized
+    // elements will always store on the first occurrance. Revent this later when printing.
+    Function *parent = b_.GetInsertBlock()->getParent();
+    call.vargs->front()->accept(*this);
+    Value *inverted = b_.CreateSub(b_.getInt64(0xffffffff), expr_);
+    BasicBlock *lt = BasicBlock::Create(module_->getContext(), "min.lt", parent);
+    BasicBlock *ge = BasicBlock::Create(module_->getContext(), "min.ge", parent);
+    b_.CreateCondBr(b_.CreateICmpSGE(inverted, oldval), ge, lt);
+
+    b_.SetInsertPoint(ge);
+    b_.CreateStore(inverted, newval);
+    b_.CreateMapUpdateElem(map, key, newval);
+    b_.CreateBr(lt);
+
+    b_.SetInsertPoint(lt);
+    b_.CreateLifetimeEnd(key);
+    b_.CreateLifetimeEnd(newval);
+    expr_ = nullptr;
+  }
+  else if (call.func == "max")
+  {
+    Map &map = *call.map;
+    AllocaInst *key = getMapKey(map);
+    Value *oldval = b_.CreateMapLookupElem(map, key);
+    AllocaInst *newval = b_.CreateAllocaBPF(map.type, map.ident + "_val");
+
+    Function *parent = b_.GetInsertBlock()->getParent();
+    call.vargs->front()->accept(*this);
+    BasicBlock *lt = BasicBlock::Create(module_->getContext(), "min.lt", parent);
+    BasicBlock *ge = BasicBlock::Create(module_->getContext(), "min.ge", parent);
+    b_.CreateCondBr(b_.CreateICmpSGE(expr_, oldval), ge, lt);
+
+    b_.SetInsertPoint(ge);
+    b_.CreateStore(expr_, newval);
+    b_.CreateMapUpdateElem(map, key, newval);
+    b_.CreateBr(lt);
+
+    b_.SetInsertPoint(lt);
+    b_.CreateLifetimeEnd(key);
+    b_.CreateLifetimeEnd(newval);
+    expr_ = nullptr;
+  }
+  else if (call.func == "avg" || call.func == "stats")
+  {
+    // avg stores the count and total in a quantize map using indexes 0 and 1
+    // respectively, and the calculation is made when printing.
+    Map &map = *call.map;
+
+    AllocaInst *count_key = getQuantizeMapKey(map, b_.getInt64(0));
+    Value *count_old = b_.CreateMapLookupElem(map, count_key);
+    AllocaInst *count_new = b_.CreateAllocaBPF(map.type, map.ident + "_num");
+    b_.CreateStore(b_.CreateAdd(count_old, b_.getInt64(1)), count_new);
+    b_.CreateMapUpdateElem(map, count_key, count_new);
+    b_.CreateLifetimeEnd(count_key);
+    b_.CreateLifetimeEnd(count_new);
+
+    AllocaInst *total_key = getQuantizeMapKey(map, b_.getInt64(1));
+    Value *total_old = b_.CreateMapLookupElem(map, total_key);
+    AllocaInst *total_new = b_.CreateAllocaBPF(map.type, map.ident + "_val");
+    call.vargs->front()->accept(*this);
+    b_.CreateStore(b_.CreateAdd(expr_, total_old), total_new);
+    b_.CreateMapUpdateElem(map, total_key, total_new);
+    b_.CreateLifetimeEnd(total_key);
+    b_.CreateLifetimeEnd(total_new);
+
+    expr_ = nullptr;
+  }
   else if (call.func == "quantize")
   {
     Map &map = *call.map;
@@ -315,6 +405,7 @@ void CodegenLLVM::visit(Call &call)
 
   else
   {
+    std::cerr << "Error: missing codegen for function \"" << call.func << "\"" << std::endl;
     abort();
   }
 }

src/ast/semantic_analyser.cpp

@@ -94,6 +94,36 @@ void SemanticAnalyser::visit(Call &call)
 
     call.type = SizedType(Type::count, 8);
   }
+  else if (call.func == "sum") {
+    check_assignment(call, true, false);
+    check_nargs(call, 1);
+
+    call.type = SizedType(Type::sum, 8);
+  }
+  else if (call.func == "min") {
+    check_assignment(call, true, false);
+    check_nargs(call, 1);
+
+    call.type = SizedType(Type::min, 8);
+  }
+  else if (call.func == "max") {
+    check_assignment(call, true, false);
+    check_nargs(call, 1);
+
+    call.type = SizedType(Type::max, 8);
+  }
+  else if (call.func == "avg") {
+    check_assignment(call, true, false);
+    check_nargs(call, 1);
+
+    call.type = SizedType(Type::avg, 8);
+  }
+  else if (call.func == "stats") {
+    check_assignment(call, true, false);
+    check_nargs(call, 1);
+
+    call.type = SizedType(Type::stats, 8);
+  }
   else if (call.func == "delete") {
     check_assignment(call, false, false);
     if (check_nargs(call, 1)) {

src/bpftrace.cpp

@@ -365,6 +365,8 @@ int BPFtrace::print_maps()
     int err;
     if (map.type_.type == Type::quantize)
       err = print_map_quantize(map, 0, 0);
+    else if (map.type_.type == Type::avg || map.type_.type == Type::stats)
+      err = print_map_stats(map);
     else
       err = print_map(map, 0, 0);
 
@@ -530,7 +532,8 @@ int BPFtrace::print_map(IMap &map, uint32_t top, uint32_t div)
   while (bpf_get_next_key(map.mapfd_, old_key.data(), key.data()) == 0)
   {
     int value_size = map.type_.size;
-    if (map.type_.type == Type::count)
+    if (map.type_.type == Type::count ||
+        map.type_.type == Type::sum || map.type_.type == Type::min || map.type_.type == Type::max)
       value_size *= ncpus_;
     auto value = std::vector<uint8_t>(value_size);
     int err = bpf_lookup_elem(map.mapfd_, key.data(), value.data());
@@ -545,13 +548,27 @@ int BPFtrace::print_map(IMap &map, uint32_t top, uint32_t div)
     old_key = key;
   }
 
-  if (map.type_.type == Type::count)
+  if (map.type_.type == Type::count || map.type_.type == Type::sum)
   {
     std::sort(values_by_key.begin(), values_by_key.end(), [&](auto &a, auto &b)
     {
       return reduce_value(a.second, ncpus_) < reduce_value(b.second, ncpus_);
     });
   }
+  else if (map.type_.type == Type::min)
+  {
+    std::sort(values_by_key.begin(), values_by_key.end(), [&](auto &a, auto &b)
+    {
+      return min_value(a.second, ncpus_) < min_value(b.second, ncpus_);
+    });
+  }
+  else if (map.type_.type == Type::max)
+  {
+    std::sort(values_by_key.begin(), values_by_key.end(), [&](auto &a, auto &b)
+    {
+      return max_value(a.second, ncpus_) < max_value(b.second, ncpus_);
+    });
+  }
   else
   {
     sort_by_key(map.key_.args_, values_by_key);
@@ -583,8 +600,12 @@ int BPFtrace::print_map(IMap &map, uint32_t top, uint32_t div)
       std::cout << resolve_usym(*(uintptr_t*)value.data());
     else if (map.type_.type == Type::string)
       std::cout << value.data() << std::endl;
-    else if (map.type_.type == Type::count)
+    else if (map.type_.type == Type::count || map.type_.type == Type::sum)
       std::cout << reduce_value(value, ncpus_) / div << std::endl;
+    else if (map.type_.type == Type::min)
+      std::cout << min_value(value, ncpus_) / div << std::endl;
+    else if (map.type_.type == Type::max)
+      std::cout << max_value(value, ncpus_) / div << std::endl;
     else
       std::cout << *(int64_t*)value.data() / div << std::endl;
   }
@@ -683,6 +704,88 @@ int BPFtrace::print_map_quantize(IMap &map, uint32_t top, uint32_t div)
   return 0;
 }
 
+int BPFtrace::print_map_stats(IMap &map)
+{
+  // A quantize-map adds an extra 8 bytes onto the end of its key for storing
+  // the bucket number.
+
+  std::vector<uint8_t> old_key;
+  try
+  {
+    old_key = find_empty_key(map, map.key_.size() + 8);
+  }
+  catch (std::runtime_error &e)
+  {
+    std::cerr << "Error getting key for map '" << map.name_ << "': "
+              << e.what() << std::endl;
+    return -2;
+  }
+  auto key(old_key);
+
+  std::map<std::vector<uint8_t>, std::vector<uint64_t>> values_by_key;
+
+  while (bpf_get_next_key(map.mapfd_, old_key.data(), key.data()) == 0)
+  {
+    auto key_prefix = std::vector<uint8_t>(map.key_.size());
+    int bucket = key.at(map.key_.size());
+
+    for (size_t i=0; i<map.key_.size(); i++)
+      key_prefix.at(i) = key.at(i);
+
+    int value_size = map.type_.size * ncpus_;
+    auto value = std::vector<uint8_t>(value_size);
+    int err = bpf_lookup_elem(map.mapfd_, key.data(), value.data());
+    if (err)
+    {
+      std::cerr << "Error looking up elem: " << err << std::endl;
+      return -1;
+    }
+
+    if (values_by_key.find(key_prefix) == values_by_key.end())
+    {
+      // New key - create a list of buckets for it
+      values_by_key[key_prefix] = std::vector<uint64_t>(2);
+    }
+    values_by_key[key_prefix].at(bucket) = reduce_value(value, ncpus_);
+
+    old_key = key;
+  }
+
+  // Sort based on sum of counts in all buckets
+  std::vector<std::pair<std::vector<uint8_t>, uint64_t>> total_counts_by_key;
+  for (auto &map_elem : values_by_key)
+  {
+    assert(map_elem.second.size() == 2);
+    uint64_t count = map_elem.second.at(0);
+    uint64_t total = map_elem.second.at(1);
+    assert(count != 0);
+    total_counts_by_key.push_back({map_elem.first, total / count});
+  }
+  std::sort(total_counts_by_key.begin(), total_counts_by_key.end(), [&](auto &a, auto &b)
+  {
+    return a.second < b.second;
+  });
+
+  for (auto &key_count : total_counts_by_key)
+  {
+    auto &key = key_count.first;
+    auto &value = values_by_key[key];
+    std::cout << map.name_ << map.key_.argument_value_list(*this, key) << ": ";
+
+    uint64_t count = value.at(0);
+    uint64_t total = value.at(1);
+
+    if (map.type_.type == Type::stats)
+      std::cout << "count " << count << ", average " << total / count << ", total " << total << std::endl;
+    else
+      std::cout << total / count << std::endl;
+  }
+
+  std::cout << std::endl;
+
+  return 0;
+}
+
 int BPFtrace::print_quantize(const std::vector<uint64_t> &values, uint32_t div) const
 {
   int max_index = -1;
@@ -767,12 +870,39 @@ uint64_t BPFtrace::reduce_value(const std::vector<uint8_t> &value, int ncpus)
   return sum;
 }
 
+uint64_t BPFtrace::max_value(const std::vector<uint8_t> &value, int ncpus)
+{
+  uint64_t val, max = 0;
+  for (int i=0; i<ncpus; i++)
+  {
+    val = *(uint64_t*)(value.data() + i*sizeof(uint64_t*));
+    if (val > max)
+      max = val;
+  }
+  return max;
+}
+
+uint64_t BPFtrace::min_value(const std::vector<uint8_t> &value, int ncpus)
+{
+  uint64_t val, max = 0;
+  for (int i=0; i<ncpus; i++)
+  {
+    val = *(uint64_t*)(value.data() + i*sizeof(uint64_t*));
+    if (val > max)
+      max = val;
+  }
+  return (0xffffffff - max);
+}
+
 std::vector<uint8_t> BPFtrace::find_empty_key(IMap &map, size_t size) const
 {
   if (size == 0) size = 8;
   auto key = std::vector<uint8_t>(size);
   int value_size = map.type_.size;
-  if (map.type_.type == Type::count || map.type_.type == Type::quantize)
+  if (map.type_.type == Type::count || map.type_.type == Type::quantize ||
+      map.type_.type == Type::sum || map.type_.type == Type::min ||
+      map.type_.type == Type::max || map.type_.type == Type::avg ||
+      map.type_.type == Type::stats)
     value_size *= ncpus_;
   auto value = std::vector<uint8_t>(value_size);
 

src/bpftrace.h

@@ -67,8 +67,11 @@ private:
   int zero_map(IMap &map);
   int print_map(IMap &map, uint32_t top, uint32_t div);
   int print_map_quantize(IMap &map, uint32_t top, uint32_t div);
+  int print_map_stats(IMap &map);
   int print_quantize(const std::vector<uint64_t> &values, uint32_t div) const;
   static uint64_t reduce_value(const std::vector<uint8_t> &value, int ncpus);
+  static uint64_t min_value(const std::vector<uint8_t> &value, int ncpus);
+  static uint64_t max_value(const std::vector<uint8_t> &value, int ncpus);
   static std::string quantize_index_label(int power);
   std::vector<uint8_t> find_empty_key(IMap &map, size_t size) const;
 };

src/map.cpp

@@ -16,13 +16,15 @@ Map::Map(const std::string &name, const SizedType &type, const MapKey &key)
   key_ = key;
 
   int key_size = key.size();
-  if (type.type == Type::quantize)
+  if (type.type == Type::quantize ||
+      type.type == Type::avg || type.type == Type::stats)
     key_size += 8;
   if (key_size == 0)
     key_size = 8;
 
   enum bpf_map_type map_type;
-  if ((type.type == Type::quantize || type.type == Type::count) &&
+  if ((type.type == Type::quantize || type.type == Type::count ||
+      type.type == Type::sum || type.type == Type::min || type.type == Type::max || type.type == Type::avg || type.type == Type::stats) &&
       (LINUX_VERSION_CODE >= KERNEL_VERSION(4, 6, 0)))
   {
       map_type = BPF_MAP_TYPE_PERCPU_HASH;

src/types.cpp

@@ -29,6 +29,11 @@ std::string typestr(Type t)
     case Type::integer:  return "integer";  break;
     case Type::quantize: return "quantize"; break;
     case Type::count:    return "count";    break;
+    case Type::sum:      return "sum";      break;
+    case Type::min:      return "min";      break;
+    case Type::max:      return "max";      break;
+    case Type::avg:      return "avg";      break;
+    case Type::stats:    return "stats";    break;
     case Type::stack:    return "stack";    break;
     case Type::ustack:   return "ustack";   break;
     case Type::string:   return "string";   break;

src/types.h

@@ -17,6 +17,11 @@ enum class Type
   integer,
   quantize,
   count,
+  sum,
+  min,
+  max,
+  avg,
+  stats,
   stack,
   ustack,
   string,

tests/parser.cpp

@@ -63,6 +63,41 @@ TEST(Parser, map_assign)
       "  =\n"
       "   map: @x\n"
       "   call: count\n");
+  test("kprobe:sys_read { @x = sum(arg2); }",
+      "Program\n"
+      " kprobe:sys_read\n"
+      "  =\n"
+      "   map: @x\n"
+      "   call: sum\n"
+      "    builtin: arg2\n");
+  test("kprobe:sys_read { @x = min(arg2); }",
+      "Program\n"
+      " kprobe:sys_read\n"
+      "  =\n"
+      "   map: @x\n"
+      "   call: min\n"
+      "    builtin: arg2\n");
+  test("kprobe:sys_read { @x = max(arg2); }",
+      "Program\n"
+      " kprobe:sys_read\n"
+      "  =\n"
+      "   map: @x\n"
+      "   call: max\n"
+      "    builtin: arg2\n");
+  test("kprobe:sys_read { @x = avg(arg2); }",
+      "Program\n"
+      " kprobe:sys_read\n"
+      "  =\n"
+      "   map: @x\n"
+      "   call: avg\n"
+      "    builtin: arg2\n");
+  test("kprobe:sys_read { @x = stats(arg2); }",
+      "Program\n"
+      " kprobe:sys_read\n"
+      "  =\n"
+      "   map: @x\n"
+      "   call: stats\n"
+      "    builtin: arg2\n");
   test("kprobe:sys_open { @x = \"mystring\" }",
       "Program\n"
       " kprobe:sys_open\n"

tests/semantic_analyser.cpp

@@ -80,6 +80,11 @@ TEST(semantic_analyser, builtin_functions)
 {
   test("kprobe:f { @x = quantize(123) }", 0);
   test("kprobe:f { @x = count() }", 0);
+  test("kprobe:f { @x = sum(pid) }", 0);
+  test("kprobe:f { @x = min(pid) }", 0);
+  test("kprobe:f { @x = max(pid) }", 0);
+  test("kprobe:f { @x = avg(pid) }", 0);
+  test("kprobe:f { @x = stats(pid) }", 0);
   test("kprobe:f { @x = 1; delete(@x) }", 0);
   test("kprobe:f { @x = 1; print(@x) }", 0);
   test("kprobe:f { @x = 1; clear(@x) }", 0);
@@ -91,6 +96,8 @@ TEST(semantic_analyser, builtin_functions)
   test("kprobe:f { sym(0xffff) }", 0);
   test("kprobe:f { usym(0xffff) }", 0);
   test("kprobe:f { reg(\"ip\") }", 0);
+  test("kprobe:f { @x = count(pid) }", 1);
+  test("kprobe:f { @x = sum(pid, 123) }", 1);
   test("kprobe:f { fake() }", 1);
 }
 
@@ -120,6 +127,7 @@ TEST(semantic_analyser, predicate_expressions)
 TEST(semantic_analyser, mismatched_call_types)
 {
   test("kprobe:f { @x = 1; @x = count(); }", 1);
+  test("kprobe:f { @x = count(); @x = sum(pid); }", 1);
   test("kprobe:f { @x = 1; @x = quantize(0); }", 1);
 }
 
@@ -137,6 +145,41 @@ TEST(semantic_analyser, call_count)
   test("kprobe:f { count(); }", 1);
 }
 
+TEST(semantic_analyser, call_sum)
+{
+  test("kprobe:f { @x = sum(); }", 1);
+  test("kprobe:f { @x = sum(123); }", 0);
+  test("kprobe:f { sum(); }", 1);
+}
+
+TEST(semantic_analyser, call_min)
+{
+  test("kprobe:f { @x = min(); }", 1);
+  test("kprobe:f { @x = min(123); }", 0);
+  test("kprobe:f { min(); }", 1);
+}
+
+TEST(semantic_analyser, call_max)
+{
+  test("kprobe:f { @x = max(); }", 1);
+  test("kprobe:f { @x = max(123); }", 0);
+  test("kprobe:f { max(); }", 1);
+}
+
+TEST(semantic_analyser, call_avg)
+{
+  test("kprobe:f { @x = avg(); }", 1);
+  test("kprobe:f { @x = avg(123); }", 0);
+  test("kprobe:f { avg(); }", 1);
+}
+
+TEST(semantic_analyser, call_stats)
+{
+  test("kprobe:f { @x = stats(); }", 1);
+  test("kprobe:f { @x = stats(123); }", 0);
+  test("kprobe:f { stats(); }", 1);
+}
+
 TEST(semantic_analyser, call_delete)
 {
   test("kprobe:f { @x = 1; delete(@x); }", 0);