extmod/modlwip: Fix TCP send and POLLOUT when out of memory. - #19705

# Non-blocking and timed sends under lwIP memory starvation, one file for both sides.
# On the host (CPython) it sends itself to the board through mpremote and plays two slow
# readers; on the board (MicroPython) it runs the two phases. The readers fill the lwIP heap
# with unacked segments, so a socket's tcp_write meets ERR_MEM while its tcp_sndbuf still
# reports room.
#   unpatched:  phase 1 a non-blocking write() blocks for seconds; phase 2 a send() on a
#               settimeout(0.5) socket blocks past its timeout
#   patched:    both phases PASS
# Usage: set WIFI_SSID / WIFI_PASS, then python3 repro.py with the board on USB.
#        python3 repro.py <board-ip> plays the readers only, for a board started by hand.
import sys

WIFI_SSID = "..."
WIFI_PASS = "..."
PORT = 8080
PHASE_S = 120
TIMEOUT_S = 0.5
# TIMEOUT_S plus the 50 ms grain of the ERR_MEM loop and scheduling slack.
MAX_TIMED_SEND_MS = 600
# A non-blocking write copies one chunk; a stall sits far above 100 ms.
MAX_NONBLOCKING_WRITE_MS = 100
READER_RATE = 20000


def board_main():
    import network, socket, select, time, errno, os

    print("firmware: %s; %s" % (os.uname().version, sys.version))
    print("board:    %s" % os.uname().machine)
    print("modlwip send under lwIP memory starvation: 2 phases of %d s, each needs 2 slow readers:" % PHASE_S)
    print("   Phase 1 NON-BLOCKING - setblocking(False) sockets: write() with no memory must return EAGAIN at once,")
    print("                          and POLLOUT must not report writable while the next write still fails")
    print("   Phase 2 TIMEOUT      - settimeout(%.1f) sockets: send() with no memory must raise ETIMEDOUT at the timeout" % TIMEOUT_S)
    print()

    wlan = network.WLAN(network.STA_IF)

    def join():
        for attempt in (1, 2, 3):
            # A radio left associated or mid-join by an earlier run refuses a new join.
            wlan.active(False)
            time.sleep_ms(500)
            wlan.active(True)
            print("joining AP(%s), attempt %d of 3..." % (WIFI_SSID, attempt), end="")
            wlan.connect(WIFI_SSID, WIFI_PASS)
            t0 = time.ticks_ms()
            while not wlan.isconnected() and time.ticks_diff(time.ticks_ms(), t0) < 20000:
                time.sleep_ms(500)
                print(".", end="")
            if wlan.isconnected():
                print(" joined")
                return True
            print(" not joined, status %d" % wlan.status())
        return False

    try:
        joined = join()
    except KeyboardInterrupt:
        print()
        print("INTERRUPTED during the Wi-Fi join")
        return
    if not joined:
        seen = [s for s in wlan.scan() if s[0] == WIFI_SSID.encode()]
        if seen:
            print("FAILED: could not join AP(%s) in 3 attempts; a scan sees it on channel %d at %d dBm,"
                  " so check the password or the AP's access list" % (WIFI_SSID, seen[0][2], seen[0][3]))
        else:
            print("FAILED: could not join AP(%s) in 3 attempts; a scan does not see it,"
                  " so the AP is out of range or off" % WIFI_SSID)
        return
    ip = wlan.ifconfig()[0]

    # Every socket lands in open_socks the moment it exists; the finally below closes them all.
    open_socks = []
    payload = memoryview(bytearray(32768))

    def listen():
        server = socket.socket()
        open_socks.append(server)
        server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        server.bind(("0.0.0.0", PORT))
        server.listen(2)
        server.settimeout(1.0)
        print("listening %s %d" % (ip, PORT))
        print()
        return server

    def accept_two(server, setup):
        print("waiting for 2 readers", end="")
        clients = []
        while len(clients) < 2:
            try:
                conn, addr = server.accept()
            except OSError:
                print(".", end="")
                continue
            open_socks.append(conn)
            if not clients:
                print()
            setup(conn)
            clients.append(conn)
            print("  Reader %d (%s:%d) connected" % (len(clients), addr[0], addr[1]))
        print()
        return clients

    def drop(clients):
        for c in clients:
            c.close()

    def verdict(failures):
        if failures:
            print("       FAIL: " + "; ".join(failures))
        else:
            print("       PASS")
        print()
        return not failures

    def phase_nonblocking(server):
        clients = accept_two(server, lambda c: c.setblocking(False))
        print("PHASE 1 NON-BLOCKING: setblocking(False) socket, write() with no memory;"
              " must return EAGAIN at once, and POLLOUT must not be followed by EAGAIN")
        pollers = {}
        for c in clients:
            pollers[c] = select.poll()
            pollers[c].register(c, select.POLLOUT)
        eagain = spins = enomem = 0
        maxwr = 0

        def nb_write(c):
            # Unpatched firmware blocks a non-blocking write in the ERR_MEM loop, then may raise ENOMEM.
            try:
                return c.write(payload)
            except OSError as e:
                if e.args[0] != errno.ENOMEM:
                    raise
                return -1

        end = time.ticks_add(time.ticks_ms(), PHASE_S * 1000)
        last = time.ticks_ms()
        while time.ticks_diff(end, time.ticks_ms()) > 0:
            if time.ticks_diff(time.ticks_ms(), last) >= 10000:
                last = time.ticks_ms()
                print("  t=%3ds  EAGAIN: %d  EAGAIN right after POLLOUT: %d  max(write) ms: %d" % (
                    time.ticks_diff(last, end) // 1000 + PHASE_S, eagain, spins, maxwr // 1000))
            for c in clients:
                t = time.ticks_us()
                n = nb_write(c)
                maxwr = max(maxwr, time.ticks_diff(time.ticks_us(), t))
                if n == -1:
                    enomem += 1
                    continue
                if n:
                    continue
                eagain += 1
                tp = time.ticks_us()
                if pollers[c].poll(1000):
                    n2 = nb_write(c)
                    if n2 == -1:
                        enomem += 1
                    elif not n2 and time.ticks_diff(time.ticks_us(), tp) < 2000:
                        spins += 1
        print("DONE:  EAGAIN: %d  EAGAIN right after POLLOUT: %d  max(write) ms: %d" % (eagain, spins, maxwr // 1000))
        drop(clients)
        failures = []
        if maxwr // 1000 > MAX_NONBLOCKING_WRITE_MS:
            failures.append("non-blocking write() blocked %d ms" % (maxwr // 1000))
        elif eagain == 0:
            failures.append("no EAGAIN, memory never ran out")
        if enomem:
            failures.append("%d writes raised ENOMEM" % enomem)
        if spins:
            failures.append("%d writes got EAGAIN right after POLLOUT" % spins)
        return verdict(failures)

    def phase_timeout(server):
        clients = accept_two(server, lambda c: c.settimeout(TIMEOUT_S))
        print("PHASE 2 TIMEOUT: settimeout(%.1f) socket, one send() with no memory;"
              " must raise ETIMEDOUT at the timeout" % TIMEOUT_S)
        timedout = enomem = 0
        maxsend = 0
        end = time.ticks_add(time.ticks_ms(), PHASE_S * 1000)
        last = time.ticks_ms()
        while time.ticks_diff(end, time.ticks_ms()) > 0:
            if time.ticks_diff(time.ticks_ms(), last) >= 10000:
                last = time.ticks_ms()
                print("  t=%3ds  ETIMEDOUT: %d  ENOMEM: %d  max(send) ms: %d" % (
                    time.ticks_diff(last, end) // 1000 + PHASE_S, timedout, enomem, maxsend // 1000))
            for c in clients:
                t = time.ticks_us()
                try:
                    # send() is one lwip_tcp_send call; write() loops over chunks, a clock each.
                    c.send(payload)
                except OSError as e:
                    if e.args[0] == errno.ETIMEDOUT:
                        timedout += 1
                    elif e.args[0] == errno.ENOMEM:
                        enomem += 1
                    else:
                        raise
                maxsend = max(maxsend, time.ticks_diff(time.ticks_us(), t))
        print("DONE:  ETIMEDOUT: %d  ENOMEM: %d  max(send) ms: %d" % (timedout, enomem, maxsend // 1000))
        drop(clients)
        failures = []
        if maxsend // 1000 > MAX_TIMED_SEND_MS:
            failures.append("send() blocked %d ms, socket timeout is %d ms"
                            % (maxsend // 1000, int(TIMEOUT_S * 1000)))
        elif timedout == 0:
            failures.append("no send blocked, memory never ran out")
        if enomem:
            failures.append("%d sends raised ENOMEM" % enomem)
        return verdict(failures)

    stage = "opening the server socket"
    try:
        server = listen()
        stage = "phase 1"
        ok1 = phase_nonblocking(server)
        print("Dropping readers, they reconnect for phase 2")
        print()
        stage = "phase 2"
        ok2 = phase_timeout(server)
        print("ALL DONE: phase 1 %s, phase 2 %s" % ("PASS" if ok1 else "FAIL", "PASS" if ok2 else "FAIL"))
    except KeyboardInterrupt:
        print()
        print("INTERRUPTED during %s: closing sockets" % stage)
    except OSError as e:
        print()
        print("FAILED: socket error %d during %s" % (e.args[0], stage))
    except Exception as e:
        print()
        print("FAILED during %s: %r" % (stage, e))
    finally:
        # One close failing must not leave the rest open.
        for s in open_socks:
            try:
                s.close()
            except OSError:
                pass


def host_main():
    import os, re, shutil, signal, socket, subprocess, threading, time

    # One line per failure mpremote reports as a traceback.
    causes = (
        ("no device found", "no board on USB: plug it in, or close the program holding its port"),
        ("could not enter raw repl", "the board does not answer on its serial port: reset or replug it"),
        ("failed to access", "the serial port is held by another program (an IDE, another mpremote)"),
        ("device disconnected", "the board dropped off USB during the run: replug it"),
    )
    stop = threading.Event()

    def reader(host, port):
        while not stop.is_set():
            s = None
            try:
                s = socket.create_connection((host, port), timeout=5)
                s.settimeout(1.0)
                while not stop.is_set():
                    try:
                        data = s.recv(4096)
                    except socket.timeout:
                        continue
                    if not data:
                        break
                    time.sleep(len(data) / READER_RATE)
            except OSError:
                pass
            finally:
                if s is not None:
                    s.close()
            time.sleep(0.5)

    def start_readers(host, port):
        for _ in range(2):
            threading.Thread(target=reader, args=(host, port), daemon=True).start()

    if len(sys.argv) > 1:
        print("2 readers to %s:%d at %d B/s each, for a board already running the test; Ctrl-C stops them"
              % (sys.argv[1], PORT, READER_RATE))
        start_readers(sys.argv[1], PORT)
        try:
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            stop.set()
            print("\nreaders stopped")
            return 0

    mpremote = ["mpremote"] if shutil.which("mpremote") else ["uvx", "mpremote"]
    try:
        # Own session: the terminal's Ctrl-C reaches this script alone, which then stops the
        # child once, instead of both dying mid-close.
        proc = subprocess.Popen(mpremote + ["run", os.path.abspath(__file__)], stdout=subprocess.PIPE,
                                stderr=subprocess.STDOUT, text=True, bufsize=1, start_new_session=True)
    except FileNotFoundError:
        print("FAILED: neither mpremote nor uvx is installed")
        return 2

    traceback_lines = []
    try:
        for line in proc.stdout:
            # mpremote prints the board's empty reply, b'', ahead of its raw-REPL traceback.
            if traceback_lines or line.startswith(("Traceback", "mpremote:")) or line.strip() == "b''":
                traceback_lines.append(line)
                continue
            sys.stdout.write(line)
            sys.stdout.flush()
            m = re.match(r"listening (\S+) (\d+)", line)
            if m:
                start_readers(m.group(1), int(m.group(2)))
        rc = proc.wait()
    except KeyboardInterrupt:
        print("\nINTERRUPTED: stopping the board")
        stop.set()
        proc.send_signal(signal.SIGINT)
        try:
            proc.wait(timeout=10)
        except subprocess.TimeoutExpired:
            proc.kill()
        return 130
    stop.set()

    if traceback_lines:
        text = "".join(traceback_lines)
        for needle, cause in causes:
            if needle in text:
                print("FAILED: " + cause)
                break
        else:
            print("FAILED: mpremote said: " + traceback_lines[-1].strip())
        return 1
    return rc


if sys.implementation.name == "micropython":
    board_main()
else:
    sys.exit(host_main())