Making Ruby go brrrr (with Codex)

You know how sometimes you just want to fix one thing, and then a week later you find yourself five levels down a rabbit hole? Yeah…

So, I just had an idea to simplify some things in Lycan. To import your liked posts, it loops over your whole history of records using the listRecords endpoint, for posts, likes, and reposts in parallel (and also in parallel fetches the original posts they're linking to from the AppView). Many, many pages, and it takes a bit of time.

Right at the point when I was launching it, I had realized I could probably make it quite a lot simpler and faster if I just fetched the whole repo .car with getRepo and parsed it locally, that would be just one big request instead of hundreds – but I didn't want to mess with something that was ready back then.

Parsing CAR repos

So I started looking at this again this month, and I realized there's one problem, that I don't have code ready to parse a .car repo (and in fact I have no idea how to parse one). I did have code for parsing the CAR structure in skyfall since forever, because this is needed to parse firehose commit messages, but a repo is something more on top of that. I remembered there was some post in Bluesky's docs about it, but it turns out it mostly just says "call these methods ReadRepoFromCar and ForEach from indigo"… So I had to dig a bit in the indigo source code to figure out how to read this whole MST tree in Ruby (it's not that hard, in the end, just not obvious).

Then, the next step was to make it read a large .car repo like mine in a reasonable time. The first problem was that this walk_all_nodes algorithm basically calls a "look up a section in the CAR by a CID we have" method once for every record in the repo and some more, which for me is a lot of times. And the method I had for this in Skyfall looked like this:

def section_with_cid(cid)
  if section = @sections.detect { |s| s.cid == cid }
    return section.body
  end

  # ... else parse more sections
end

So… this is a bit suboptimal for this use case (in case you can't tell, it checks all previously loaded sections sequentially, every time). So I added a Hash @section_map there, mapping CIDs to loaded sections for quick lookup.

This improved things significantly, except… it made firehose processing slower. Because now the CARs loaded from commit events were building up this section_map index too, even though they don't need it, because they don't really look up sections this way. Also, when I ran a benchmark with ruby-prof, encoding the CID data with Base32 now showed up pretty high on the list, because it was being used to put the CIDs in the hash map.

Base32

I showed the profile output to Codex and told it to maybe look at the base32 gem code if some specific things could be optimized there. It came back 5 minutes later saying basically "here you go, I rewrote the whole thing from scratch lol" 😅

The existing base32 gem wasn't written with performance in mind – it's just normal Ruby code written in a readable way. It slices the string into 8- or 5-byte "chunks", wraps them in Chunk objects, and then those chunks encode and merge their parts using various Enumerator methods like take_while, inject, collect, flatten, join. It works and makes sense, but it creates quite a lot of objects on the way, and if this is called thousands of times, that becomes a problem.

The new code looks a bit more like C written in Ruby… it does everything in one method call, operates on a single mutable string instance and on numbers with byte values, and uses a plain while loop instead of the more functional coding style. And the inner loop is completely "unwrapped" like this to avoid iteration overhead:

v0 = table[data.getbyte(offset)]
v1 = table[data.getbyte(offset + 1)]
v2 = table[data.getbyte(offset + 2)]
v3 = table[data.getbyte(offset + 3)]
v4 = table[data.getbyte(offset + 4)]
v5 = table[data.getbyte(offset + 5)]
v6 = table[data.getbyte(offset + 6)]
v7 = table[data.getbyte(offset + 7)]

value = (v0 << 35) | (v1 << 30) | (v2 << 25) | (v3 << 20) |
        (v4 << 15) | (v5 << 10) | (v6 << 5) | v7

It also does some neat things with pre-generated lookup tables stored in constants: the decoder has a 256-item array, where the value under n is the 5 bits encoded by character with ascii code n, or 255 for "invalid":

BASE32_DECODE_TABLE = begin
  table = Array.new(256, 255)

  BASE32_ALPHABET.each_byte.with_index do |byte, val|
    table[byte] = val
    table[byte - 32] = val if byte >= 97 && byte <= 122
  end

  table.freeze
end

# 'b'.ord == 98
# BASE32_DECODE_TABLE[98] == 1

And the encoder has a 32 x 32 "flattened" table, which it uses to look up 2-character or 2-byte values using 10-bit slices, which lets us do half as many lookups in the main loop there:

BASE32_ENCODE_TABLE = Array.new(1024) { |i|
  (BASE32_ALPHABET.getbyte(i >> 5).chr +
   BASE32_ALPHABET.getbyte(i & 31).chr).freeze
}.freeze

(I also tested a version of the code using String#append_as_bytes added in Ruby 3.4, but I couldn't see any noticeable difference vs. appending characters using <<.)

Optimizing

I ran the benchmarks again, saved new profile output, and gave it to Codex again to see what easy wins it can find there. And this turned out to be a really good approach in general, because apparently it's much better at this than I am! It can notice patterns and details, and has knowledge about what changes in code can improve performance, which is sometimes not obvious at all.

I'm not gonna go through everything step by step because it was basically a full week of hacking, switching between: running benchmarks and tracking all numbers in a spreadsheet, which looked like this (yes I'm neurotypical, why are you asking):

… generating profile output again and giving it to Codex to explore, discussing with Codex how best to design some APIs, reading the IPLD / DASL / ATProto specs to check minor details, telling Codex to implement some change or doing it myself, tweaking the unit tests, re-running all benchmarks or adding new ones, and so on.

I've managed to do a series of several different improvements in different parts, though in the end it ended up a bit underwhelming, because I realized that if I had only used the binary CID data for the hash map instead of the encoded base32 version, that would have brought me most of the way there… but the end result is still quite a lot faster than that. Also some unrelated parts got optimized on the way, which should also speed up the firehose processing or some other use cases:

Processing my .car repo:

• first "map" version: 9.601s just keys, 16.599s with records

• after all optimizations: 1.139s just keys, 2.278s with records

• "quick win" version: 2.327s just keys, 3.648s with records

Encoding CIDs to JSON/Base32 form:

• original version (base32 gem): 9.990s

• updated version (Codex's code): 3.810s

Parsing firehose messages (CBOR + CAR):

• original version: 6.225s

• after optimizations: 4.383s
  (in the real app like my feed service the difference is much smaller, because it also does a lot of other things apart from this specific parsing, like reading websocket frames, analyzing the post text with each feed's regexps, building AR models, saving records to Postgres, etc.)

Takeaways

I also learned some general things about what to look for in the code I'm trying to optimize.

The main theme of the various optimizations (apart from the obvious "try to avoid doing work you don't need to do") was: make less allocations everywhere on the hot path, especially String allocations; refactor things so existing objects are reused instead where possible. This was probably the biggest win by far.

• One no-brainer is to have frozen_string_literal: true everywhere at least in library code, but I already had that one – apart from making string literals in code frozen (immutable), this also tells Ruby to only allocate the literal strings once instead of every time it gets to this line (a bit like a &'static str in Rust?…). This makes a lot of difference e.g. for things like Hash keys used over and over.

• There are a lot of other places where strings are copied in non-obvious ways: every time you slice a string like s[1..-1], or gsub(/=$/, ''), or chomp(), and so on, or add two strings together with + . A lot of those can be replaced with a variant that modifies the string in-place if it makes sense, like chomp!, <<, prepend etc. (but of course avoid this if the string can be shared with other parts of the code).

- { '$bytes' => Base64.encode64(data).chomp.gsub(/=+$/, '') }
+ string = Base64.strict_encode64(data)
+ string.chomp!('=') while string.getbyte(-1) == 61
+
+ { '$bytes' => string }

  data = section_data.byteslice(0, 36)
- cid = CID.new("\x00" + data)
+ data.prepend("\x00")
+ cid = CID.new(data)

• Two string concatenations like a + b + c also end up creating two new strings, so changing that to "#{a}#{b}#{c}" avoids one allocation again.

- cid = CID.new("\x00" + prefix + cid_data, true, true)
+ cid = CID.new("\x00#{prefix}#{cid_data}", true, true)

• Even a character lookup like s[0] == 'b' creates a new one-character string instance! Because there's no "Char" type in Ruby. Changing that to use getbyte instead, which returns a number, avoids the string allocation. (Numbers are technically also objects, but they seem to optimize much better; also getbyte is specifically listed as having been optimized further under YJIT in Ruby 3.4.)

- raise DecodeError unless str[0] == 'b'
+ raise DecodeError unless str.getbyte(0) == JSON_PREFIX_CODE

Here's what Codex wrote about the Base32 optimization:

Classic Ruby performance trap: the code is elegant, readable, and maps neatly onto the algorithm—but every intermediate concept becomes a real heap object. For occasional Base32 encoding, the original is perfectly reasonable. At 35,000 CIDs, however, “make a chunk, map characters, flatten everything” becomes millions of tiny operations.

So the real problem wasn’t bit shifting itself. It was representing a straightforward byte-to-text conversion as hundreds of thousands of objects, nested arrays, and millions of temporary one-character strings. 🤖

Some APIs were reorganized to optionally allow for less string copying or to skip some processing or wrapping data in objects on the hot code path.

For example, the base32 decoder/encoder is modified to allow passing a starting prefix and starting offset, avoiding two extra string allocations, first to slice the prefix off the input, and second to append a prefix to the output:

- @json ||= 'b' + Base32.encode(@binary_data[1..-1])
+ @json ||= Base32.encode(@binary_data, 1, 'b')

And the section_with_cid function was changed to allow lookup by CIDs passed as as raw binary data instead of a CID object, which avoids creating the objects that are only used to wrap a string and then immediately extract that string two lines later. The caller still gets wrapped CID objects in the end, but the iteration through the tree skips them.

Same with avoiding the recursive conversion of decoded objects to wrap CIDs with $link or binary strings with $bytes as used in ATProto. section_with_cid previously returned converted JSON body immediately; it was changed to return a CARSection, on which you call either decoded_body or json_body depending on whether you want the version with converted $links or not. So iterating through the tree skips the conversion, until the part where we return data to the caller.

Some other things:

• There was a fragment that checked the beginning of each section by reading multiple "varints" (variable-length numbers) by reading and adding up byte by byte, in order to verify that the first few bytes are an expected CID prefix – this was replaced with a constant string header like "\x01\x71\x12\x20" and checking start_with?.

read_section calls read_varint five times for each of 319,788 sections: (…). The last four values are required constants and all have canonical one-byte varint encodings: \x01\x71\x12\x20. That means 1,279,152 of the 1,598,940 section-level varint calls are parsing four bytes whose values are already known. 🤖

• Some strings are frozen with .freeze to prevent modification – it turns out this also has an effect on performance in some cases: when strings are used as Hash keys, Hash makes a copy of the string first if it's mutable, but uses it as is if it's frozen.

• Avoiding extra method calls in some places, like:

  def raw_record_for_operation(op)
-   op.cid && blocks.section_with_cid(op.cid)
+   cid = op.cid
+   cid && blocks.section_with_cid(cid)

To be continued

I'm confident now that we (me & 🤖) can optimize some things there even further – for example, I want to replace the CBOR gem (written in C and built to handle all CBOR, not only the DAG-CBOR/DRISL subset) with pure Ruby code built to only handle the ATProto CBOR; I think it should be possible to make that at least as fast or faster, at least with YJIT on. This will also let me integrate it better with the rest of the code, e.g. avoiding the intermediate step of CIDs encoded as CBOR::Tagged objects and then converting them to CID objects (and instead build them as CID objects immediately), and avoiding the second pass of converting those to $link if needed and instead doing that in the same loop.

I also extracted these lower-level data decoding pieces from Skyfall – CID, CARArchive, CARRepo, Base32 – to a new gem I've named "oxygene" – because oxygen is a key component of the Atmosphere, and because I love Jarre's "Oxygène" album :) So it can be used for things like decoding .car repos, without pulling in the whole Skyfall with eventmachine, faye-websocket etc.