Add global vars which can be set from userspace

This adds the ability to create const volatile global variables in bpf and allow them to be updated in userspace after program opening but before program loading. This is useful for things like AOT, where we don't know know the number of CPUs on the machine where the bpftrace program will run at the time of codegen.
delphix · Jul 13, 2024 · 36c4b72 · 36c4b72
1 parent 6cc2798
commit 36c4b72
Show file tree

Hide file tree

Showing 38 changed files with 932 additions and 732 deletions.
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -29,6 +29,7 @@ add_library(runtime
   disasm.cpp
   dwarf_parser.cpp
   format_string.cpp
+  globalvars.cpp
   log.cpp
   mapkey.cpp
   output.cpp

diff --git a/src/aot/aot.cpp b/src/aot/aot.cpp
@@ -341,7 +341,7 @@ int load(BPFtrace &bpftrace, const std::string &in)
 
   bpftrace.bytecode_ = BpfBytecode(btaot_section + hdr->elf_off,
                                    hdr->elf_len,
-                                   bpftrace.config_);
+                                   bpftrace);
   if (err)
     goto out;
 

diff --git a/src/aot/aot_main.cpp b/src/aot/aot_main.cpp
@@ -117,15 +117,6 @@ int main(int argc, char* argv[])
 
   BPFtrace bpftrace(std::move(output));
 
-  // TODO: remove this once we move to libbpf or move to open-coded iterators
-  auto num_cpus = bpftrace.get_num_possible_cpus();
-  if (num_cpus > 1024) {
-    LOG(WARNING) << "Detected " << num_cpus
-                 << " cpus. For ahead-of-time compilation there is a max of "
-                    "1024 cpus so there may be incorrect data for 'count' "
-                    "and 'sum' aggregations.";
-  }
-
   int err = aot::load(bpftrace, argv[0]);
   if (err) {
     LOG(ERROR) << "Failed to load AOT script";

diff --git a/src/ast/dibuilderbpf.cpp b/src/ast/dibuilderbpf.cpp
@@ -240,5 +240,12 @@ DIGlobalVariableExpression *DIBuilderBPF::createMapEntry(
       file, name, "global", file, 0, map_entry_type, false);
 }
 
+DIGlobalVariableExpression *DIBuilderBPF::createGlobalInt64(
+    std::string_view name)
+{
+  return createGlobalVariableExpression(
+      file, name, "global", file, 0, getInt64Ty(), false);
+}
+
 } // namespace ast
 } // namespace bpftrace
diff --git a/src/ast/dibuilderbpf.h b/src/ast/dibuilderbpf.h
@@ -44,6 +44,7 @@ class DIBuilderBPF : public DIBuilder {
                                              uint64_t max_entries,
                                              const MapKey &key,
                                              const SizedType &value_type);
+  DIGlobalVariableExpression *createGlobalInt64(std::string_view name);
 
   DIFile *file = nullptr;
 

diff --git a/src/ast/irbuilderbpf.cpp b/src/ast/irbuilderbpf.cpp
@@ -12,6 +12,7 @@
 #include "ast/codegen_helper.h"
 #include "bpfmap.h"
 #include "bpftrace.h"
+#include "globalvars.h"
 #include "log.h"
 #include "utils.h"
 
@@ -576,8 +577,7 @@ Value *IRBuilderBPF::CreatePerCpuMapAggElems(Value *ctx,
                                              Map &map,
                                              Value *key,
                                              const SizedType &type,
-                                             const location &loc,
-                                             bool is_aot)
+                                             const location &loc)
 {
   /*
    * int ret = 0;
@@ -604,10 +604,6 @@ Value *IRBuilderBPF::CreatePerCpuMapAggElems(Value *ctx,
   AllocaInst *ret = CreateAllocaBPF(getInt64Ty(), "ret");
   AllocaInst *i = CreateAllocaBPF(getInt32Ty(), "i");
 
-  // Set a large upper bound if we don't know the number of cpus
-  // when generating the instructions
-  int nr_cpus = is_aot ? 1024 : bpftrace_.get_num_possible_cpus();
-
   CreateStore(getInt32(0), i);
   CreateStore(getInt64(0), ret);
 
@@ -623,11 +619,12 @@ Value *IRBuilderBPF::CreatePerCpuMapAggElems(Value *ctx,
                                              parent);
   CreateBr(while_cond);
   SetInsertPoint(while_cond);
-  // TODO: after full libbpf support update the number of cpus from userspace
-  // dynamically using a global mutable variable and the skeleton
+
   auto *cond = CreateICmp(CmpInst::ICMP_ULT,
                           CreateLoad(getInt32Ty(), i),
-                          getInt32(nr_cpus),
+                          CreateLoad(getInt32Ty(),
+                                     module_.getGlobalVariable(
+                                         bpftrace::globalvars::NUM_CPUS)),
                           "num_cpu.cmp");
   CreateCondBr(cond, while_body, while_end);
 

diff --git a/src/ast/irbuilderbpf.h b/src/ast/irbuilderbpf.h
@@ -91,8 +91,7 @@ class IRBuilderBPF : public IRBuilder<> {
                                  Map &map,
                                  Value *key,
                                  const SizedType &type,
-                                 const location &loc,
-                                 bool is_aot);
+                                 const location &loc);
   void CreateMapUpdateElem(Value *ctx,
                            const std::string &map_ident,
                            Value *key,

diff --git a/src/ast/passes/codegen_llvm.cpp b/src/ast/passes/codegen_llvm.cpp
@@ -39,6 +39,7 @@
 #include "ast/signal_bt.h"
 #include "bpfmap.h"
 #include "collect_nodes.h"
+#include "globalvars.h"
 #include "log.h"
 #include "tracepoint_format_parser.h"
 #include "types.h"
@@ -47,14 +48,13 @@
 namespace bpftrace {
 namespace ast {
 
-CodegenLLVM::CodegenLLVM(Node *root, BPFtrace &bpftrace, bool is_aot)
-    : CodegenLLVM(root, bpftrace, is_aot, std::make_unique<USDTHelper>())
+CodegenLLVM::CodegenLLVM(Node *root, BPFtrace &bpftrace)
+    : CodegenLLVM(root, bpftrace, std::make_unique<USDTHelper>())
 {
 }
 
 CodegenLLVM::CodegenLLVM(Node *root,
                          BPFtrace &bpftrace,
-                         bool is_aot,
                          std::unique_ptr<USDTHelper> usdt_helper)
     : root_(root),
       bpftrace_(bpftrace),
@@ -63,8 +63,7 @@ CodegenLLVM::CodegenLLVM(Node *root,
       module_(std::make_unique<Module>("bpftrace", *context_)),
       async_ids_(AsyncIds()),
       b_(*context_, *module_, bpftrace, async_ids_),
-      debug_(*module_),
-      is_aot_(is_aot)
+      debug_(*module_)
 {
   llvm::InitializeAllTargets();
   llvm::InitializeAllTargetMCs();
@@ -1399,17 +1398,16 @@ void CodegenLLVM::visit(Map &map)
     if (val_type.IsAvgTy()) {
       AllocaInst *count_key = getHistMapKey(map, b_.getInt64(0));
       Value *count_val = b_.CreatePerCpuMapAggElems(
-          ctx_, map, count_key, val_type, map.loc, is_aot_);
+          ctx_, map, count_key, val_type, map.loc);
       b_.CreateLifetimeEnd(count_key);
 
       AllocaInst *total_key = getHistMapKey(map, b_.getInt64(1));
       Value *total_val = b_.CreatePerCpuMapAggElems(
-          ctx_, map, total_key, val_type, map.loc, is_aot_);
+          ctx_, map, total_key, val_type, map.loc);
       b_.CreateLifetimeEnd(total_key);
       value = b_.CreateUDiv(total_val, count_val);
     } else {
-      value = b_.CreatePerCpuMapAggElems(
-          ctx_, map, key, val_type, map.loc, is_aot_);
+      value = b_.CreatePerCpuMapAggElems(ctx_, map, key, val_type, map.loc);
     }
   } else {
     value = b_.CreateMapLookupElem(ctx_, map, key, map.loc);
@@ -3480,6 +3478,8 @@ void CodegenLLVM::generate_ir()
 {
   assert(state_ == State::INIT);
   generate_maps(bpftrace_.resources);
+  generate_global_vars(bpftrace_.resources);
+
   auto scoped_del = accept(root_);
   debug_.finalize();
   state_ = State::IR;
@@ -3652,6 +3652,20 @@ void CodegenLLVM::generate_maps(const RequiredResources &resources)
                       CreateInt(loss_cnt_val_size));
 }
 
+void CodegenLLVM::generate_global_vars(const RequiredResources &resources)
+{
+  for (const auto &name : resources.needed_global_vars) {
+    auto var = llvm::dyn_cast<GlobalVariable>(
+        module_->getOrInsertGlobal(name, b_.getInt64Ty()));
+    var->setInitializer(b_.getInt64(1));
+    var->setConstant(true);
+    var->setSection(bpftrace::globalvars::SECTION_NAME);
+    var->setExternallyInitialized(true);
+    var->setDSOLocal(true);
+    var->addDebugInfo(debug_.createGlobalInt64(name));
+  }
+}
+
 void CodegenLLVM::emit_elf(const std::string &filename)
 {
   assert(state_ == State::OPT);
@@ -3759,7 +3773,7 @@ BpfBytecode CodegenLLVM::emit(void)
   assert(!output.empty());
 
   state_ = State::DONE;
-  return BpfBytecode(output.data(), output.size(), bpftrace_.config_);
+  return BpfBytecode(output.data(), output.size(), bpftrace_);
 }
 
 BpfBytecode CodegenLLVM::compile(void)
@@ -4231,8 +4245,7 @@ Function *CodegenLLVM::createForEachMapCallback(const For &f, llvm::Type *ctx_t)
                                              "lookup_key");
     b_.CreateStore(key, key_ptr);
 
-    val = b_.CreatePerCpuMapAggElems(
-        ctx_, map, key_ptr, map_val_type, map.loc, is_aot_);
+    val = b_.CreatePerCpuMapAggElems(ctx_, map, key_ptr, map_val_type, map.loc);
   } else if (!onStack(val_type)) {
     val = b_.CreateLoad(b_.GetType(val_type), val, "val");
   }

diff --git a/src/ast/passes/codegen_llvm.h b/src/ast/passes/codegen_llvm.h
@@ -29,10 +29,9 @@ using CallArgs = std::vector<std::tuple<FormatString, std::vector<Field>>>;
 
 class CodegenLLVM : public Visitor {
 public:
-  explicit CodegenLLVM(Node *root, BPFtrace &bpftrace, bool is_aot = false);
+  explicit CodegenLLVM(Node *root, BPFtrace &bpftrace);
   explicit CodegenLLVM(Node *root,
                        BPFtrace &bpftrace,
-                       bool is_aot,
                        std::unique_ptr<USDTHelper> usdt_helper);
 
   void visit(Integer &integer) override;
@@ -99,6 +98,7 @@ class CodegenLLVM : public Visitor {
   libbpf::bpf_map_type get_map_type(const SizedType &val_type,
                                     const MapKey &key);
   void generate_maps(const RequiredResources &resources);
+  void generate_global_vars(const RequiredResources &resources);
   void optimize(void);
   bool verify(void);
   BpfBytecode emit(void);
@@ -252,7 +252,6 @@ class CodegenLLVM : public Visitor {
   IRBuilderBPF b_;
 
   DIBuilderBPF debug_;
-  bool is_aot_;
 
   const DataLayout &datalayout() const
   {

diff --git a/src/ast/passes/resource_analyser.cpp b/src/ast/passes/resource_analyser.cpp
@@ -1,6 +1,7 @@
 #include "resource_analyser.h"
 
 #include "bpftrace.h"
+#include "globalvars.h"
 #include "log.h"
 #include "struct.h"
 
@@ -115,6 +116,9 @@ void ResourceAnalyser::visit(Call &call)
                                         : " ";
     resources_.join_args.push_back(delim);
     resources_.needs_join_map = true;
+  } else if (call.func == "count" || call.func == "sum" || call.func == "min" ||
+             call.func == "max" || call.func == "avg") {
+    resources_.needed_global_vars.insert(bpftrace::globalvars::NUM_CPUS);
   } else if (call.func == "hist") {
     auto &map_info = resources_.maps_info[call.map->ident];
     int bits = static_cast<Integer *>(call.vargs->at(1))->n;

diff --git a/src/bpfbytecode.cpp b/src/bpfbytecode.cpp
@@ -1,6 +1,7 @@
 #include "bpfbytecode.h"
 
 #include "bpftrace.h"
+#include "globalvars.h"
 #include "log.h"
 #include "utils.h"
 
@@ -11,8 +12,8 @@
 
 namespace bpftrace {
 
-BpfBytecode::BpfBytecode(const void *elf, size_t elf_size, const Config &config)
-    : log_size_(config.get(ConfigKeyInt::log_size))
+BpfBytecode::BpfBytecode(const void *elf, size_t elf_size, BPFtrace &bpftrace)
+    : log_size_(bpftrace.config_.get(ConfigKeyInt::log_size))
 {
   int log_level = 0;
   // In debug mode, show full verifier log.
@@ -31,12 +32,33 @@ BpfBytecode::BpfBytecode(const void *elf, size_t elf_size, const Config &config)
   if (!bpf_object_)
     LOG(BUG) << "The produced ELF is not a valid BPF object";
 
+  struct bpf_map *global_vars_map = nullptr;
+  bool needs_global_vars = !bpftrace.resources.needed_global_vars.empty();
+
   // Discover maps
   struct bpf_map *m;
   bpf_map__for_each (m, bpf_object_.get()) {
+    if (needs_global_vars) {
+      std::string_view name = bpf_map__name(m);
+      // there are some random chars in the beginning of the map name
+      if (name.npos != name.find(globalvars::SECTION_NAME)) {
+        global_vars_map = m;
+        continue;
+      }
+    }
     maps_.emplace(bpftrace_map_name(bpf_map__name(m)), m);
   }
 
+  if (needs_global_vars) {
+    if (!global_vars_map) {
+      LOG(BUG) << "No map found for " << globalvars::SECTION_NAME
+               << " which is needed to set global variables";
+    }
+    globalvars::update_global_vars(bpf_object_.get(),
+                                   global_vars_map,
+                                   bpftrace);
+  }
+
   // Discover programs
   struct bpf_program *p;
   bpf_object__for_each_program (p, bpf_object_.get()) {

diff --git a/src/bpfbytecode.h b/src/bpfbytecode.h
@@ -23,7 +23,7 @@ class BpfBytecode {
   BpfBytecode()
   {
   }
-  BpfBytecode(const void *elf, size_t elf_size, const Config &config);
+  BpfBytecode(const void *elf, size_t elf_size, BPFtrace &bpftrace);
 
   BpfBytecode(const BpfBytecode &) = delete;
   BpfBytecode &operator=(const BpfBytecode &) = delete;

diff --git a/src/bpftrace.cpp b/src/bpftrace.cpp
@@ -2180,11 +2180,6 @@ struct bcc_symbol_option &BPFtrace::get_symbol_opts()
   return symopts;
 }
 
-int BPFtrace::get_num_possible_cpus() const
-{
-  return libbpf_num_possible_cpus();
-}
-
 /*
  * This prevents an ABBA deadlock when attaching to spin lock internal
  * functions e.g. "kfunc:queued_spin_lock_slowpath".

diff --git a/src/bpftrace.h b/src/bpftrace.h
@@ -140,7 +140,6 @@ class BPFtrace {
   std::optional<int64_t> get_int_literal(const ast::Expression *expr) const;
   std::optional<std::string> get_watchpoint_binary_path() const;
   virtual bool is_traceable_func(const std::string &func_name) const;
-  virtual int get_num_possible_cpus() const;
   virtual std::unordered_set<std::string> get_func_modules(
       const std::string &func_name) const;
   int create_pcaps(void);