throughput_protocols.py 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  1. #!/usr/bin/python3
  2. #
  3. import basic_protocols
  4. import logging
  5. import enum
  6. import time
  7. import socket
  8. import struct
  9. #
  10. class ClientProtocol(basic_protocols.Protocol):
  11. def __init__(self, socket, total_bytes, send_buffer_len=None, use_acceleration=None, custom_data=b'', push_start_cb=None, push_done_cb=None):
  12. self.socket = socket
  13. self.total_bytes = total_bytes
  14. self.send_buffer_len = send_buffer_len
  15. self.use_acceleration = use_acceleration
  16. self.custom_data = custom_data
  17. self.push_start_cb = push_start_cb
  18. self.push_done_cb = push_done_cb
  19. #
  20. self.states = enum.Enum('CLIENT_CONN_STATES', 'READY_TO_BEGIN SEND_CUSTOM_DATA PUSH_DATA SEND_RESULTS DONE')
  21. self.state = self.states.READY_TO_BEGIN
  22. #
  23. self.sub_protocol = None
  24. #
  25. def _run_iteration(self):
  26. if self.state is self.states.READY_TO_BEGIN:
  27. self.sub_protocol = basic_protocols.SendDataProtocol(self.socket, self.custom_data)
  28. self.state = self.states.SEND_CUSTOM_DATA
  29. #
  30. if self.state is self.states.SEND_CUSTOM_DATA:
  31. if self.sub_protocol.run():
  32. self.sub_protocol = basic_protocols.PushDataProtocol(self.socket, self.total_bytes,
  33. send_buffer_len=self.send_buffer_len,
  34. use_acceleration=self.use_acceleration,
  35. push_start_cb=self.push_start_cb,
  36. push_done_cb=self.push_done_cb)
  37. self.state = self.states.PUSH_DATA
  38. #
  39. #
  40. if self.state is self.states.PUSH_DATA:
  41. if self.sub_protocol.run():
  42. self.sub_protocol = basic_protocols.SendDataProtocol(self.socket,
  43. struct.pack('d', self.sub_protocol.time_started_push))
  44. self.state = self.states.SEND_RESULTS
  45. #
  46. #
  47. if self.state is self.states.SEND_RESULTS:
  48. if self.sub_protocol.run():
  49. self.state = self.states.DONE
  50. #
  51. #
  52. if self.state is self.states.DONE:
  53. return True
  54. #
  55. return False
  56. #
  57. #
  58. class ServerProtocol(basic_protocols.Protocol):
  59. def __init__(self, socket, results_callback=None, use_acceleration=None):
  60. self.socket = socket
  61. self.results_callback = results_callback
  62. self.use_acceleration = use_acceleration
  63. #
  64. self.states = enum.Enum('SERVER_CONN_STATES', 'READY_TO_BEGIN RECV_CUSTOM_DATA PULL_DATA RECV_RESULTS DONE')
  65. self.state = self.states.READY_TO_BEGIN
  66. #
  67. self.sub_protocol = None
  68. self.results = {}
  69. #
  70. def _run_iteration(self):
  71. if self.state is self.states.READY_TO_BEGIN:
  72. self.sub_protocol = basic_protocols.ReceiveDataProtocol(self.socket)
  73. self.state = self.states.RECV_CUSTOM_DATA
  74. #
  75. if self.state is self.states.RECV_CUSTOM_DATA:
  76. if self.sub_protocol.run():
  77. self.results['custom_data'] = self.sub_protocol.received_data
  78. #
  79. self.sub_protocol = basic_protocols.PullDataProtocol(self.socket, use_acceleration=self.use_acceleration)
  80. self.state = self.states.PULL_DATA
  81. #
  82. #
  83. if self.state is self.states.PULL_DATA:
  84. if self.sub_protocol.run():
  85. self.results['data_size'] = self.sub_protocol.data_size
  86. self.results['time_of_first_byte'] = self.sub_protocol.time_of_first_byte
  87. self.results['time_of_last_byte'] = self.sub_protocol.time_of_last_byte
  88. self.results['transfer_rate'] = self.sub_protocol.calc_transfer_rate()
  89. self.results['deltas'] = self.sub_protocol.deltas
  90. #
  91. self.sub_protocol = basic_protocols.ReceiveDataProtocol(self.socket)
  92. self.state = self.states.RECV_RESULTS
  93. #
  94. #
  95. if self.state is self.states.RECV_RESULTS:
  96. if self.sub_protocol.run():
  97. time_started_push = struct.unpack('d', self.sub_protocol.received_data)[0]
  98. self.results['time_started_push'] = time_started_push
  99. #
  100. if self.results_callback:
  101. self.results_callback(self.results)
  102. #
  103. self.state = self.states.DONE
  104. #
  105. #
  106. if self.state is self.states.DONE:
  107. return True
  108. #
  109. return False
  110. #
  111. #