ESP32/control_iperf.py

#!/usr/bin/env python3
import asyncio
import argparse
import serial_asyncio
import sys
import re
import glob

# --- 1. Protocol Class (Async Logic) ---
class SerialController(asyncio.Protocol):
    def __init__(self, port_name, cmd_type, value, loop, completion_future, sync_event):
        self.port_name = port_name
        self.cmd_type = cmd_type  # 'start', 'stop', 'pps', 'status'
        self.value = value
        self.loop = loop
        self.transport = None
        self.buffer = ""
        self.completion_future = completion_future
        self.sync_event = sync_event # Global trigger for tight sync
        self.status_data = {}

        # Command Construction
        if self.cmd_type == 'pps':
            self.cmd_str = f"iperf pps {self.value}\n"
            self.target_key = "IPERF_PPS_UPDATED"
        elif self.cmd_type == 'status':
            self.cmd_str = "iperf status\n"
            self.target_key = "IPERF_STATUS"
        elif self.cmd_type == 'start':
            self.cmd_str = "iperf start\n"
            self.target_key = "IPERF_STARTED"
        elif self.cmd_type == 'stop':
            self.cmd_str = "iperf stop\n"
            self.target_key = "IPERF_STOPPED"

    def connection_made(self, transport):
        self.transport = transport
        # 1. Wake up the device immediately
        transport.write(b'\n')
        # 2. Schedule the command to wait for the global trigger
        self.loop.create_task(self.await_trigger_and_send())

    async def await_trigger_and_send(self):
        # 3. Wait here until trigger is fired (allows sync start/stop)
        if self.sync_event:
            await self.sync_event.wait()

        # 4. FIRE!
        self.transport.write(self.cmd_str.encode())

    def data_received(self, data):
        self.buffer += data.decode(errors='ignore')

        while '\n' in self.buffer:
            line, self.buffer = self.buffer.split('\n', 1)
            line = line.strip()

            # --- Status Parsing ---
            if self.cmd_type == 'status':
                if "IPERF_STATUS" in line:
                    m = re.search(r'Src=([\d\.]+), Dst=([\d\.]+), Running=(\d+), Config=(\d+), Actual=(\d+), Err=([-\d\.]+)%, Pkts=(\d+), AvgBW=([\d\.]+) Mbps', line)
                    if m:
                        self.status_data['main'] = {
                            'src': m.group(1), 'dst': m.group(2),
                            'run': "Run" if m.group(3) == '1' else "Stop",
                            'cfg': m.group(4), 'act': m.group(5),
                            'err': m.group(6), 'pkts': m.group(7), 'bw': m.group(8)
                        }

                elif "IPERF_STATES" in line:
                    m = re.search(r'TX=([\d\.]+)s/([\d\.]+)% \((\d+)\), SLOW=([\d\.]+)s/([\d\.]+)% \((\d+)\), STALLED=([\d\.]+)s/([\d\.]+)% \((\d+)\)', line)
                    if m:
                        self.status_data['states'] = {
                            'tx_t': m.group(1), 'tx_p': m.group(2), 'tx_c': m.group(3),
                            'sl_t': m.group(4), 'sl_p': m.group(5), 'sl_c': m.group(6),
                            'st_t': m.group(7), 'st_p': m.group(8), 'st_c': m.group(9)
                        }

                if 'main' in self.status_data and 'states' in self.status_data:
                     if not self.completion_future.done():
                         d = self.status_data['main']
                         s = self.status_data['states']
                         output = (f"{d['src']} -> {d['dst']} | {d['run']}, "
                                   f"Cfg:{d['cfg']}, Act:{d['act']}, Err:{d['err']}%, Pkts:{d['pkts']}, BW:{d['bw']}M | "
                                   f"TX:{s['tx_t']}s/{s['tx_p']}%({s['tx_c']}) "
                                   f"SL:{s['sl_t']}s/{s['sl_p']}%({s['sl_c']}) "
                                   f"ST:{s['st_t']}s/{s['st_p']}%({s['st_c']})")
                         self.completion_future.set_result(output)
                     self.transport.close()
                     return

            # --- Simple Command Parsing ---
            else:
                if self.target_key in line:
                    if not self.completion_future.done():
                        self.completion_future.set_result(True)
                    self.transport.close()
                    return

    def connection_lost(self, exc):
        if not self.completion_future.done():
            if self.cmd_type == 'status' and 'main' in self.status_data:
                 d = self.status_data['main']
                 output = (f"{d['src']} -> {d['dst']} | {d['run']}, "
                           f"Cfg:{d['cfg']}, Act:{d['act']}, BW:{d['bw']}M (Partial)")
                 self.completion_future.set_result(output)
            else:
                 self.completion_future.set_exception(Exception("Closed"))

# --- 2. Helper Functions ---
def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('action', choices=['start', 'stop', 'pps', 'status', 'step-all'])
    parser.add_argument('value_arg', nargs='?', type=int, help='Value for PPS')
    parser.add_argument('--value', type=int, help='Value for PPS')
    parser.add_argument('--devices', required=True, help="List (e.g. /dev/ttyUSB0, /dev/ttyUSB1), Range (/dev/ttyUSB0-29), or 'all'")

    args = parser.parse_args()
    if args.value_arg is not None: args.value = args.value_arg
    if args.action == 'pps' and args.value is None:
        print("Error: 'pps' action requires a value")
        sys.exit(1)
    return args

def natural_sort_key(s):
    return [int(text) if text.isdigit() else text.lower()
            for text in re.split('([0-9]+)', s)]

def expand_devices(device_str):
    if device_str.lower() == 'all':
        devices = glob.glob('/dev/ttyUSB*')
        devices.sort(key=natural_sort_key)
        if not devices:
            print("Error: No /dev/ttyUSB* devices found!")
            sys.exit(1)
        return devices

    devices = []
    parts = [d.strip() for d in device_str.split(',')]
    for part in parts:
        match = re.match(r'^(.*?)(\d+)-(\d+)$', part)
        if match:
            prefix, start, end = match.group(1), int(match.group(2)), int(match.group(3))
            step = 1 if end >= start else -1
            for i in range(start, end + step, step):
                devices.append(f"{prefix}{i}")
        else:
            devices.append(part)
    return devices

# --- 3. Async Execution Core ---
async def run_single_cmd(port, cmd_type, value, sync_event=None):
    """Runs a single command on a single port."""
    loop = asyncio.get_running_loop()
    fut = loop.create_future()
    try:
        await serial_asyncio.create_serial_connection(
            loop,
            lambda: SerialController(port, cmd_type, value, loop, fut, sync_event),
            port,
            baudrate=115200
        )
        return await asyncio.wait_for(fut, timeout=5.0)
    except asyncio.TimeoutError:
        return None
    except Exception:
        return None

async def run_parallel_action(devices, action, value):
    """Runs the specified action on all devices in parallel."""
    print(f"Initializing {len(devices)} devices for '{action}'...")
    sync_event = asyncio.Event()
    tasks = [run_single_cmd(d, action, value, sync_event) for d in devices]

    # Allow connections to settle
    await asyncio.sleep(0.5)

    # Fire all commands at once
    sync_event.set()
    results = await asyncio.gather(*tasks)

    print("\nResults:")
    for dev, res in zip(devices, results):
        if action == 'status':
            print(f"{dev}: {res if res else 'TIMEOUT'}")
        else:
            status = "OK" if res is True else "FAIL"
            print(f"{dev}: {status}")

async def run_step_all(devices):
    """Stops all, then starts/stops devices one by one."""
    print("\n>>> STEP-ALL PHASE 1: STOPPING ALL DEVICES <<<")
    await run_parallel_action(devices, 'stop', None)

    print("\n>>> STEP-ALL PHASE 2: SEQUENTIAL TEST <<<")
    for i, dev in enumerate(devices):
        print(f"\n[{i+1}/{len(devices)}] Testing {dev}...")

        # Start
        print(f"  -> Starting {dev}...")
        res_start = await run_single_cmd(dev, 'start', None, None) # No sync needed for single
        if res_start is not True:
            print(f"  -> FAILED to start {dev}. Skipping.")
            continue

        # Wait (Run Traffic)
        print("  -> Running traffic (5 seconds)...")
        await asyncio.sleep(5)

        # Stop
        print(f"  -> Stopping {dev}...")
        res_stop = await run_single_cmd(dev, 'stop', None, None)
        if res_stop is not True:
            print(f"  -> Warning: Failed to stop {dev}")
        else:
            print(f"  -> {dev} OK")

        # Wait (Gap between devices)
        print("  -> Waiting 1 second...")
        await asyncio.sleep(1)

async def async_main(args, devices):
    if args.action == 'step-all':
        await run_step_all(devices)
    else:
        await run_parallel_action(devices, args.action, args.value)

# --- 4. Main Execution Block ---
if __name__ == '__main__':
    args = parse_arguments()
    dev_list = expand_devices(args.devices)

    if sys.platform == 'win32':
        asyncio.set_event_loop(asyncio.ProactorEventLoop())

    asyncio.run(async_main(args, dev_list))