Allow definition of a load and unload NIF callback (#16)

Co-authored-by: Jonatan Kłosko <[email protected]>

Author: brodeuralexis

README.md

@@ -604,6 +604,40 @@ Attempting to decode STL containers making use of `std::pmr::polymorphic_allocat
 will result in the `std::pmr::get_default_resource()` memory resource being
 used.
 
+## Global Callbacks
+
+The Erlang NIF API allows the creation of global callbacks. This section
+describes how to leverage these callbacks while using Fine.
+
+### Load
+
+The NIF load callback is called by the ERTS when the NIFs are being loaded by
+`:erlang.load_nif/2`.  Fine allows customizing the behavior of the load callback
+using the `fine::Registration::register_load` function:
+
+```c++
+static std::unique_ptr<ThreadPool> s_pool;
+
+static auto load_registration = fine::Registration::register_load([](ErlNifEnv *env, void **priv_data, fine::Term load_info) {
+  const auto thread_count = fine::decode<std::uint64_t>(caller_env, load_info);
+  s_pool = std::make_unique<FixedThreadPool>(thread_count);
+});
+```
+
+### Unload
+
+The NIF unload callback is called by the ERTS when the NIFs are being unloaded
+from the runtime.  Fine allows customizing the behavior of the unload callback
+using the `fine::Registration::register_unload` function:
+
+```c++
+static std::unique_ptr<ThreadPool> s_pool;
+
+static auto unload_registration = fine::Registration::register_unload([](ErlNifEnv *env, void *priv_data) noexcept {
+  s_pool.stop();
+});
+```
+
 <!-- Docs -->
 
 ## Prior work

c_include/fine.hpp

@@ -5,6 +5,7 @@
 #include <cstdint>
 #include <cstdlib>
 #include <cstring>
+#include <functional>
 #include <map>
 #include <memory>
 #include <optional>
@@ -47,7 +48,10 @@ template <typename T, typename SFINAE = void> struct Encoder;
 
 namespace __private__ {
 std::vector<ErlNifFunc> &get_erl_nif_funcs();
-int load(ErlNifEnv *env, void **priv_data, ERL_NIF_TERM load_info);
+void init_atoms(ErlNifEnv *env);
+bool init_resources(ErlNifEnv *env);
+int load(ErlNifEnv *, void **, ERL_NIF_TERM) noexcept;
+void unload(ErlNifEnv *, void *) noexcept;
 } // namespace __private__
 
 // Definitions
@@ -97,8 +101,7 @@ class Atom {
   friend struct Decoder<Atom>;
   friend struct ::std::hash<Atom>;
 
-  friend int __private__::load(ErlNifEnv *env, void **priv_data,
-                               ERL_NIF_TERM load_info);
+  friend void __private__::init_atoms(ErlNifEnv *env);
 
   // We accumulate all globally defined atom objects and create the
   // terms upfront as part of init (called from the NIF load callback).
@@ -1187,6 +1190,12 @@ template <typename T> void raise(ErlNifEnv *env, const T &value) {
 // Mechanism for accumulating information via static object definitions.
 class Registration {
 public:
+  // A function compatible with the load callback of Erlang's NIFs.
+  using LoadCallback = std::function<void(ErlNifEnv *, void **, fine::Term)>;
+
+  // A function compatible with the unload callback of Erlang's NIFs.
+  using UnloadCallback = std::function<void(ErlNifEnv *, void *)>;
+
   template <typename T>
   static Registration register_resource(const char *name) {
     Registration::resources.push_back({&fine::ResourcePtr<T>::resource_type,
@@ -1200,6 +1209,26 @@ class Registration {
     return {};
   }
 
+  // Registers a load callback.
+  static Registration register_load(LoadCallback callback) {
+    if (erl_nif_load_callback) {
+      throw std::logic_error("load callback already registered");
+    }
+
+    Registration::erl_nif_load_callback = callback;
+    return {};
+  }
+
+  // Registers an unload callback.
+  static Registration register_unload(UnloadCallback callback) {
+    if (erl_nif_unload_callback) {
+      throw std::logic_error("unload callback already registered");
+    }
+
+    Registration::erl_nif_unload_callback = callback;
+    return {};
+  }
+
 private:
   static bool init_resources(ErlNifEnv *env) {
     for (const auto &[resource_type_ptr, name, dtor] :
@@ -1220,15 +1249,18 @@ class Registration {
   }
 
   friend std::vector<ErlNifFunc> &__private__::get_erl_nif_funcs();
+  friend int __private__::load(ErlNifEnv *, void **, ERL_NIF_TERM) noexcept;
+  friend void __private__::unload(ErlNifEnv *, void *) noexcept;
 
-  friend int __private__::load(ErlNifEnv *env, void **priv_data,
-                               ERL_NIF_TERM load_info);
+  friend bool __private__::init_resources(ErlNifEnv *env);
 
   inline static std::vector<std::tuple<ErlNifResourceType **, const char *,
                                        void (*)(ErlNifEnv *, void *)>>
       resources = {};
 
   inline static std::vector<ErlNifFunc> erl_nif_funcs = {};
+  inline static LoadCallback erl_nif_load_callback = {};
+  inline static UnloadCallback erl_nif_unload_callback = {};
 };
 
 // NIF definitions
@@ -1295,23 +1327,49 @@ constexpr unsigned int nif_arity(Ret (*)(Args...)) {
 }
 
 namespace __private__ {
+void init_atoms(ErlNifEnv *env) { fine::Atom::init_atoms(env); }
+
+bool init_resources(ErlNifEnv *env) {
+  return fine::Registration::init_resources(env);
+}
+
 inline std::vector<ErlNifFunc> &get_erl_nif_funcs() {
   return Registration::erl_nif_funcs;
 }
 
-inline int load(ErlNifEnv *env, void **, ERL_NIF_TERM) {
-  Atom::init_atoms(env);
+inline int load(ErlNifEnv *caller_env, void **priv_data,
+                ERL_NIF_TERM load_info) noexcept {
+  init_atoms(caller_env);
 
-  if (!Registration::init_resources(env)) {
+  if (!init_resources(caller_env)) {
+    return -1;
+  }
+
+  try {
+    if (fine::Registration::erl_nif_load_callback) {
+      std::invoke(fine::Registration::erl_nif_load_callback, caller_env,
+                  priv_data, load_info);
+    }
+  } catch (const std::exception &e) {
+    enif_fprintf(stderr, "unhandled exception: %s\n", e.what());
+    return -1;
+  } catch (...) {
+    enif_fprintf(stderr, "unhandled throw\n");
     return -1;
   }
 
   return 0;
 }
+
+inline void unload(ErlNifEnv *caller_env, void *priv_data) noexcept {
+  if (fine::Registration::erl_nif_unload_callback) {
+    std::invoke(fine::Registration::erl_nif_unload_callback, caller_env,
+                priv_data);
+  }
+}
 } // namespace __private__
 
 // Macros
-
 #define FINE_NIF(name, flags)                                                  \
   static ERL_NIF_TERM name##_nif(ErlNifEnv *env, int argc,                     \
                                  const ERL_NIF_TERM argv[]) {                  \
@@ -1345,16 +1403,16 @@ inline int load(ErlNifEnv *env, void **, ERL_NIF_TERM) {
     auto &nif_funcs = fine::__private__::get_erl_nif_funcs();                  \
     auto num_funcs = static_cast<int>(nif_funcs.size());                       \
     auto funcs = nif_funcs.data();                                             \
-    auto load = fine::__private__::load;                                       \
+                                                                               \
     static ErlNifEntry entry = {ERL_NIF_MAJOR_VERSION,                         \
                                 ERL_NIF_MINOR_VERSION,                         \
                                 name,                                          \
                                 num_funcs,                                     \
                                 funcs,                                         \
-                                load,                                          \
-                                NULL,                                          \
+                                fine::__private__::load,                       \
                                 NULL,                                          \
                                 NULL,                                          \
+                                fine::__private__::unload,                     \
                                 ERL_NIF_VM_VARIANT,                            \
                                 1,                                             \
                                 sizeof(ErlNifResourceTypeInit),                \

example/c_src/example.cpp

@@ -1,8 +1,6 @@
 #include <fine.hpp>
 
-int64_t add(ErlNifEnv *env, int64_t x, int64_t y) {
-  return x + y;
-}
+int64_t add(ErlNifEnv *env, int64_t x, int64_t y) { return x + y; }
 
 FINE_NIF(add, 0);
 FINE_INIT("Elixir.Example");

test/c_src/finest.cpp

@@ -462,6 +462,16 @@ std::uint64_t hash_atom(ErlNifEnv *, fine::Atom atom) noexcept {
 }
 FINE_NIF(hash_atom, 0);
 
+static bool s_loaded = false;
+
+static auto load_registration = fine::Registration::register_load(
+    [](ErlNifEnv *, void **, ERL_NIF_TERM) { s_loaded = true; });
+
+static auto unload_registration = fine::Registration::register_unload(
+    [](ErlNifEnv *, void *) noexcept { s_loaded = false; });
+
+bool is_loaded(ErlNifEnv *) { return s_loaded; }
+FINE_NIF(is_loaded, 0);
 } // namespace finest
 
 FINE_INIT("Elixir.Finest.NIF");

test/lib/finest/nif.ex

@@ -80,5 +80,7 @@ defmodule Finest.NIF do
   def hash_term(_term), do: err!()
   def hash_atom(_term), do: err!()
 
+  def is_loaded(), do: err!()
+
   defp err!(), do: :erlang.nif_error(:not_loaded)
 end

test/test/finest_test.exs

@@ -502,4 +502,10 @@ defmodule FinestTest do
       end
     end
   end
+
+  describe "callbacks" do
+    test "load" do
+      assert NIF.is_loaded()
+    end
+  end
 end