#!/usr/bin/env python3
"""
Bidirectional TCP proxy for the MDEV-38190 SHUTDOWN-hang repro -- no DBUG
hook, no debug build required on the server side.

Sits between a semi-sync slave and its real master:

    slave  <--listen_port-->  PROXY  <--upstream_host:upstream_port-->  master

- master -> slave (the binlog dump stream) is ALWAYS forwarded immediately.
- slave -> master (which carries the semi-sync ACK packets, interleaved on
  the same connection) is forwarded immediately only while the gate file's
  content is "OPEN". While it says "CLOSED", those bytes are queued in
  memory instead of being sent -- genuinely withholding the ACK over the
  wire. Flipping the gate file back to OPEN flushes anything queued and
  resumes normal passthrough within one poll interval (0.1s).

Usage:
  ./semisync_ack_proxy.py <listen_port> <upstream_host> <upstream_port> <gate_file>

The gate file is created with "OPEN" if it doesn't already exist.
"""
import asyncio
import sys


def gate_open(path):
    try:
        with open(path) as f:
            return f.read().strip() == "OPEN"
    except FileNotFoundError:
        return True


async def pump(reader, writer, gate_file, label):
    """Copy bytes reader -> writer. If gate_file is None, forward every
    read immediately (the master->slave direction). If gate_file is set,
    buffer whatever arrives while the gate says CLOSED, and flush it as
    soon as the gate says OPEN -- checked on a timer, not just on new
    reads, so a gate flip takes effect even if the peer has gone quiet."""
    buffered = bytearray()
    while True:
        try:
            data = await asyncio.wait_for(reader.read(65536), timeout=0.1)
        except asyncio.TimeoutError:
            data = None
        else:
            if data == b"":
                break  # EOF
            if data and gate_file is None:
                writer.write(data)
                await writer.drain()
                continue
            if data:
                buffered.extend(data)
                print(f"[proxy] {label}: withheld {len(data)} bytes "
                      f"(buffered={len(buffered)})", file=sys.stderr)

        if gate_file is not None and buffered and gate_open(gate_file):
            writer.write(bytes(buffered))
            await writer.drain()
            print(f"[proxy] {label}: gate OPEN, flushed {len(buffered)} bytes",
                  file=sys.stderr)
            buffered.clear()


async def handle_client(client_reader, client_writer, upstream_host, upstream_port, gate_file):
    try:
        upstream_reader, upstream_writer = await asyncio.open_connection(upstream_host, upstream_port)
    except Exception as e:
        print(f"[proxy] could not connect upstream: {e}", file=sys.stderr)
        client_writer.close()
        return

    down = asyncio.create_task(
        pump(upstream_reader, client_writer, gate_file=None, label="master->slave"))
    up = asyncio.create_task(
        pump(client_reader, upstream_writer, gate_file=gate_file, label="slave->master(ACK)"))

    done, pending = await asyncio.wait([down, up], return_when=asyncio.FIRST_COMPLETED)
    for t in pending:
        t.cancel()
    client_writer.close()
    upstream_writer.close()


async def main():
    if len(sys.argv) != 5:
        print(__doc__, file=sys.stderr)
        sys.exit(1)
    listen_port, upstream_host, upstream_port, gate_file = sys.argv[1:5]
    listen_port = int(listen_port)
    upstream_port = int(upstream_port)

    try:
        with open(gate_file, "x") as f:
            f.write("OPEN")
    except FileExistsError:
        pass

    server = await asyncio.start_server(
        lambda r, w: handle_client(r, w, upstream_host, upstream_port, gate_file),
        "127.0.0.1", listen_port)
    addr = server.sockets[0].getsockname()
    print(f"[proxy] listening on {addr}, forwarding to {upstream_host}:{upstream_port}, "
          f"gate file={gate_file}", file=sys.stderr)
    async with server:
        await server.serve_forever()


if __name__ == "__main__":
    asyncio.run(main())