Merge pull request #13219 from xokdvium/eval-profiler

libexpr: Add `EvalProfiler` and use it for `FunctionCallTrace`

src/libexpr/eval-profiler.cc

@@ -0,0 +1,48 @@
+#include "nix/expr/eval-profiler.hh"
+#include "nix/expr/nixexpr.hh"
+
+namespace nix {
+
+void EvalProfiler::preFunctionCallHook(
+    const EvalState & state, const Value & v, std::span<Value *> args, const PosIdx pos)
+{
+}
+
+void EvalProfiler::postFunctionCallHook(
+    const EvalState & state, const Value & v, std::span<Value *> args, const PosIdx pos)
+{
+}
+
+void MultiEvalProfiler::preFunctionCallHook(
+    const EvalState & state, const Value & v, std::span<Value *> args, const PosIdx pos)
+{
+    for (auto & profiler : profilers) {
+        if (profiler->getNeededHooks().test(Hook::preFunctionCall))
+            profiler->preFunctionCallHook(state, v, args, pos);
+    }
+}
+
+void MultiEvalProfiler::postFunctionCallHook(
+    const EvalState & state, const Value & v, std::span<Value *> args, const PosIdx pos)
+{
+    for (auto & profiler : profilers) {
+        if (profiler->getNeededHooks().test(Hook::postFunctionCall))
+            profiler->postFunctionCallHook(state, v, args, pos);
+    }
+}
+
+EvalProfiler::Hooks MultiEvalProfiler::getNeededHooksImpl() const
+{
+    Hooks hooks;
+    for (auto & p : profilers)
+        hooks |= p->getNeededHooks();
+    return hooks;
+}
+
+void MultiEvalProfiler::addProfiler(ref<EvalProfiler> profiler)
+{
+    profilers.push_back(profiler);
+    invalidateNeededHooks();
+}
+
+}

src/libexpr/eval.cc

@@ -372,6 +372,10 @@ EvalState::EvalState(
     );
 
     createBaseEnv(settings);
+
+    /* Register function call tracer. */
+    if (settings.traceFunctionCalls)
+        profiler.addProfiler(make_ref<FunctionCallTrace>());
 }
 
 
@@ -1526,9 +1530,14 @@ void EvalState::callFunction(Value & fun, std::span<Value *> args, Value & vRes,
 {
     auto _level = addCallDepth(pos);
 
-    auto trace = settings.traceFunctionCalls
-        ? std::make_unique<FunctionCallTrace>(positions[pos])
-        : nullptr;
+    auto neededHooks = profiler.getNeededHooks();
+    if (neededHooks.test(EvalProfiler::preFunctionCall)) [[unlikely]]
+        profiler.preFunctionCallHook(*this, fun, args, pos);
+
+    Finally traceExit_{[&](){
+        if (profiler.getNeededHooks().test(EvalProfiler::postFunctionCall)) [[unlikely]]
+            profiler.postFunctionCallHook(*this, fun, args, pos);
+    }};
 
     forceValue(fun, pos);
 

src/libexpr/function-trace.cc

@@ -3,16 +3,20 @@
 
 namespace nix {
 
-FunctionCallTrace::FunctionCallTrace(const Pos & pos) : pos(pos) {
+void FunctionCallTrace::preFunctionCallHook(
+    const EvalState & state, const Value & v, std::span<Value *> args, const PosIdx pos)
+{
     auto duration = std::chrono::high_resolution_clock::now().time_since_epoch();
     auto ns = std::chrono::duration_cast<std::chrono::nanoseconds>(duration);
-    printMsg(lvlInfo, "function-trace entered %1% at %2%", pos, ns.count());
+    printMsg(lvlInfo, "function-trace entered %1% at %2%", state.positions[pos], ns.count());
 }
 
-FunctionCallTrace::~FunctionCallTrace() {
+void FunctionCallTrace::postFunctionCallHook(
+    const EvalState & state, const Value & v, std::span<Value *> args, const PosIdx pos)
+{
     auto duration = std::chrono::high_resolution_clock::now().time_since_epoch();
     auto ns = std::chrono::duration_cast<std::chrono::nanoseconds>(duration);
-    printMsg(lvlInfo, "function-trace exited %1% at %2%", pos, ns.count());
+    printMsg(lvlInfo, "function-trace exited %1% at %2%", state.positions[pos], ns.count());
 }
 
 }

src/libexpr/include/nix/expr/eval-profiler.hh

@@ -0,0 +1,113 @@
+#pragma once
+/**
+ * @file
+ *
+ * Evaluation profiler interface definitions and builtin implementations.
+ */
+
+#include "nix/util/ref.hh"
+
+#include <vector>
+#include <span>
+#include <bitset>
+#include <optional>
+
+namespace nix {
+
+class EvalState;
+class PosIdx;
+struct Value;
+
+class EvalProfiler
+{
+public:
+    enum Hook {
+        preFunctionCall,
+        postFunctionCall,
+    };
+
+    static constexpr std::size_t numHooks = Hook::postFunctionCall + 1;
+    using Hooks = std::bitset<numHooks>;
+
+private:
+    std::optional<Hooks> neededHooks;
+
+protected:
+    /** Invalidate the cached neededHooks. */
+    void invalidateNeededHooks()
+    {
+        neededHooks = std::nullopt;
+    }
+
+    /**
+     * Get which hooks need to be called.
+     *
+     * This is the actual implementation which has to be defined by subclasses.
+     * Public API goes through the needsHooks, which is a
+     * non-virtual interface (NVI) which caches the return value.
+     */
+    virtual Hooks getNeededHooksImpl() const
+    {
+        return Hooks{};
+    }
+
+public:
+    /**
+     * Hook called in the EvalState::callFunction preamble.
+     * Gets called only if (getNeededHooks().test(Hook::preFunctionCall)) is true.
+     *
+     * @param state Evaluator state.
+     * @param v Function being invoked.
+     * @param args Function arguments.
+     * @param pos Function position.
+     */
+    virtual void
+    preFunctionCallHook(const EvalState & state, const Value & v, std::span<Value *> args, const PosIdx pos);
+
+    /**
+     * Hook called on EvalState::callFunction exit.
+     * Gets called only if (getNeededHooks().test(Hook::postFunctionCall)) is true.
+     *
+     * @param state Evaluator state.
+     * @param v Function being invoked.
+     * @param args Function arguments.
+     * @param pos Function position.
+     */
+    virtual void
+    postFunctionCallHook(const EvalState & state, const Value & v, std::span<Value *> args, const PosIdx pos);
+
+    virtual ~EvalProfiler() = default;
+
+    /**
+     * Get which hooks need to be invoked for this EvalProfiler instance.
+     */
+    Hooks getNeededHooks()
+    {
+        if (neededHooks.has_value())
+            return *neededHooks;
+        return *(neededHooks = getNeededHooksImpl());
+    }
+};
+
+/**
+ * Profiler that invokes multiple profilers at once.
+ */
+class MultiEvalProfiler : public EvalProfiler
+{
+    std::vector<ref<EvalProfiler>> profilers;
+
+    [[gnu::noinline]] Hooks getNeededHooksImpl() const override;
+
+public:
+    MultiEvalProfiler() = default;
+
+    /** Register a profiler instance. */
+    void addProfiler(ref<EvalProfiler> profiler);
+
+    [[gnu::noinline]] void
+    preFunctionCallHook(const EvalState & state, const Value & v, std::span<Value *> args, const PosIdx pos) override;
+    [[gnu::noinline]] void
+    postFunctionCallHook(const EvalState & state, const Value & v, std::span<Value *> args, const PosIdx pos) override;
+};
+
+}

src/libexpr/include/nix/expr/eval.hh

@@ -3,6 +3,7 @@
 
 #include "nix/expr/attr-set.hh"
 #include "nix/expr/eval-error.hh"
+#include "nix/expr/eval-profiler.hh"
 #include "nix/util/types.hh"
 #include "nix/expr/value.hh"
 #include "nix/expr/nixexpr.hh"
@@ -903,6 +904,9 @@ private:
     typedef std::map<ExprLambda *, size_t> FunctionCalls;
     FunctionCalls functionCalls;
 
+    /** Evaluation/call profiler. */
+    MultiEvalProfiler profiler;
+
     void incrFunctionCall(ExprLambda * fun);
 
     typedef std::map<PosIdx, size_t> AttrSelects;

src/libexpr/include/nix/expr/function-trace.hh

@@ -2,15 +2,24 @@
 ///@file
 
 #include "nix/expr/eval.hh"
-
-#include <chrono>
+#include "nix/expr/eval-profiler.hh"
 
 namespace nix {
 
-struct FunctionCallTrace
+class FunctionCallTrace : public EvalProfiler
 {
-    const Pos pos;
-    FunctionCallTrace(const Pos & pos);
-    ~FunctionCallTrace();
+    Hooks getNeededHooksImpl() const override
+    {
+        return Hooks().set(preFunctionCall).set(postFunctionCall);
+    }
+
+public:
+    FunctionCallTrace() = default;
+
+    [[gnu::noinline]] void
+    preFunctionCallHook(const EvalState & state, const Value & v, std::span<Value *> args, const PosIdx pos) override;
+    [[gnu::noinline]] void
+    postFunctionCallHook(const EvalState & state, const Value & v, std::span<Value *> args, const PosIdx pos) override;
 };
+
 }

src/libexpr/include/nix/expr/meson.build

@@ -14,6 +14,7 @@ headers = [config_pub_h] + files(
   'eval-error.hh',
   'eval-gc.hh',
   'eval-inline.hh',
+  'eval-profiler.hh',
   'eval-settings.hh',
   'eval.hh',
   'function-trace.hh',

src/libexpr/meson.build

@@ -140,6 +140,7 @@ sources = files(
   'eval-cache.cc',
   'eval-error.cc',
   'eval-gc.cc',
+  'eval-profiler.cc',
   'eval-settings.cc',
   'eval.cc',
   'function-trace.cc',