putbench.py

putbench.py

import swift.common.utils
import StringIO
import pprint
import random
import sys, os
import requests
from commandr import command, Run

from gevent import monkey
monkey.patch_all()

import gevent
from gevent.queue import JoinableQueue

# Monkey patch swift.common.utils code so we don't need
# to create /etc/swift/swif.conf with the hashes
swift.common.utils.HASH_PATH_SUFFIX = '0d85cf6346c7086d'
swift.common.utils.HASH_PATH_PREFIX = '64543f1f5108d509'

from swift.common.ring import Ring
from time import time

def generate_names(ntype, times):
    return ["%s_%s" % (ntype,n) for n in xrange(times)]

def do_put():
    while True:
        next_item = queue.get()
        node_ip = next_item.target_node
        account = next_item.account
        container = next_item.container
        obj = next_item.obj
        data = StringIO.StringIO(next_item.data)

        # Get available nodes in the ring for this object/container
        # and the target partition for the object
        opart, onodes = object_ring.get_nodes(account, container, obj)
        cpart, cnodes = container_ring.get_nodes(account, container)

        # We're only interested in benchmarking target_node,
        # discard the rest
        node = [n for n in onodes if n['ip'] == node_ip][0]

        # Choose a random container server
        cnode = cnodes[int(random.random() * 3)]

        # container related headers we will send
        x_container_host = '%s:%s' % (cnode['ip'], cnode['port'])
        x_container_partition = cpart
        x_container_device = '%s' % cnode['device']

        # save container host/device distribution stats
        cdkey = x_container_host + ":" + x_container_device
        if not cdevices.has_key(cdkey):
            cdevices[cdkey] = 0
        cdevices[cdkey] += 1
        
        # save object host/device distribution stats
        odkey = "%s:%s:%s" % (node_ip, node['port'], node['device'])
        if not odevices.has_key(odkey):
            odevices[odkey] = 0
        odevices[odkey] += 1

        # create the request
        path = "%s/%s/%s/%s/%s" % (node['device'], opart, account, container, obj)
        url = "http://%s:6000/%s" % (node_ip, path)
        headers = { 
            'X-Container-Host': x_container_host,
            'X-Container-Partition': x_container_partition,
            'X-Container-Device': x_container_device,
            'X-Timestamp': time(),
            'Content-Type': 'application/octet-stream'
        }

        # send the request
        response = requests.put(url, headers = headers, data = data)

        # 201 if PUT was successful
        rcode = response.status_code

        # some more stats
        if not rcodes.has_key(rcode):
            rcodes[rcode] = 0
        rcodes[rcode] += 1

        # tell gevent we're done with the task
        queue.task_done()

class Data(object):
    def __init__(self, target_node, account, container, obj, data):
        self.target_node = target_node
        self.account = account
        self.container = container
        self.obj = obj
        self.data = data

@command('run')
def run(target_node, account, count = 1000, object_prefix = 'object_',
        container_prefix = 'test', container_count = 1, data_size = 1024,
        concurrency = 100):
    
    print "Target node:       %s" % target_node
    print "Concurrency:       %s" % concurrency
    print "Account:           %s" % account
    print "PUT count:         %s" % count
    print "Container prefix:  %s" % container_prefix
    print "Container count:   %s" % container_count
    print "Data size:         %s" % data_size
    print

    tstart = time()
    # We want a file or StringIO here
    # otherwise memory usage goes sky high when sending
    # the put request
    data = 'x' * data_size

    for n in range(concurrency):
        gevent.spawn(do_put)

    if container_count > 1:
        containers = generate_names(container_prefix, container_count) 

    for obj in xrange(count):
        o = "%s_%s" % (object_prefix, obj)
        if container_count > 1:
            # choose a random container
            r = int(round(random.random()*(len(containers) - 1)))
            c = containers[r]
        else:
            c = container_prefix
        queue.put(Data(target_node, account, c, o, data))

    queue.join()
    ttime = time() - tstart

    printer = pprint.PrettyPrinter(indent = 4)
    print "Took %.2f seconds, %.2f req/s" % (ttime, sum(rcodes.itervalues())/ttime)
    print "Response codes distribution: %s" % rcodes

    print "Container distribution:"
    printer.pprint(cdevices)

    print "Object distribution:"
    printer.pprint(odevices)

# Ring objects
object_ring    = Ring('./', ring_name = 'object')
container_ring = Ring('./', ring_name = 'container')
account_ring = Ring('./', ring_name = 'account')
queue = JoinableQueue()
rcodes = {}
cdevices = {}
odevices = {}

if __name__ == '__main__':
      Run()