Filtering messages was overloading our servers

When Aria delivers messages to a client via TCP, it keeps the recent stream, what we call the stream tip, in an in-memory ring buffer. If a client falls behind, it can request recent messages from this ring buffer to get up to date. We call this process tip recovery.

One challenge with tip recovery is that Aria stores the entire stream of messages, when clients often only care about a much smaller subset of the messages. Aria divides the message stream into topics, which form a hierarchical namespace like a file system. Clients can subscribe to an individual topic or to a topic subtree, which consists of all topics under a given topic, similar to a globstar. Aria filters the entire stream down to only the messages from the subscribed topics and delivers these messages to the client.

Originally we were fine doing this filtering as a simple linear pass because, while algorithmically inefficient, looping over the stream is CPU cache friendly and therefore decently fast. However, as the number of clients doing tip recovery grew, we noticed that our servers were struggling with the increased load. Some servers reached 100% CPU utilization, resulting in clients “falling off the tip” and being unable to catch up. We added more servers as a temporary fix, but it became clear that we desperately needed to re-think our tip recovery code.

Theodor decided to solve this problem by adding index data structures that Aria could use to efficiently filter messages from the stream. But what should be indexed? A naive approach would be to make an index for each topic. However, Aria instances can have almost a million topics, making this approach untenable. Instead, Theodor created an index for each topic partition. A topic partition consists of all topics with the same two-segment prefix, a segment being a single part of the topic name separated by a slash. For instance, app/codestore/commits and app/codestore/features would both be under the app/codestore topic partition.

Each topic partition gets an index of its messages’ location in the Aria stream:

When a client requests messages, Aria does an n-way merge of the indexes using a min-heap. This allows Aria to efficiently reconstruct the message stream, in order, for just these specific topic partitions:

Prototyping and benchmarking the index

As with any performance work, we had to benchmark in order to validate that we were actually improving the system. We also wanted extensive testing, since the message publishing logic is in the critical path for Aria.

Theodor started by writing a tool to profile various recovery scenarios. That way, we could test out different configurations, such as the number of topic partitions, the degree of interleaving, the number of readers, etc., and see how changing these variables affected the performance of the system.

With the benchmark in hand, he implemented the initial version of the topic partition index. When we profiled this implementation, we found that the min-heap was the biggest bottleneck in our code. Before AI, this is probably where we’d pick a single implementation that appeared to be better. But these days, running experiments is cheap: Theodor prompted an agent to spin up five different heap implementations and profile them overnight. The next day, we had our answer: fast_heap_unboxed gave us a 2x performance improvement to our indexed recovery code.

Theodor tested the new index with plenty of expect tests, exercised the new logic in Antithesis, and finally had a swarm of agents analyze the code with high scrutiny.

Using a block pool to store the indexes

Once we devised the index data structure, we needed to figure out an efficient representation. We didn’t want to make a ring buffer for each index, since ring buffers are not really dynamically resizable. They would therefore have to be sized to accommodate the worst case scenario. Specifically, index size is inversely proportional to message size (smaller messages mean more messages can fit in the tip store, which means you need a larger index). An Aria message can be as small as 32 bytes. If the entire tip store consisted of messages that small, the corresponding index would be 2GB. Keeping a 2GB index for each topic partition would waste too much memory.

Instead, we wanted an index to dynamically shrink and grow with the number of messages in its topic partition.

What we settled on was using a shared pool of blocks that each contain 1024 entries. The indexes point to a block and insert values into it. Once all of the messages within a block have left the ring buffer, blocks can be popped off from the index and re-used for other indexes:

Reading old messages was taking too long

Theodor’s work on indexing fixed our tip recovery woes, but we had another message delivery problem: Clients that needed messages from before the stream tip were waiting way too long.

When a client reboots, it often requests all messages since the beginning of the week from its subscribed topics. We call this process initial recovery. Initial recovery can involve millions, even billions of messages. It’s absolutely critical that Aria read these messages and deliver them to clients as quickly as possible. Also, clients tend to reboot all around the same time, so it’s especially important that this process scales well.

In one particular incident, an initial recovery that usually took under 2.5 seconds was taking over 13 minutes. When we dug into why, we found the same issue that plagued the tip store: Aria was reading and filtering 10x more data than it actually needed to send. It was clear that we needed to rethink how we were storing messages on disk.

Aria persists messages in subtree stores

Aria initially stores the messages on disk in chronological segments for each topic partition. Then, a separate process takes these segments and splits them into multiple subtree stores that each contain messages for a topic subtree. These subtree stores are intended to strike a balance between writing to many files, which results in efficient filtering but requires more work to recombine into a message stream, and writing to fewer files, which is easier to recombine but less efficient to filter.

To do this splitting, Aria had been using a simple heuristic: it chose the subtree store by reading the first three segments of the topic. So if we had topics app/options/orders/created and app/options/orders/cancelled, those would both get put into a subtree store for app/options/orders. If the topic had fewer than three segments, it was put into its own subtree store; for instance, app/options would get its own store. Effectively, this heuristic gave each direct child of a topic partition a separate subtree store.

However, this heuristic didn’t work well when one topic in a store had significantly more messages than the other. If the client were to subscribe to just the less active topic, Aria would have to filter out all of the other topic’s messages. Like with tip recovery, this filtering created a lot of work.

More broadly, we weren’t happy that users had to understand this specific aspect of Aria’s internal design when designing their topic structure. We wanted users to structure their topics in a way that made sense to them and have them trust that Aria would do the splitting intelligently.

Theodor was tasked with developing an algorithm that could avoid this issue by intelligently splitting the topic tree based on each topic’s message volume. Aria is able to calculate the message volume per topic, since it does this segment processing step after persisting the message stream.

Gathering data and testing different algorithms

Like with the tip recovery implementation, Theodor’s work on adaptive splitting required exploring different options. He started by gathering data from various Aria sessions that could be used to test out different algorithms. He then implemented a few different algorithms for splitting and ran them on the collected data. Theodor analyzed these results from a few different lenses, such as:

• How many bytes do we have to ignore from a store if we read a topic?
• What is the sum of these “wasted bytes” over all the topics?
• What topic, if read, would result in the most wasted bytes?
• What topic, if read, would result in the worst ratio of total bytes to useful bytes?

As before, this process of exploration was aided by LLMs. Theodor was able to quickly vibe-code a web UI that visualized the different splitting algorithms across real data sets.

The final algorithm

Eventually Theodor settled on a solution that combined a greedy clustering algorithm with a binary search. This gave us the properties that we were looking for: Topics with a lot of messages receive their own subtree store, while smaller topics are folded into a single store.

Here’s an example of a topic tree using the previous heuristic:

Notice how the 2.72GB topic is in the same store as the 10MB and 28.5MB topics. This means that if someone wants to read the 10MB topic, they have to filter out all those other messages. Meanwhile, the smaller topics are split between three stores, even though they’re orders of magnitude smaller than the large topics.

With the new adaptive splitting algorithm, this is how the same tree looks:

Now the 2.72GB topic has its own store separate from the 28.5MB and 10MB topics, while the smaller topics are all bundled together into a shared store.

Using multiple testing strategies on the algorithm

Since the adaptive splitting directly affects how messages are stored in Aria, we wanted to ensure with a high degree of certainty that there were no bugs. Dropping or re-ordering a message would have serious consequences.

We started by using expect tests to check various properties of the algorithm, such as large topics getting their own store, small siblings sharing stores, and so on.

For message consistency, we already had an existing property-based test that inserted a sequence of randomly generated messages on randomly selected topics and then ran the splitting algorithm. The test read out random pieces of the topic tree and confirmed that the messages had the correct ordering and content. Theodor expanded the test to randomize the tree splitting, which confirmed that message content stayed the same regardless of how the topic tree was split.

Finally, we did multiple runs in Antithesis to confirm that the changes didn’t break anything else in Aria.

How did the optimizations do?

The tip indexing code has been shipped to production. We’ve seen preliminary results in our staging environment that showed 30% less CPU utilization in some real-world scenarios. Overall latency across clients went down significantly during those times. We’ve noticed some multi-second tail latencies in servers that didn’t implement the indexing, while these issues didn’t appear in any servers running the new code.

Adaptive splitting is deployed to our staging environment and will be running in production shortly.