~exarkun/+junk/twisted-benchmarks

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
"""
This benchmarks runs a trivial Twisted TCP echo server and a client pumps as
much data to it as it can in a fixed period of time.

The size of the string passed to each write call may play a significant
factor in the performance of this benchmark.
"""

from twisted.internet.defer import Deferred
from twisted.internet.endpoints import TCP4ClientEndpoint
from twisted.internet.protocol import ServerFactory, Factory, Protocol
from twisted.protocols.wire import Echo

from benchlib import driver


class Counter(Protocol):
    count = 0

    def dataReceived(self, bytes):
        self.count += len(bytes)



class Client(object):
    _finished = None

    def __init__(self, reactor, server):
        self._reactor = reactor
        self._server = server


    def run(self, duration, chunkSize):
        self._duration = duration
        self._bytes = 'x' * chunkSize
        # Set up a connection
        factory = Factory()
        factory.protocol = Counter
        d = self._server.connect(factory)
        d.addCallback(self._connected)
        return d


    def _connected(self, client):
        self._client = client
        self._stopCall = self._reactor.callLater(self._duration, self._stop)
        client.transport.registerProducer(self, False)
        self._finished = Deferred()
        return self._finished


    def _stop(self):
        self.stopProducing()
        self._client.transport.unregisterProducer()
        self._finish(self._client.count)


    def _finish(self, value):
        if self._finished is not None:
            finished = self._finished
            self._finished = None
            finished.callback(value)


    def resumeProducing(self):
        self._client.transport.write(self._bytes)


    def stopProducing(self):
        self._client.transport.loseConnection()


    def connectionLost(self, reason):
        self._finish(reason)



def main(reactor, duration):
    chunkSize = 16384

    server = ServerFactory()
    server.protocol = Echo
    serverPort = reactor.listenTCP(0, server)
    client = Client(
        reactor,
        TCP4ClientEndpoint(
            reactor, '127.0.0.1', serverPort.getHost().port))
    d = client.run(duration, chunkSize)
    return d


if __name__ == '__main__':
    import sys
    import tcp
    driver(tcp.main, sys.argv)