Load-testing for exchange to exchange bindings in RabbitMQ

bindingload.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Note: Requires https://github.com/skarab/pika/tree/rabbitmq-ext
# for RabbitMQ exchange binding extensions
import csv
import getopt
import os
import random
import sys
import time

import pika

from pika.exceptions import AMQPChannelError
from pika.adapters.blocking_connection import BlockingConnection


def pad(num, places=32):
    "Pad with zeroes"
    str = '%x' % num
    padding = '0' * (places - len(str))
    return '%s%s' % (padding, str)


def random_key():
    "Random topic routing key"
    return '%s.%s' % (
        pad(random.randint(1, 200000)),
        pad(random.randint(1, 540000)),
    )


def main(filename, mode, sleep):
    """
    Example which tests exchange to exchange bindings and their
    impact on unrelated topic exchanges
    """
    cparams = pika.ConnectionParameters(
        host='127.0.0.1',
        port=5672,
        virtual_host='/',
        credentials=pika.PlainCredentials('guest', 'guest'))
    connection = BlockingConnection(cparams)
    print 'Getting AMQP channel'
    channel = connection.channel()
    if filename:
        print 'Reading CSV file'
        channel.exchange_declare(
            exchange='f1', type='fanout', durable=False,
            auto_delete=False)
        channel.exchange_declare(
            exchange='d1', type='direct', durable=False,
            auto_delete=False)
        channel.exchange_declare(
            exchange='d2', type='direct', durable=False,
            auto_delete=False)
        for s_id, d1_id, d2_id in csv.reader(
            open(os.path.join(base_dir, 'bindings.csv'), 'rb')):
            # Convert to integers
            exchange = 's%d' % int(s_id)
            if mode == 'declare':
                print 'Declaring and binding %s' % exchange
                channel.exchange_declare(
                    exchange=exchange,
                    type='fanout',
                    durable=True)
                channel.exchange_bind(
                    source='f1',
                    destination=exchange,
                    routing_key='')
                channel.exchange_bind(
                    source='d1',
                    destination=exchange,
                    routing_key=d1_id)
                channel.exchange_bind(
                    source='d2',
                    destination=exchange,
                    routing_key=d2_id)
            elif mode == 'delete':
                print 'Deleting %s' % exchange
                channel.exchange_delete(exchange=exchange)
    # Declare *unrelated* topic exchange
    channel.exchange_declare(
        exchange='t1',
        type='topic',
        durable=True,
        auto_delete=False)
    i = 0
    while True:
        i = i + 1
        if i % 5000 == 0:
            print 'Sent %d messages' % i
        if sleep:
            time.sleep(sleep)
        channel.basic_publish(
            exchange='t1',
            routing_key=random_key(),
            body='null',
            properties=pika.BasicProperties(content_type='application/json'))
    connection.close()
    print 'Channel closed!'


if __name__ == '__main__':
    try:
        base_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)))
        filename = os.path.join(base_dir, 'bindings.csv')
        opts, _ = getopt.getopt(sys.argv[1:], "m:f:s:", ["mode=", "filename=", "sleep="])
        mode = 'skip'
        sleep = 0
        for opt, arg in opts:
            if opt == '-f':
                filename = arg
            elif opt == '-s':
                sleep = float(arg)
            elif opt == '-m':
                mode = arg
    except getopt.GetoptError:
        pass
    main(filename, mode, sleep)

bindings.csv