pounder_test/scripts/stream_throughput.py

157 lines
4.2 KiB
Python
Raw Normal View History

2021-06-15 19:18:16 +08:00
#!/usr/bin/python3
"""
Author: Ryan Summers
Description: Provides a mechanism for measuring Stabilizer stream data throughput.
"""
2021-07-23 21:14:30 +08:00
import argparse
2021-06-15 19:18:16 +08:00
import socket
import collections
import struct
import time
import logging
# Representation of a single data batch transmitted by Stabilizer.
Packet = collections.namedtuple('Packet', ['index', 'data'])
2021-06-15 19:18:16 +08:00
2021-07-26 19:07:07 +08:00
# The magic header half-word at the start of each packet.
MAGIC_HEADER = 0x057B
# The struct format of the header.
HEADER_FORMAT = '<HBBI'
2021-07-22 20:45:58 +08:00
# All supported formats by this reception script.
2021-07-26 19:07:07 +08:00
#
2021-07-27 19:12:57 +08:00
# The items in this dict are functions that will be provided the sample batch size and will return
2021-07-26 19:07:07 +08:00
# the struct deserialization code to unpack a single batch.
2021-07-22 20:45:58 +08:00
FORMAT = {
2021-07-26 19:07:07 +08:00
1: lambda batch_size: f'<{batch_size}H{batch_size}H{batch_size}H{batch_size}H'
2021-07-22 20:45:58 +08:00
}
2021-07-26 19:07:07 +08:00
def parse_packet(buf):
""" Attempt to parse packets from the received buffer. """
# Attempt to parse a block from the buffer.
if len(buf) < struct.calcsize(HEADER_FORMAT):
2021-07-27 19:12:57 +08:00
return
2021-07-26 19:07:07 +08:00
# Parse out the packet header
magic, format_id, batch_size, sequence_number = struct.unpack_from(HEADER_FORMAT, buf)
buf = buf[struct.calcsize(HEADER_FORMAT):]
if magic != MAGIC_HEADER:
logging.warning('Encountered bad magic header: %s', hex(magic))
2021-07-27 19:12:57 +08:00
return
2021-07-26 19:07:07 +08:00
frame_format = FORMAT[format_id](batch_size)
2021-07-26 19:47:03 +08:00
batch_count = int(len(buf) / struct.calcsize(frame_format))
2021-07-26 19:07:07 +08:00
packets = []
for offset in range(batch_count):
data = struct.unpack_from(frame_format, buf)
buf = buf[struct.calcsize(frame_format):]
2021-07-27 19:12:57 +08:00
yield Packet(sequence_number + offset, data)
2021-07-26 19:07:07 +08:00
2021-06-15 19:18:16 +08:00
class Timer:
""" A basic timer for measuring elapsed time periods. """
def __init__(self, period=1.0):
""" Create the timer with the provided period. """
self.start_time = time.time()
self.trigger_time = self.start_time + period
self.period = period
self.started = False
def is_triggered(self):
""" Check if the timer period has elapsed. """
now = time.time()
return now >= self.trigger_time
def start(self):
""" Start the timer. """
self.start_time = time.time()
self.started = True
def is_started(self):
""" Check if the timer has started. """
return self.started
def arm(self):
""" Arm the timer trigger. """
self.trigger_time = time.time() + self.period
def elapsed(self):
""" Get the elapsed time since the timer was started. """
now = time.time()
return now - self.start_time
2021-07-26 19:07:07 +08:00
def sequence_delta(previous_sequence, next_sequence):
""" Check the number of items between two sequence numbers. """
if previous_sequence is None:
return 0
2021-06-15 19:18:16 +08:00
2021-07-26 19:07:07 +08:00
delta = next_sequence - (previous_sequence + 1)
return delta & 0xFFFFFFFF
2021-06-15 19:18:16 +08:00
def main():
""" Main program. """
2021-07-23 21:14:30 +08:00
parser = argparse.ArgumentParser(description='Measure Stabilizer livestream quality')
parser.add_argument('--port', default=1111, help='The port that stabilizer is streaming to')
args = parser.parse_args()
2021-06-15 19:18:16 +08:00
connection = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
2021-07-23 21:14:30 +08:00
connection.bind(("", args.port))
2021-06-15 19:18:16 +08:00
logging.basicConfig(level=logging.INFO,
format='%(asctime)s.%(msecs)03d %(levelname)-8s %(message)s')
2021-07-26 19:07:07 +08:00
last_index = None
2021-06-15 19:18:16 +08:00
drop_count = 0
good_blocks = 0
2021-07-26 19:07:07 +08:00
total_bytes = 0
2021-06-15 19:18:16 +08:00
timer = Timer()
while True:
# Receive any data over UDP and parse it.
2021-07-27 19:12:57 +08:00
data = connection.recv(4096)
2021-06-15 19:18:16 +08:00
if data and not timer.is_started():
timer.start()
# Handle any received packets.
2021-07-26 19:07:07 +08:00
total_bytes += len(data)
2021-07-26 19:47:03 +08:00
for packet in parse_packet(data):
2021-06-15 19:18:16 +08:00
# Handle any dropped packets.
2021-07-26 19:07:07 +08:00
drop_count += sequence_delta(last_index, packet.index)
2021-06-15 19:18:16 +08:00
last_index = packet.index
good_blocks += 1
# Report the throughput periodically.
if timer.is_triggered():
2021-07-26 19:07:07 +08:00
drate = total_bytes * 8 / 1e6 / timer.elapsed()
2021-06-15 19:18:16 +08:00
print(f'''
Data Rate: {drate:.3f} Mbps
Received Blocks: {good_blocks}
Dropped blocks: {drop_count}
2021-07-26 19:07:07 +08:00
Metadata: {total_bytes / 1e6:.3f} MB in {timer.elapsed():.2f} s
2021-06-15 19:18:16 +08:00
----
''')
timer.arm()
if __name__ == '__main__':
main()