pounder_test/scripts/stream_throughput.py

#!/usr/bin/python3
"""
Author: Ryan Summers

Description: Provides a mechanism for measuring Stabilizer stream data throughput.
"""
import argparse
import logging
import sys
import time

#from stabilizer.stream import StabilizerStream

import socket
import struct

class StabilizerStream:
    """ Provides access to Stabilizer's livestreamed data. """

    # The magic header half-word at the start of each packet.
    MAGIC_HEADER = 0x057B

    # The struct format of the header.
    HEADER_FORMAT = '<HBBI'

    # All supported formats by this reception script.
    #
    # The items in this dict are functions that will be provided the sample batch size and will
    # return the struct deserialization code to unpack a single batch.
    FORMAT = {
        1: lambda batch_size: f'<{batch_size}H{batch_size}H{batch_size}H{batch_size}H'
    }

    def __init__(self, port, timeout=None):
        """ Initialize the stream.

        Args:
            port: The UDP port to receive the stream from.
            timeout: The timeout to set on the UDP socket.
        """
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.socket.bind(("", port))
        self.total_bytes = 0

        if timeout is not None:
            self.socket.settimeout(timeout)


    def clear(self, duration=1):
        """ Clear the socket RX buffer by reading all available data.

        Args:
            duration: The maximum duration in seconds to read data for.
        """
        time.sleep(duration)

        try:
            while self.socket.recv(4096):
                pass
        except socket.timeout:
            pass


    def get_rx_bytes(self):
        """ Get the number of bytes read from the stream. """
        return self.total_bytes


    def read_frame(self):
        """ Read a single frame from the stream.

        Returns:
            Yields the (seqnum, data) of the batches available in the frame.
        """
        buf = self.socket.recv(4096)
        self.total_bytes += len(buf)

        # Attempt to parse a block from the buffer.
        if len(buf) < struct.calcsize(self.HEADER_FORMAT):
            return

        # Parse out the packet header
        magic, format_id, batch_size, sequence_number = struct.unpack_from(self.HEADER_FORMAT, buf)
        buf = buf[struct.calcsize(self.HEADER_FORMAT):]

        if magic != self.MAGIC_HEADER:
            logging.warning('Encountered bad magic header: %s', hex(magic))
            return

        frame_format = self.FORMAT[format_id](batch_size)

        batch_count = int(len(buf) / struct.calcsize(frame_format))

        for offset in range(batch_count):
            data = struct.unpack_from(frame_format, buf)
            buf = buf[struct.calcsize(frame_format):]
            yield (sequence_number + offset, data)



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


def sequence_delta(previous_sequence, next_sequence):
    """ Check the number of items between two sequence numbers. """
    if previous_sequence is None:
        return 0

    delta = next_sequence - (previous_sequence + 1)
    return delta & 0xFFFFFFFF


def main():
    """ Main program. """
    parser = argparse.ArgumentParser(description='Measure Stabilizer livestream quality')
    parser.add_argument('--port', type=int, default=1883,
                        help='The port that stabilizer is streaming to')

    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO,
                        format='%(asctime)s.%(msecs)03d %(levelname)-8s %(message)s')

    last_index = None

    drop_count = 0
    good_blocks = 0
    total_bytes = 0

    timer = Timer()

    stream = StabilizerStream(args.port)

    while True:
        # Receive any data over UDP and parse it.
        for (seqnum, _data) in stream.read_frame():
            if not timer.is_started():
                timer.start()

            # Handle any dropped packets.
            drop_count += sequence_delta(last_index, seqnum)
            last_index = seqnum
            good_blocks += 1

            if 1000 < good_blocks and good_blocks < 1020 :
                print(_data[:8])

        # Report the throughput periodically.
        if timer.is_triggered():
            drate = stream.get_rx_bytes() * 8 / 1e6 / timer.elapsed()

            print(f'''
Data Rate:       {drate:.3f} Mbps
Received Blocks: {good_blocks}
Dropped blocks:  {drop_count}

Metadata: {total_bytes / 1e6:.3f} MB in {timer.elapsed():.2f} s
----
''')
            sys.stdout.flush()
            timer.arm()


if __name__ == '__main__':
    main()