python/ovstest/tcp.py

   1 # Copyright (c) 2011, 2012 Nicira, Inc.
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at:
   6 #
   7 #     http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14
  15 """
  16 tcp module contains listener and sender classes for TCP protocol
  17 """
  18
  19 from twisted.internet.protocol import Factory, ClientFactory, Protocol
  20 from twisted.internet import interfaces
  21 from zope.interface import implements
  22 import time
  23
  24
  25 class TcpListenerConnection(Protocol):
  26     """
  27     This per-connection class is instantiated each time sender connects
  28     """
  29     def __init__(self):
  30         self.stats = 0
  31
  32     def dataReceived(self, data):
  33         self.stats += len(data)
  34
  35     def connectionLost(self, reason):
  36         self.factory.stats += self.stats
  37
  38
  39 class TcpListenerFactory(Factory):
  40     """
  41     This per-listening socket class is used to
  42     instantiate TcpListenerConnections
  43     """
  44     protocol = TcpListenerConnection
  45
  46     def __init__(self):
  47         self.stats = 0
  48
  49     def getResults(self):
  50         """ returns the number of bytes received as string"""
  51         # XML RPC does not support 64bit int (http://bugs.python.org/issue2985)
  52         # so we have to convert the amount of bytes into a string
  53         return str(self.stats)
  54
  55
  56 class Producer(object):
  57     implements(interfaces.IPushProducer)
  58     """
  59     This producer class generates infinite byte stream for a specified time
  60     duration
  61     """
  62     def __init__(self, proto, duration):
  63         self.proto = proto
  64         self.start = time.time()
  65         self.produced = 0
  66         self.paused = False
  67         self.data = "X" * 65535
  68         self.duration = duration
  69
  70     def pauseProducing(self):
  71         """This function is called whenever write() to socket would block"""
  72         self.paused = True
  73
  74     def resumeProducing(self):
  75         """This function is called whenever socket becomes writable"""
  76         self.paused = False
  77         current = time.time()
  78         while (not self.paused) and (current < self.start + self.duration):
  79             self.proto.transport.write(self.data)
  80             self.produced += len(self.data)
  81             current = time.time()
  82         if current >= self.start + self.duration:
  83             self.proto.factory.stats += self.produced
  84             self.proto.transport.unregisterProducer()
  85             self.proto.transport.loseConnection()
  86
  87     def stopProducing(self):
  88         pass
  89
  90
  91 class TcpSenderConnection(Protocol):
  92     """
  93     TCP connection instance class that sends all traffic at full speed.
  94     """
  95
  96     def connectionMade(self):
  97         producer = Producer(self, self.factory.duration)
  98         self.transport.registerProducer(producer, True)
  99         producer.resumeProducing()
 100
 101     def dataReceived(self, data):
 102         self.transport.loseConnection()
 103
 104
 105 class TcpSenderFactory(ClientFactory):
 106     """
 107     This factory is responsible to instantiate TcpSenderConnection classes
 108     each time sender initiates connection
 109     """
 110     protocol = TcpSenderConnection
 111
 112     def __init__(self, duration):
 113         self.duration = duration
 114         self.stats = 0
 115
 116     def getResults(self):
 117         """Returns amount of bytes sent to the Listener (as a string)"""
 118         return str(self.stats)