Support disabling the Erlang allocator

README.md

@@ -143,6 +143,16 @@ Once you have `exqlite` configured, you can use the `:key` option in the databas
 config :exqlite, key: "super-secret'
 ```
 
+### Disabling the Erlang allocator
+
+By default, we configure SQLite to use allocators provided by Erlang ([`enif_alloc`](https://www.erlang.org/doc/apps/erts/erl_nif.html#enif_alloc)). This allows SQLite's memory usage to be reported by `:erlang.memory()`. However, in some cases this may perform worse than using SQLite's default allocator, in which case you can disable the Erlang allocator:
+
+```elixir
+config :exqlite, disable_erlang_allocator: true
+```
+
+Note that this setting is read when the NIF is loaded and subsequent changes will have no effect unless you restart the BEAM VM.
+
 ## Usage
 
 The `Exqlite.Sqlite3` module usage is fairly straight forward.

c_src/sqlite3_nif.c

@@ -10,6 +10,8 @@
 #include <erl_nif.h>
 #include <sqlite3.h>
 
+static ERL_NIF_TERM am_false;
+static ERL_NIF_TERM am_true;
 static ERL_NIF_TERM am_ok;
 static ERL_NIF_TERM am_error;
 static ERL_NIF_TERM am_badarg;
@@ -41,6 +43,8 @@ static ERL_NIF_TERM am_update;
 static ERL_NIF_TERM am_invalid_pid;
 static ERL_NIF_TERM am_log;
 
+static int erlang_allocator_enabled = 0;
+
 static ErlNifResourceType* connection_type       = NULL;
 static ErlNifResourceType* statement_type        = NULL;
 static sqlite3_mem_methods default_alloc_methods = {0};
@@ -1435,9 +1439,8 @@ on_load(ErlNifEnv* env, void** priv, ERL_NIF_TERM info)
       exqlite_mem_shutdown,
       0};
 
-    sqlite3_config(SQLITE_CONFIG_GETMALLOC, &default_alloc_methods);
-    sqlite3_config(SQLITE_CONFIG_MALLOC, &methods);
-
+    am_true                                = enif_make_atom(env, "true");
+    am_false                               = enif_make_atom(env, "false");
     am_ok                                  = enif_make_atom(env, "ok");
     am_error                               = enif_make_atom(env, "error");
     am_badarg                              = enif_make_atom(env, "badarg");
@@ -1469,6 +1472,17 @@ on_load(ErlNifEnv* env, void** priv, ERL_NIF_TERM info)
     am_invalid_pid                         = enif_make_atom(env, "invalid_pid");
     am_log                                 = enif_make_atom(env, "log");
 
+    ERL_NIF_TERM disable_erlang_allocator;
+    if (!enif_get_map_value(env, info, enif_make_atom(env, "disable_erlang_allocator"), &disable_erlang_allocator)) {
+        return -1;
+    }
+    erlang_allocator_enabled = enif_is_identical(disable_erlang_allocator, am_false);
+
+    sqlite3_config(SQLITE_CONFIG_GETMALLOC, &default_alloc_methods);
+    if (erlang_allocator_enabled) {
+        sqlite3_config(SQLITE_CONFIG_MALLOC, &methods);
+    }
+
     connection_type = enif_open_resource_type(
       env,
       NULL,
@@ -2069,6 +2083,21 @@ exqlite_errstr(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
     return make_binary(env, msg, strlen(msg));
 }
 
+//
+// This is only used in tests to verify whether the Erlang allocator is being used.
+//
+ERL_NIF_TERM
+exqlite_erlang_allocator_enabled(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
+{
+    assert(env);
+
+    if (argc != 0) {
+        return enif_make_badarg(env);
+    }
+
+    return erlang_allocator_enabled ? am_true : am_false;
+}
+
 //
 // Most of our nif functions are going to be IO bounded
 //
@@ -2105,6 +2134,7 @@ static ErlNifFunc nif_funcs[] = {
   {"cancel", 1, exqlite_cancel, 0},
   {"errmsg", 1, exqlite_errmsg},
   {"errstr", 1, exqlite_errstr},
+  {"erlang_allocator_enabled", 0, exqlite_erlang_allocator_enabled},
 };
 
 ERL_NIF_INIT(Elixir.Exqlite.Sqlite3NIF, nif_funcs, on_load, NULL, on_upgrade, on_unload)

lib/exqlite/sqlite3_nif.ex

@@ -20,7 +20,15 @@ defmodule Exqlite.Sqlite3NIF do
 
   def load_nif() do
     path = :filename.join(:code.priv_dir(:exqlite), ~c"sqlite3_nif")
-    :erlang.load_nif(path, 0)
+
+    disable_erlang_allocator =
+      Application.get_env(:exqlite, :disable_erlang_allocator, false)
+
+    if not is_boolean(disable_erlang_allocator) do
+      raise ArgumentError, ":disable_erlang_allocator must be a boolean"
+    end
+
+    :erlang.load_nif(path, %{disable_erlang_allocator: disable_erlang_allocator})
   end
 
   @spec open(String.t(), integer()) :: {:ok, db()} | {:error, reason()}
@@ -117,5 +125,9 @@ defmodule Exqlite.Sqlite3NIF do
   @spec errstr(integer) :: String.t()
   def errstr(_rc), do: :erlang.nif_error(:not_loaded)
 
+  @doc false
+  @spec erlang_allocator_enabled() :: boolean()
+  def erlang_allocator_enabled(), do: :erlang.nif_error(:not_loaded)
+
   # add statement inspection tooling https://sqlite.org/c3ref/expanded_sql.html
 end

test/exqlite/allocator_test.exs

@@ -0,0 +1,8 @@
+defmodule Exqlite.AllocatorTest do
+  use ExUnit.Case
+
+  test "erlang allocator is used by default" do
+    assert Application.get_env(:exqlite, :disable_erlang_allocator) == nil
+    assert Exqlite.Sqlite3NIF.erlang_allocator_enabled()
+  end
+end