Week 38 note: Load testing preparations
This week we started doing load testing prep work using k6. Previously I had been looking at k6 through rose-tinted glasses, but this week I finally learned some new things about it.
Our problem scenario starts with a need to authenticate a group of users to enable testing of actual endpoints. The initial plan was to open up a file and then do a bunch of authenticating HTTP requests before running the tests. As we needed to perform this for 5000 users, we wanted to do this concurrently (quickly enough) without breaking the API. Client-side rate limiting was essential here.
However, we didn't really understand how k6 data initialization should be done. We also learned that await was not supported, we were not able to use a node library, and that HTTP requests were only possible inside the setup stage (and not during the init stage).
Stopping to regroup what would be the best way forward, we ended up circumventing k6 for the authentication part. We wrote a python script to perform authentication for all users, then wrote the access tokens to a file, after which we could load them up in k6 (during the open phase). This meant that our load tests would have a fixed maximum session, but we were okay with that.
To support concurrency, we used ThreadPoolExecutor, urllib3's PoolManager as well as a token bucket approach. During every second, we would allow n amount of requests to be made while replenishing the token bucket using a fill rate & would block if that capacity had been used.
requirements.txt:
urllib3==2.2.3
simple_concurrency_example.py:
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
import json
import time
import urllib3
class TokenBucket:
def __init__(self, tokens, fill_rate):
"""
tokens is the total tokens in the bucket.
fill_rate is the rate in tokens/second that the bucket will be refilled.
"""
self.capacity = tokens
self._tokens = tokens
self.fill_rate = fill_rate
self.timestamp = time.time()
def consume(self, tokens):
if tokens <= self.get_tokens():
self._tokens -= tokens
return True
return False
def get_tokens(self):
now = time.time()
delta = self.fill_rate * (now - self.timestamp)
self._tokens = min(self.capacity, self._tokens + delta)
self.timestamp = now
return self._tokens
def fetch_token(user_id, cmanager):
# perform IO-expensive http authentication here
print(f"{datetime.now()} {user_id}")
return f"<token:{user_id}>"
connection_manager = urllib3.PoolManager(maxsize=5, block=True)
thread_pool = ThreadPoolExecutor(max_workers=5)
bucket = TokenBucket(5, 5)
user_ids = [n for n in range(15)] # just numbers as an example here
futures = []
user_access_tokens = []
for user_id in user_ids:
if not bucket.consume(1):
sleep_time = 1 - (bucket.get_tokens() - 1)
time.sleep(sleep_time)
future = thread_pool.submit(fetch_token, user_id, connection_manager)
futures.append(future)
for future in futures:
result = future.result()
print(f"{datetime.now()} {result}")
user_access_tokens.append(result)
with open("./access-tokens.json", "w") as f:
json.dump({"user_access_tokens": user_access_tokens}, f, ensure_ascii=False, indent=4)
Example output:
2024-09-22 21:08:44.362819 0
2024-09-22 21:08:44.362906 1
2024-09-22 21:08:44.363118 2
2024-09-22 21:08:44.363139 3
2024-09-22 21:08:44.363147 4
2024-09-22 21:08:46.366548 5
2024-09-22 21:08:46.366825 6
2024-09-22 21:08:46.367540 8
2024-09-22 21:08:46.367703 9
2024-09-22 21:08:46.367955 10
2024-09-22 21:08:46.366950 7
2024-09-22 21:08:48.364249 <token:0>
2024-09-22 21:08:48.364406 11
2024-09-22 21:08:48.364602 12
2024-09-22 21:08:48.364732 13
2024-09-22 21:08:48.364489 <token:1>
2024-09-22 21:08:48.365148 <token:2>
2024-09-22 21:08:48.365181 <token:3>
2024-09-22 21:08:48.365205 <token:4>
2024-09-22 21:08:48.365224 <token:5>
2024-09-22 21:08:48.365242 <token:6>
2024-09-22 21:08:48.365259 <token:7>
2024-09-22 21:08:48.365276 <token:8>
2024-09-22 21:08:48.364827 14
2024-09-22 21:08:48.365291 <token:9>
2024-09-22 21:08:48.365508 <token:10>
2024-09-22 21:08:48.365543 <token:11>
2024-09-22 21:08:48.365562 <token:12>
2024-09-22 21:08:48.365578 <token:13>
2024-09-22 21:08:48.365593 <token:14>
After that, we had needed to ensure that the CI environment had all of the dependencies installed before running the script:
python3 -m ensurepip --default-pip
pip install -r requirements.txt
Yes, we ran this globally without virtualenv on purpose.
Another nice way (I later learned) would have been to run these using uv, which is a virtualenv-less package & project manager tool to run Python software. Quickly trying it out:
uv init simple-concurrency-example
cd simple-concurrency-example
uv add urllib3
# copy `simple_concurrency_example.py` under this directory
uv run simple_concurrency_example.py
In the end, we had our authenticating script in place and we were free to focus on our load testing scenarios.