Max Bernstein on Rails at Scale shows how Rails enum predicate methods were made dramatically faster by generating static code at runtime instead of dynamic public_send dispatch, with benchmarks and lessons about fixing performance upstream.

Rails is a JIT compiler

Or, you wouldn’t believe how Aaron made enums faster with one weird trick.

The title was bait to get you to open this article. But fear not! We will be generating code at run-time today. (See the PR for a sneak peek.)

Sam popped into the JIT channel asking how we might optimize Rails enums in Roundhouse, motivated by the following code snippet:

class Order < ActiveRecord::Base
  enum :status, { pending: 0, paid: 1, shipped: 2, refunded: 3 }

  def shippable?
    paid? && !refunded?
  end
end

In particular, there is a lot of dynamic behavior going on with the code generated by Rails enum and Sam wanted to know if inlining it all away in an (external-to-Rails) build step would help.

That didn’t sound right to me because I thought that:

Rails got a lot better in recent years about removing these dynamic lookups
in favor of creating and caching methods
ZJIT should be able to handle dispatch to methods (generated or otherwise)
just fine

So I inquired and Sam presented to us the implementation of enum which contained this gem:

value = 0
define_method("#{value_method_name}?") {
  public_send(:"#{name}_for_database") == value
}

Now fortunately, this is not the fully dynamic lookup I was worried about! No hashes, etc. This is Rails using the enum to generate a method on the ActiveRecord class. So in our case, for example, it turns pending into a method on Order:

value = 0
Order.define_method("pending?") {
  public_send(:status_for_database) == value
}

Unfortunately, it is still using very dynamic behavior. There is no reason to be using public_send for method dispatch if you know the name statically: it doesn’t get cached and has a lot more machinery than we need. There is also no reason to use a bmethod (define_method with a block that captures the environment) and look up the value dynamically; the enum values are burned in at enum definition time.

I wanted to make it use module_eval with a string instead to generate a normal looking method. I don’t know anything about Rails so I nerd sniped Aaron into fixing it.

Nerd sniping Aaron

The change to the the code started off fairly mechanical: we should be able to do module_eval to generate a similar looking method:

code = []
code << "def #{value_method_name}?; #{name}_for_database == #{value.inspect}; end"
# ...
module_eval(code.join("\n"), __FILE__, __LINE__) unless code.empty?

The differences here being:

Direct dispatch to the method instead of going through public_send
Burning the value into the method instead of looking it up from the parent
environment of the block

It passed almost all the tests but then we ran into a couple of weird cases:

People have the darndest enum names
1. The Rails test suite, for example, uses fake timezones: enum :timezone, [:"Etc/GMT+1", :"Etc/GMT-1"]
2. Our application, has an instance with a capital letter: enum :thing, [:Foo, :Bar]
3. Probably other stuff that wouldn’t be a identifier in Ruby
Floats. Naturally. While all the enum values are required to be serializable
(and therefore reifiable in Ruby source code), the roundtrippability gets a little weird with floats. So we have this irritating check: !(Float === value && !value.finite?)

In the end it wasn’t quite a one-line change because Ufuk also pointed out that we should be using ActiveSupport::CodeGenerator instead of yolo-ing it. But it was pretty close to it.

Check out the PR.

Performance numbers

Go check out the PR! But you’re still here, so I’ll repeat the numbers here. The following microbenchmark of the code at the beginning of the article:

IPS.ips do |x|
  x.report("shippable?") { order.shippable? }
end

got a lot better.

Here are the performance numbers for main:

aaron@Aarons-MBP ~/g/rails (main)> ruby -I~/git/ips/lib test.rb
-- create_table(:orders)
   -> 0.0073s
          shippable?:     1.721M i/s (± 2.3%, GC  4.4%)
aaron@Aarons-MBP ~/g/rails (main)> ruby --yjit -I~/git/ips/lib test.rb
-- create_table(:orders)
   -> 0.0048s
          shippable?:     3.568M i/s (± 2.6%, GC  8.2%)
aaron@Aarons-MBP ~/g/rails (main)> ruby --zjit -I~/git/ips/lib test.rb
-- create_table(:orders)
   -> 0.0029s
          shippable?:     3.812M i/s (± 2.8%, GC  8.9%)

and for Aaron’s PR:

aaron@Aarons-MBP ~/g/rails (fix-enum)> ruby -I~/git/ips/lib test.rb
-- create_table(:orders)
   -> 0.0073s
          shippable?:     3.145M i/s (± 1.5%)
aaron@Aarons-MBP ~/g/rails (fix-enum)> ruby --yjit -I~/git/ips/lib test.rb
-- create_table(:orders)
   -> 0.0065s
          shippable?:    16.274M i/s (± 1.3%)
aaron@Aarons-MBP ~/g/rails (fix-enum)> ruby --zjit -I~/git/ips/lib test.rb
-- create_table(:orders)
   -> 0.0027s
          shippable?:    18.509M i/s (± 2.2%)

Wahoo!

Lessons

Static is better than dynamic. Burn in those constants.
In particular, public_send is slow!
We can fix performance issues upstream! Especially if the upstream libraries
already do code generation. Not all problems require a custom compiler.