"""Does closing with unread data reset the connection, and is the reply lost?

The server reads only part of a large upload, answers, and closes in one of
three ways while the client is still sending. The client reports how much of
the reply arrived and what errors it saw.
"""
import socket, sys, threading, time

REPLY = b"R" * 65536
UPLOAD = 1500000


def server(lsock, variant):
    c, _ = lsock.accept()
    got = 0
    while got < 16384:
        d = c.recv(16384 - got)
        if not d:
            break
        got += len(d)
    c.sendall(REPLY)
    if variant == "close":
        c.close()
    elif variant == "shutrdwr":
        c.shutdown(socket.SHUT_RDWR)
        c.close()
    elif variant == "drain":
        c.shutdown(socket.SHUT_WR)
        c.settimeout(5)
        try:
            while c.recv(65536):
                pass
        except OSError:
            pass
        c.close()


def run(variant):
    lsock = socket.socket()
    lsock.bind(("127.0.0.1", 0))
    lsock.listen(1)
    port = lsock.getsockname()[1]
    t = threading.Thread(target=server, args=(lsock, variant))
    t.start()
    cl = socket.create_connection(("127.0.0.1", port))
    errs = {}

    def send():
        try:
            cl.sendall(b"x" * UPLOAD)
            cl.shutdown(socket.SHUT_WR)
        except OSError as e:
            errs["send"] = type(e).__name__

    st = threading.Thread(target=send)
    st.start()
    # read slowly at first, as a client busy uploading would
    time.sleep(0.3)
    rx = 0
    cl.settimeout(5)
    try:
        while True:
            d = cl.recv(65536)
            if not d:
                break
            rx += len(d)
    except OSError as e:
        errs["recv"] = type(e).__name__
    st.join()
    t.join()
    cl.close()
    lsock.close()
    return rx, errs


for variant in sys.argv[1:] or ["close", "shutrdwr", "drain"]:
    results = [run(variant) for _ in range(5)]
    full = sum(1 for rx, _ in results if rx == len(REPLY))
    errs = sorted({k + ":" + v for _, e in results for k, v in e.items()})
    print("%-9s reply complete %d/5, received %s, errors %s"
          % (variant, full, [rx for rx, _ in results], errs or "none"))