add join() function

README.md

@@ -235,5 +235,6 @@ Functions:
 - `sym(void *p)` - Resolve kernel address
 - `usym(void *p)` - Resolve user space address (incomplete)
 - `reg(char *name)` - Returns the value stored in the named register
+- `join(char *arr[])` - Prints the string array
 - `time(char *fmt)` - Print the current time
 - `exit()` - Quit bpftrace

src/ast/codegen_llvm.cpp

@@ -244,6 +244,40 @@ void CodegenLLVM::visit(Call &call)
     b_.CreateProbeReadStr(buf, call.type.size, expr_);
     expr_ = buf;
   }
+  else if (call.func == "join")
+  {
+    call.vargs->front()->accept(*this);
+    AllocaInst *first = b_.CreateAllocaBPF(SizedType(Type::integer, 8), call.func + "_first");
+    AllocaInst *second = b_.CreateAllocaBPF(b_.getInt64Ty(), call.func+"_second");
+    Value *perfdata = b_.CreateGetJoinMap(ctx_);
+    Function *parent = b_.GetInsertBlock()->getParent();
+    BasicBlock *zero = BasicBlock::Create(module_->getContext(), "joinzero", parent);
+    BasicBlock *notzero = BasicBlock::Create(module_->getContext(), "joinnotzero", parent);
+    b_.CreateCondBr(b_.CreateICmpNE(perfdata, ConstantExpr::getCast(Instruction::IntToPtr, b_.getInt64(0), b_.getInt8PtrTy()), "joinzerocond"), notzero, zero);
+
+    // arg0
+    b_.SetInsertPoint(notzero);
+    b_.CreateStore(b_.getInt64(asyncactionint(AsyncAction::join)), perfdata);
+    AllocaInst *arr = b_.CreateAllocaBPF(b_.getInt64Ty(), call.func+"_r0");
+    b_.CreateProbeRead(arr, 8, expr_);
+    b_.CreateProbeReadStr(b_.CreateAdd(perfdata, b_.getInt64(8)), bpftrace_.join_argsize_, b_.CreateLoad(arr));
+
+    for (int i = 1; i < bpftrace_.join_argnum_; i++) {
+      // argi
+      b_.CreateStore(b_.CreateAdd(expr_, b_.getInt64(8 * i)), first);
+      b_.CreateProbeRead(second, 8, b_.CreateLoad(first));
+      b_.CreateProbeReadStr(b_.CreateAdd(perfdata, b_.getInt64(8 + i * bpftrace_.join_argsize_)), bpftrace_.join_argsize_, b_.CreateLoad(second));
+    }
+
+    // emit
+    b_.CreatePerfEventOutput(ctx_, perfdata, 8 + bpftrace_.join_argnum_ * bpftrace_.join_argsize_);
+
+    b_.CreateBr(zero);
+
+    // done
+    b_.SetInsertPoint(zero);
+    expr_ = nullptr;
+  }
   else if (call.func == "sym" || call.func == "usym")
   {
     // We want expr_ to just pass through from the child node - don't set it here

src/ast/irbuilderbpf.cpp

@@ -90,6 +90,26 @@ CallInst *IRBuilderBPF::CreateBpfPseudoCall(Map &map)
   return CreateBpfPseudoCall(mapfd);
 }
 
+CallInst *IRBuilderBPF::CreateGetJoinMap(Value *ctx)
+{
+  Value *map_ptr = CreateBpfPseudoCall(bpftrace_.join_map_->mapfd_);
+  AllocaInst *key = CreateAllocaBPF(getInt32Ty(), "key");
+  Value *keyv = getInt32(0);
+  CreateStore(keyv, key);
+
+  FunctionType *lookup_func_type = FunctionType::get(
+      getInt8PtrTy(),
+      {getInt8PtrTy(), getInt8PtrTy()},
+      false);
+  PointerType *lookup_func_ptr_type = PointerType::get(lookup_func_type, 0);
+  Constant *lookup_func = ConstantExpr::getCast(
+      Instruction::IntToPtr,
+      getInt64(BPF_FUNC_map_lookup_elem),
+      lookup_func_ptr_type);
+  CallInst *call = CreateCall(lookup_func, {map_ptr, key}, "join_elem");
+  return call;
+}
+
 Value *IRBuilderBPF::CreateMapLookupElem(Map &map, AllocaInst *key)
 {
   Value *map_ptr = CreateBpfPseudoCall(map);
@@ -194,7 +214,22 @@ void IRBuilderBPF::CreateProbeRead(AllocaInst *dst, size_t size, Value *src)
   CallInst *call = CreateCall(proberead_func, {dst, getInt64(size), src}, "probe_read");
 }
 
-void IRBuilderBPF::CreateProbeReadStr(AllocaInst *dst, size_t size, Value *src)
+CallInst *IRBuilderBPF::CreateProbeReadStr(AllocaInst *dst, size_t size, Value *src)
+{
+  // int bpf_probe_read_str(void *dst, int size, const void *unsafe_ptr)
+  FunctionType *probereadstr_func_type = FunctionType::get(
+      getInt64Ty(),
+      {getInt8PtrTy(), getInt64Ty(), getInt8PtrTy()},
+      false);
+  PointerType *probereadstr_func_ptr_type = PointerType::get(probereadstr_func_type, 0);
+  Constant *probereadstr_func = ConstantExpr::getCast(
+      Instruction::IntToPtr,
+      getInt64(BPF_FUNC_probe_read_str),
+      probereadstr_func_ptr_type);
+  return CreateCall(probereadstr_func, {dst, getInt64(size), src}, "probe_read_str");
+}
+
+CallInst *IRBuilderBPF::CreateProbeReadStr(Value *dst, size_t size, Value *src)
 {
   // int bpf_probe_read_str(void *dst, int size, const void *unsafe_ptr)
   FunctionType *probereadstr_func_type = FunctionType::get(
@@ -206,7 +241,7 @@ void IRBuilderBPF::CreateProbeReadStr(AllocaInst *dst, size_t size, Value *src)
       Instruction::IntToPtr,
       getInt64(BPF_FUNC_probe_read_str),
       probereadstr_func_ptr_type);
-  CallInst *call = CreateCall(probereadstr_func, {dst, getInt64(size), src}, "probe_read_str");
+  return CreateCall(probereadstr_func, {dst, getInt64(size), src}, "map_read_str");
 }
 
 CallInst *IRBuilderBPF::CreateGetNs()

src/ast/irbuilderbpf.h

@@ -28,12 +28,14 @@ public:
   void        CreateMapUpdateElem(Map &map, AllocaInst *key, Value *val);
   void        CreateMapDeleteElem(Map &map, AllocaInst *key);
   void        CreateProbeRead(AllocaInst *dst, size_t size, Value *src);
-  void        CreateProbeReadStr(AllocaInst *dst, size_t size, Value *src);
+  CallInst   *CreateProbeReadStr(AllocaInst *dst, size_t size, Value *src);
+  CallInst   *CreateProbeReadStr(Value *dst, size_t size, Value *src);
   CallInst   *CreateGetNs();
   CallInst   *CreateGetPidTgid();
   CallInst   *CreateGetUidGid();
   CallInst   *CreateGetCpuId();
   CallInst   *CreateGetStackId(Value *ctx, bool ustack);
+  CallInst   *CreateGetJoinMap(Value *ctx);
   void        CreateGetCurrentComm(AllocaInst *buf, size_t size);
   void        CreatePerfEventOutput(Value *ctx, Value *data, size_t size);
 

src/ast/semantic_analyser.cpp

@@ -146,6 +146,13 @@ void SemanticAnalyser::visit(Call &call)
     else if (call.func == "usym")
       call.type = SizedType(Type::usym, 8);
   }
+  else if (call.func == "join") {
+    check_assignment(call, false, false);
+    check_nargs(call, 1);
+    check_arg(call, Type::integer, 0);
+    call.type = SizedType(Type::none, 0);
+    needs_join_map_ = true;
+  }
   else if (call.func == "reg") {
     if (check_nargs(call, 1)) {
       if (check_arg(call, Type::string, 0, true)) {
@@ -675,12 +682,28 @@ int SemanticAnalyser::create_maps(bool debug)
   {
     if (needs_stackid_map_)
       bpftrace_.stackid_map_ = std::make_unique<bpftrace::FakeMap>(BPF_MAP_TYPE_STACK_TRACE);
+    if (needs_join_map_)
+    {
+      // join uses map storage as we'd like to process data larger than can fit on the BPF stack.
+      std::string map_ident = "join";
+      SizedType type = SizedType(Type::join, 8 + bpftrace_.join_argnum_ * bpftrace_.join_argsize_);
+      MapKey key;
+      bpftrace_.join_map_ = std::make_unique<bpftrace::FakeMap>(map_ident, type, key);
+    }
     bpftrace_.perf_event_map_ = std::make_unique<bpftrace::FakeMap>(BPF_MAP_TYPE_PERF_EVENT_ARRAY);
   }
   else
   {
     if (needs_stackid_map_)
       bpftrace_.stackid_map_ = std::make_unique<bpftrace::Map>(BPF_MAP_TYPE_STACK_TRACE);
+    if (needs_join_map_)
+    {
+      // join uses map storage as we'd like to process data larger than can fit on the BPF stack.
+      std::string map_ident = "join";
+      SizedType type = SizedType(Type::join, 8 + bpftrace_.join_argnum_ * bpftrace_.join_argsize_);
+      MapKey key;
+      bpftrace_.join_map_ = std::make_unique<bpftrace::Map>(map_ident, type, key);
+    }
     bpftrace_.perf_event_map_ = std::make_unique<bpftrace::Map>(BPF_MAP_TYPE_PERF_EVENT_ARRAY);
   }
 

src/ast/semantic_analyser.h

@@ -61,6 +61,7 @@ private:
   std::map<std::string, SizedType> map_val_;
   std::map<std::string, MapKey> map_key_;
   bool needs_stackid_map_ = false;
+  bool needs_join_map_ = false;
   bool has_begin_probe_ = false;
   bool has_end_probe_ = false;
 };

src/bpftrace.cpp

@@ -181,6 +181,20 @@ void perf_event_printer(void *cb_cookie, void *data, int size)
     printf("%s", timestr);
     return;
   }
+  else if (printf_id == asyncactionint(AsyncAction::join))
+  {
+     const char *joinstr = " ";
+     for (int i = 0; i < bpftrace->join_argnum_; i++) {
+       auto *arg = arg_data + i * bpftrace->join_argsize_;
+       if (arg[0] == 0)
+         break;
+       if (i)
+         printf("%s", joinstr);
+       printf("%s", arg);
+     }
+     printf("\n");
+     return;
+  }
 
   // printf
   auto fmt = std::get<0>(bpftrace->printf_args_[printf_id]).c_str();
@@ -318,7 +332,7 @@ int BPFtrace::setup_perf_events()
   online_cpus_ = cpus.size();
   for (int cpu : cpus)
   {
-    int page_cnt = 8;
+    int page_cnt = 64;
     void *reader = bpf_open_perf_buffer(&perf_event_printer, &perf_event_lost, this, -1, cpu, page_cnt);
     if (reader == nullptr)
     {

src/bpftrace.h

@@ -44,7 +44,10 @@ public:
   std::vector<std::tuple<std::string, std::vector<SizedType>>> printf_args_;
   std::vector<std::string> time_args_;
   std::unique_ptr<IMap> stackid_map_;
+  std::unique_ptr<IMap> join_map_;
   std::unique_ptr<IMap> perf_event_map_;
+  int join_argnum_;
+  int join_argsize_;
 
   static void sort_by_key(std::vector<SizedType> key_args,
       std::vector<std::pair<std::vector<uint8_t>, std::vector<uint8_t>>> &values_by_key);

src/main.cpp

@@ -113,6 +113,10 @@ int main(int argc, char *argv[])
 
   BPFtrace bpftrace;
 
+  // defaults
+  bpftrace.join_argnum_ = 16;
+  bpftrace.join_argsize_ = 1024;
+
   // PID is currently only used for USDT probes that need enabling. Future work:
   // - make PID a filter for all probe types: pass to perf_event_open(), etc.
   // - provide PID in USDT probe specification as a way to override -p.

src/map.cpp

@@ -22,6 +22,7 @@ Map::Map(const std::string &name, const SizedType &type, const MapKey &key)
   if (key_size == 0)
     key_size = 8;
 
+  int max_entries = 128;
   enum bpf_map_type map_type;
   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) &&
@@ -29,11 +30,16 @@ Map::Map(const std::string &name, const SizedType &type, const MapKey &key)
   {
       map_type = BPF_MAP_TYPE_PERCPU_HASH;
   }
+  else if (type.type == Type::join)
+  {
+    map_type = BPF_MAP_TYPE_PERCPU_ARRAY;
+    max_entries = 1;
+    key_size = 4;
+  }
   else
     map_type = BPF_MAP_TYPE_HASH;
 
   int value_size = type.size;
-  int max_entries = 128;
   int flags = 0;
   mapfd_ = bpf_create_map(map_type, name.c_str(), key_size, value_size, max_entries, flags);
   if (mapfd_ < 0)

src/types.h

@@ -28,6 +28,7 @@ enum class Type
   sym,
   usym,
   cast,
+  join,
 };
 
 std::ostream &operator<<(std::ostream &os, Type type);
@@ -85,6 +86,7 @@ enum class AsyncAction
   clear,
   zero,
   time,
+  join,
 };
 
 uint64_t asyncactionint(AsyncAction a);

tests/codegen.cpp

@@ -1476,6 +1476,8 @@ attributes #1 = { argmemonly nounwind }
 )EXPECTED");
 }
 
+// TODO: add a join() test. It gets stuck in codegen.compile().
+
 TEST(codegen, int_propagation)
 {
   test("kprobe:f { @x = 1234; @y = @x }",

tests/semantic_analyser.cpp

@@ -93,6 +93,7 @@ TEST(semantic_analyser, builtin_functions)
   test("kprobe:f { exit() }", 0);
   test("kprobe:f { str(0xffff) }", 0);
   test("kprobe:f { printf(\"hello\\n\") }", 0);
+  test("kprobe:f { join(0) }", 0);
   test("kprobe:f { sym(0xffff) }", 0);
   test("kprobe:f { usym(0xffff) }", 0);
   test("kprobe:f { reg(\"ip\") }", 0);
@@ -399,6 +400,16 @@ TEST(semantic_analyser, printf_format_multi)
   test("kprobe:f { printf(\"%d %s %d\", 1, 2, \"mystr\") }", 10);
 }
 
+TEST(semantic_analyser, join)
+{
+  test("kprobe:f { join(arg0) }", 0);
+  test("kprobe:f { printf(\"%s\", join(arg0)) }", 10);
+  test("kprobe:f { join() }", 1);
+  test("kprobe:f { $fmt = \"mystring\"; join($fmt) }", 10);
+  test("kprobe:f { @x = join(arg0) }", 1);
+  test("kprobe:f { $x = join(arg0) }", 1);
+}
+
 TEST(semantic_analyser, kprobe)
 {
   test("kprobe:f { 1 }", 0);