/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.kafka.common.network;

import org.apache.kafka.common.MetricName;
import org.apache.kafka.common.memory.MemoryPool;
import org.apache.kafka.common.memory.SimpleMemoryPool;
import org.apache.kafka.common.metrics.KafkaMetric;
import org.apache.kafka.common.metrics.Metrics;
import org.apache.kafka.common.security.auth.SecurityProtocol;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.common.utils.Time;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.test.TestUtils;

import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;
import org.mockito.MockedConstruction;

import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.lang.reflect.Field;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;
import java.util.stream.Collectors;

import static org.apache.kafka.test.TestUtils.waitForCondition;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.atLeastOnce;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockConstruction;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

/**
 * A set of tests for the selector. These use a test harness that runs a simple socket server that echos back responses.
 */
@Timeout(240)
public class SelectorTest {
    protected static final int BUFFER_SIZE = 4 * 1024;
    private static final String METRIC_GROUP = "MetricGroup";
    private static final long CONNECTION_MAX_IDLE_MS = 5_000;

    protected EchoServer server;
    protected Time time;
    protected Selector selector;
    protected ChannelBuilder channelBuilder;
    protected Metrics metrics;

    @BeforeEach
    public void setUp() throws Exception {
        Map<String, Object> configs = new HashMap<>();
        this.server = new EchoServer(SecurityProtocol.PLAINTEXT, configs);
        this.server.start();
        this.time = new MockTime();
        this.channelBuilder = new PlaintextChannelBuilder(ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT));
        this.channelBuilder.configure(clientConfigs());
        this.metrics = new Metrics();
        this.selector = new Selector(CONNECTION_MAX_IDLE_MS, this.metrics, time, METRIC_GROUP, channelBuilder, new LogContext());
    }

    @AfterEach
    public void tearDown() throws Exception {
        try {
            verifySelectorEmpty();
        } finally {
            this.selector.close();
            this.server.close();
            this.metrics.close();
        }
    }

    protected Map<String, Object> clientConfigs() {
        return new HashMap<>();
    }

    /**
     * Validate that when the server disconnects, a client send ends up with that node in the disconnected list.
     */
    @Test
    public void testServerDisconnect() throws Exception {
        final String node = "0";

        // connect and do a simple request
        blockingConnect(node);
        assertEquals("hello", blockingRequest(node, "hello"));

        KafkaChannel channel = selector.channel(node);

        // disconnect
        this.server.closeConnections();
        waitForCondition(() -> {
            try {
                selector.poll(1000L);
                return selector.disconnected().containsKey(node);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }, 5000, "Failed to observe disconnected node in disconnected set");

        assertNull(channel.selectionKey().attachment());

        // reconnect and do another request
        blockingConnect(node);
        assertEquals("hello", blockingRequest(node, "hello"));
    }

    /**
     * Sending a request with one already in flight should result in an exception
     */
    @Test
    public void testCantSendWithInProgress() throws Exception {
        String node = "0";
        blockingConnect(node);
        selector.send(createSend(node, "test1"));
        try {
            selector.send(createSend(node, "test2"));
            fail("IllegalStateException not thrown when sending a request with one in flight");
        } catch (IllegalStateException e) {
            // Expected exception
        }
        selector.poll(0);
        assertTrue(selector.disconnected().containsKey(node), "Channel not closed");
        assertEquals(ChannelState.FAILED_SEND, selector.disconnected().get(node));
    }

    /**
     * Sending a request to a node without an existing connection should result in an exception
     */
    @Test
    public void testSendWithoutConnecting() {
        assertThrows(IllegalStateException.class, () -> selector.send(createSend("0", "test")));
    }

    /**
     * Sending a request to a node with a bad hostname should result in an exception during connect
     */
    @Test
    public void testNoRouteToHost() {
        assertThrows(IOException.class,
            () -> selector.connect("0", new InetSocketAddress("some.invalid.hostname.foo.bar.local", server.port), BUFFER_SIZE, BUFFER_SIZE));
    }

    /**
     * Sending a request to a node not listening on that port should result in disconnection
     */
    @Test
    public void testConnectionRefused() throws Exception {
        String node = "0";
        ServerSocket nonListeningSocket = new ServerSocket(0);
        int nonListeningPort = nonListeningSocket.getLocalPort();
        selector.connect(node, new InetSocketAddress("localhost", nonListeningPort), BUFFER_SIZE, BUFFER_SIZE);
        while (selector.disconnected().containsKey(node)) {
            assertEquals(ChannelState.NOT_CONNECTED, selector.disconnected().get(node));
            selector.poll(1000L);
        }
        nonListeningSocket.close();
    }

    /**
     * Send multiple requests to several connections in parallel. Validate that responses are received in the order that
     * requests were sent.
     */
    @Test
    public void testNormalOperation() throws Exception {
        int conns = 5;
        int reqs = 500;

        // create connections
        InetSocketAddress addr = new InetSocketAddress("localhost", server.port);
        for (int i = 0; i < conns; i++)
            connect(Integer.toString(i), addr);
        // send echo requests and receive responses
        Map<String, Integer> requests = new HashMap<>();
        Map<String, Integer> responses = new HashMap<>();
        int responseCount = 0;
        for (int i = 0; i < conns; i++) {
            String node = Integer.toString(i);
            selector.send(createSend(node, node + "-0"));
        }

        // loop until we complete all requests
        while (responseCount < conns * reqs) {
            // do the i/o
            selector.poll(0L);

            assertEquals(0, selector.disconnected().size(), "No disconnects should have occurred.");

            // handle any responses we may have gotten
            for (NetworkReceive receive : selector.completedReceives()) {
                String[] pieces = asString(receive).split("-");
                assertEquals(2, pieces.length, "Should be in the form 'conn-counter'");
                assertEquals(receive.source(), pieces[0], "Check the source");
                assertEquals(0, receive.payload().position(), "Check that the receive has kindly been rewound");
                if (responses.containsKey(receive.source())) {
                    assertEquals((int) responses.get(receive.source()), Integer.parseInt(pieces[1]), "Check the request counter");
                    responses.put(receive.source(), responses.get(receive.source()) + 1);
                } else {
                    assertEquals(0, Integer.parseInt(pieces[1]), "Check the request counter");
                    responses.put(receive.source(), 1);
                }
                responseCount++;
            }

            // prepare new sends for the next round
            for (NetworkSend send : selector.completedSends()) {
                String dest = send.destinationId();
                if (requests.containsKey(dest))
                    requests.put(dest, requests.get(dest) + 1);
                else
                    requests.put(dest, 1);
                if (requests.get(dest) < reqs)
                    selector.send(createSend(dest, dest + "-" + requests.get(dest)));
            }
        }
        if (channelBuilder instanceof PlaintextChannelBuilder) {
            assertEquals(0, cipherMetrics(metrics).size());
        } else {
            waitForCondition(() -> cipherMetrics(metrics).size() == 1,
                "Waiting for cipher metrics to be created.");
            assertEquals(5, cipherMetrics(metrics).get(0).metricValue());
        }
    }

    static List<KafkaMetric> cipherMetrics(Metrics metrics) {
        return metrics.metrics().entrySet().stream().
            filter(e -> e.getKey().description().
                contains("The number of connections with this SSL cipher and protocol.")).
            map(Map.Entry::getValue).
            collect(Collectors.toList());
    }

    /**
     * Validate that we can send and receive a message larger than the receive and send buffer size
     */
    @Test
    public void testSendLargeRequest() throws Exception {
        String node = "0";
        blockingConnect(node);
        String big = TestUtils.randomString(10 * BUFFER_SIZE);
        assertEquals(big, blockingRequest(node, big));
    }

    @Test
    public void testPartialSendAndReceiveReflectedInMetrics() throws Exception {
        // We use a large payload to attempt to trigger the partial send and receive logic.
        int payloadSize = 20 * BUFFER_SIZE;
        String payload = TestUtils.randomString(payloadSize);
        String nodeId = "0";
        blockingConnect(nodeId);
        ByteBufferSend send = ByteBufferSend.sizePrefixed(ByteBuffer.wrap(payload.getBytes()));
        NetworkSend networkSend = new NetworkSend(nodeId, send);

        selector.send(networkSend);
        KafkaChannel channel = selector.channel(nodeId);

        KafkaMetric outgoingByteTotal = findUntaggedMetricByName("outgoing-byte-total");
        KafkaMetric incomingByteTotal = findUntaggedMetricByName("incoming-byte-total");

        waitForCondition(() -> {
            long bytesSent = send.size() - send.remaining();
            assertEquals(bytesSent, ((Double) outgoingByteTotal.metricValue()).longValue());

            NetworkReceive currentReceive = channel.currentReceive();
            if (currentReceive != null) {
                assertEquals(currentReceive.bytesRead(), ((Double) incomingByteTotal.metricValue()).intValue());
            }

            selector.poll(50);
            return !selector.completedReceives().isEmpty();
        }, "Failed to receive expected response");

        KafkaMetric requestTotal = findUntaggedMetricByName("request-total");
        assertEquals(1, ((Double) requestTotal.metricValue()).intValue());

        KafkaMetric responseTotal = findUntaggedMetricByName("response-total");
        assertEquals(1, ((Double) responseTotal.metricValue()).intValue());
    }

    @Test
    public void testLargeMessageSequence() throws Exception {
        int bufferSize = 512 * 1024;
        String node = "0";
        int reqs = 50;
        InetSocketAddress addr = new InetSocketAddress("localhost", server.port);
        connect(node, addr);
        String requestPrefix = TestUtils.randomString(bufferSize);
        sendAndReceive(node, requestPrefix, 0, reqs);
    }

    @Test
    public void testEmptyRequest() throws Exception {
        String node = "0";
        blockingConnect(node);
        assertEquals("", blockingRequest(node, ""));
    }

    @Test
    public void testClearCompletedSendsAndReceives() throws Exception {
        int bufferSize = 1024;
        String node = "0";
        InetSocketAddress addr = new InetSocketAddress("localhost", server.port);
        connect(node, addr);
        String request = TestUtils.randomString(bufferSize);
        selector.send(createSend(node, request));
        boolean sent = false;
        boolean received = false;
        while (!sent || !received) {
            selector.poll(1000L);
            assertEquals(0, selector.disconnected().size(), "No disconnects should have occurred.");
            if (!selector.completedSends().isEmpty()) {
                assertEquals(1, selector.completedSends().size());
                selector.clearCompletedSends();
                assertEquals(0, selector.completedSends().size());
                sent = true;
            }

            if (!selector.completedReceives().isEmpty()) {
                assertEquals(1, selector.completedReceives().size());
                assertEquals(request, asString(selector.completedReceives().iterator().next()));
                selector.clearCompletedReceives();
                assertEquals(0, selector.completedReceives().size());
                received = true;
            }
        }
    }

    @Test
    public void testExistingConnectionId() throws IOException {
        blockingConnect("0");
        assertThrows(IllegalStateException.class, () -> blockingConnect("0"));
    }

    @Test
    public void testMute() throws Exception {
        blockingConnect("0");
        blockingConnect("1");

        selector.send(createSend("0", "hello"));
        selector.send(createSend("1", "hi"));

        selector.mute("1");

        while (selector.completedReceives().isEmpty())
            selector.poll(5);
        assertEquals(1, selector.completedReceives().size(), "We should have only one response");
        assertEquals("0", selector.completedReceives().iterator().next().source(),
            "The response should not be from the muted node");

        selector.unmute("1");
        do {
            selector.poll(5);
        } while (selector.completedReceives().isEmpty());
        assertEquals(1, selector.completedReceives().size(), "We should have only one response");
        assertEquals("1", selector.completedReceives().iterator().next().source(), "The response should be from the previously muted node");
    }

    @Test
    public void testCloseAllChannels() throws Exception {
        AtomicInteger closedChannelsCount = new AtomicInteger(0);
        ChannelBuilder channelBuilder = new PlaintextChannelBuilder(null) {
            private int channelIndex = 0;
            @Override
            KafkaChannel buildChannel(String id, TransportLayer transportLayer, Supplier<Authenticator> authenticatorCreator,
                                      int maxReceiveSize, MemoryPool memoryPool, ChannelMetadataRegistry metadataRegistry) {
                return new KafkaChannel(id, transportLayer, authenticatorCreator, maxReceiveSize, memoryPool, metadataRegistry) {
                    private final int index = channelIndex++;
                    @Override
                    public void close() throws IOException {
                        closedChannelsCount.getAndIncrement();
                        super.close();
                        if (index == 0) throw new RuntimeException("you should fail");
                    }
                };
            }
        };
        channelBuilder.configure(clientConfigs());
        Selector selector = new Selector(CONNECTION_MAX_IDLE_MS, new Metrics(), new MockTime(), "MetricGroup", channelBuilder, new LogContext());
        selector.connect("0", new InetSocketAddress("localhost", server.port), BUFFER_SIZE, BUFFER_SIZE);
        selector.connect("1", new InetSocketAddress("localhost", server.port), BUFFER_SIZE, BUFFER_SIZE);
        assertThrows(RuntimeException.class, selector::close);
        assertEquals(2, closedChannelsCount.get());
    }

    @Test
    public void registerFailure() throws Exception {
        final String channelId = "1";

        final ChannelBuilder channelBuilder = mock(ChannelBuilder.class);

        when(channelBuilder.buildChannel(eq(channelId), any(SelectionKey.class), anyInt(), any(MemoryPool.class),
                any(ChannelMetadataRegistry.class))).thenThrow(new RuntimeException("Test exception"));

        try (MockedConstruction<Selector.SelectorChannelMetadataRegistry> mockedMetadataRegistry =
                     mockConstruction(Selector.SelectorChannelMetadataRegistry.class)) {
            Selector selector = new Selector(CONNECTION_MAX_IDLE_MS, new Metrics(), new MockTime(), "MetricGroup", channelBuilder, new LogContext());
            final SocketChannel socketChannel = SocketChannel.open();
            socketChannel.configureBlocking(false);
            IOException e = assertThrows(IOException.class, () -> selector.register(channelId, socketChannel));
            assertTrue(e.getCause().getMessage().contains("Test exception"), "Unexpected exception: " + e);
            assertFalse(socketChannel.isOpen(), "Socket not closed");
            // Ideally, metadataRegistry is closed by the KafkaChannel but if the KafkaChannel is not created due to
            // an error such as in a case like this, the Selector should be closing the metadataRegistry instead.
            verify(mockedMetadataRegistry.constructed().get(0)).close();
            selector.close();
        }
    }

    @Test
    public void testCloseOldestConnection() throws Exception {
        String id = "0";
        selector.connect(id, new InetSocketAddress("localhost", server.port), BUFFER_SIZE, BUFFER_SIZE);
        NetworkTestUtils.waitForChannelConnected(selector, id);
        time.sleep(CONNECTION_MAX_IDLE_MS + 1_000);
        selector.poll(0);

        assertTrue(selector.disconnected().containsKey(id), "The idle connection should have been closed");
        assertEquals(ChannelState.EXPIRED, selector.disconnected().get(id));
    }

    @Test
    public void testIdleExpiryWithoutReadyKeys() throws IOException {
        String id = "0";
        selector.connect(id, new InetSocketAddress("localhost", server.port), BUFFER_SIZE, BUFFER_SIZE);
        KafkaChannel channel = selector.channel(id);
        channel.selectionKey().interestOps(0);

        time.sleep(CONNECTION_MAX_IDLE_MS + 1_000);
        selector.poll(0);
        assertTrue(selector.disconnected().containsKey(id), "The idle connection should have been closed");
        assertEquals(ChannelState.EXPIRED, selector.disconnected().get(id));
    }

    @Test
    public void testImmediatelyConnectedCleaned() throws Exception {
        Metrics metrics = new Metrics(); // new metrics object to avoid metric registration conflicts
        Selector selector = new ImmediatelyConnectingSelector(CONNECTION_MAX_IDLE_MS, metrics, time, "MetricGroup", channelBuilder, new LogContext());

        try {
            testImmediatelyConnectedCleaned(selector, true);
            testImmediatelyConnectedCleaned(selector, false);
        } finally {
            selector.close();
            metrics.close();
        }
    }

    private static class ImmediatelyConnectingSelector extends Selector {
        public ImmediatelyConnectingSelector(long connectionMaxIdleMS,
                                             Metrics metrics,
                                             Time time,
                                             String metricGrpPrefix,
                                             ChannelBuilder channelBuilder,
                                             LogContext logContext) {
            super(connectionMaxIdleMS, metrics, time, metricGrpPrefix, channelBuilder, logContext);
        }

        @Override
        protected boolean doConnect(SocketChannel channel, InetSocketAddress address) throws IOException {
            // Use a blocking connect to trigger the immediately connected path
            channel.configureBlocking(true);
            boolean connected = super.doConnect(channel, address);
            channel.configureBlocking(false);
            return connected;
        }
    }

    private void testImmediatelyConnectedCleaned(Selector selector, boolean closeAfterFirstPoll) throws Exception {
        String id = "0";
        selector.connect(id, new InetSocketAddress("localhost", server.port), BUFFER_SIZE, BUFFER_SIZE);
        verifyNonEmptyImmediatelyConnectedKeys(selector);
        if (closeAfterFirstPoll) {
            selector.poll(0);
            verifyEmptyImmediatelyConnectedKeys(selector);
        }
        selector.close(id);
        verifySelectorEmpty(selector);
    }

    /**
     * Verify that if Selector#connect fails and throws an Exception, all related objects
     * are cleared immediately before the exception is propagated.
     */
    @Test
    public void testConnectException() throws Exception {
        Metrics metrics = new Metrics();
        AtomicBoolean throwIOException = new AtomicBoolean();
        Selector selector = new ImmediatelyConnectingSelector(CONNECTION_MAX_IDLE_MS, metrics, time, "MetricGroup", channelBuilder, new LogContext()) {
            @Override
            protected SelectionKey registerChannel(String id, SocketChannel socketChannel, int interestedOps) throws IOException {
                SelectionKey key = super.registerChannel(id, socketChannel, interestedOps);
                key.cancel();
                if (throwIOException.get())
                    throw new IOException("Test exception");
                return key;
            }
        };

        try {
            verifyImmediatelyConnectedException(selector, "0");
            throwIOException.set(true);
            verifyImmediatelyConnectedException(selector, "1");
        } finally {
            selector.close();
            metrics.close();
        }
    }

    private void verifyImmediatelyConnectedException(Selector selector, String id) throws Exception {
        try {
            selector.connect(id, new InetSocketAddress("localhost", server.port), BUFFER_SIZE, BUFFER_SIZE);
            fail("Expected exception not thrown");
        } catch (Exception e) {
            verifyEmptyImmediatelyConnectedKeys(selector);
            assertNull(selector.channel(id), "Channel not removed");
            ensureEmptySelectorFields(selector);
        }
    }

    /*
     * Verifies that a muted connection is expired on idle timeout even if there are pending
     * receives on the socket.
     */
    @Test
    public void testExpireConnectionWithPendingReceives() throws Exception {
        KafkaChannel channel = createConnectionWithPendingReceives(5);
        verifyChannelExpiry(channel);
    }

    /**
     * Verifies that a muted connection closed by peer is expired on idle timeout even if there are pending
     * receives on the socket.
     */
    @Test
    public void testExpireClosedConnectionWithPendingReceives() throws Exception {
        KafkaChannel channel = createConnectionWithPendingReceives(5);
        server.closeConnections();
        verifyChannelExpiry(channel);
    }

    private void verifyChannelExpiry(KafkaChannel channel) throws Exception {
        String id = channel.id();
        selector.mute(id); // Mute to allow channel to be expired even if more data is available for read
        time.sleep(CONNECTION_MAX_IDLE_MS + 1_000);
        selector.poll(0);
        assertNull(selector.channel(id), "Channel not expired");
        assertNull(selector.closingChannel(id), "Channel not removed from closingChannels");
        assertEquals(ChannelState.EXPIRED, channel.state());
        assertNull(channel.selectionKey().attachment());
        assertTrue(selector.disconnected().containsKey(id), "Disconnect not notified");
        assertEquals(ChannelState.EXPIRED, selector.disconnected().get(id));
        verifySelectorEmpty();
    }

    /**
     * Verifies that sockets with incoming data available are not expired.
     * For PLAINTEXT, pending receives are always read from socket without any buffering, so this
     * test is only verifying that channels are not expired while there is data to read from socket.
     * For SSL, pending receives may also be in SSL netReadBuffer or appReadBuffer. So the test verifies
     * that connection is not expired when data is available from buffers or network.
     */
    @Test
    public void testCloseOldestConnectionWithMultiplePendingReceives() throws Exception {
        int expectedReceives = 5;
        KafkaChannel channel = createConnectionWithPendingReceives(expectedReceives);
        int completedReceives = selector.completedReceives().size();

        while (selector.disconnected().isEmpty()) {
            time.sleep(CONNECTION_MAX_IDLE_MS + 1_000);
            selector.poll(completedReceives == expectedReceives ? 0 : 1_000);
            completedReceives += selector.completedReceives().size();
        }

        assertEquals(expectedReceives, completedReceives);
        assertNull(selector.channel(channel.id()), "Channel not expired");
        assertNull(selector.closingChannel(channel.id()), "Channel not expired");
        assertTrue(selector.disconnected().containsKey(channel.id()), "Disconnect not notified");
        assertTrue(selector.completedReceives().isEmpty(), "Unexpected receive");
    }

    /**
     * Tests that graceful close of channel processes remaining data from socket read buffers.
     * Since we cannot determine how much data is available in the buffers, this test verifies that
     * multiple receives are completed after server shuts down connections, with retries to tolerate
     * cases where data may not be available in the socket buffer.
     */
    @Test
    public void testGracefulClose() throws Exception {
        int maxReceiveCountAfterClose = 0;
        for (int i = 6; i <= 100 && maxReceiveCountAfterClose < 5; i++) {
            int receiveCount = 0;
            KafkaChannel channel = createConnectionWithPendingReceives(i);
            // Poll until one or more receives complete and then close the server-side connection
            waitForCondition(() -> {
                selector.poll(1000);
                return !selector.completedReceives().isEmpty();
            }, 5000, "Receive not completed");
            server.closeConnections();
            while (selector.disconnected().isEmpty()) {
                selector.poll(1);
                receiveCount += selector.completedReceives().size();
                assertTrue(selector.completedReceives().size() <= 1, "Too many completed receives in one poll");
            }
            assertEquals(channel.id(), selector.disconnected().keySet().iterator().next());
            maxReceiveCountAfterClose = Math.max(maxReceiveCountAfterClose, receiveCount);
        }
        assertTrue(maxReceiveCountAfterClose >= 5, "Too few receives after close: " + maxReceiveCountAfterClose);
    }

    /**
     * Tests that graceful close is not delayed if only part of an incoming receive is
     * available in the socket buffer.
     */
    @Test
    public void testPartialReceiveGracefulClose() throws Exception {
        String id = "0";
        blockingConnect(id);
        KafkaChannel channel = selector.channel(id);
        // Inject a NetworkReceive into Kafka channel with a large size
        injectNetworkReceive(channel, 100000);
        sendNoReceive(channel, 2); // Send some data that gets received as part of injected receive
        selector.poll(1000); // Wait until some data arrives, but not a completed receive
        assertEquals(0, selector.completedReceives().size());
        server.closeConnections();
        waitForCondition(() -> {
            try {
                selector.poll(100);
                return !selector.disconnected().isEmpty();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }, 10000, "Channel not disconnected");
        assertEquals(1, selector.disconnected().size());
        assertEquals(channel.id(), selector.disconnected().keySet().iterator().next());
        assertEquals(0, selector.completedReceives().size());
    }

    @Test
    public void testMuteOnOOM() throws Exception {
        //clean up default selector, replace it with one that uses a finite mem pool
        selector.close();
        MemoryPool pool = new SimpleMemoryPool(900, 900, false, null);
        selector = new Selector(NetworkReceive.UNLIMITED, CONNECTION_MAX_IDLE_MS, metrics, time, "MetricGroup",
            new HashMap<>(), true, false, channelBuilder, pool, new LogContext());

        try (ServerSocketChannel ss = ServerSocketChannel.open()) {
            ss.bind(new InetSocketAddress(0));

            InetSocketAddress serverAddress = (InetSocketAddress) ss.getLocalAddress();

            Thread sender1 = createSender(serverAddress, randomPayload(900));
            Thread sender2 = createSender(serverAddress, randomPayload(900));
            sender1.start();
            sender2.start();

            //wait until everything has been flushed out to network (assuming payload size is smaller than OS buffer size)
            //this is important because we assume both requests' prefixes (1st 4 bytes) have made it.
            sender1.join(5000);
            sender2.join(5000);

            SocketChannel channelX = ss.accept(); //not defined if its 1 or 2
            channelX.configureBlocking(false);
            SocketChannel channelY = ss.accept();
            channelY.configureBlocking(false);
            selector.register("clientX", channelX);
            selector.register("clientY", channelY);

            Collection<NetworkReceive> completed = Collections.emptyList();
            long deadline = System.currentTimeMillis() + 5000;
            while (System.currentTimeMillis() < deadline && completed.isEmpty()) {
                selector.poll(1000);
                completed = selector.completedReceives();
            }
            assertEquals(1, completed.size(), "could not read a single request within timeout");
            NetworkReceive firstReceive = completed.iterator().next();
            assertEquals(0, pool.availableMemory());
            assertTrue(selector.isOutOfMemory());

            selector.poll(10);
            assertTrue(selector.completedReceives().isEmpty());
            assertEquals(0, pool.availableMemory());
            assertTrue(selector.isOutOfMemory());

            firstReceive.close();
            assertEquals(900, pool.availableMemory()); //memory has been released back to pool

            completed = Collections.emptyList();
            deadline = System.currentTimeMillis() + 5000;
            while (System.currentTimeMillis() < deadline && completed.isEmpty()) {
                selector.poll(1000);
                completed = selector.completedReceives();
            }
            assertEquals(1, selector.completedReceives().size(), "could not read a single request within timeout");
            assertEquals(0, pool.availableMemory());
            assertFalse(selector.isOutOfMemory());
        }
    }

    private Thread createSender(InetSocketAddress serverAddress, byte[] payload) {
        return new PlaintextSender(serverAddress, payload);
    }

    protected byte[] randomPayload(int sizeBytes) throws Exception {
        Random random = new Random();
        byte[] payload = new byte[sizeBytes + 4];
        random.nextBytes(payload);
        ByteArrayOutputStream prefixOs = new ByteArrayOutputStream();
        DataOutputStream prefixDos = new DataOutputStream(prefixOs);
        prefixDos.writeInt(sizeBytes);
        prefixDos.flush();
        prefixDos.close();
        prefixOs.flush();
        prefixOs.close();
        byte[] prefix = prefixOs.toByteArray();
        System.arraycopy(prefix, 0, payload, 0, prefix.length);
        return payload;
    }

    /**
     * Tests that a connect and disconnect in a single poll invocation results in the channel id being
     * in `disconnected`, but not `connected`.
     */
    @Test
    public void testConnectDisconnectDuringInSinglePoll() throws Exception {
        // channel is connected, not ready and it throws an exception during prepare
        KafkaChannel kafkaChannel = mock(KafkaChannel.class);
        when(kafkaChannel.id()).thenReturn("1");
        when(kafkaChannel.socketDescription()).thenReturn("");
        when(kafkaChannel.state()).thenReturn(ChannelState.NOT_CONNECTED);
        when(kafkaChannel.finishConnect()).thenReturn(true);
        when(kafkaChannel.isConnected()).thenReturn(true);
        when(kafkaChannel.ready()).thenReturn(false);
        doThrow(new IOException()).when(kafkaChannel).prepare();

        SelectionKey selectionKey = mock(SelectionKey.class);
        when(kafkaChannel.selectionKey()).thenReturn(selectionKey);
        SocketChannel socket = SocketChannel.open();
        when(selectionKey.channel()).thenReturn(socket);
        when(selectionKey.readyOps()).thenReturn(SelectionKey.OP_CONNECT);
        when(selectionKey.attachment()).thenReturn(kafkaChannel);

        Set<SelectionKey> selectionKeys = Set.of(selectionKey);
        selector.pollSelectionKeys(selectionKeys, false, System.nanoTime());

        assertFalse(selector.connected().contains(kafkaChannel.id()));
        assertTrue(selector.disconnected().containsKey(kafkaChannel.id()));

        verify(kafkaChannel, atLeastOnce()).ready();
        verify(kafkaChannel).disconnect();
        verify(kafkaChannel).close();
        verify(selectionKey).cancel();
        socket.close();
    }

    @Test
    public void testOutboundConnectionsCountInConnectionCreationMetric() throws Exception {
        // create connections
        int expectedConnections = 5;
        InetSocketAddress addr = new InetSocketAddress("localhost", server.port);
        for (int i = 0; i < expectedConnections; i++)
            connect(Integer.toString(i), addr);

        // Poll continuously, as we cannot guarantee that the first call will see all connections
        int seenConnections = 0;
        for (int i = 0; i < 10; i++) {
            selector.poll(100L);
            seenConnections += selector.connected().size();
            if (seenConnections == expectedConnections)
                break;
        }

        assertEquals((double) expectedConnections, getMetric("connection-creation-total").metricValue());
        assertEquals((double) expectedConnections, getMetric("connection-count").metricValue());
    }

    @Test
    public void testInboundConnectionsCountInConnectionCreationMetric() throws Exception {
        int conns = 5;

        try (ServerSocketChannel ss = ServerSocketChannel.open()) {
            ss.bind(new InetSocketAddress(0));
            InetSocketAddress serverAddress = (InetSocketAddress) ss.getLocalAddress();

            for (int i = 0; i < conns; i++) {
                Thread sender = createSender(serverAddress, randomPayload(1));
                sender.start();
                try (SocketChannel channel = ss.accept()) {
                    channel.configureBlocking(false);
                    selector.register(Integer.toString(i), channel);
                } finally {
                    sender.join();
                }
            }
        }

        assertEquals((double) conns, getMetric("connection-creation-total").metricValue());
        assertEquals((double) conns, getMetric("connection-count").metricValue());
    }

    @Test
    public void testConnectionsByClientMetric() throws Exception {
        String node = "0";
        Map<String, String> unknownNameAndVersion = softwareNameAndVersionTags(
            ClientInformation.UNKNOWN_NAME_OR_VERSION, ClientInformation.UNKNOWN_NAME_OR_VERSION);
        Map<String, String> knownNameAndVersion = softwareNameAndVersionTags("A", "B");

        try (ServerSocketChannel ss = ServerSocketChannel.open()) {
            ss.bind(new InetSocketAddress(0));
            InetSocketAddress serverAddress = (InetSocketAddress) ss.getLocalAddress();

            Thread sender = createSender(serverAddress, randomPayload(1));
            sender.start();
            SocketChannel channel = ss.accept();
            channel.configureBlocking(false);

            // Metric with unknown / unknown should be there
            selector.register(node, channel);
            assertEquals(1,
                getMetric("connections", unknownNameAndVersion).metricValue());
            assertEquals(ClientInformation.EMPTY,
                selector.channel(node).channelMetadataRegistry().clientInformation());

            // Metric with unknown / unknown should not be there, metric with A / B should be there
            ClientInformation clientInformation = new ClientInformation("A", "B");
            selector.channel(node).channelMetadataRegistry()
                .registerClientInformation(clientInformation);
            assertEquals(clientInformation,
                selector.channel(node).channelMetadataRegistry().clientInformation());
            assertEquals(0, getMetric("connections", unknownNameAndVersion).metricValue());
            assertEquals(1, getMetric("connections", knownNameAndVersion).metricValue());

            // Metric with A / B should not be there,
            selector.close(node);
            assertEquals(0, getMetric("connections", knownNameAndVersion).metricValue());
        }
    }

    private Map<String, String> softwareNameAndVersionTags(String clientSoftwareName, String clientSoftwareVersion) {
        Map<String, String> tags = new HashMap<>(2);
        tags.put("clientSoftwareName", clientSoftwareName);
        tags.put("clientSoftwareVersion", clientSoftwareVersion);
        return tags;
    }

    private KafkaMetric getMetric(String name, Map<String, String> tags) throws Exception {
        Optional<Map.Entry<MetricName, KafkaMetric>> metric = metrics.metrics().entrySet().stream()
            .filter(entry ->
                entry.getKey().name().equals(name) && entry.getKey().tags().equals(tags))
            .findFirst();
        if (metric.isEmpty())
            throw new Exception(String.format("Could not find metric called %s with tags %s", name, tags.toString()));

        return metric.get().getValue();
    }

    @SuppressWarnings("unchecked")
    @Test
    public void testLowestPriorityChannel() throws Exception {
        int conns = 5;
        InetSocketAddress addr = new InetSocketAddress("localhost", server.port);
        for (int i = 0; i < conns; i++) {
            connect(String.valueOf(i), addr);
        }
        assertNotNull(selector.lowestPriorityChannel());
        for (int i = conns - 1; i >= 0; i--) {
            if (i != 2) 
                assertEquals("", blockingRequest(String.valueOf(i), ""));
            time.sleep(10);
        }
        assertEquals("2", selector.lowestPriorityChannel().id());

        Field field = Selector.class.getDeclaredField("closingChannels");
        field.setAccessible(true);
        Map<String, KafkaChannel> closingChannels = (Map<String, KafkaChannel>) field.get(selector);
        closingChannels.put("3", selector.channel("3"));
        assertEquals("3", selector.lowestPriorityChannel().id());
        closingChannels.remove("3");

        for (int i = 0; i < conns; i++) {
            selector.close(String.valueOf(i));
        }
        assertNull(selector.lowestPriorityChannel());
    }

    @Test
    public void testMetricsCleanupOnSelectorClose() throws Exception {
        Metrics metrics = new Metrics();
        Selector selector = new ImmediatelyConnectingSelector(CONNECTION_MAX_IDLE_MS, metrics, time, "MetricGroup", channelBuilder, new LogContext()) {
            @Override
            public void close(String id) {
                super.close(id);
                throw new RuntimeException();
            }
        };
        assertTrue(metrics.metrics().size() > 1);
        String id = "0";
        selector.connect(id, new InetSocketAddress("localhost", server.port), BUFFER_SIZE, BUFFER_SIZE);

        // Close the selector and ensure a RuntimeException has been throw
        assertThrows(RuntimeException.class, selector::close);

        // We should only have one remaining metric for kafka-metrics-count, which is a global metric
        assertEquals(1, metrics.metrics().size());
    }

    @Test
    public void testWriteCompletesSendWithNoBytesWritten() throws IOException {
        KafkaChannel channel = mock(KafkaChannel.class);
        when(channel.id()).thenReturn("1");
        when(channel.write()).thenReturn(0L);
        NetworkSend send = new NetworkSend("destination", new ByteBufferSend(ByteBuffer.allocate(0)));
        when(channel.maybeCompleteSend()).thenReturn(send);
        selector.write(channel);
        assertEquals(Collections.singletonList(send), selector.completedSends());
    }

    /**
     * Ensure that no errors are thrown if channels are closed while processing multiple completed receives
     */
    @Test
    public void testChannelCloseWhileProcessingReceives() throws Exception {
        int numChannels = 4;
        Map<String, KafkaChannel> channels = TestUtils.fieldValue(selector, Selector.class, "channels");
        Set<SelectionKey> selectionKeys = new HashSet<>();
        for (int i = 0; i < numChannels; i++) {
            String id = String.valueOf(i);
            KafkaChannel channel = mock(KafkaChannel.class);
            channels.put(id, channel);
            when(channel.id()).thenReturn(id);
            when(channel.state()).thenReturn(ChannelState.READY);
            when(channel.isConnected()).thenReturn(true);
            when(channel.ready()).thenReturn(true);
            when(channel.read()).thenReturn(1L);

            SelectionKey selectionKey = mock(SelectionKey.class);
            when(channel.selectionKey()).thenReturn(selectionKey);
            when(selectionKey.isValid()).thenReturn(true);
            when(selectionKey.isReadable()).thenReturn(true);
            when(selectionKey.readyOps()).thenReturn(SelectionKey.OP_READ);
            when(selectionKey.attachment())
                    .thenReturn(channel)
                    .thenReturn(null);
            selectionKeys.add(selectionKey);

            NetworkReceive receive = mock(NetworkReceive.class);
            when(receive.source()).thenReturn(id);
            when(receive.size()).thenReturn(10);
            when(receive.bytesRead()).thenReturn(1);
            when(receive.payload()).thenReturn(ByteBuffer.allocate(10));
            when(channel.maybeCompleteReceive()).thenReturn(receive);
        }

        selector.pollSelectionKeys(selectionKeys, false, System.nanoTime());
        assertEquals(numChannels, selector.completedReceives().size());
        Set<KafkaChannel> closed = new HashSet<>();
        Set<KafkaChannel> notClosed = new HashSet<>();
        for (NetworkReceive receive : selector.completedReceives()) {
            KafkaChannel channel = selector.channel(receive.source());
            assertNotNull(channel);
            if (closed.size() < 2) {
                selector.close(channel.id());
                closed.add(channel);
            } else
                notClosed.add(channel);
        }
        assertEquals(notClosed, new HashSet<>(selector.channels()));
        closed.forEach(channel -> assertNull(selector.channel(channel.id())));

        selector.poll(0);
        assertEquals(0, selector.completedReceives().size());
    }

    private String blockingRequest(String node, String s) throws IOException {
        selector.send(createSend(node, s));
        while (true) {
            selector.poll(1000L);
            for (NetworkReceive receive : selector.completedReceives())
                if (receive.source().equals(node))
                    return asString(receive);
        }
    }

    protected void connect(String node, InetSocketAddress serverAddr) throws IOException {
        selector.connect(node, serverAddr, BUFFER_SIZE, BUFFER_SIZE);
    }

    /* connect and wait for the connection to complete */
    private void blockingConnect(String node) throws IOException {
        blockingConnect(node, new InetSocketAddress("localhost", server.port));
    }

    protected void blockingConnect(String node, InetSocketAddress serverAddr) throws IOException {
        selector.connect(node, serverAddr, BUFFER_SIZE, BUFFER_SIZE);
        NetworkTestUtils.waitForChannelReady(selector, node);
    }

    protected final NetworkSend createSend(String node, String payload) {
        return new NetworkSend(node, ByteBufferSend.sizePrefixed(ByteBuffer.wrap(payload.getBytes())));
    }

    protected String asString(NetworkReceive receive) {
        return new String(Utils.toArray(receive.payload()));
    }

    private void sendAndReceive(String node, String requestPrefix, int startIndex, int endIndex) throws Exception {
        int requests = startIndex;
        int responses = startIndex;
        selector.send(createSend(node, requestPrefix + "-" + startIndex));
        requests++;
        while (responses < endIndex) {
            // do the i/o
            selector.poll(0L);
            assertEquals(0, selector.disconnected().size(), "No disconnects should have occurred.");
            // handle requests and responses of the fast node
            for (NetworkReceive receive : selector.completedReceives()) {
                assertEquals(requestPrefix + "-" + responses, asString(receive));
                responses++;
            }

            for (int i = 0; i < selector.completedSends().size() && requests < endIndex; i++, requests++) {
                selector.send(createSend(node, requestPrefix + "-" + requests));
            }
        }
    }

    private void verifyNonEmptyImmediatelyConnectedKeys(Selector selector) throws Exception {
        Field field = Selector.class.getDeclaredField("immediatelyConnectedKeys");
        field.setAccessible(true);
        Collection<?> immediatelyConnectedKeys = (Collection<?>) field.get(selector);
        assertFalse(immediatelyConnectedKeys.isEmpty());
    }

    private void verifyEmptyImmediatelyConnectedKeys(Selector selector) throws Exception {
        Field field = Selector.class.getDeclaredField("immediatelyConnectedKeys");
        ensureEmptySelectorField(selector, field);
    }

    protected void verifySelectorEmpty() throws Exception {
        verifySelectorEmpty(this.selector);
    }

    public void verifySelectorEmpty(Selector selector) throws Exception {
        for (KafkaChannel channel : selector.channels()) {
            selector.close(channel.id());
            assertNull(channel.selectionKey().attachment());
        }
        selector.poll(0);
        selector.poll(0); // Poll a second time to clear everything
        ensureEmptySelectorFields(selector);
    }

    private void ensureEmptySelectorFields(Selector selector) throws Exception {
        for (Field field : Selector.class.getDeclaredFields()) {
            ensureEmptySelectorField(selector, field);
        }
    }

    private void ensureEmptySelectorField(Selector selector, Field field) throws Exception {
        field.setAccessible(true);
        Object obj = field.get(selector);
        if (obj instanceof Collection)
            assertTrue(((Collection<?>) obj).isEmpty(), "Field not empty: " + field + " " + obj);
        else if (obj instanceof Map)
            assertTrue(((Map<?, ?>) obj).isEmpty(), "Field not empty: " + field + " " + obj);
    }

    private KafkaMetric getMetric(String name) throws Exception {
        Optional<Map.Entry<MetricName, KafkaMetric>> metric = metrics.metrics().entrySet().stream()
                .filter(entry -> entry.getKey().name().equals(name))
                .findFirst();
        if (metric.isEmpty())
            throw new Exception(String.format("Could not find metric called %s", name));

        return metric.get().getValue();
    }

    private KafkaMetric findUntaggedMetricByName(String name) {
        MetricName metricName = new MetricName(name, METRIC_GROUP + "-metrics", "", new HashMap<>());
        KafkaMetric metric = metrics.metrics().get(metricName);
        assertNotNull(metric);
        return metric;
    }

    /**
     * Creates a connection, sends the specified number of requests and returns without reading
     * any incoming data. Some of the incoming data may be in the socket buffers when this method
     * returns, but there is no guarantee that all the data from the server will be available
     * immediately. 
     */
    private KafkaChannel createConnectionWithPendingReceives(int pendingReceives) throws Exception {
        String id = "0";
        blockingConnect(id);
        KafkaChannel channel = selector.channel(id);
        sendNoReceive(channel, pendingReceives);
        return channel;
    }

    /**
     * Sends the specified number of requests and waits for the requests to be sent.
     * The channel is muted during polling to ensure that incoming data is not received.
     */
    private void sendNoReceive(KafkaChannel channel, int numRequests) throws Exception {
        selector.mute(channel.id());
        for (int i = 0; i < numRequests; i++) {
            selector.send(createSend(channel.id(), String.valueOf(i)));
            do {
                selector.poll(10);
            } while (selector.completedSends().isEmpty());
        }
        selector.unmute(channel.id());
    }

    /**
     * Injects a NetworkReceive for channel with size buffer filled in with the provided size
     * and a payload buffer allocated with that size, but no data in the payload buffer.
     */
    private void injectNetworkReceive(KafkaChannel channel, int size) throws Exception {
        NetworkReceive receive = new NetworkReceive();
        TestUtils.setFieldValue(channel, "receive", receive);
        ByteBuffer sizeBuffer = TestUtils.fieldValue(receive, NetworkReceive.class, "size");
        sizeBuffer.putInt(size);
        TestUtils.setFieldValue(receive, "buffer", ByteBuffer.allocate(size));
    }
}
