Skip to content

Instantly share code, notes, and snippets.

@tkaemming
Last active May 4, 2016 00:06
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save tkaemming/5f5ce3cbb37b4b0c44a58f70db2236d2 to your computer and use it in GitHub Desktop.
Save tkaemming/5f5ce3cbb37b4b0c44a58f70db2236d2 to your computer and use it in GitHub Desktop.
#!/usr/bin/env python
import click
import functools
import signal
import uuid
from multiprocessing import Pool
def run(keys, iterations, retry_failures):
from redis import (
StrictRedis,
WatchError,
)
result = [0, 0]
def increment():
"""Update counters."""
with client.pipeline(transaction=True) as transaction:
transaction.watch(*keys)
counts = transaction.mget(keys)
try:
transaction.multi()
updates = {}
for key, count in zip(keys, counts):
updates[key] = (int(count) if count is not None else 0) + 1
transaction.mset(updates)
transaction.execute()
except WatchError:
result[1] += 1
if retry_failures:
increment()
else:
result[0] += 1
try:
client = StrictRedis()
for i in xrange(iterations):
increment()
except KeyboardInterrupt:
pass
return result
@click.command()
@click.option('-c', '--concurrency', default=4)
@click.option('-n', '--num-keys', default=2)
@click.option('-i', '--iterations', default=500)
@click.option('--retry-failures', default=False, is_flag=True)
def __main__(concurrency, num_keys, iterations, retry_failures):
namespace = uuid.uuid1().hex
keys = ['{}:{}'.format(namespace, i) for i in xrange(num_keys)]
pool = Pool(concurrency)
results = []
for _ in xrange(concurrency):
results.append(pool.apply_async(run, (keys, iterations, retry_failures)))
signal.signal(signal.SIGINT, signal.SIG_IGN)
successes = 0
failures = 0
for result in results:
data = result.get()
successes += data[0]
failures += data[1]
stdout = click.get_text_stream('stdout')
stdout.write(
'{} {} {:%}\n'.format(
successes,
failures,
float(successes) / (successes + failures)
)
)
if __name__ == '__main__':
__main__()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment