test-mass-agent.py

#!/usr/bin/env python
# -*- mode: python; coding: utf-8 -*-

from __future__ import absolute_import

import random, itertools
from twisted.internet import reactor, protocol, defer
from twisted.web import client, http_headers

# ********** BEGIN CONFIGURE CONSTANTS **********

DEBUG__USE_GC_COLLECT = True
DEBUG__GC_COLLECT_PERIOD = 1000

URL_TEMPLATE = 'http://example.org/?random=%s'
THREAD_COUNT = 100
URL_COUNT = 'inf' # integer or 'inf'
REQUEST_TIMEOUT = 120.0

# ********** END CONFIGURE CONSTANTS **********

def new_url():
    rnd = '%s-%s-%s-%s' % (
            random.randrange(1000000),
            random.randrange(1000000),
            random.randrange(1000000),
            random.randrange(1000000),
            )
    
    return URL_TEMPLATE % rnd

def new_url_iter():
    while True:
        yield new_url()

class DeliverCloser(protocol.Protocol, object):
    def __init__(self, deferred):
        self._deferred = deferred
    
    def dataReceived(self, data):
        self.transport.loseConnection()
    
    def connectionLost(self, reason):
        self._deferred.callback(None)

@defer.inlineCallbacks
def request_thread(thread_name, url_iter):
    for url in url_iter:
        print u'<%s> %r: opening...' % (thread_name, url)
        try:
            agent = client.Agent(reactor, connectTimeout=REQUEST_TIMEOUT)
            response = yield agent.request(
                    'GET',
                    url,
                    headers=http_headers.Headers({
                            'User-Agent': ['Twisted Web Client for ExampleOrg'],
                            }),
                    )
            response_deferred = defer.Deferred()
            
            response.deliverBody(DeliverCloser(response_deferred))
            yield response_deferred
            print u'<%s> %r: done! [success; code is %r]' % (thread_name, url, response.code)
        except Exception as e:
            print u'<%s> %r: done! [error; %r]' % (thread_name, url, e)
        
        if DEBUG__USE_GC_COLLECT and not random.randrange(DEBUG__GC_COLLECT_PERIOD):
            from gc import collect as gc_collect
            
            print u'<%s> gc_collect() is %r' % (thread_name, gc_collect())

@defer.inlineCallbacks
def bulk_request(url_iter):
    deferred_list = []
    
    print u'bulk: opening...'
    
    for thread_i in xrange(THREAD_COUNT):
        thread_name = 'thread-%r' % thread_i
        deferred = request_thread(thread_name, url_iter)
        deferred_list.append(deferred)
    
    print u'bulk: all opened. waiting...'
    
    for deferred in deferred_list:
        yield deferred
    
    print u'bulk: done!'

@defer.inlineCallbacks
def main():
    if URL_COUNT is None or URL_COUNT in ('inf', 'infinite'):
        url_iter = new_url_iter()
    else:
        url_iter = itertools.islice(new_url_iter(), URL_COUNT)
    
    yield bulk_request(url_iter)
    
    reactor.stop()

if __name__ == '__main__':
    reactor.callLater(0, main)
    reactor.run()