"""
TCP transport classes

Wire protocol: "len(payload) msgpack({'head': SOMEHEADER, 'body': SOMEBODY})"


"""

import asyncio
import errno
import logging
import multiprocessing
import os
import queue
import socket
import threading
import urllib
import uuid
import warnings
import weakref

import salt.ext.tornado
import salt.ext.tornado.concurrent
import salt.ext.tornado.gen
import salt.ext.tornado.iostream
import salt.ext.tornado.netutil
import salt.ext.tornado.tcpclient
import salt.ext.tornado.tcpserver
import salt.master
import salt.payload
import salt.transport.frame
import salt.transport.ipc
import salt.utils.asynchronous
import salt.utils.files
import salt.utils.msgpack
import salt.utils.platform
import salt.utils.versions
from salt.exceptions import SaltClientError, SaltReqTimeoutError
from salt.utils.network import ip_bracket
from salt.utils.process import SignalHandlingProcess

if salt.utils.platform.is_windows():
    USE_LOAD_BALANCER = True
else:
    USE_LOAD_BALANCER = False

if USE_LOAD_BALANCER:
    import salt.ext.tornado.util

log = logging.getLogger(__name__)


class ClosingError(Exception):
    """ """


def _get_socket(opts):
    family = socket.AF_INET
    if opts.get("ipv6", False):
        family = socket.AF_INET6
    return socket.socket(family, socket.SOCK_STREAM)


def _get_bind_addr(opts, port_type):
    return (
        ip_bracket(opts["interface"], strip=True),
        int(opts[port_type]),
    )


def _set_tcp_keepalive(sock, opts):
    """
    Ensure that TCP keepalives are set for the socket.
    """
    if hasattr(socket, "SO_KEEPALIVE"):
        if opts.get("tcp_keepalive", False):
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
            if hasattr(socket, "SOL_TCP"):
                if hasattr(socket, "TCP_KEEPIDLE"):
                    tcp_keepalive_idle = opts.get("tcp_keepalive_idle", -1)
                    if tcp_keepalive_idle > 0:
                        sock.setsockopt(
                            socket.SOL_TCP, socket.TCP_KEEPIDLE, int(tcp_keepalive_idle)
                        )
                if hasattr(socket, "TCP_KEEPCNT"):
                    tcp_keepalive_cnt = opts.get("tcp_keepalive_cnt", -1)
                    if tcp_keepalive_cnt > 0:
                        sock.setsockopt(
                            socket.SOL_TCP, socket.TCP_KEEPCNT, int(tcp_keepalive_cnt)
                        )
                if hasattr(socket, "TCP_KEEPINTVL"):
                    tcp_keepalive_intvl = opts.get("tcp_keepalive_intvl", -1)
                    if tcp_keepalive_intvl > 0:
                        sock.setsockopt(
                            socket.SOL_TCP,
                            socket.TCP_KEEPINTVL,
                            int(tcp_keepalive_intvl),
                        )
            if hasattr(socket, "SIO_KEEPALIVE_VALS"):
                # Windows doesn't support TCP_KEEPIDLE, TCP_KEEPCNT, nor
                # TCP_KEEPINTVL. Instead, it has its own proprietary
                # SIO_KEEPALIVE_VALS.
                tcp_keepalive_idle = opts.get("tcp_keepalive_idle", -1)
                tcp_keepalive_intvl = opts.get("tcp_keepalive_intvl", -1)
                # Windows doesn't support changing something equivalent to
                # TCP_KEEPCNT.
                if tcp_keepalive_idle > 0 or tcp_keepalive_intvl > 0:
                    # Windows defaults may be found by using the link below.
                    # Search for 'KeepAliveTime' and 'KeepAliveInterval'.
                    # https://technet.microsoft.com/en-us/library/bb726981.aspx#EDAA
                    # If one value is set and the other isn't, we still need
                    # to send both values to SIO_KEEPALIVE_VALS and they both
                    # need to be valid. So in that case, use the Windows
                    # default.
                    if tcp_keepalive_idle <= 0:
                        tcp_keepalive_idle = 7200
                    if tcp_keepalive_intvl <= 0:
                        tcp_keepalive_intvl = 1
                    # The values expected are in milliseconds, so multiply by
                    # 1000.
                    sock.ioctl(
                        socket.SIO_KEEPALIVE_VALS,
                        (
                            1,
                            int(tcp_keepalive_idle * 1000),
                            int(tcp_keepalive_intvl * 1000),
                        ),
                    )
        else:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 0)


class LoadBalancerServer(SignalHandlingProcess):
    """
    Raw TCP server which runs in its own process and will listen
    for incoming connections. Each incoming connection will be
    sent via multiprocessing queue to the workers.
    Since the queue is shared amongst workers, only one worker will
    handle a given connection.
    """

    # TODO: opts!
    # Based on default used in salt.ext.tornado.netutil.bind_sockets()
    backlog = 128

    def __init__(self, opts, socket_queue, **kwargs):
        super().__init__(**kwargs)
        self.opts = opts
        self.socket_queue = socket_queue
        self._socket = None

    def close(self):
        if self._socket is not None:
            self._socket.shutdown(socket.SHUT_RDWR)
            self._socket.close()
            self._socket = None

    # pylint: disable=W1701
    def __del__(self):
        self.close()

    # pylint: enable=W1701

    def run(self):
        """
        Start the load balancer
        """
        self._socket = _get_socket(self.opts)
        self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        _set_tcp_keepalive(self._socket, self.opts)
        self._socket.setblocking(1)
        self._socket.bind(_get_bind_addr(self.opts, "ret_port"))
        self._socket.listen(self.backlog)

        while self._socket and not self._socket._closed:
            try:
                # Wait for a connection to occur since the socket is
                # blocking.
                connection, address = self._socket.accept()
                # Wait for a free slot to be available to put
                # the connection into.
                # Sockets are picklable on Windows in Python 3.
                self.socket_queue.put((connection, address), True, None)
            except OSError as e:
                # ECONNABORTED indicates that there was a connection
                # but it was closed while still in the accept queue.
                # (observed on FreeBSD).
                sock = self._socket
                if sock is None:
                    break
                name = sock.getsockname()
                if isinstance(name, tuple):
                    name = name[0]
                if salt.ext.tornado.util.errno_from_exception(e) == errno.ECONNABORTED:
                    continue
                elif e.errno == errno.EINVAL:
                    # This can happen after socket.shutdown but before socket.close
                    log.trace("Socket shutdown: %s", name)
                    break
                elif e.errno == errno.EBADF:
                    # This can happen after socket.close
                    log.trace("Socket closed: %s", name)
                    break
                raise


class Resolver:

    _resolver_configured = False

    @classmethod
    def _config_resolver(cls, num_threads=10):
        salt.ext.tornado.netutil.Resolver.configure(
            "salt.ext.tornado.netutil.ThreadedResolver", num_threads=num_threads
        )
        cls._resolver_configured = True

    def __init__(self, *args, **kwargs):
        if not self._resolver_configured:
            # TODO: add opt to specify number of resolver threads
            self._config_resolver()


class TCPPubClient(salt.transport.base.PublishClient):
    """
    Tornado based TCP Pub Client
    """

    ttype = "tcp"

    def __init__(self, opts, io_loop, **kwargs):  # pylint: disable=W0231
        super().__init__(opts, io_loop, **kwargs)
        self.opts = opts
        self.io_loop = io_loop
        self.message_client = None
        self.connected = False
        self._closing = False
        self.resolver = Resolver()

    def close(self):
        if self._closing:
            return
        self._closing = True
        if self.message_client is not None:
            self.message_client.close()
            self.message_client = None

    @salt.ext.tornado.gen.coroutine
    def connect(self, publish_port, connect_callback=None, disconnect_callback=None):
        self._connect_called = True
        self.publish_port = publish_port
        self.message_client = MessageClient(
            self.opts,
            self.opts["master_ip"],
            int(self.publish_port),
            io_loop=self.io_loop,
            connect_callback=connect_callback,
            disconnect_callback=disconnect_callback,
            source_ip=self.opts.get("source_ip"),
            source_port=self.opts.get("source_publish_port"),
        )
        yield self.message_client.connect()  # wait for the client to be connected
        self.connected = True

    @salt.ext.tornado.gen.coroutine
    def _decode_messages(self, messages):
        if not isinstance(messages, dict):
            # TODO: For some reason we need to decode here for things
            #       to work. Fix this.
            body = salt.utils.msgpack.loads(messages)
            body = salt.transport.frame.decode_embedded_strs(body)
        else:
            body = messages
        raise salt.ext.tornado.gen.Return(body)

    @salt.ext.tornado.gen.coroutine
    def send(self, msg):
        yield self.message_client._stream.write(msg)

    def on_recv(self, callback):
        """
        Register an on_recv callback
        """
        return self.message_client.on_recv(callback)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


class TCPReqServer(salt.transport.base.DaemonizedRequestServer):
    """
    Tornado based TCP Request/Reply Server

    :param dict opts: Salt master config options.
    """

    # TODO: opts!
    backlog = 5

    def __init__(self, opts):  # pylint: disable=W0231
        self.opts = opts
        self._socket = None
        self.req_server = None

    @property
    def socket(self):
        return self._socket

    def close(self):
        if self._socket is not None:
            try:
                self._socket.shutdown(socket.SHUT_RDWR)
            except OSError as exc:
                if exc.errno == errno.ENOTCONN:
                    # We may try to shutdown a socket which is already disconnected.
                    # Ignore this condition and continue.
                    pass
                else:
                    raise
            if self.req_server is None:
                # We only close the socket if we don't have a req_server instance.
                # If we did, because the req_server is also handling this socket, when we call
                # req_server.stop(), tornado will give us an AssertionError because it's trying to
                # match the socket.fileno() (after close it's -1) to the fd it holds on it's _sockets cache
                # so it can remove the socket from the IOLoop handlers
                self._socket.close()
            self._socket = None
        if self.req_server is not None:
            try:
                self.req_server.close()
            except OSError as exc:
                if exc.errno != 9:
                    raise
                log.exception(
                    "TCPReqServerChannel close generated an exception: %s", str(exc)
                )
            self.req_server = None

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()

    def pre_fork(self, process_manager):
        """
        Pre-fork we need to create the zmq router device
        """
        if USE_LOAD_BALANCER:
            self.socket_queue = multiprocessing.Queue()
            process_manager.add_process(
                LoadBalancerServer,
                args=(self.opts, self.socket_queue),
                name="LoadBalancerServer",
            )
        elif not salt.utils.platform.is_windows():
            self._socket = _get_socket(self.opts)
            self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            _set_tcp_keepalive(self._socket, self.opts)
            self._socket.setblocking(0)
            self._socket.bind(_get_bind_addr(self.opts, "ret_port"))

    def post_fork(self, message_handler, io_loop):
        """
        After forking we need to create all of the local sockets to listen to the
        router

        message_handler: function to call with your payloads
        """
        self.message_handler = message_handler

        with salt.utils.asynchronous.current_ioloop(io_loop):
            if USE_LOAD_BALANCER:
                self.req_server = LoadBalancerWorker(
                    self.socket_queue,
                    self.handle_message,
                    ssl_options=self.opts.get("ssl"),
                )
            else:
                if salt.utils.platform.is_windows():
                    self._socket = _get_socket(self.opts)
                    self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                    _set_tcp_keepalive(self._socket, self.opts)
                    self._socket.setblocking(0)
                    self._socket.bind(_get_bind_addr(self.opts, "ret_port"))
                self.req_server = SaltMessageServer(
                    self.handle_message,
                    ssl_options=self.opts.get("ssl"),
                    io_loop=io_loop,
                )
                self.req_server.add_socket(self._socket)
                self._socket.listen(self.backlog)

    @salt.ext.tornado.gen.coroutine
    def handle_message(self, stream, payload, header=None):
        payload = self.decode_payload(payload)
        reply = yield self.message_handler(payload)
        stream.write(salt.transport.frame.frame_msg(reply, header=header))

    def decode_payload(self, payload):
        return payload


class SaltMessageServer(salt.ext.tornado.tcpserver.TCPServer):
    """
    Raw TCP server which will receive all of the TCP streams and re-assemble
    messages that are sent through to us
    """

    def __init__(self, message_handler, *args, **kwargs):
        io_loop = (
            kwargs.pop("io_loop", None) or salt.ext.tornado.ioloop.IOLoop.current()
        )
        self._closing = False
        super().__init__(*args, **kwargs)
        self.io_loop = io_loop
        self.clients = []
        self.message_handler = message_handler

    @salt.ext.tornado.gen.coroutine
    def handle_stream(  # pylint: disable=arguments-differ
        self,
        stream,
        address,
        _StreamClosedError=salt.ext.tornado.iostream.StreamClosedError,
    ):
        """
        Handle incoming streams and add messages to the incoming queue
        """
        log.trace("Req client %s connected", address)
        self.clients.append((stream, address))
        unpacker = salt.utils.msgpack.Unpacker()
        try:
            while True:
                wire_bytes = yield stream.read_bytes(4096, partial=True)
                unpacker.feed(wire_bytes)
                for framed_msg in unpacker:
                    framed_msg = salt.transport.frame.decode_embedded_strs(framed_msg)
                    header = framed_msg["head"]
                    self.io_loop.spawn_callback(
                        self.message_handler, stream, framed_msg["body"], header
                    )
        except _StreamClosedError:
            log.trace("req client disconnected %s", address)
            unpacker = salt.utils.msgpack.Unpacker()
            self.remove_client((stream, address))
        except Exception as e:  # pylint: disable=broad-except
            log.trace("other master-side exception: %s", e, exc_info=True)
            unpacker = salt.utils.msgpack.Unpacker()
            self.remove_client((stream, address))
            stream.close()

    def remove_client(self, client):
        try:
            self.clients.remove(client)
        except ValueError:
            log.trace("Message server client was not in list to remove")

    def close(self):
        """
        Close the server
        """
        if self._closing:
            return
        self._closing = True
        for item in list(self.clients):
            client, address = item
            client.close()
            self.remove_client(item)
        try:
            self.stop()
        except OSError as exc:
            if exc.errno != 9:
                raise


class LoadBalancerWorker(SaltMessageServer):
    """
    This will receive TCP connections from 'LoadBalancerServer' via
    a multiprocessing queue.
    Since the queue is shared amongst workers, only one worker will handle
    a given connection.
    """

    def __init__(self, socket_queue, message_handler, *args, **kwargs):
        super().__init__(message_handler, *args, **kwargs)
        self.socket_queue = socket_queue
        self._stop = threading.Event()
        self.thread = threading.Thread(target=self.socket_queue_thread)
        self.thread.start()

    def close(self):
        self._stop.set()
        self.thread.join()
        super().close()

    def socket_queue_thread(self):
        try:
            while True:
                try:
                    client_socket, address = self.socket_queue.get(True, 1)
                except queue.Empty:
                    if self._stop.is_set():
                        break
                    continue
                # 'self.io_loop' initialized in super class
                # 'salt.ext.tornado.tcpserver.TCPServer'.
                # 'self._handle_connection' defined in same super class.
                self.io_loop.spawn_callback(
                    self._handle_connection, client_socket, address
                )
        except (KeyboardInterrupt, SystemExit):
            pass


class TCPClientKeepAlive(salt.ext.tornado.tcpclient.TCPClient):
    """
    Override _create_stream() in TCPClient to enable keep alive support.
    """

    def __init__(self, opts, resolver=None):
        self.opts = opts
        super().__init__(resolver=resolver)

    def _create_stream(
        self, max_buffer_size, af, addr, **kwargs
    ):  # pylint: disable=unused-argument,arguments-differ
        """
        Override _create_stream() in TCPClient.

        Tornado 4.5 added the kwargs 'source_ip' and 'source_port'.
        Due to this, use **kwargs to swallow these and any future
        kwargs to maintain compatibility.
        """
        # Always connect in plaintext; we'll convert to ssl if necessary
        # after one connection has completed.
        sock = _get_socket(self.opts)
        _set_tcp_keepalive(sock, self.opts)
        stream = salt.ext.tornado.iostream.IOStream(
            sock, max_buffer_size=max_buffer_size
        )
        return stream.connect(addr)


# TODO consolidate with IPCClient
# TODO: limit in-flight messages.
# TODO: singleton? Something to not re-create the tcp connection so much
class MessageClient:
    """
    Low-level message sending client
    """

    def __init__(
        self,
        opts,
        host,
        port,
        io_loop=None,
        resolver=None,
        connect_callback=None,
        disconnect_callback=None,
        source_ip=None,
        source_port=None,
    ):
        self.opts = opts
        self.host = host
        self.port = port
        self.source_ip = source_ip
        self.source_port = source_port
        self.connect_callback = connect_callback
        self.disconnect_callback = disconnect_callback
        self.io_loop = io_loop or salt.ext.tornado.ioloop.IOLoop.current()
        with salt.utils.asynchronous.current_ioloop(self.io_loop):
            self._tcp_client = TCPClientKeepAlive(opts, resolver=resolver)
        # TODO: max queue size
        self.send_future_map = {}  # mapping of request_id -> Future

        self._read_until_future = None
        self._on_recv = None
        self._closing = False
        self._closed = False
        self._connecting_future = salt.ext.tornado.concurrent.Future()
        self._stream_return_running = False
        self._stream = None

        self.backoff = opts.get("tcp_reconnect_backoff", 1)

    # TODO: timeout inflight sessions
    def close(self):
        if self._closing or self._closed:
            return
        self._closing = True
        self.io_loop.add_timeout(1, self.check_close)

    @salt.ext.tornado.gen.coroutine
    def check_close(self):
        if not self.send_future_map:
            self._tcp_client.close()
            if self._stream:
                self._stream.close()
            self._stream = None
            self._closed = True
            self._closing = False
        else:
            self.io_loop.add_timeout(1, self.check_close)

    # pylint: disable=W1701
    def __del__(self):
        self.close()

    # pylint: enable=W1701

    @salt.ext.tornado.gen.coroutine
    def getstream(self, **kwargs):
        if self.source_ip or self.source_port:
            kwargs = {
                "source_ip": self.source_ip,
                "source_port": self.source_port,
            }
        stream = None
        while stream is None and (not self._closed and not self._closing):
            try:
                stream = yield self._tcp_client.connect(
                    ip_bracket(self.host, strip=True),
                    self.port,
                    ssl_options=self.opts.get("ssl"),
                    **kwargs,
                )
            except Exception as exc:  # pylint: disable=broad-except
                log.warning(
                    "TCP Message Client encountered an exception while connecting to"
                    " %s:%s: %r, will reconnect in %d seconds",
                    self.host,
                    self.port,
                    exc,
                    self.backoff,
                )
                yield salt.ext.tornado.gen.sleep(self.backoff)
        raise salt.ext.tornado.gen.Return(stream)

    @salt.ext.tornado.gen.coroutine
    def connect(self):
        if self._stream is None:
            self._stream = yield self.getstream()
            if self._stream:
                self._closing = False
                self._closed = False
                if not self._stream_return_running:
                    self.io_loop.spawn_callback(self._stream_return)
                if self.connect_callback:
                    self.connect_callback(True)

    @salt.ext.tornado.gen.coroutine
    def _stream_return(self):
        self._stream_return_running = True
        unpacker = salt.utils.msgpack.Unpacker()
        while not self._closed and not self._closing:
            try:
                wire_bytes = yield self._stream.read_bytes(4096, partial=True)
                unpacker.feed(wire_bytes)
                for framed_msg in unpacker:
                    framed_msg = salt.transport.frame.decode_embedded_strs(framed_msg)
                    header = framed_msg["head"]
                    body = framed_msg["body"]
                    message_id = header.get("mid")

                    if message_id in self.send_future_map:
                        self.send_future_map.pop(message_id).set_result(body)
                        # self.remove_message_timeout(message_id)
                    else:
                        if self._on_recv is not None:
                            self.io_loop.spawn_callback(self._on_recv, header, body)
                        else:
                            log.error(
                                "Got response for message_id %s that we are not"
                                " tracking",
                                message_id,
                            )
            except salt.ext.tornado.iostream.StreamClosedError as e:
                log.debug(
                    "tcp stream to %s:%s closed, unable to recv",
                    self.host,
                    self.port,
                )
                for future in self.send_future_map.values():
                    future.set_exception(e)
                self.send_future_map = {}
                if self._closing or self._closed:
                    return
                if self.disconnect_callback:
                    self.disconnect_callback()
                stream = self._stream
                self._stream = None
                if stream:
                    stream.close()
                unpacker = salt.utils.msgpack.Unpacker()
                yield self.connect()
            except TypeError:
                # This is an invalid transport
                if "detect_mode" in self.opts:
                    log.info(
                        "There was an error trying to use TCP transport; "
                        "attempting to fallback to another transport"
                    )
                else:
                    raise SaltClientError
            except Exception as e:  # pylint: disable=broad-except
                log.error("Exception parsing response", exc_info=True)
                for future in self.send_future_map.values():
                    future.set_exception(e)
                self.send_future_map = {}
                if self._closing or self._closed:
                    return
                if self.disconnect_callback:
                    self.disconnect_callback()
                stream = self._stream
                self._stream = None
                if stream:
                    stream.close()
                unpacker = salt.utils.msgpack.Unpacker()
                yield self.connect()
        self._stream_return_running = False

    def _message_id(self):
        return str(uuid.uuid4())

    # TODO: return a message object which takes care of multiplexing?
    def on_recv(self, callback):
        """
        Register a callback for received messages (that we didn't initiate)
        """
        if callback is None:
            self._on_recv = callback
        else:

            def wrap_recv(header, body):
                callback(body)

            self._on_recv = wrap_recv

    def remove_message_timeout(self, message_id):
        if message_id not in self.send_timeout_map:
            return
        timeout = self.send_timeout_map.pop(message_id)
        self.io_loop.remove_timeout(timeout)

    def timeout_message(self, message_id, msg):
        if message_id not in self.send_future_map:
            return
        future = self.send_future_map.pop(message_id)
        if future is not None:
            future.set_exception(SaltReqTimeoutError("Message timed out"))

    @salt.ext.tornado.gen.coroutine
    def send(self, msg, timeout=None, callback=None, raw=False):
        if self._closing:
            raise ClosingError()
        message_id = self._message_id()
        header = {"mid": message_id}

        future = salt.ext.tornado.concurrent.Future()

        if callback is not None:

            def handle_future(future):
                response = future.result()
                self.io_loop.add_callback(callback, response)

            future.add_done_callback(handle_future)
        # Add this future to the mapping
        self.send_future_map[message_id] = future

        if self.opts.get("detect_mode") is True:
            timeout = 1

        if timeout is not None:
            self.io_loop.call_later(timeout, self.timeout_message, message_id, msg)

        item = salt.transport.frame.frame_msg(msg, header=header)

        @salt.ext.tornado.gen.coroutine
        def _do_send():
            yield self.connect()
            # If the _stream is None, we failed to connect.
            if self._stream:
                yield self._stream.write(item)

        # Run send in a callback so we can wait on the future, in case we time
        # out before we are able to connect.
        self.io_loop.add_callback(_do_send)
        recv = yield future
        raise salt.ext.tornado.gen.Return(recv)


class Subscriber:
    """
    Client object for use with the TCP publisher server
    """

    def __init__(self, stream, address):
        self.stream = stream
        self.address = address
        self._closing = False
        self._read_until_future = None
        self.id_ = None

    def close(self):
        if self._closing:
            return
        self._closing = True
        if not self.stream.closed():
            self.stream.close()
            if self._read_until_future is not None and self._read_until_future.done():
                # This will prevent this message from showing up:
                # '[ERROR   ] Future exception was never retrieved:
                # StreamClosedError'
                # This happens because the logic is always waiting to read
                # the next message and the associated read future is marked
                # 'StreamClosedError' when the stream is closed.
                self._read_until_future.exception()

    # pylint: disable=W1701
    def __del__(self):
        self.close()

    # pylint: enable=W1701


class PubServer(salt.ext.tornado.tcpserver.TCPServer):
    """
    TCP publisher
    """

    def __init__(
        self, opts, io_loop=None, presence_callback=None, remove_presence_callback=None
    ):
        super().__init__(ssl_options=opts.get("ssl"))
        self.io_loop = io_loop
        self.opts = opts
        self._closing = False
        self.clients = set()
        self.presence_events = False
        if presence_callback:
            self.presence_callback = presence_callback
        else:
            self.presence_callback = lambda subscriber, msg: msg
        if remove_presence_callback:
            self.remove_presence_callback = remove_presence_callback
        else:
            self.remove_presence_callback = lambda subscriber: subscriber

    def close(self):
        if self._closing:
            return
        self._closing = True
        for client in list(self.clients):
            client.close()
        self.clients.clear()

    # pylint: disable=W1701
    def __del__(self):
        self.close()

    # pylint: enable=W1701

    def _discard_slow_client(self, client, reason=""):
        """
        Close and forget a subscriber whose write future didn't drain in
        the ``publish_drain_timeout``.  Idempotent -- ``client.close``
        tolerates double-close, and ``set.discard`` is a no-op on absent
        entries.
        """
        if client not in self.clients and getattr(client, "_slow_closed", False):
            return
        client._slow_closed = True
        log.warning(
            "Publisher discarding slow subscriber %s (%s)",
            client.address,
            reason,
        )
        try:
            self.remove_presence_callback(client)
        except Exception:  # pylint: disable=broad-except
            pass
        self.clients.discard(client)
        try:
            client.close()
        except Exception:  # pylint: disable=broad-except
            pass

    @salt.ext.tornado.gen.coroutine
    def _stream_read(
        self, client, _StreamClosedError=salt.ext.tornado.iostream.StreamClosedError
    ):
        unpacker = salt.utils.msgpack.Unpacker()
        while not self._closing:
            try:
                client._read_until_future = client.stream.read_bytes(4096, partial=True)
                wire_bytes = yield client._read_until_future
                unpacker.feed(wire_bytes)
                for framed_msg in unpacker:
                    framed_msg = salt.transport.frame.decode_embedded_strs(framed_msg)
                    body = framed_msg["body"]
                    if self.presence_callback:
                        self.presence_callback(client, body)
            except _StreamClosedError as e:
                log.debug("tcp stream to %s closed, unable to recv", client.address)
                client.close()
                self.remove_presence_callback(client)
                self.clients.discard(client)
                break
            except Exception as e:  # pylint: disable=broad-except
                log.error(
                    "Exception parsing response from %s", client.address, exc_info=True
                )
                continue

    def handle_stream(self, stream, address):
        log.debug("Subscriber at %s connected", address)
        client = Subscriber(stream, address)
        self.clients.add(client)
        self.io_loop.spawn_callback(self._stream_read, client)

    @salt.ext.tornado.gen.coroutine
    def _drain_write(self, client, future, drain_timeout):
        """
        PATCH: drain a single subscriber's write future with a bounded
        timeout, and evict the subscriber if it can't keep up.

        Previously ``publish_payload`` awaited each client's write future
        (concurrently, per commit 73c6970351b, but still unbounded).  A
        single slow subscriber (kernel TCP send buffer full) made
        ``yield future`` never resolve; the ``publish_payload`` coroutine
        held its reference to the payload, and every subsequent broadcast
        piled up more pending ``publish_payload`` coroutines all blocked
        on the same subscriber, wedging the EventPublisher io_loop.

        With EP not draining ``master_event_pull.ipc``, MWorker's
        ``fire_event`` -> ``stream.write`` blocked in the kernel;
        SyncWrapper's ``thread.join()`` never returned and every MWorker
        deadlocked, which in turn wedged MWQ's DEALER send() and cascaded
        down to minion request timeouts and TCP churn.

        Fire-and-forget (via ``io_loop.spawn_callback``) with a per-write
        timeout means ``publish_payload`` returns immediately after
        queueing writes.  Slow subscribers drain in their own tasks.  If
        a subscriber can't drain in ``publish_drain_timeout`` seconds, it
        is closed and removed from ``self.clients`` -- fixes the wedge;
        the peer can reconnect and try again.
        """
        try:
            yield salt.ext.tornado.gen.with_timeout(
                self.io_loop.time() + drain_timeout, future
            )
        except salt.ext.tornado.iostream.StreamClosedError:
            self._discard_slow_client(client, reason="stream closed")
        except salt.ext.tornado.gen.TimeoutError:
            self._discard_slow_client(client, reason=f"drain timeout {drain_timeout}s")
        except Exception as exc:  # pylint: disable=broad-except
            log.warning("Publisher drain to %s failed: %s", client.address, exc)
            self._discard_slow_client(client, reason=str(exc))

    # TODO: ACK the publish through IPC
    @salt.ext.tornado.gen.coroutine
    def publish_payload(self, package, topic_list=None):
        log.trace("TCP PubServer sending payload: %s \n\n %r", package, topic_list)
        payload = salt.transport.frame.frame_msg(package)
        to_remove = []
        # Start writes to every targeted client concurrently so a single
        # slow subscriber can't stall delivery to the rest of the fleet.
        # See https://github.com/saltstack/salt/issues/66282 — commit
        # 73c6970351b made the writes concurrent; this patch adds a
        # per-subscriber timeout that evicts wedged subscribers so their
        # unresolved write futures can no longer pin the io_loop.
        drain_timeout = self.opts.get("publish_drain_timeout", 60.0)
        if topic_list:
            for topic in topic_list:
                sent = False
                for client in list(self.clients):
                    if topic == client.id_:
                        try:
                            fut = client.stream.write(payload)
                            self.io_loop.spawn_callback(
                                self._drain_write, client, fut, drain_timeout
                            )
                            sent = True
                        except salt.ext.tornado.iostream.StreamClosedError:
                            to_remove.append(client)
                if not sent:
                    log.debug("Publish target %s not connected %r", topic, self.clients)
        else:
            for client in list(self.clients):
                try:
                    fut = client.stream.write(payload)
                    self.io_loop.spawn_callback(
                        self._drain_write, client, fut, drain_timeout
                    )
                except salt.ext.tornado.iostream.StreamClosedError:
                    to_remove.append(client)
        for client in to_remove:
            log.debug(
                "Subscriber at %s has disconnected from publisher", client.address
            )
            client.close()
            self.remove_presence_callback(client)
            self.clients.discard(client)
        log.trace("TCP PubServer finished publishing payload")


class TCPPublishServer(salt.transport.base.DaemonizedPublishServer):
    """
    Tornado based TCP PublishServer
    """

    # TODO: opts!
    # Based on default used in salt.ext.tornado.netutil.bind_sockets()
    backlog = 128

    def __init__(self, opts):
        self.opts = opts
        self.pub_sock = None
        self.pub_server = None
        self.io_loop = None
        self._closing = False

    @property
    def topic_support(self):
        return not self.opts.get("order_masters", False)

    def __setstate__(self, state):
        self.__init__(state["opts"])

    def __getstate__(self):
        return {"opts": self.opts}

    def publish_daemon(
        self,
        publish_payload,
        presence_callback=None,
        remove_presence_callback=None,
    ):
        """
        Bind to the interface specified in the configuration file
        """
        io_loop = salt.ext.tornado.ioloop.IOLoop()
        io_loop.make_current()
        self.io_loop = io_loop

        # Spin up the publisher
        self.pub_server = pub_server = PubServer(
            self.opts,
            io_loop=io_loop,
            presence_callback=presence_callback,
            remove_presence_callback=remove_presence_callback,
        )
        sock = _get_socket(self.opts)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        _set_tcp_keepalive(sock, self.opts)
        sock.setblocking(0)
        sock.bind(_get_bind_addr(self.opts, "publish_port"))
        sock.listen(self.backlog)
        # pub_server will take ownership of the socket
        pub_server.add_socket(sock)

        # Set up Salt IPC server
        if self.opts.get("ipc_mode", "") == "tcp":
            pull_uri = int(self.opts.get("tcp_master_publish_pull", 4514))
        else:
            pull_uri = os.path.join(self.opts["sock_dir"], "publish_pull.ipc")
        self.pub_server = pub_server
        pull_sock = salt.transport.ipc.IPCMessageServer(
            pull_uri,
            io_loop=io_loop,
            payload_handler=publish_payload,
        )

        # Securely create socket
        log.warning("Starting the Salt Puller on %s", pull_uri)
        with salt.utils.files.set_umask(0o177):
            pull_sock.start()

        # run forever
        try:
            io_loop.start()
        except (KeyboardInterrupt, SystemExit):
            pass
        finally:
            pull_sock.close()

    def pre_fork(self, process_manager):
        """
        Do anything necessary pre-fork. Since this is on the master side this will
        primarily be used to create IPC channels and create our daemon process to
        do the actual publishing
        """
        process_manager.add_process(self.publish_daemon, name=self.__class__.__name__)

    @salt.ext.tornado.gen.coroutine
    def publish_payload(self, payload, *args):
        ret = yield self.pub_server.publish_payload(payload, *args)
        raise salt.ext.tornado.gen.Return(ret)

    def publish(self, payload, **kwargs):
        """
        Publish "load" to minions
        """
        if self.opts.get("ipc_mode", "") == "tcp":
            pull_uri = int(self.opts.get("tcp_master_publish_pull", 4514))
        else:
            pull_uri = os.path.join(self.opts["sock_dir"], "publish_pull.ipc")

        # PATCH: avoid the nested-SyncWrapper deadlock in the
        # ``fire_event`` -> ``TCPPublishServer.publish`` -> ``pub_sock.send``
        # chain.  ``self.pub_sock`` is a ``SyncWrapper(IPCMessageClient)``.
        # When ``publish`` is invoked while an asyncio/Tornado io_loop is
        # running on the current thread (as is the case in every
        # ``MWorker._handle_payload`` coroutine -> ``_handle_clear`` ->
        # ``_send_pub`` -> ``chan.publish`` path), the outer SyncWrapper
        # around ``SaltEvent.pusher`` (or the MWorker's own io_loop) has
        # already spawned a worker thread that runs this method, then
        # ``self.pub_sock.send`` invokes SyncWrapper *again* -- it detects
        # the running io_loop, spawns yet another thread, and the two
        # threads can deadlock on ``threading.Thread.join()``.  All
        # MWorkers wedge, MWQ's DEALER send() blocks (queue backlog),
        # minions time out and reconnect, dead-peer TCP conns pile up.
        #
        # Fix: when we're already in an async context, bypass the outer
        # SyncWrapper entirely and schedule an ``IPCMessageClient.send``
        # coroutine directly on the running loop (fire-and-forget, which
        # matches the ``fire_event`` precedent at
        # ``salt/utils/event.py``: ``self.io_loop.spawn_callback(
        # self.pusher.send, msg)``).  Cache the raw
        # ``IPCMessageClient`` per running loop because Tornado
        # ``IOStream`` instances (and any locks bound to a specific loop)
        # cannot be safely shared across loops -- this
        # ``TCPPublishServer`` is used by both the sync-mode SyncWrapper
        # thread's io_loop and any coroutine-mode io_loop.  A per-loop
        # ``asyncio.Lock`` serializes concurrent ``fire_event`` tasks so
        # their length-prefixed frames don't interleave on the shared
        # stream (framing corruption would otherwise surface as bogus
        # ~GB length prefixes on the puller side).
        try:
            loop = asyncio.get_running_loop()
            in_async = True
        except RuntimeError:
            in_async = False

        if in_async:
            per_loop = getattr(self, "_async_pub_by_loop", None)
            if per_loop is None:
                # PATCH: WeakKeyDictionary so entries drop when the loop
                # is GC'd.  Keying on ``id(loop)`` would be unsafe --
                # CPython recycles integer ids after GC and a fresh
                # ``SyncWrapper.asyncio_loop`` could land on the same id
                # as a dead one and inherit that dead loop's cached
                # (dead) publisher.
                per_loop = self._async_pub_by_loop = weakref.WeakKeyDictionary()

            entry = per_loop.get(loop)
            if entry is not None:
                pub, _lock = entry
                # PATCH: invalidate on a dead stream.  If the puller
                # side went away (slow-subscriber discard, subscriber
                # process restart) the stream is closed but the entry
                # is still cached -- next ``send`` raises
                # ``StreamClosedError`` forever until we rebuild.  A
                # closed stream is unrecoverable in tornado's
                # ``IOStream``; drop the entry so we reconnect below.
                stream = getattr(pub, "stream", None)
                if stream is None or stream.closed():
                    del per_loop[loop]
                    entry = None

            if entry is None:
                # ``IPCMessageClient`` expects a Tornado ``IOLoop`` for
                # ``add_callback``/``add_future`` scheduling.  In Tornado
                # 6.x ``IOLoop.current()`` is a thin wrapper around the
                # currently-running asyncio loop, so constructing it
                # here (on the loop's own thread) binds the client to
                # the correct loop.
                tio_loop = salt.ext.tornado.ioloop.IOLoop.current()
                pub = salt.transport.ipc.IPCMessageClient(pull_uri, io_loop=tio_loop)
                lock = asyncio.Lock()
                entry = (pub, lock)
                per_loop[loop] = entry
            pub, lock = entry

            async def _send_async():
                async with lock:
                    # Re-resolve the active publisher inside the lock so
                    # tasks that captured a stale ``pub`` from the outer
                    # scope before a concurrent StreamClosedError retry
                    # replaced ``per_loop[loop]`` don't cascade evict
                    # the healthy replacement.  Falls back to the
                    # outer-scope ``pub`` if the cache is empty (e.g.
                    # close() ran between capture and lock acquisition).
                    current_entry = per_loop.get(loop)
                    active_pub = current_entry[0] if current_entry else pub
                    try:
                        if not active_pub.connected():
                            await active_pub.connect()
                        await active_pub.send(payload)
                    except salt.ext.tornado.iostream.StreamClosedError:
                        # PATCH: puller closed on us mid-send.  Only
                        # drop the cache entry if it still points at
                        # the publisher we tried -- otherwise a peer
                        # task already replaced it with a healthy new
                        # one and we must not evict that.
                        current_entry = per_loop.get(loop)
                        if current_entry is not None and current_entry[0] is active_pub:
                            per_loop.pop(loop, None)
                        tio_loop = salt.ext.tornado.ioloop.IOLoop.current()
                        new_pub = salt.transport.ipc.IPCMessageClient(
                            pull_uri, io_loop=tio_loop
                        )
                        per_loop[loop] = (new_pub, lock)
                        try:
                            await new_pub.connect()
                            await new_pub.send(payload)
                        except Exception:  # pylint: disable=broad-except
                            log.exception("TCPPublishServer async publish retry failed")
                    except Exception:  # pylint: disable=broad-except
                        log.exception("TCPPublishServer async publish failed")

            # Schedule on the running loop.  Fire-and-forget matches the
            # sync-branch semantics (SyncWrapper.send returns after the
            # local IPC write completes; callers do not observe an ack
            # from the puller side).
            loop.create_task(_send_async())
            return

        if not self.pub_sock:
            self.pub_sock = salt.utils.asynchronous.SyncWrapper(
                salt.transport.ipc.IPCMessageClient,
                (pull_uri,),
                loop_kwarg="io_loop",
            )
            self.pub_sock.connect()
        self.pub_sock.send(payload)

    def close(self):
        self._closing = True
        if self.pub_sock:
            self.pub_sock.close()
            self.pub_sock = None
        if self.pub_server:
            self.pub_server.close()
            self.pub_server = None
        if self.io_loop:
            self.io_loop.stop()
            self.io_loop.close(all_fds=True)
            self.io_loop = None

    # pylint: disable=W1701
    def __del__(self):
        if not self._closing:
            warnings.warn(
                f"unclosed publish server {self!r}", ResourceWarning, source=self
            )

    # pylint: enable=W1701


class TCPReqClient(salt.transport.base.RequestClient):
    """
    Tornado based TCP RequestClient
    """

    ttype = "tcp"

    def __init__(self, opts, io_loop, **kwargs):  # pylint: disable=W0231
        super().__init__(opts, io_loop, **kwargs)
        self.opts = opts
        self.io_loop = io_loop
        parse = urllib.parse.urlparse(self.opts["master_uri"])
        master_host, master_port = parse.netloc.rsplit(":", 1)
        master_addr = (master_host, int(master_port))
        # self.resolver = Resolver()
        resolver = kwargs.get("resolver")
        self.message_client = salt.transport.tcp.MessageClient(
            opts,
            master_host,
            int(master_port),
            io_loop=io_loop,
            resolver=resolver,
            source_ip=opts.get("source_ip"),
            source_port=opts.get("source_ret_port"),
        )
        self._closing = False

    @salt.ext.tornado.gen.coroutine
    def connect(self):
        self._connect_called = True
        yield self.message_client.connect()

    @salt.ext.tornado.gen.coroutine
    def send(self, load, timeout=60):
        ret = yield self.message_client.send(load, timeout=timeout)
        raise salt.ext.tornado.gen.Return(ret)

    def close(self):
        if self._closing:
            return
        self._closing = True
        self.message_client.close()
