1
#!/usr/bin/env python3
2
"""Exercise two S3 DataChannels, SFU fanout, and per-viewer reply permission."""
3
import asyncio
4
import json
5
from pathlib import Path
6
import runpy
7
import time
8
 
9
from aiortc import RTCPeerConnection, RTCConfiguration, RTCSessionDescription
10
 
11
ROOT = Path(__file__).resolve().parents[2]
12
OUT = ROOT / "artifacts/archive/sfu-bringup"
13
base = runpy.run_path(str(Path(__file__).with_name("probe-sfu.py")))
14
PROFILES = {
15
    "robot": {"ordered": True},
16
    "spectrum": {"ordered": False, "maxRetransmits": 0},
17
}
18
 
19
 
20
async def until(test, device, timeout=15):
21
    deadline = time.monotonic() + timeout
22
    while time.monotonic() < deadline:
23
        device.poll()
24
        if test():
25
            return
26
        await asyncio.sleep(0.01)
27
    raise TimeoutError("Condition not met before deadline")
28
 
29
 
30
async def main():
31
    api = base["SFU"]()
32
    device = base["Device"]()
33
    hosts = []
34
    resources = []
35
    local_channels = {}
36
    result = {"profiles": PROFILES, "messages_sent_per_profile": 100}
37
    try:
38
        await device.handshake()
39
        device.send(cmd="peer_init", initiator=True)
40
        response = await device.wait(lambda e: e.get("cmd") == "peer_init")
41
        if response["result"] != 0:
42
            raise RuntimeError("Peer initialization failed")
43
        offer = await device.wait(lambda e: e.get("event") == "sdp")
44
        source = await api.call("/sessions/new", {
45
            "sessionDescription": {"type": "offer", "sdp": offer["text"]}})
46
        source_id = source["sessionId"]
47
        device.send(cmd="sdp", text=source["sessionDescription"]["sdp"])
48
        await device.wait(lambda e: e.get("event") == "peer_state" and e.get("state") == 9)
49
        result["s3_connected"] = True
50
 
51
        for label, profile in PROFILES.items():
52
            created = await api.call(f"/sessions/{source_id}/datachannels/new", {
53
                "dataChannels": [{"location": "local", "dataChannelName": label, **profile}]})
54
            item = created["dataChannels"][0]
55
            if item.get("errorCode"):
56
                raise RuntimeError(item["errorCode"])
57
            local_channels[label] = item["id"]
58
            resources.append((source_id, item["id"]))
59
            # Register local send state too; enabling the SCTP transport alone
60
            # does not create a channel when manual_ch_create is enabled.
61
            command = {"cmd": "create_channel", "label": label, "ordered": profile["ordered"]}
62
            if "maxRetransmits" in profile:
63
                command["max_retransmits"] = profile["maxRetransmits"]
64
            device.send(**command)
65
            created_local = await device.wait(lambda e: e.get("cmd") == "create_channel")
66
            if created_local["result"] != 0:
67
                raise RuntimeError(f"Local {label} creation: {created_local['result']}")
68
        result["source_stream_ids"] = local_channels
69
 
70
        for viewer in range(2):
71
            peer = RTCPeerConnection(RTCConfiguration(iceServers=[]))
72
            endpoint = {"peer": peer, "channels": {}, "received": {label: [] for label in PROFILES}}
73
            hosts.append(endpoint)
74
            sid, transport = await api.establish()
75
            endpoint["session_id"] = sid
76
            await peer.setRemoteDescription(RTCSessionDescription(**transport["sessionDescription"]))
77
            await peer.setLocalDescription(await peer.createAnswer())
78
            await api.call(f"/sessions/{sid}/renegotiate", {
79
                "sessionDescription": {"type": "answer", "sdp": peer.localDescription.sdp}}, method="PUT")
80
            await until(lambda: peer.connectionState == "connected", device)
81
            for label, profile in PROFILES.items():
82
                pulled = await api.call(f"/sessions/{sid}/datachannels/new", {
83
                    "dataChannels": [{"location": "remote", "sessionId": source_id, "dataChannelName": label,
84
                        "waitForAck": True, "canReply": viewer == 0 and label == "robot", **profile}]})
85
                item = pulled["dataChannels"][0]
86
                if item.get("errorCode"):
87
                    raise RuntimeError(item["errorCode"])
88
                resources.append((sid, item["id"]))
89
                channel = peer.createDataChannel(label, negotiated=True, id=item["id"], **profile)
90
                endpoint["channels"][label] = channel
91
                channel.on("message", endpoint["received"][label].append)
92
                await until(lambda: channel.readyState == "open", device)
93
                channel.send("subscriber-ready")
94
        print("Both viewers connected; sending two profiles", flush=True)
95
        start = time.monotonic()
96
        for sequence in range(100):
97
            for label in PROFILES:
98
                payload = {"kind": label, "sequence": sequence}
99
                if label == "robot":
100
                    payload["value"] = sequence * 17 % 101
101
                else:
102
                    payload["bands"] = [(sequence * 7 + band * 11) % 256 for band in range(32)]
103
                device.send(cmd="send", stream_id=local_channels[label], text=json.dumps(payload))
104
                sent = await device.wait(lambda e: e.get("cmd") == "send")
105
                if sent["result"] != 0:
106
                    raise RuntimeError(f"Send on {label}: {sent['result']}")
107
            await asyncio.sleep(max(0, start + (sequence + 1) * 0.05 - time.monotonic()))
108
        result["send_duration_seconds"] = round(time.monotonic() - start, 3)
109
        await until(lambda: all(len(host["received"]["robot"]) == 100 for host in hosts), device)
110
        await asyncio.sleep(0.5)
111
        result["receivers"] = []
112
        for host in hosts:
113
            metrics = {}
114
            for label, messages in host["received"].items():
115
                decoded = [json.loads(msg) for msg in messages]
116
                sequence = [msg["sequence"] for msg in decoded]
117
                metrics[label] = {"received": len(messages), "unique": len(set(sequence)),
118
                    "all_in_order": sequence == list(range(100)),
119
                    "payloads_valid": all(msg.get("kind") == label and
120
                        (msg.get("value") == msg["sequence"] * 17 % 101 if label == "robot" else
121
                         msg.get("bands") == [(msg["sequence"] * 7 + band * 11) % 256 for band in range(32)])
122
                        for msg in decoded)}
123
            result["receivers"].append(metrics)
124
 
125
        async def command(viewer, command_id, color):
126
            started = time.monotonic()
127
            hosts[viewer]["channels"]["robot"].send(json.dumps({"led": color, "command_id": command_id}))
128
            received = await device.wait(lambda e: e.get("event") == "data_received" and
129
                json.loads(e.get("text", "{}" )).get("command_id") == command_id, timeout=5)
130
            led = await device.wait(lambda e: e.get("event") == "led" and
131
                [e.get("r"), e.get("g"), e.get("b")] == color, timeout=2)
132
            return {"received": bool(received), "led_result": led["result"],
133
                "observed_ms_including_usb": round((time.monotonic() - started) * 1000, 1)}
134
 
135
        result["controller_command"] = await command(0, 101, [16, 0, 16])
136
        hosts[1]["channels"]["robot"].send('{"led":[16,0,0],"command_id":999}')
137
        deadline = time.monotonic() + 1
138
        while time.monotonic() < deadline:
139
            device.poll()
140
            await asyncio.sleep(0.02)
141
        result["spectator_command_blocked"] = not any(e.get("event") == "data_received" and
142
            json.loads(e.get("text", "{}" )).get("command_id") == 999 for e in device.events)
143
        await api.call(f"/sessions/{hosts[1]['session_id']}/datachannels/update", {
144
            "dataChannels": [{"location": "remote", "sessionId": source_id,
145
                "dataChannelName": "robot", "canReply": True}]}, method="PUT")
146
        result["new_controller_command"] = await command(1, 102, [0, 16, 0])
147
    except Exception as exc:
148
        result["error"] = f"{type(exc).__name__}: {exc}"
149
    finally:
150
        result["local_channel_close_results"] = {}
151
        for label in reversed(local_channels):
152
            try:
153
                device.send(cmd="close_channel", label=label)
154
                closed = await device.wait(lambda e: e.get("cmd") == "close_channel", timeout=3)
155
                result["local_channel_close_results"][label] = closed["result"]
156
            except Exception:
157
                result["local_channel_close_results"][label] = "timeout"
158
        # Keep polling the protocol loop while stream-reset acknowledgments arrive.
159
        end = time.monotonic() + 1
160
        while time.monotonic() < end:
161
            device.poll()
162
            await asyncio.sleep(0.02)
163
        result["channel_closed_events"] = [e for e in device.events if e.get("event") == "channel_closed"]
164
        for sid, cid in reversed(resources):
165
            try:
166
                await api.call(f"/sessions/{sid}/datachannels/close", {"dataChannels": [{"id": cid}]}, method="PUT")
167
            except Exception:
168
                result["cleanup_incomplete"] = True
169
        for host in hosts:
170
            await host["peer"].close()
171
        try:
172
            device.send(cmd="peer_close")
173
            closed = await device.wait(lambda e: e.get("cmd") == "peer_close", timeout=5)
174
            result["peer_close_result"] = closed["result"]
175
        except Exception:
176
            result["peer_close_result"] = "timeout"
177
        device.close()
178
        (OUT / "datachannels-result.json").write_text(json.dumps(result, indent=2) + "\n")
179
        print(json.dumps(result, indent=2), flush=True)
180
 
181
 
182
if __name__ == "__main__":
183
    asyncio.run(main())