aboutsummaryrefslogtreecommitdiffstats
path: root/packages/meshbay-node/tests/test_proactor_udp.py
blob: 0b22b4eb9eb53e29489a08a53048e92be115fe7f (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
"""
On Windows the node's datagrams go straight through the socket, not one
overlapped operation per datagram and per turn of the loop (`proactor_udp.py`).

Run on a real ProactorEventLoop, which only Windows has.
"""

import asyncio
import socket
import sys
import time

import pytest
from meshbay_node.transport.webrtc import proactor_udp

pytestmark = pytest.mark.skipif(sys.platform != "win32", reason="the proactor is Windows'")

N = 2000


class _Collect(asyncio.DatagramProtocol):
    def __init__(self, raise_on: int = -1):
        self.got: list[int] = []
        self.raise_on = raise_on
        self.all = asyncio.get_running_loop().create_future()

    def datagram_received(self, data, addr):
        self.got.append(int.from_bytes(data[:4], "big"))
        if len(self.got) == N and not self.all.done():
            self.all.set_result(None)
        if len(self.got) == self.raise_on:
            raise RuntimeError("a protocol's own bug")


class _Refusing(socket.socket):
    """A socket that refuses a send every so often, as a full buffer does, and
    can answer one read with the reset an ICMP port unreachable leaves. The
    overlapped path works by handle and never calls these."""

    def __init__(self, every: int = 0, reset_once: bool = False):
        super().__init__(socket.AF_INET, socket.SOCK_DGRAM)
        self.every, self.calls, self.reset_once = every, 0, reset_once

    def sendto(self, data, addr):
        self.calls += 1
        if self.every and self.calls % self.every == 0:
            raise BlockingIOError
        return super().sendto(data, addr)

    def recvfrom(self, n):
        if self.reset_once:
            self.reset_once = False
            raise ConnectionResetError(10054, "port unreachable")
        return super().recvfrom(n)


def _busy_turns(ms: float):
    """A loop with other work to do: every turn takes `ms` longer."""
    async def busy():
        while True:
            end = time.perf_counter() + ms / 1000
            while time.perf_counter() < end:
                pass
            await asyncio.sleep(0)
    return asyncio.ensure_future(busy())


async def _exchange(*, refuse_every=0, reset_once=False, raise_on=-1, busy_ms=0.0):
    """N datagrams from one proactor transport to another on loopback: what
    arrived in what order, and how many overlapped sends and reads it took."""
    loop = asyncio.get_running_loop()
    rsock = _Refusing(reset_once=reset_once)
    rsock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 8 << 20)
    rsock.bind(("127.0.0.1", 0))
    rtr, rx = await loop.create_datagram_endpoint(lambda: _Collect(raise_on), sock=rsock)
    tsock = _Refusing(refuse_every)
    tsock.bind(("127.0.0.1", 0))
    tx, _ = await loop.create_datagram_endpoint(asyncio.DatagramProtocol, sock=tsock)
    counts = {"sendto": 0, "recvfrom": 0}
    originals = {name: getattr(loop._proactor, name) for name in counts}

    def counting(name):
        def call(*a, **kw):
            counts[name] += 1
            return originals[name](*a, **kw)
        return call
    for name in counts:
        setattr(loop._proactor, name, counting(name))
    errors = []
    loop.set_exception_handler(lambda _loop, ctx: errors.append(ctx))
    busy = _busy_turns(busy_ms) if busy_ms else None
    try:
        dest = rsock.getsockname()
        for i in range(N):
            tx.sendto(i.to_bytes(4, "big") * 50, dest)
            if i % 50 == 49:
                await asyncio.sleep(0)
        await asyncio.wait_for(rx.all, 20)
        return rx.got, counts, errors
    finally:
        if busy:
            busy.cancel()
        for name, fn in originals.items():
            setattr(loop._proactor, name, fn)
        tx.close()
        rtr.close()


def _run(**kw):
    return asyncio.run(_exchange(**kw), loop_factory=asyncio.ProactorEventLoop)


def test_the_transport_is_patched():
    cls = asyncio.proactor_events._ProactorDatagramTransport
    assert cls.sendto is proactor_udp.sendto
    assert cls._loop_reading is proactor_udp._loop_reading


def test_datagrams_go_straight_out_in_order():
    got, counts, _ = _run()
    assert got == list(range(N))
    assert counts["sendto"] == 0


def test_what_the_socket_refuses_is_queued_and_order_is_kept():
    """A refusal sends that datagram the overlapped way, and the ones after it
    wait behind it rather than overtake it."""
    got, counts, _ = _run(refuse_every=97)
    assert got == list(range(N))
    assert counts["sendto"] >= N // 97


def test_a_busy_loop_reads_everything_waiting_in_order():
    """Unpatched, this loop read one datagram a turn: about 2,000 a second
    whatever was waiting, and the kernel dropped the rest."""
    got, counts, _ = _run(busy_ms=2)
    assert got == list(range(N))
    assert counts["recvfrom"] < N // 10


def test_a_port_unreachable_on_a_direct_read_is_skipped():
    got, _, errors = _run(reset_once=True, busy_ms=2)
    assert got == list(range(N))
    assert not errors


def test_a_protocol_that_raises_does_not_end_the_reading():
    got, _, errors = _run(raise_on=10)
    assert got == list(range(N))
    assert [str(e["exception"]) for e in errors] == ["a protocol's own bug"]