1 files changed, 128 insertions, 13 deletions
diff --git a/packages/meshbay-node/src/meshbay_node/transport/webrtc/sctp_fix.py b/packages/meshbay-node/src/meshbay_node/transport/webrtc/sctp_fix.py
index 0d45a7d..a7e6f86 100644
--- a/packages/meshbay-node/src/meshbay_node/transport/webrtc/sctp_fix.py
+++ b/packages/meshbay-node/src/meshbay_node/transport/webrtc/sctp_fix.py
@@ -1,12 +1,13 @@
-"""A fix to aiortc's SCTP sender, applied on import (aiortc 1.15).
+"""Two fixes to aiortc's SCTP stack, applied on import (aiortc 1.15).
 
-aiortc counts the bytes in flight and sends nothing new while that count is at
-or above the congestion window. A chunk the peer acknowledged in a gap block is
-marked `_acked` and taken out of the count. If the T3 timer then expires, every
-chunk still queued is sent again, that one included, and goes back into the
-count with `_acked` still set. When the cumulative ack finally covers it, the
-SACK handler skips it because it reads as acknowledged already, and its bytes
-stay counted for the life of the association.
+**The sender leaked bytes in flight.** aiortc counts the bytes in flight and
+sends nothing new while that count is at or above the congestion window. A chunk
+the peer acknowledged in a gap block is marked `_acked` and taken out of the
+count. If the T3 timer then expires, every chunk still queued is sent again,
+that one included, and goes back into the count with `_acked` still set. When
+the cumulative ack finally covers it, the SACK handler skips it because it reads
+as acknowledged already, and its bytes stay counted for the life of the
+association.
 
 Each such loss leaks a chunk. On a lossy path (a phone on Wi-Fi) the leak grows
 past the window, which T3 itself keeps small. From then on the node sends
@@ -14,14 +15,49 @@ nothing on that connection: nothing is outstanding, so no timer runs to unstick
 it, and the association still reads as established. The peer keeps talking and
 the node keeps answering into a queue that never moves. Live symptom: the photo
 backup stalled after a few dozen photos, its slot granted and the grant never
-delivered (docs/MESHBAY_DESIGN.md §9.12).
+delivered (docs/MESHBAY_DESIGN.md §9.12). The fix is the invariant the code
+assumed: a chunk being sent again is in flight, so it is not acknowledged.
 
-The fix is the invariant the code assumed: a chunk being sent again is in
-flight, so it is not acknowledged. `test_sctp_fix.py` replays the sequence
-against aiortc itself.
+**The receiver reassembled in quadratic time.** After every DATA chunk,
+`InboundStream.pop_messages` walks the reassembly list from its start, through
+every fragment of the message still being assembled. A message of n fragments
+costs n²/2 steps, and a 1 MiB `file_chunk` is about 870 fragments: measured on a
+receiving aiortc, 153 million steps for 400 MB, two thirds of its CPU, and the
+whole of why it could not take more than 20 MB/s on loopback. No message can
+come out of a reassembly holding no last fragment, so the walk is skipped until
+one is there; when one is, the walk is aiortc's own, unchanged.
+
+The node receives uploads in 48 KB messages, where the walk is short; the
+command-line client receives 1 MiB ones and carries the same fix
+(`meshbay_cli/client/sctp_fix.py`).
+
+**The receiver acknowledged every packet.** RFC 4960 §6.2 asks for one SACK
+every second packet, and Chromium does that; aiortc sent one at the end of
+every packet of data. Each SACK costs its receiver a DTLS decryption and a parse
+on the one core that also does everything else — on a node sending a download
+to the command-line client, more time than sending the data. In an upload the
+node is the receiver and pays the other half: the encryption and the send of a
+SACK per packet. A SACK still goes at once when something is out of order or
+duplicated (the loss signal the sender is waiting for) or the packet carried
+anything other than data; a lone packet is acknowledged within `SACK_DELAY_S`.
+
+`test_sctp_fix.py` replays all three against aiortc itself.
 """
 
-from aiortc.rtcsctptransport import RTCSctpTransport
+import asyncio
+import contextlib
+
+from aiortc.rtcsctptransport import (
+    SCTP_DATA_LAST_FRAG,
+    DataChunk,
+    InboundStream,
+    InitChunk,
+    RTCSctpTransport,
+    parse_packet,
+)
+
+# RFC 4960 §6.2: at most 500 ms; 200 is what implementations use.
+SACK_DELAY_S = 0.2
 
 
 def _flight_size_increase(self, chunk) -> None:
@@ -30,3 +66,82 @@ def _flight_size_increase(self, chunk) -> None:
 
 
 RTCSctpTransport._flight_size_increase = _flight_size_increase
+
+
+_add_chunk = InboundStream.add_chunk
+_pop_messages = InboundStream.pop_messages
+_prune_chunks = InboundStream.prune_chunks
+
+
+def add_chunk(self, chunk) -> None:
+    _add_chunk(self, chunk)
+    if chunk.flags & SCTP_DATA_LAST_FRAG:
+        self._last_frags = getattr(self, "_last_frags", 0) + 1
+
+
+def pop_messages(self):
+    # Only ever too high (a pruned last fragment), never too low: a message
+    # leaves the reassembly with exactly one last fragment, and only here.
+    if not getattr(self, "_last_frags", 0):
+        return
+    for message in _pop_messages(self):
+        self._last_frags -= 1
+        yield message
+
+
+def prune_chunks(self, tsn: int) -> int:
+    size = _prune_chunks(self, tsn)
+    self._last_frags = sum(1 for c in self.reassembly if c.flags & SCTP_DATA_LAST_FRAG)
+    return size
+
+
+InboundStream.add_chunk = add_chunk
+InboundStream.pop_messages = pop_messages
+InboundStream.prune_chunks = prune_chunks
+
+
+async def _handle_data(self, data: bytes) -> None:
+    """aiortc's own, down to the SACK at the end: that one may wait for a
+    second packet."""
+    try:
+        _, _, verification_tag, chunks = parse_packet(data)
+    except ValueError:
+        return
+    if any(isinstance(x, InitChunk) for x in chunks):
+        assert len(chunks) == 1
+        expected_tag = 0
+    else:
+        expected_tag = self._local_verification_tag
+    if verification_tag != expected_tag:
+        return
+    for chunk in chunks:
+        await self._receive_chunk(chunk)
+    if not self._sack_needed:
+        return
+    self._sack_packets = getattr(self, "_sack_packets", 0) + 1
+    if (self._sack_packets >= 2 or self._sack_misordered or self._sack_duplicates
+            or not all(isinstance(c, DataChunk) for c in chunks)):
+        await _sack_now(self)
+    elif getattr(self, "_sack_timer", None) is None:
+        self._sack_timer = self._loop.call_later(SACK_DELAY_S, _sack_later, self)
+
+
+async def _sack_now(self) -> None:
+    timer = getattr(self, "_sack_timer", None)
+    if timer is not None:
+        timer.cancel()
+        self._sack_timer = None
+    self._sack_packets = 0
+    if self._sack_needed and self._association_state != self.State.CLOSED:
+        # The DTLS transport may be gone under a timer that fired late.
+        with contextlib.suppress(ConnectionError):
+            await self._send_sack()
+
+
+def _sack_later(self) -> None:
+    self._sack_timer = None
+    # Held: a task nobody references can be collected before it runs.
+    self._sack_task = asyncio.ensure_future(_sack_now(self))
+
+
+RTCSctpTransport._handle_data = _handle_data