silenthacks
Back to writing

When "Roughly Similar" Became a Scaling Problem

Scaling a weighted random selection feature from a naive SQL RANDOM() query through Redis and shared memory, then back to Postgres.

14 min read

Sacrifice a few items, get one back.

What even goes into making something like this? When we were brainstorming features for the bot, this concept was something we were all keen on. The idea itself is not new; plenty of systems implement some variation of it, and perhaps that’s what drew our attention towards it.

Originally, there were a few contenders for how it would work. Should the inputs have to be related to each other somehow, or could they be anything at all? Was the feature even a good idea? Our first discussions were, looking back, extremely unrefined, but we toyed with various alterations and executions.

A concern came up early on: if you’re putting in more than you’re taking out, the pool becomes a soup of random items. To counter this, we gave the inputs themselves a “cost”. Based on usage stats and a few other numbers we tracked, we’d compute a value for each item that became its contribution to the operation. Put in a few high-value items, or more low-value ones, and you’d land roughly the same result.

By that point, we had a basic idea. We wanted a system where you put in 5 items to receive 1 out. After much more discussion, we reduced the input count to 3 and built the operation around it.

First Try

SELECT * FROM items WHERE in_pool = true ORDER BY RANDOM() LIMIT 1;

Roughly, that’s all it was at first. When there were so few items, and fewer still in the pool, pretty much everything had the same score. This did the job more or less during the initial stages of alpha and beta testing, but it became apparent quickly that this was not going to scale well. As the number of items being stored grew, the command would return slower and slower. We didn’t know it at the time, but this little problem would become the biggest scaling issue we faced.

Another Shot

We were now fighting a battle on two fronts:

  1. We needed to make the command faster
  2. Scores were starting to diverge and we needed to return an item that actually was similar in value

Calculating scores was wildly inefficient and processing more than a handful of items would take a while. We wanted to cache the score values so that the calculation wasn’t being run each time, but of course didn’t want to retain stale data. The idea came about moving the pooled items to a second table, adding a column to keep its score. We could then periodically refresh this by running the calculation again and updating the values.

However, perhaps in our inexperience, we could not find an efficient solution to the second problem. Fetching a row using RANDOM() was slow but on top of that, we needed to weigh it against items of similar score. Postgres is more than powerful enough for this, but we had no clue how to do it. We could, however, solve this in Python fairly easily, so we decided to move the problem there. Perhaps one day down the line, we’d be good enough to figure it out…

As a quick recap, the operation promises an item within the same “10s” score range. For example, if your inputs averaged a score of 74, you could receive an item ranging from [70, 80). It would also (on average) try to give you back an item with a similar rank (a secondary attribute) to what you put in. We took a naive and simple approach to it at first:

  1. Create an empty dictionary
  2. Fetch all the items in the pool from the database at startup
  3. Iterate over them, calculating their scores
  4. Calculate their key in the dictionary with score // 10
  5. If the key wasn’t present in the dictionary, insert it with a list
  6. Fetch the value at this key and append the item to it

In Python, it looked something like this:

pool = {}
for item in items:
    score = calculate_score(item)
    key = score // 10
    group = pool.setdefault(key, [])
    group.append(item)

This meant we could grow the pool dynamically as the max score values grew. To actually find an item, we just look into the list corresponding to the right range and pick one, weighted against the average input rank:

key = avg_score // 10
group = pool[key]
weights = [1 / (1 + abs(avg_rank - item.rank)) for item in group]
item = random.choices(group, weights=weights, k=1)[0]

(Not quite exactly it but illustrates the idea.)

This worked for quite a while. Every day or so, we’d refresh the cache to prevent the values getting too stale. When an item was consumed, we could easily add the new one to the cache too. Anything unclaimed just got included during the refresh. Seemed like things were working well enough, so why didn’t we keep it like this?

As we continued to add new features and grow our playerbase, we began running into performance bottlenecks. We upgraded to a server with more cores, but we weren’t making any use of them. We were, however, making use of sharding by this point, so splitting the bot across multiple cores wasn’t too drastic of a change. There was one drawback to this though: we would have to rework the cache.

Suppose we split the bot across 8 different processes. If we kept the code as is, each process would have to maintain a copy of the cache, which not only eats up RAM but without any synchronisation primitives, just leaves the cache prone to being out of sync across processes. Python does not make managing inter-process synchronisation easy at all, so that idea was gone. Maybe instead, we could cache it once and share it across all processes?

Introducing Redis

What it said on the tin was a perfect match for what we were looking for: an in-memory cache to be used across all the processes. In reality, it turned out to be a huge frustration.

At the time, the most popular Redis library for Python didn’t have stable async bindings. We turned to aioredis (which has now merged with redis-py). All we had to do now was change the cache target from a dictionary to the Redis cache.

Since Redis is ultimately just a key-value store, we restructured our item format to fit with it. We reduced it to a binary serialisation of a JSON with the minimal necessary data. Once we got Redis working though, we found that it took a long time to cache the pool. We assigned the caching task to just one process and while we could have potentially split it up, it was still a non-negligible issue. The easiest fix was to just reduce the number of items we were caching, which would also help with our second issue.

If the bot was asked to take out an item in the 70s range, it would have to fetch the whole 70s pool from Redis, convert the binary back into something Python can understand, then run the weights calculations. This whole thing took time and reducing the cache size made the most sense.

We played around with some numbers and tried to find a good balance between variety and size. The bot would randomly sample this number of items from the pool and cache it on startup. As the bot grew, and so did the pool, we continued to upgrade our server’s RAM just to accommodate for the cache.

It wasn’t all great though. As I mentioned, working with Redis gave us some very frustrating moments. We found the library unpredictably dropping connection to Redis whilst the bot was running, leaving the pool empty for periods of time. We used its pub/sub handling for some coordination between processes, which would also fail and leave us requiring restarts. We tried our best working around these, introducing many failsafe mechanics and patches to reduce the likelihood of connection being lost. Eventually, we managed to reach a point we were happy to keep it at. It took a lot of time and effort whilst we were working on other features, so we decided to focus our attention away from it whilst it and the pool were stable.

No Longer Stable

Redis was taking longer to cache than we liked every time we restarted the bot and the operation was just too intensive as we had more players using it. It was time for another rework. After some rough prototyping with Redis alternatives, we realised we couldn’t solve the problem by just switching engines - this was going to need a whole redesign. We started thinking, how can we get faster than an in-memory data store? Long story short, you can’t really. How else then can we cut down on time?

The real bottleneck we were facing was the whole “load and convert” process every time an item was pulled out of the pool. If we could optimise that, we’d already be halfway there. The easiest way to achieve this was to keep the pipeline as close to the bot as it could be i.e. in Python. We settled on Python’s multiprocessing.shared_memory, specifically the SharedList. This would allocate a block of shared memory and let multiple processes access it - exactly what we needed. We would need to change our data format again, this time a tuple serialised into binary.

There was still one bit missing though and it was finally time to put my CS degree to good use. Appending items to the shared list arbitrarily would mean having to scan over the whole thing and keeping track of which ones lie in our range. We originally used a dictionary to keep mini groups of 10s, meaning we could easily grab what range we needed to look at. We would have to convert this convenience into the list, so we came up with a fairly simple way:

  1. Run analysis on the pool to see the score distribution
  2. Allocate specific ranges of indexes to each score range
  3. When pulling an item, only search between the indexes for the correct range

This was essentially re-inventing the dictionary from first principles. Now, with some heuristic based indexes, we could cache the pool with some sense of order. Although we did need to read the binary into Python during pulling, there was much less of an overhead compared to our Redis implementation. We kept a good buffer in each range, giving us breathing room as the distribution changed. Overall, we cut down on cache and pull times, allowing us to even cache the pool much more frequently. Once more, things were looking good.

Until They Weren’t Again

There’s quite a lot of items in the pool and we were still caching only a random subset of them. It seemed that the pool was less and less consistently able to return something that fit the score required. It was time to take another look.

The natural response to this problem was to try and cache the whole pool. In implementing the shared list, we’d cut down our RAM footprint a lot, so we actually had the capacity to cache everything now. Just to be on the safe side though, we added a custom serialisation format so that the binary contained only as much as was needed. We parallelised database fetches where we could and completely overhauled our list allocation implementation. Now, it would perform a first pass to find exactly how much it needed to allocate to each range and dynamically store these indexes. Pulling an item would then sample these ranges and do its calculation. This kept resources manageable since some ranges had millions of items in them, and some as few as just 1.

The pool took longer to cache, so we changed our refresh time back to daily. Now, if the bot returned an item with a lower score than the average input, it truly meant there were no items in the pool with that range. The change in structure also meant we could handle special items differently now, instead of including them in the same pool. By this point, we thought we’d really cracked it. We had the whole pool cached, freed up so much RAM and improved speeds. Even if we came up with another implementation, we’d already achieved the exact specifications we only hoped to months ago. Surely, now, we could finally leave it alone?

One Step Forward, Two Steps Back

Things were absolutely not fine. There seemed to be more score discrepancy and the pool seemed to have developed its own mind: sometimes it would return only items with one particular attribute value, other times skewing entirely to one rank. Despite no code being changed, its daily behaviour was wildly inconsistent. When we thought we were finally done, we realised we were just back to the drawing board.

Whilst we were working on other features, we kept it in the back of our minds. Sometimes that’s the best thing to do - give the problem some breathing room while you work on other things. Eventually, inspiration struck and only writing this up now did I realise that our very first idea was right on the money all along.

One part we accepted without choice was having to fetch all the items in the pool from the database. What if we completely cut that out? We could move the problem back to the database and completely remove any complications in the bot’s code. We definitely had the experience to do so now, so we set out for one more redesign.

The first step was to figure out how to store the data. We could calculate scores for the whole pool on each query, or store them and periodically refresh them. I hope it’s obvious to see why the only real option was the latter. We created a new table just for the pool with a column for storing the score. Postgres even comes with table partitioning on range, meaning we can specifically optimise the table to be split up according to the 10-point ranges we use.

Moving the problem to the database also meant we would have to convert the calculation logic to PL/pgSQL, which eliminates the round trip needed to query and fetch the data. It even came with the side benefit of easily being able to filter a user’s items against their score. A few stored procedures and functions later, and we can fetch a random row weighted against rank. Finally, with the help of a cron job, we can also refresh the pool daily.

When a user triggers the operation, we just query the database to atomically insert and pop an item. No more need for synchronisation primitives. All the logic is handled database-side so the bot only needs to make one trip. No more Redis, no more SharedMemory.

Are We There Yet?

That was probably as good as it was going to get - the thing we originally aimed for and eventually came back to. It leveraged the power of Postgres to keep things speedy, while letting us easily query and analyse the data. It removed all of our blocking startup processes, took the feature out of the bot’s memory footprint entirely, and gave us much more flexibility in maintaining it. For the first time, instead of pouring all our effort into getting a good, working backend, we could actually look at the higher-level details of what the feature meant for our users.

It was never perfect, but by then we no longer had the shaky foundations we’d spent so long trying to balance everything on. I’d hoped that would be the last redesign - I wasn’t sure I had enough caffeine or sanity left for another round. I’ve since moved on from the project, so whether it ever needed another rework after me, who knows. But for the stretch I was there, it finally felt like we had it.