#!/usr/bin/env python3
"""A PostgreSQL client that sends one query and then stops reading.

Speaks just enough of the v3 protocol (startup, trust auth, simple Query) to
make a backend start streaming a large result, then never reads the socket
again.  The kernel buffers fill and the backend blocks in secure_write()
(wait_event ClientWrite).  The process stays alive and keeps ACKing TCP
zero-window probes, which is what a paused or wedged application looks like.

usage: stall_client.py PORT DBNAME SQL [--rcvbuf BYTES]
Prints "SENT <backend pid> <local port>" once the query is on the wire, then sleeps until killed.
"""
import argparse
import os
import socket
import struct
import sys
import time

ap = argparse.ArgumentParser()
ap.add_argument("port", type=int)
ap.add_argument("dbname")
ap.add_argument("sql")
ap.add_argument("--host", default="127.0.0.1")
ap.add_argument("--user", default=os.environ.get("USER", "postgres"))
ap.add_argument("--rcvbuf", type=int, default=0, help="0 = kernel default")
a = ap.parse_args()

s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
# Optional: shrink the receive buffer.  Do NOT use this for the repro: with a tiny
# buffer the kernel keeps reopening the window a few bytes at a time, the backend
# completes a flush now and then, reaches CHECK_FOR_INTERRUPTS() and absorbs the
# barrier.  The hang needs a window that stays at zero (default buffers do that).
if a.rcvbuf:
    s.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, a.rcvbuf)
s.connect((a.host, a.port))


def recv_exact(n):
    buf = b""
    while len(buf) < n:
        chunk = s.recv(n - len(buf))
        if not chunk:
            sys.exit("server closed the connection")
        buf += chunk
    return buf


def read_msg():
    t = recv_exact(1)
    (ln,) = struct.unpack("!I", recv_exact(4))
    return t, recv_exact(ln - 4)


params = b"user\0" + a.user.encode() + b"\0database\0" + a.dbname.encode() + b"\0\0"
s.sendall(struct.pack("!II", 8 + len(params), 196608) + params)
backend_pid = None
while True:
    t, body = read_msg()
    if t == b"R" and struct.unpack("!I", body[:4])[0] != 0:
        sys.exit("only trust auth is supported")
    if t == b"K":
        backend_pid = struct.unpack("!I", body[:4])[0]
    if t == b"E":
        sys.exit("error during startup: %r" % body)
    if t == b"Z":
        break

q = a.sql.encode() + b"\0"
s.sendall(b"Q" + struct.pack("!I", 4 + len(q)) + q)
print("SENT", backend_pid, s.getsockname()[1], flush=True)
while True:          # alive, connected, never reading
    time.sleep(3600)