savannah-1

Supercharging Elixir with Rust

I’ve been writing quite a bit of Elixir over the past 6 months. There’s so much about the language that I really enjoy, it feels disingenuous of me to be poking at its weaknesses. But if we’re being honest, there are plenty of people out there singing the praises of Elixir and the BEAM in general. So let’s talk about this reality: if you are going to be doing a lot of computation, the BEAM might not be where you want to do that.

A while back I decided to write a solution to Project Euler: Problem 50 in Elixir. It’s a fun problem and I recommend it, but the relevant bit here is that it required me to generate all of the prime numbers less than 1 million. Now, like Elixir, I’m not particularly good at the maths and didn’t feel especially equipped or inclined to implement the Sieve of Eratosthenes myself, so I just went out in search of code snippets that I thought might do the job.

I started with Krzystof Wende’s prime number generator:

def is_prime(x), do: (2..x |> Enum.filter(fn a -> rem(x, a) == 0 end) |> length()) == 1
def prime(n), do: Stream.interval(1) |> Stream.drop(2) |> Stream.filter(&is_prime/1) |> Enum.take(n)

This generator worked well for smaller numbers, but was prohibitively slow for generating all the primes below 1 million.

Next I found Andrew Haines’ implementation of the Sieve of Eratosthenes. Its performance was an improvement, finishing in 1.8 seconds. But I was curious what kind of improvements we might see if we were using a native extension instead.

It had been a couple of years since I had played with Rust, but I had recently heard about an interesting project by Hans Elias Bukholm Josephsen called Rustler, which allows us to write Erlang NIFs in Rust.

Similar to FFIs in other languages, NIFs are a way to interact with native extensions on the BEAM. They have been around for a long time and can be useful, but historically they’ve also been really dangerous. A segfault in an NIF doesn’t just kill the process that called it. It can bring down the whole VM.

Avoiding this pitfall is what makes Rust and Rustler so interesting. Where code written in Rust is safe by default, any code written in Rust that isn’t explicitly unsafe shouldn’t be able to crash the BEAM. This could open up the door to some serious performance gains for computationally intense work, without the dangers traditionally associated with NIFs.

Ultimately I ended up writing a library called primal_ex, which is wrapper around the excellent Rust library, Primal. This tutorial will track quite closely with what primal_ex came to look like, so if you get stuck, it might be helpful to take a look at that repo.

The first thing we need to do is to install Rust. And if you don’t have Elixir installed already, you can find instructions here.

Now we can get started with mix new!


$ mix new primes_please
$ cd primes_please

Then we need to add rustler as a dependency in mix.exs

  defp deps do
    [
      {:rustler, "~> 0.9.0"},
    ]
  end

and run

$ mix deps.get

Once Rustler is installed, we run:

$ mix rustler.new

At this point, we’ll be asked for the Elixir module name we want to use. We want to keep the Elixir module in the PrimesPlease namespace, so we’ll enter:

PrimesPlease.Generate

We’ll then be asked for the Rust crate name. The default will suffice, so we’ll just hit “enter”.

Cool, now we have a Rust NIF ready to go! Well almost. In the resulting output of the rustler.new command, we see that there’s a README.md for further instruction. It tells us to do a few things.

First, we need to add the Rustler complier to our project, in our Mixfile. We also need to add a rustler_crates attribute which we’ll just point to a separate function. The project function in mix.exs now looks like this:

def project do
  [
		app: :primes_please,
  	version: "0.1.0",
  	elixir: "~> 1.4",
  	compilers: [:rustler] ++ Mix.compilers(),
  	rustler_crates: rustler_crates(),
  	build_embedded: Mix.env == :prod,
  	start_permanent: Mix.env == :prod,
  	deps: deps()
	]
end

And we’ll add that rustler_crates function, which will return a keyword list. We’re only adding one keyword, which will be the snake-cased version of our Elixir module name, Generate. And the path is the folder that mix rustler.new just created. We’ll also set the mode as :debug, unless the ENV is :prod, in which case it will be set to :release. Be mindful that debug mode is not optimized, so you may see some nice performance gains on a full release compile.

defp rustler_crates do 
  [generate: [
    path: "native/primesplease_generate",
    mode: (if Mix.env == :prod, do: :release, else: :debug),
  ]]
end

Awesome. Now our NIF will build with our project.

Next we need to hook up the NIF to the Elixir module that we’ll use to access the Rust code, which in our case is PrimePlease.Generate, which will live in lib/primes_please/generate.ex.

In this module, we also need to define the function we’re going to hook up: primes/1.

defmodule PrimesPlease.Generate do

  use Rustler, otp_app: :primes_please, crate: :primesplease_generate

  def err(), do: throw :nif_not_loaded

  # When your NIF is loaded, it will override this function.
  def primes(_a), do: err()    

end

If we run

$ iex -S mix

everything gets compiled, but we get an error that tells us that the NIF can’t be loaded, because we’re not implementing the auto-generated add/2. On the one hand, that’s ok, because we’re not trying to implement add/2, we’re trying to implement primes/1. Still, if we open up native/primesplease_generate/scr/lib.rs, we learn a bit about the NIF interface by looking at how they implemented add/2 as an example.

The first thing we should notice is rustler_export_nifs. Here we see a string representation of the Elixir module that we are trying to map the functions to. And then we notice an array of tuples that actual handle the mapping. The first element of the tuple is a string representation of the Elixir function being mapped. The second element is the arity of that function, and third element is the Rust function we’re trying to map to. So we’ll go ahead and replace

rustler_export_nifs! {
  "Elixir.PrimesPlease.Generate",
  [("add", 2, add)],
  None
}

with

rustler_export_nifs! {
  "Elixir.PrimesPlease.Generate",
  [("primes", 1, primes)],
  None
}

Now we’re ready to write the primes function.

Well, almost. Before we do that, we need to install the Primal crate. So, we’ll hop into native/primesplease_generate/Cargo.toml and paste primal = "0.2.3" under the [dependencies] header. Rustler will handle downloading the crate for us when we compile the NIF later on.

Now that we have the Primal crate available to us, we’ll need to include it in native/primesplease_generate/scr/lib.rs. We also need to include FromIterator from the StdLib. So toward the top of the file, we’ll include these lines:

extern crate primal;
use std::iter::FromIterator;

Then we’ll take the autogenerated add function and modify it into the primes that we want.
Here’s the before:

fn add<'a>(env: NifEnv<'a>, args: &[NifTerm<'a>]) -> NifResult<NifTerm<'a>> {
    let num1: i64 = try!(args[0].decode());
    let num2: i64 = try!(args[1].decode());

    Ok((atoms::ok(), num1 + num2).encode(env))
}

And the after:

fn primes<'a>(env: NifEnv<'a>, args: &[NifTerm<'a>]) -> NifResult<NifTerm<'a>> {
    let num: usize = try!(args[0].decode());
    let primes  = primal::Primes::all().take_while(|p| *p < num);
    let response = Vec::from_iter(primes);
    Ok((atoms::ok(), response).encode(env))
}

You’ll note that we made the following changes:
1. We changed the function name from add to primes.
2. We only need to first incoming argument, so we got rid of the second one.
3. We also needed to cast that argument to a uname rather than an i64, primary because that’s what the Primal library returns and we need the same type for comparison.
4. We need to return a vector rather than an iterator. Fortunately, Vec::from_iter allows us to make that conversion pretty easily.

Cool. Now we should be ready to go.

$ iex -S mix
iex(1)> {:ok, primes} = PrimesPlease.Generate.primes(1_000_000)
{:ok, [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, ...]}

Congratulations! You are now throwing some serious math computation down to a low level language which handles it beautifully.

I ran some preliminary benchmarks comparing primal_ex (with optimized binaries) to Andrew Haines’ version of the Sieve of Eratosthenes, when returning the first million primes.

As you might expect the native code performed really well.
I ran the benchmarks in a batch of 20.
On the first run, the NIF outperformed the straight Elixir code by a factor of 400:
Rust: 17484µs
Elixir: 6986807µs

I should probably just stop there. A 400x increase in performance is pretty awesome. But I do want to mention the results of the subsequent 19 runs, even if they are not especially representative of a typical Rust/Elixir comparison. From what I can gather, the Rust seive is able to use CPU level caching in a way that the Elixir sieve simply can’t. So after the initial run of grabbing all the primes below 1 million, on each subsequent run, the Rust sieve was able to leverage that cache to outperform the Elixir code by a factor of ~3,000.
These are the results of run 5, which are largely representative:
Rust: 2270µs
Elixir: 6962072µs

I was curious if the caching was somehow tied to the inputs, but after running a fair number of tests I concluded that the Rust sieve was able to leverage the cache even when the input values changed.

At the end of the day, I certainly don’t want to set unreasonable expectations for what a Rust NIF can do. In a sense, generating prime numbers is the perfect use case for illustrating when a NIF can have a real leg up on Elixir code. Whether a prime number generator in Elixir has any real value is its own question, but what this exercise does illustrate is that we are now able to leverage type-safe libraries from a low level ecosystem in Elixir. Use cases that have often been a good fit for native extensions are starting to see implementations using Rust and Rustler. html5ever_elixir, for example, provides bindings for Mozilla’s html5ever HTML parser. But I’m particular interested see to what other use cases will emerge. In the past, performance gains would have to be truly substantial to just the cost associated with adding a native extension. Rust and Rustler could go a long way toward bringing those costs down.

Rustler’s is still a relatively young tool. But if it does mature and provide us with reliable, seamless integration between Elixir and Rust, that could be really exciting. We have in Elixir, a language that does concurrency as well as anybody, and in Rust, a language that does computation as well as anybody. Together they could be quite the 1 – 2 punch.