relay-throughput-testing/src/throughput_protocols.py

#!/usr/bin/python3
#
import basic_protocols
import logging
import enum
import time
import socket
import struct
#
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 SEND_RESULTS 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.sub_protocol = basic_protocols.SendDataProtocol(self.socket,
				                                                     struct.pack('d', self.sub_protocol.time_started_push))
				self.state = self.states.SEND_RESULTS
			#
		#
		if self.state is self.states.SEND_RESULTS:
			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 RECV_RESULTS DONE')
		self.state = self.states.READY_TO_BEGIN
		#
		self.sub_protocol = None
		self.results = {}
	#
	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.results['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():
				self.results['data_size'] = self.sub_protocol.data_size
				self.results['time_of_first_byte'] = self.sub_protocol.time_of_first_byte
				self.results['time_of_last_byte'] = self.sub_protocol.time_of_last_byte
				self.results['transfer_rate'] = self.sub_protocol.calc_transfer_rate()
				self.results['deltas'] = self.sub_protocol.deltas
				#
				self.sub_protocol = basic_protocols.ReceiveDataProtocol(self.socket)
				self.state = self.states.RECV_RESULTS
			#
		#
		if self.state is self.states.RECV_RESULTS:
			if self.sub_protocol.run():
				time_started_push = struct.unpack('d', self.sub_protocol.received_data)[0]
				self.results['time_started_push'] = time_started_push
				#
				if self.results_callback:
					self.results_callback(self.results)
				#
				self.state = self.states.DONE
			#
		#
		if self.state is self.states.DONE:
			return True
		#
		return False
	#
#