Attempt at adapting KuzuDB's Rust crate into Elixir via NIF, any tips/thoughts?

I got a very basic NIF working for KuzuDB: GitHub - bgoosmanviz/kuzu_nif: Adapting KuzuDB's Rust crate to Elixir using Rustler.

KuzuDB is an embedded graph database, in the same sense sqlite is an embedded sql database.

Posting here in case anyone is interested in helping flesh it out. I don’t know Rust very well.

9 Likes

Hi, great job with starting the project, @bgoosman :slight_smile:

I’m wondering what your plans are in regard to the repo’s future; thanks!

I think integrating it with the Explorer library would be nice. i.e. pass dataframes from Elixir to/from the NIF, to be loaded into kuzudb, or fetched from kuzudb. I tried without much success.

WDYT?

Thanks for reminding me about this. I just updated the kuzu crate from 0.6.0 to 0.7.0 and improved the Elixir unit test.

We can do some pretty useful things with just run_query, like create a kuzudb using COPY FROM commands.

1 Like

Besides that, feature parity with the API.

Thanks for updating your repo @bgoosman . Yes passing dataframes would be a good idea.

I’m wondering how the BEAM process management maps onto Rust threads, i.e. can we achieve proper query concurrency using KuzuDB via Elixir?

Are you using KuzuDB in production with your integration? What is the use case, without getting into any trade secrets :wink:

Not sure about BEAM & Rust threads. Here’s KuzuDB’s concurrency model Connections & Concurrency | Kùzu.

You got me thinking about the connection to the db. Right now run_query will create a new connection every time it’s called, which doesn’t sound ideal. I wonder how to manage a persistent connection.

We’re a graph viz company, so having an embeddable graph engine like KuzuDB would be an excellent improvement over a server based graph engine like Neo4j.

That’s my goal as well - being able to do in-mem Cypher queries.

Maybe @filmor @hansihe or @evnu could point us in the right direction in terms of how the mapping between BEAM processes and threads in Rust works.

:wave:

I was curious so I took a peek, and I think this function might benefit from being scheduled on a dirty scheduler: kuzu_nif/native/kuzu_ex/src/lib.rs at f94c2e243983650db26c4b19bb0f73835fd52ec5 · bgoosmanviz/kuzu_nif · GitHub and re: “mapping processes to threads”, I think usually libraries for wrappers around embedded databases create a NIF resource per database connection and then Elixir / Erlang manage them however they want, no 3rd-party threads involved. Looking at kuzu_nif/native/kuzu_ex/src/lib.rs at f94c2e243983650db26c4b19bb0f73835fd52ec5 · bgoosmanviz/kuzu_nif · GitHub, it seems like the API might be similar to DuckDB wrappers, since they also have a notion of “databases” and “connections” (SQLite only has “databases”). Two example DuckDB wrappers in Elixir:

So I think the API could be:

# NIF resource, DB is closed when the underlying reference refcount goes to 0
# can be a Dirty IO NIF
db = Kuzu.open() 

conns =
  for _ <- 1..:erlang.system_info(:dirty_io_schedulers) do
    # also NIF resources, since Kuzu is C++
    # they probably reference and keep DB alive too,
    # so it's easy to just wrap them info resources individually.
    # When none of them are referenced anymore, Erlang would "garbage collect"
    # them and call NIF resource destructors, which would destroy Kuzu connections.
    # once Kuzu connections are destroyed, Kuzu would probably close DB as well
    # since it would no longer be referenced either.
    # can be a Dirty CPU NIF or a plain NIF, depending on what Kuzu does here
    Kuzu.connect(db)
  end

{:ok, pool} = GenServer.start_link(SomePool, conns)
{ref, conn} = GenServer.call(pool, :checkout)

result =
  try do
    # since it might trigger disk access, it would be safer if it was a Dirty IO NIF
    Kuzu.query(conn, query)
    # or it can be a Kuzu.prepare followed by a Kuzu.execute
    # this would allow statement caching and query params
    # Kuzu.prepare can be a Dirty CPU or a plain NIF
    # Kuzu.execute would be a Dirty IO NIF for the same reasons as Kuzu.query
  after
    GenServer.cast(pool, {:checkin, ref, conn})
  end

# can probably be a plain NIF
Kuzu.to_arrow(result)

I think this would be useful in the upcoming AoC :slight_smile:

1 Like

NIFs are executed by regular scheduled threads of the VM. See [the Erlang docs] (erl_nif — erts v15.1.2) for the big fat warning on long-running work, special care should be taken when a call takes more time than about 1ms. I usually just scheduled NIFs on an async scheduler (dirty CPU mostly) when the NIF was supposed to work for some longer time.

Note that Rustler doesn’t implement all of erl_nif, for example creating threads in the VM is not implemented as rust threads can be spawned if necessary.

3 Likes

Very good point!

I think if you go with the route of handling concurrency on rust side, using NIFs is a clear mistake from the start, they are explicitly designed in the way for concurrency to be handled from elixir side.

One could argue that usage of dirty schedulers makes it possible, however I think this is a misuse of the tool, as dirty schedulers were designed primarily to have some tolerance when it comes to NIFs that are not behaving nicely.

I think in such cases if handling of concurrency on rust side cannot be avoided, it would be smarter to have 2 separate programs communicating between each-other using Port, that will give a better tolerance and flexibility.

Not sure if that is correct, as Erlang has had a function to create threads from a NIF for a long time. I think it is just another tool in the toolbox, and it is the responsibility of the developer to pick the right mechanism.

Might be true, but I think the description of NIFs makes it pretty clear about their use-case:

NIFs are a simpler and more efficient way of calling C-code than using port drivers. NIFs are most suitable for synchronous functions, such as foo and bar in the example, that do some relatively short calculations without side effects and return the result.

I am personally not a fan on relying on third-party code to not crash the erlang VM, no matter how well it was written, so I would be much more happier to have them isolated where it’s possible. This is especially true when we talk about long-running things.

Excluding actual lower-level (kernel) crashes on physically faulty devices, using Rust and making sure certain functions are never used in the code (like Option’s unwrap or expect) is 100% safe. This has been proven in kernel development (though in the interest of full objectivity, only a subset of Rust is used there) as well.

The real issue is making sure you respond within 1ms or less. Parallelism is not scary if your native code is written well (famous last words, right? “Just don’t write bugs”. :003: )

The thread creation stuff in enif only exists to provide a platform-independent API for threading in C. They don’t have (to my knowledge) any deeper “integration” with the BEAM.

One thing to keep in mind is that not all work has to happen in the actual NIF call. You can easily spawn a thread (or a pool of threads) and pass the PID of the caller that you then use to send the result of the query back to. This is what I use in my XML processing faster_xml/rust_src/faster_xml_nif/src/lib.rs at master · filmor/faster_xml · GitHub and it works fine in production, even without an actual thread pool.

So, simple suggestion to get started: Have a Database stored in a resource object, and on every query, spawn a thread that creates a Connection and sends the result back.

Later you can refine this to use a thread pool (with some retry logic on the Elixir side).

That way, the NIF will be quick enough that you don’t need to worry about scheduling and you get proper concurrent processing.

5 Likes

This is fantastic, thank you, I was looking for a ready idiom about something very similar and your snippet puts me on the right path.

Did you happen to have saved an entire tokio future that way and later resume it?

After 3 days of pain, I successfully added rustler_precompiled to kuzu_nif and published to hex.pm. Users of kuzu_nif on supported targets (see below) no longer need to compile the Rust crate.

Why? This saves us a few minutes on every build. For me, I’m using kuzu_nif in a Phoenix Framework project. Having a precompiled version of kuzu_nif means 1) I don’t need to install a Rust compiler in my Phoenix Framework dockerfile and 2) I save precious time compiling my Phoenix app.

The supported targets are indicated in the packaged release files. TLDR; Apple arm/x86 and Debian Linux arm/x64.

  • libkuzu_ex-v0.7.0-nif-2.17-aarch64-apple-darwin.so.tar.gz
  • libkuzu_ex-v0.7.0-nif-2.17-aarch64-unknown-linux-gnu.so.tar.gz
  • libkuzu_ex-v0.7.0-nif-2.17-x86_64-apple-darwin.so.tar.gz
  • libkuzu_ex-v0.7.0-nif-2.17-x86_64-unknown-linux-gnu.so.tar.gz

Figuring out how to successfully compile on aarch64 linux was the hardest part. In the end, I had to compile a custom Docker image just for building aarch64-unknown-linux-gnu.

Enjoy! :tada:

6 Likes

I released a new library called kuzu_pyport_ex. I realized I didn’t need the performance gains from using the Rust NIF, so I made kuzu_pyport_ex, which uses the Python lib via a Port, which is safer and simpler than using a NIF.

I also made a change to kuzu_nif to use the DirtyCpu scheduler, which I think was necessary to avoid instability during long running queries.

1 Like