relay-throughput-testing/src/throughput_protocols.py

#!/usr/bin/python3
#
import basic_protocols
import logging
import enum
import time
import socket
#
class ClientProtocol(basic_protocols.Protocol):
	def __init__(self, socket, total_bytes, send_buffer_len=None, use_acceleration=None, custom_data=b'', push_start_cb=None, push_done_cb=None):
		self.socket = socket
		self.total_bytes = total_bytes
		self.send_buffer_len = send_buffer_len
		self.use_acceleration = use_acceleration
		self.custom_data = custom_data
		self.push_start_cb = push_start_cb
		self.push_done_cb = push_done_cb
		#
		self.states = enum.Enum('CLIENT_CONN_STATES', 'READY_TO_BEGIN SEND_CUSTOM_DATA PUSH_DATA DONE')
		self.state = self.states.READY_TO_BEGIN
		#
		self.sub_protocol = None
	#
	def _run_iteration(self):
		if self.state is self.states.READY_TO_BEGIN:
			self.sub_protocol = basic_protocols.SendDataProtocol(self.socket, self.custom_data)
			self.state = self.states.SEND_CUSTOM_DATA
		#
		if self.state is self.states.SEND_CUSTOM_DATA:
			if self.sub_protocol.run():
				self.sub_protocol = basic_protocols.PushDataProtocol(self.socket, self.total_bytes,
				                                                     send_buffer_len=self.send_buffer_len,
				                                                     use_acceleration=self.use_acceleration,
				                                                     push_start_cb=self.push_start_cb,
				                                                     push_done_cb=self.push_done_cb)
				self.state = self.states.PUSH_DATA
			#
		#
		if self.state is self.states.PUSH_DATA:
			if self.sub_protocol.run():
				self.state = self.states.DONE
			#
		#
		if self.state is self.states.DONE:
			return True
		#
		return False
	#
#
class ServerProtocol(basic_protocols.Protocol):
	def __init__(self, socket, results_callback=None, use_acceleration=None):
		self.socket = socket
		self.results_callback = results_callback
		self.use_acceleration = use_acceleration
		#
		self.states = enum.Enum('SERVER_CONN_STATES', 'READY_TO_BEGIN RECV_CUSTOM_DATA PULL_DATA DONE')
		self.state = self.states.READY_TO_BEGIN
		#
		self.sub_protocol = None
		self.custom_data = None
	#
	def _run_iteration(self):
		if self.state is self.states.READY_TO_BEGIN:
			self.sub_protocol = basic_protocols.ReceiveDataProtocol(self.socket)
			self.state = self.states.RECV_CUSTOM_DATA
		#
		if self.state is self.states.RECV_CUSTOM_DATA:
			if self.sub_protocol.run():
				self.custom_data = self.sub_protocol.received_data
				#
				self.sub_protocol = basic_protocols.PullDataProtocol(self.socket, use_acceleration=self.use_acceleration)
				self.state = self.states.PULL_DATA
			#
		#
		if self.state is self.states.PULL_DATA:
			if self.sub_protocol.run():
				if self.results_callback:
					results = {}
					results['custom_data'] = self.custom_data
					results['data_size'] = self.sub_protocol.data_size
					results['time_of_first_byte'] = self.sub_protocol.time_of_first_byte
					results['time_of_last_byte'] = self.sub_protocol.time_of_last_byte
					results['transfer_rate'] = self.sub_protocol.calc_transfer_rate()
					results['deltas'] = self.sub_protocol.deltas
					#
					self.results_callback(results)
				#
				self.state = self.states.DONE
			#
		#
		if self.state is self.states.DONE:
			return True
		#
		return False
	#
#