content: /*
   * 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.clients;
  
  import org.apache.kafka.common.Cluster;
  import org.apache.kafka.common.KafkaException;
  import org.apache.kafka.common.Node;
  import org.apache.kafka.common.errors.AuthenticationException;
  import org.apache.kafka.common.errors.RebootstrapRequiredException;
  import org.apache.kafka.common.errors.UnsupportedVersionException;
  import org.apache.kafka.common.internals.ClusterResourceListeners;
  import org.apache.kafka.common.message.ApiMessageType;
  import org.apache.kafka.common.message.ApiVersionsResponseData;
  import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersion;
  import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionCollection;
  import org.apache.kafka.common.message.GetTelemetrySubscriptionsRequestData;
  import org.apache.kafka.common.message.GetTelemetrySubscriptionsResponseData;
  import org.apache.kafka.common.message.ProduceRequestData;
  import org.apache.kafka.common.message.ProduceResponseData;
  import org.apache.kafka.common.message.PushTelemetryRequestData;
  import org.apache.kafka.common.message.PushTelemetryResponseData;
  import org.apache.kafka.common.network.NetworkReceive;
  import org.apache.kafka.common.protocol.ApiKeys;
  import org.apache.kafka.common.protocol.Errors;
  import org.apache.kafka.common.requests.AbstractResponse;
  import org.apache.kafka.common.requests.ApiVersionsResponse;
  import org.apache.kafka.common.requests.GetTelemetrySubscriptionsRequest;
  import org.apache.kafka.common.requests.GetTelemetrySubscriptionsResponse;
  import org.apache.kafka.common.requests.MetadataRequest;
  import org.apache.kafka.common.requests.MetadataResponse;
  import org.apache.kafka.common.requests.ProduceRequest;
  import org.apache.kafka.common.requests.ProduceResponse;
  import org.apache.kafka.common.requests.PushTelemetryRequest;
  import org.apache.kafka.common.requests.PushTelemetryResponse;
  import org.apache.kafka.common.requests.RequestHeader;
  import org.apache.kafka.common.requests.RequestTestUtils;
  import org.apache.kafka.common.security.authenticator.SaslClientAuthenticator;
  import org.apache.kafka.common.telemetry.internals.ClientTelemetrySender;
  import org.apache.kafka.common.utils.LogContext;
  import org.apache.kafka.common.utils.MockTime;
  import org.apache.kafka.test.DelayedReceive;
  import org.apache.kafka.test.MockSelector;
  import org.apache.kafka.test.TestUtils;
  
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.Test;
  
  import java.net.InetAddress;
  import java.net.InetSocketAddress;
  import java.net.UnknownHostException;
  import java.nio.ByteBuffer;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Optional;
  import java.util.Set;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.stream.Collectors;
  import java.util.stream.IntStream;
  
  import static org.apache.kafka.common.protocol.ApiKeys.PRODUCE;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertNotEquals;
  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.anyLong;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.times;
  import static org.mockito.Mockito.verify;
  import static org.mockito.Mockito.when;
  
  public class NetworkClientTest {
  
      protected final int defaultRequestTimeoutMs = 1000;
      protected final MockTime time = new MockTime();
      protected final MockSelector selector = new MockSelector(time);
      protected final Node node = TestUtils.singletonCluster().nodes().iterator().next();
      protected final long reconnectBackoffMsTest = 10 * 1000;
      protected final long reconnectBackoffMaxMsTest = 10 * 10000;
      protected final long connectionSetupTimeoutMsTest = 5 * 1000;
      protected final long connectionSetupTimeoutMaxMsTest = 127 * 1000;
      private final int reconnectBackoffExpBase = ClusterConnectionStates.RECONNECT_BACKOFF_EXP_BASE;
      private final double reconnectBackoffJitter = ClusterConnectionStates.RECONNECT_BACKOFF_JITTER;
      private final TestMetadataUpdater metadataUpdater = new TestMetadataUpdater(Collections.singletonList(node));
      private final NetworkClient client = createNetworkClient(reconnectBackoffMaxMsTest);
      private final NetworkClient clientWithNoExponentialBackoff = createNetworkClient(reconnectBackoffMsTest);
      private final NetworkClient clientWithStaticNodes = createNetworkClientWithStaticNodes();
      private final NetworkClient clientWithNoVersionDiscovery = createNetworkClientWithNoVersionDiscovery();
  
      private static ArrayList<InetAddress> initialAddresses;
      private static ArrayList<InetAddress> newAddresses;
  
      static {
          try {
              initialAddresses = new ArrayList<>(Arrays.asList(
                      InetAddress.getByName("10.200.20.100"),
                      InetAddress.getByName("10.200.20.101"),
                      InetAddress.getByName("10.200.20.102")
              ));
              newAddresses = new ArrayList<>(Arrays.asList(
                      InetAddress.getByName("10.200.20.103"),
                      InetAddress.getByName("10.200.20.104"),
                      InetAddress.getByName("10.200.20.105")
              ));
          } catch (UnknownHostException e) {
              fail("Attempted to create an invalid InetAddress, this should not happen");
          }
      }
  
      private NetworkClient createNetworkClient(long reconnectBackoffMaxMs) {
          return new NetworkClient(selector, metadataUpdater, "mock", Integer.MAX_VALUE,
                  reconnectBackoffMsTest, reconnectBackoffMaxMs, 64 * 1024, 64 * 1024,
                  defaultRequestTimeoutMs, connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest, time, true, new ApiVersions(), new LogContext(),
                  MetadataRecoveryStrategy.NONE);
      }
  
      private NetworkClient createNetworkClientWithMaxInFlightRequestsPerConnection(
              int maxInFlightRequestsPerConnection, long reconnectBackoffMaxMs) {
          return new NetworkClient(selector, metadataUpdater, "mock", maxInFlightRequestsPerConnection,
                  reconnectBackoffMsTest, reconnectBackoffMaxMs, 64 * 1024, 64 * 1024,
                  defaultRequestTimeoutMs, connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest, time, true, new ApiVersions(), new LogContext(),
                  MetadataRecoveryStrategy.NONE);
      }
  
      private NetworkClient createNetworkClientWithMultipleNodes(long reconnectBackoffMaxMs, long connectionSetupTimeoutMsTest, int nodeNumber) {
          List<Node> nodes = TestUtils.clusterWith(nodeNumber).nodes();
          TestMetadataUpdater metadataUpdater = new TestMetadataUpdater(nodes);
          return new NetworkClient(selector, metadataUpdater, "mock", Integer.MAX_VALUE,
                  reconnectBackoffMsTest, reconnectBackoffMaxMs, 64 * 1024, 64 * 1024,
                  defaultRequestTimeoutMs, connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest, time, true, new ApiVersions(), new LogContext(),
                  MetadataRecoveryStrategy.NONE);
      }
  
      private NetworkClient createNetworkClientWithStaticNodes() {
          return new NetworkClient(selector, metadataUpdater,
                  "mock-static", Integer.MAX_VALUE, 0, 0, 64 * 1024, 64 * 1024, defaultRequestTimeoutMs,
                  connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest, time, true, new ApiVersions(), new LogContext(),
                  MetadataRecoveryStrategy.NONE);
      }
  
      private NetworkClient createNetworkClientWithNoVersionDiscovery(Metadata metadata) {
          return new NetworkClient(selector, metadata, "mock", Integer.MAX_VALUE,
                  reconnectBackoffMsTest, 0, 64 * 1024, 64 * 1024,
                  defaultRequestTimeoutMs, connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest, time, false, new ApiVersions(), new LogContext(),
                  MetadataRecoveryStrategy.NONE);
      }
  
      private NetworkClient createNetworkClientWithNoVersionDiscovery() {
          return new NetworkClient(selector, metadataUpdater, "mock", Integer.MAX_VALUE,
                  reconnectBackoffMsTest, reconnectBackoffMaxMsTest,
                  64 * 1024, 64 * 1024, defaultRequestTimeoutMs,
                  connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest, time, false, new ApiVersions(), new LogContext(),
                  MetadataRecoveryStrategy.NONE);
      }
  
      @BeforeEach
      public void setup() {
          selector.reset();
      }
  
      @Test
      public void testSendToUnreadyNode() {
          MetadataRequest.Builder builder = new MetadataRequest.Builder(Collections.singletonList("test"), true);
          long now = time.milliseconds();
          ClientRequest request = client.newClientRequest("5", builder, now, false);
          assertThrows(IllegalStateException.class, () -> client.send(request, now));
      }
  
      @Test
      public void testSimpleRequestResponse() {
          checkSimpleRequestResponse(client);
      }
  
      @Test
      public void testSimpleRequestResponseWithStaticNodes() {
          checkSimpleRequestResponse(clientWithStaticNodes);
      }
  
      @Test
      public void testSimpleRequestResponseWithNoBrokerDiscovery() {
          checkSimpleRequestResponse(clientWithNoVersionDiscovery);
      }
  
      @Test
      public void testDnsLookupFailure() {
          /* Fail cleanly when the node has a bad hostname */
          assertFalse(client.ready(new Node(1234, "badhost", 1234), time.milliseconds()));
      }
  
      @Test
      public void testClose() {
          client.ready(node, time.milliseconds());
          awaitReady(client, node);
          client.poll(1, time.milliseconds());
          assertTrue(client.isReady(node, time.milliseconds()), "The client should be ready");
  
          ProduceRequest.Builder builder = ProduceRequest.builder(new ProduceRequestData()
                  .setTopicData(new ProduceRequestData.TopicProduceDataCollection())
                  .setAcks((short) 1)
                  .setTimeoutMs(1000));
          ClientRequest request = client.newClientRequest(node.idString(), builder, time.milliseconds(), true);
          client.send(request, time.milliseconds());
          assertEquals(1, client.inFlightRequestCount(node.idString()),
              "There should be 1 in-flight request after send");
          assertTrue(client.hasInFlightRequests(node.idString()));
          assertTrue(client.hasInFlightRequests());
  
          client.close(node.idString());
          assertEquals(0, client.inFlightRequestCount(node.idString()), "There should be no in-flight request after close");
          assertFalse(client.hasInFlightRequests(node.idString()));
          assertFalse(client.hasInFlightRequests());
          assertFalse(client.isReady(node, 0), "Connection should not be ready after close");
      }
  
      @Test
      public void testUnsupportedVersionDuringInternalMetadataRequest() {
          List<String> topics = Collections.singletonList("topic_1");
  
          // disabling auto topic creation for versions less than 4 is not supported
          MetadataRequest.Builder builder = new MetadataRequest.Builder(topics, false, (short) 3);
          client.sendInternalMetadataRequest(builder, node.idString(), time.milliseconds());
          assertEquals(UnsupportedVersionException.class, metadataUpdater.getAndClearFailure().getClass());
      }
  
      @Test
      public void testRebootstrap() {
          long rebootstrapTriggerMs = 1000;
          AtomicInteger rebootstrapCount = new AtomicInteger();
          Metadata metadata = new Metadata(50, 50, 5000, new LogContext(), new ClusterResourceListeners()) {
              @Override
              public synchronized void rebootstrap() {
                  super.rebootstrap();
                  rebootstrapCount.incrementAndGet();
              }
          };
  
          NetworkClient client = new NetworkClient(selector, metadata, "mock", Integer.MAX_VALUE,
                  reconnectBackoffMsTest, 0, 64 * 1024, 64 * 1024,
                  defaultRequestTimeoutMs, connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest, time, false, new ApiVersions(), new LogContext(),
                  rebootstrapTriggerMs,
                  MetadataRecoveryStrategy.REBOOTSTRAP);
          MetadataUpdater metadataUpdater = TestUtils.fieldValue(client, NetworkClient.class, "metadataUpdater");
          metadata.bootstrap(Collections.singletonList(new InetSocketAddress("localhost", 9999)));
  
          metadata.requestUpdate(true);
          client.poll(0, time.milliseconds());
          time.sleep(rebootstrapTriggerMs + 1);
          client.poll(0, time.milliseconds());
          assertEquals(1, rebootstrapCount.get());
          time.sleep(1);
          client.poll(0, time.milliseconds());
          assertEquals(1, rebootstrapCount.get());
  
          metadata.requestUpdate(true);
          client.poll(0, time.milliseconds());
          assertEquals(1, rebootstrapCount.get());
          metadataUpdater.handleFailedRequest(time.milliseconds(), Optional.of(new KafkaException()));
          client.poll(0, time.milliseconds());
          assertEquals(1, rebootstrapCount.get());
          time.sleep(rebootstrapTriggerMs);
          client.poll(0, time.milliseconds());
          assertEquals(2, rebootstrapCount.get());
  
          metadata.requestUpdate(true);
          client.poll(0, time.milliseconds());
          assertEquals(2, rebootstrapCount.get());
  
          MetadataRequest.Builder builder = new MetadataRequest.Builder(Collections.emptyList(), true);
          ClientRequest request = client.newClientRequest(node.idString(), builder, time.milliseconds(), true);
          MetadataResponse rebootstrapResponse = (MetadataResponse) builder.build().getErrorResponse(0, new RebootstrapRequiredException("rebootstrap"));
          metadataUpdater.handleSuccessfulResponse(request.makeHeader(builder.latestAllowedVersion()), time.milliseconds(), rebootstrapResponse);
          assertEquals(2, rebootstrapCount.get());
          time.sleep(50);
          client.poll(0, time.milliseconds());
          assertEquals(3, rebootstrapCount.get());
      }
  
      @Test
      public void testInflightRequestsDuringRebootstrap() {
          long refreshBackoffMs = 50;
          long rebootstrapTriggerMs = 1000;
          int defaultRequestTimeoutMs = 5000;
          AtomicInteger rebootstrapCount = new AtomicInteger();
          Metadata metadata = new Metadata(refreshBackoffMs, refreshBackoffMs, 5000, new LogContext(), new ClusterResourceListeners())  {
              @Override
              public synchronized void rebootstrap() {
                  super.rebootstrap();
                  rebootstrapCount.incrementAndGet();
              }
          };
          metadata.bootstrap(Collections.singletonList(new InetSocketAddress("localhost", 9999)));
          NetworkClient client = new NetworkClient(selector, metadata, "mock", Integer.MAX_VALUE,
                  reconnectBackoffMsTest, 0, 64 * 1024, 64 * 1024,
                  defaultRequestTimeoutMs, connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest, time, false, new ApiVersions(), new LogContext(),
                  rebootstrapTriggerMs, MetadataRecoveryStrategy.REBOOTSTRAP);
  
          MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith(2, Collections.emptyMap());
          metadata.updateWithCurrentRequestVersion(metadataResponse, false, time.milliseconds());
          List<Node> nodes = metadata.fetch().nodes();
          nodes.forEach(node -> {
              client.ready(node, time.milliseconds());
              awaitReady(client, node);
          });
  
          // Queue a request
          sendEmptyProduceRequest(client, nodes.get(0).idString());
          List<ClientResponse> responses = client.poll(0, time.milliseconds());
          assertEquals(0, responses.size());
          assertEquals(1, client.inFlightRequestCount());
  
          // Trigger rebootstrap
          metadata.requestUpdate(true);
          time.sleep(refreshBackoffMs);
          responses = client.poll(0, time.milliseconds());
          assertEquals(0, responses.size());
          assertEquals(2, client.inFlightRequestCount());
          time.sleep(rebootstrapTriggerMs + 1);
          responses = client.poll(0, time.milliseconds());
  
          // Verify that inflight produce request was aborted with disconnection
          assertEquals(1, responses.size());
          assertEquals(PRODUCE, responses.get(0).requestHeader().apiKey());
          assertTrue(responses.get(0).wasDisconnected());
          assertEquals(0, client.inFlightRequestCount());
          assertEquals(Collections.emptySet(), nodes.stream().filter(node -> !client.connectionFailed(node)).collect(Collectors.toSet()));
      }
  
      private void checkSimpleRequestResponse(NetworkClient networkClient) {
          awaitReady(networkClient, node); // has to be before creating any request, as it may send ApiVersionsRequest and its response is mocked with correlation id 0
          short requestVersion = PRODUCE.latestVersion();
          ProduceRequest.Builder builder = new ProduceRequest.Builder(
                  requestVersion,
                  requestVersion,
                  new ProduceRequestData()
                      .setAcks((short) 1)
                      .setTimeoutMs(1000));
          TestCallbackHandler handler = new TestCallbackHandler();
          ClientRequest request = networkClient.newClientRequest(node.idString(), builder, time.milliseconds(),
              true, defaultRequestTimeoutMs, handler);
          networkClient.send(request, time.milliseconds());
          networkClient.poll(1, time.milliseconds());
          assertEquals(1, networkClient.inFlightRequestCount());
          ProduceResponse produceResponse = new ProduceResponse(new ProduceResponseData());
          ByteBuffer buffer = RequestTestUtils.serializeResponseWithHeader(produceResponse, requestVersion, request.correlationId());
          selector.completeReceive(new NetworkReceive(node.idString(), buffer));
          List<ClientResponse> responses = networkClient.poll(1, time.milliseconds());
          assertEquals(1, responses.size());
          assertTrue(handler.executed, "The handler should have executed.");
          assertTrue(handler.response.hasResponse(), "Should have a response body.");
          assertEquals(request.correlationId(), handler.response.requestHeader().correlationId(),
              "Should be correlated to the original request");
      }
  
      private void delayedApiVersionsResponse(int correlationId, short version, ApiVersionsResponse response) {
          ByteBuffer buffer = RequestTestUtils.serializeResponseWithHeader(response, version, correlationId);
          selector.delayedReceive(new DelayedReceive(node.idString(), new NetworkReceive(node.idString(), buffer)));
      }
  
      private void setExpectedApiVersionsResponse(ApiVersionsResponse response) {
          short apiVersionsResponseVersion = response.apiVersion(ApiKeys.API_VERSIONS.id).maxVersion();
          delayedApiVersionsResponse(0, apiVersionsResponseVersion, response);
      }
  
      private void awaitReady(NetworkClient client, Node node) {
          if (client.discoverBrokerVersions()) {
              setExpectedApiVersionsResponse(TestUtils.defaultApiVersionsResponse(
                  ApiMessageType.ListenerType.BROKER));
          }
          while (!client.ready(node, time.milliseconds()))
              client.poll(1, time.milliseconds());
          selector.clear();
      }
  
      @Test
      public void testInvalidApiVersionsRequest() {
          // initiate the connection
          client.ready(node, time.milliseconds());
  
          // handle the connection, send the ApiVersionsRequest
          client.poll(0, time.milliseconds());
  
          // check that the ApiVersionsRequest has been initiated
          assertTrue(client.hasInFlightRequests(node.idString()));
  
          // prepare response
          delayedApiVersionsResponse(0, ApiKeys.API_VERSIONS.latestVersion(),
              new ApiVersionsResponse(
                  new ApiVersionsResponseData()
                      .setErrorCode(Errors.INVALID_REQUEST.code())
                      .setThrottleTimeMs(0)
              ));
  
          // handle completed receives
          client.poll(0, time.milliseconds());
  
          // the ApiVersionsRequest is gone
          assertFalse(client.hasInFlightRequests(node.idString()));
  
          // various assertions
          assertFalse(client.isReady(node, time.milliseconds()));
      }
  
      @Test
      public void testApiVersionsRequest() {
          // initiate the connection
          client.ready(node, time.milliseconds());
  
          // handle the connection, send the ApiVersionsRequest
          client.poll(0, time.milliseconds());
  
          // check that the ApiVersionsRequest has been initiated
          assertTrue(client.hasInFlightRequests(node.idString()));
  
          // prepare response
          delayedApiVersionsResponse(0, ApiKeys.API_VERSIONS.latestVersion(), defaultApiVersionsResponse());
  
          // handle completed receives
          client.poll(0, time.milliseconds());
  
          // the ApiVersionsRequest is gone
          assertFalse(client.hasInFlightRequests(node.idString()));
  
          // various assertions
          assertTrue(client.isReady(node, time.milliseconds()));
      }
  
      @Test
      public void testUnsupportedApiVersionsRequestWithVersionProvidedByTheBroker() {
          // initiate the connection
          client.ready(node, time.milliseconds());
  
          // handle the connection, initiate first ApiVersionsRequest
          client.poll(0, time.milliseconds());
  
          // ApiVersionsRequest is in flight but not sent yet
          assertTrue(client.hasInFlightRequests(node.idString()));
  
          // completes initiated sends
          client.poll(0, time.milliseconds());
          assertEquals(1, selector.completedSends().size());
  
          ByteBuffer buffer = selector.completedSendBuffers().get(0).buffer();
          RequestHeader header = parseHeader(buffer);
          assertEquals(ApiKeys.API_VERSIONS, header.apiKey());
          assertEquals(4, header.apiVersion());
  
          // prepare response
          ApiVersionCollection apiKeys = new ApiVersionCollection();
          apiKeys.add(new ApiVersion()
              .setApiKey(ApiKeys.API_VERSIONS.id)
              .setMinVersion((short) 0)
              .setMaxVersion((short) 2));
          delayedApiVersionsResponse(0, (short) 0,
              new ApiVersionsResponse(
                  new ApiVersionsResponseData()
                      .setErrorCode(Errors.UNSUPPORTED_VERSION.code())
                      .setApiKeys(apiKeys)
              ));
  
          // handle ApiVersionResponse, initiate second ApiVersionRequest
          client.poll(0, time.milliseconds());
  
          // ApiVersionsRequest is in flight but not sent yet
          assertTrue(client.hasInFlightRequests(node.idString()));
  
          // ApiVersionsResponse has been received
          assertEquals(1, selector.completedReceives().size());
  
          // clean up the buffers
          selector.completedSends().clear();
          selector.completedSendBuffers().clear();
          selector.completedReceives().clear();
  
          // completes initiated sends
          client.poll(0, time.milliseconds());
  
          // ApiVersionsRequest has been sent
          assertEquals(1, selector.completedSends().size());
  
          buffer = selector.completedSendBuffers().get(0).buffer();
          header = parseHeader(buffer);
          assertEquals(ApiKeys.API_VERSIONS, header.apiKey());
          assertEquals(2, header.apiVersion());
  
          // prepare response
          delayedApiVersionsResponse(1, (short) 0, defaultApiVersionsResponse());
  
          // handle completed receives
          client.poll(0, time.milliseconds());
  
          // the ApiVersionsRequest is gone
          assertFalse(client.hasInFlightRequests(node.idString()));
          assertEquals(1, selector.completedReceives().size());
  
          // the client is ready
          assertTrue(client.isReady(node, time.milliseconds()));
      }
  
      @Test
      public void testUnsupportedApiVersionsRequestWithoutVersionProvidedByTheBroker() {
          // initiate the connection
          client.ready(node, time.milliseconds());
  
          // handle the connection, initiate first ApiVersionsRequest
          client.poll(0, time.milliseconds());
  
          // ApiVersionsRequest is in flight but not sent yet
          assertTrue(client.hasInFlightRequests(node.idString()));
  
          // completes initiated sends
          client.poll(0, time.milliseconds());
          assertEquals(1, selector.completedSends().size());
  
          ByteBuffer buffer = selector.completedSendBuffers().get(0).buffer();
          RequestHeader header = parseHeader(buffer);
          assertEquals(ApiKeys.API_VERSIONS, header.apiKey());
          assertEquals(4, header.apiVersion());
  
          // prepare response
          delayedApiVersionsResponse(0, (short) 0,
              new ApiVersionsResponse(
                  new ApiVersionsResponseData()
                      .setErrorCode(Errors.UNSUPPORTED_VERSION.code())
              ));
  
          // handle ApiVersionResponse, initiate second ApiVersionRequest
          client.poll(0, time.milliseconds());
  
          // ApiVersionsRequest is in flight but not sent yet
          assertTrue(client.hasInFlightRequests(node.idString()));
  
          // ApiVersionsResponse has been received
          assertEquals(1, selector.completedReceives().size());
  
          // clean up the buffers
          selector.completedSends().clear();
          selector.completedSendBuffers().clear();
          selector.completedReceives().clear();
  
          // completes initiated sends
          client.poll(0, time.milliseconds());
  
          // ApiVersionsRequest has been sent
          assertEquals(1, selector.completedSends().size());
  
          buffer = selector.completedSendBuffers().get(0).buffer();
          header = parseHeader(buffer);
          assertEquals(ApiKeys.API_VERSIONS, header.apiKey());
          assertEquals(0, header.apiVersion());
  
          // prepare response
          delayedApiVersionsResponse(1, (short) 0, defaultApiVersionsResponse());
  
          // handle completed receives
          client.poll(0, time.milliseconds());
  
          // the ApiVersionsRequest is gone
          assertFalse(client.hasInFlightRequests(node.idString()));
          assertEquals(1, selector.completedReceives().size());
  
          // the client is ready
          assertTrue(client.isReady(node, time.milliseconds()));
      }
  
      @Test
      public void testRequestTimeout() {
          testRequestTimeout(defaultRequestTimeoutMs + 5000);
      }
  
      @Test
      public void testDefaultRequestTimeout() {
          testRequestTimeout(defaultRequestTimeoutMs);
      }
  
      /**
       * This is a helper method that will execute two produce calls. The first call is expected to work and the
       * second produce call is intentionally made to emulate a request timeout. In the case that a timeout occurs
       * during a request, we want to ensure that we {@link Metadata#requestUpdate(boolean) request a metadata update} so that
       * on a subsequent invocation of {@link NetworkClient#poll(long, long) poll}, the metadata request will be sent.
       *
       * <p/>
       *
       * The {@link MetadataUpdater} has a specific method to handle
       * {@link NetworkClient.DefaultMetadataUpdater#handleServerDisconnect(long, String, Optional) server disconnects}
       * which is where we {@link Metadata#requestUpdate(boolean) request a metadata update}. This test helper method ensures
       * that is invoked by checking {@link Metadata#updateRequested()} after the simulated timeout.
       *
       * @param requestTimeoutMs Timeout in ms
       */
      private void testRequestTimeout(int requestTimeoutMs) {
          Metadata metadata = new Metadata(50, 50, 5000, new LogContext(), new ClusterResourceListeners());
          MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith(2, Collections.emptyMap());
          metadata.updateWithCurrentRequestVersion(metadataResponse, false, time.milliseconds());
  
          NetworkClient client = createNetworkClientWithNoVersionDiscovery(metadata);
  
          // Send first produce without any timeout.
          ClientResponse clientResponse = produce(client, requestTimeoutMs, false);
          assertEquals(node.idString(), clientResponse.destination());
          assertFalse(clientResponse.wasDisconnected(), "Expected response to succeed and not disconnect");
          assertFalse(clientResponse.wasTimedOut(), "Expected response to succeed and not time out");
          assertFalse(metadata.updateRequested(), "Expected NetworkClient to not need to update metadata");
  
          // Send second request, but emulate a timeout.
          clientResponse = produce(client, requestTimeoutMs, true);
          assertEquals(node.idString(), clientResponse.destination());
          assertTrue(clientResponse.wasDisconnected(), "Expected response to fail due to disconnection");
          assertTrue(clientResponse.wasTimedOut(), "Expected response to fail due to timeout");
          assertTrue(metadata.updateRequested(), "Expected NetworkClient to have called requestUpdate on metadata on timeout");
      }
  
      private ClientResponse produce(NetworkClient client, int requestTimeoutMs, boolean shouldEmulateTimeout) {
          awaitReady(client, node); // has to be before creating any request, as it may send ApiVersionsRequest and its response is mocked with correlation id 0
          ProduceRequest.Builder builder = ProduceRequest.builder(new ProduceRequestData()
                  .setTopicData(new ProduceRequestData.TopicProduceDataCollection())
                  .setAcks((short) 1)
                  .setTimeoutMs(1000));
          TestCallbackHandler handler = new TestCallbackHandler();
          ClientRequest request = client.newClientRequest(node.idString(), builder, time.milliseconds(), true,
                  requestTimeoutMs, handler);
          client.send(request, time.milliseconds());
  
          if (shouldEmulateTimeout) {
              // For a delay of slightly more than our timeout threshold to emulate the request timing out.
              time.sleep(requestTimeoutMs + 1);
          } else {
              ProduceResponse produceResponse = new ProduceResponse(new ProduceResponseData());
              ByteBuffer buffer = RequestTestUtils.serializeResponseWithHeader(produceResponse, PRODUCE.latestVersion(), request.correlationId());
              selector.completeReceive(new NetworkReceive(node.idString(), buffer));
          }
  
          List<ClientResponse> responses = client.poll(0, time.milliseconds());
          assertEquals(1, responses.size());
          return responses.get(0);
      }
  
      @Test
      public void testConnectionSetupTimeout() {
          // Use two nodes to ensure that the logic iterate over a set of more than one
          // element. ConcurrentModificationException is not triggered otherwise.
          final Cluster cluster = TestUtils.clusterWith(2);
          final Node node0 = cluster.nodeById(0);
          final Node node1 = cluster.nodeById(1);
  
          client.ready(node0, time.milliseconds());
          selector.serverConnectionBlocked(node0.idString());
  
          client.ready(node1, time.milliseconds());
          selector.serverConnectionBlocked(node1.idString());
  
          client.poll(0, time.milliseconds());
          assertFalse(client.connectionFailed(node),
              "The connections should not fail before the socket connection setup timeout elapsed");
  
          time.sleep((long) (connectionSetupTimeoutMsTest * 1.2) + 1);
          client.poll(0, time.milliseconds());
          assertTrue(client.connectionFailed(node),
              "Expected the connections to fail due to the socket connection setup timeout");
      }
  
      @Test
      public void testConnectionTimeoutAfterThrottling() {
          awaitReady(client, node);
          short requestVersion = PRODUCE.latestVersion();
          int timeoutMs = 1000;
          ProduceRequest.Builder builder = new ProduceRequest.Builder(
              requestVersion,
              requestVersion,
              new ProduceRequestData()
                  .setAcks((short) 1)
                  .setTimeoutMs(timeoutMs));
          TestCallbackHandler handler = new TestCallbackHandler();
          ClientRequest r1 = client.newClientRequest(node.idString(), builder, time.milliseconds(), true,
                  defaultRequestTimeoutMs, handler);
  
          client.send(r1, time.milliseconds());
          client.poll(0, time.milliseconds());
  
          // Throttle long enough to ensure other inFlight requests timeout.
          ProduceResponse pr = new ProduceResponse(new ProduceResponseData().setThrottleTimeMs(timeoutMs));
          ByteBuffer buffer = RequestTestUtils.serializeResponseWithHeader(pr, requestVersion, r1.correlationId());
          selector.delayedReceive(new DelayedReceive(node.idString(), new NetworkReceive(node.idString(), buffer)));
          ClientRequest r2 = client.newClientRequest(node.idString(), builder, time.milliseconds(), true,
                  defaultRequestTimeoutMs, handler);
          client.send(r2, time.milliseconds());
          time.sleep(timeoutMs);
          client.poll(0, time.milliseconds());
  
          assertEquals(1, client.inFlightRequestCount(node.idString()));
          assertFalse(client.connectionFailed(node), "Connection should not have failed due to the extra time spent throttling.");
      }
  
      @Test
      public void testConnectionThrottling() {
          // Instrument the test to return a response with a 100ms throttle delay.
          awaitReady(client, node);
          short requestVersion = PRODUCE.latestVersion();
          ProduceRequest.Builder builder = new ProduceRequest.Builder(
              requestVersion,
              requestVersion,
              new ProduceRequestData()
                  .setAcks((short) 1)
                  .setTimeoutMs(1000));
          TestCallbackHandler handler = new TestCallbackHandler();
          ClientRequest request = client.newClientRequest(node.idString(), builder, time.milliseconds(), true,
                  defaultRequestTimeoutMs, handler);
          client.send(request, time.milliseconds());
          client.poll(1, time.milliseconds());
          int throttleTime = 100;
          ProduceResponse produceResponse = new ProduceResponse(new ProduceResponseData().setThrottleTimeMs(throttleTime));
          ByteBuffer buffer = RequestTestUtils.serializeResponseWithHeader(produceResponse, requestVersion, request.correlationId());
          selector.completeReceive(new NetworkReceive(node.idString(), buffer));
          client.poll(1, time.milliseconds());
  
          // The connection is not ready due to throttling.
          assertFalse(client.ready(node, time.milliseconds()));
          assertEquals(100, client.throttleDelayMs(node, time.milliseconds()));
  
          // After 50ms, the connection is not ready yet.
          time.sleep(50);
          assertFalse(client.ready(node, time.milliseconds()));
          assertEquals(50, client.throttleDelayMs(node, time.milliseconds()));
  
          // After another 50ms, the throttling is done and the connection becomes ready again.
          time.sleep(50);
          assertTrue(client.ready(node, time.milliseconds()));
          assertEquals(0, client.throttleDelayMs(node, time.milliseconds()));
      }
  
      private int sendEmptyProduceRequest() {
          return sendEmptyProduceRequest(client, node.idString());
      }
  
      private int sendEmptyProduceRequest(NetworkClient client, String nodeId) {
          ProduceRequest.Builder builder = ProduceRequest.builder(new ProduceRequestData()
                  .setTopicData(new ProduceRequestData.TopicProduceDataCollection())
                  .setAcks((short) 1)
                  .setTimeoutMs(1000));
          TestCallbackHandler handler = new TestCallbackHandler();
          ClientRequest request = client.newClientRequest(nodeId, builder, time.milliseconds(), true,
                  defaultRequestTimeoutMs, handler);
          client.send(request, time.milliseconds());
          return request.correlationId();
      }
  
      private void sendResponse(AbstractResponse response, short version, int correlationId) {
          ByteBuffer buffer = RequestTestUtils.serializeResponseWithHeader(response, version, correlationId);
          selector.completeReceive(new NetworkReceive(node.idString(), buffer));
      }
  
      private void sendThrottledProduceResponse(int correlationId, int throttleMs, short version) {
          ProduceResponse response = new ProduceResponse(new ProduceResponseData().setThrottleTimeMs(throttleMs));
          sendResponse(response, version, correlationId);
      }
  
      @Test
      public void testLeastLoadedNode() {
          client.ready(node, time.milliseconds());
          assertFalse(client.isReady(node, time.milliseconds()));
          LeastLoadedNode leastLoadedNode = client.leastLoadedNode(time.milliseconds());
          assertEquals(node, leastLoadedNode.node());
          assertTrue(leastLoadedNode.hasNodeAvailableOrConnectionReady());
  
          awaitReady(client, node);
          client.poll(1, time.milliseconds());
          assertTrue(client.isReady(node, time.milliseconds()), "The client should be ready");
  
          // leastloadednode should be our single node
          leastLoadedNode = client.leastLoadedNode(time.milliseconds());
          assertTrue(leastLoadedNode.hasNodeAvailableOrConnectionReady());
          Node leastNode = leastLoadedNode.node();
          assertEquals(leastNode.id(), node.id(), "There should be one leastloadednode");
  
          // sleep for longer than reconnect backoff
          time.sleep(reconnectBackoffMsTest);
  
          // CLOSE node
          selector.serverDisconnect(node.idString());
  
          client.poll(1, time.milliseconds());
          assertFalse(client.ready(node, time.milliseconds()), "After we forced the disconnection the client is no longer ready.");
          leastLoadedNode = client.leastLoadedNode(time.milliseconds());
          assertFalse(leastLoadedNode.hasNodeAvailableOrConnectionReady());
          assertNull(leastLoadedNode.node(), "There should be NO leastloadednode");
      }
  
      @Test
      public void testHasNodeAvailableOrConnectionReady() {
          NetworkClient client = createNetworkClientWithMaxInFlightRequestsPerConnection(1, reconnectBackoffMaxMsTest);
          awaitReady(client, node);
  
          long now = time.milliseconds();
          LeastLoadedNode leastLoadedNode = client.leastLoadedNode(now);
          assertEquals(node, leastLoadedNode.node());
          assertTrue(leastLoadedNode.hasNodeAvailableOrConnectionReady());
  
          MetadataRequest.Builder builder = new MetadataRequest.Builder(Collections.emptyList(), true);
          ClientRequest request = client.newClientRequest(node.idString(), builder, now, true);
          client.send(request, now);
          client.poll(defaultRequestTimeoutMs, now);
  
          leastLoadedNode = client.leastLoadedNode(now);
          assertNull(leastLoadedNode.node());
          assertTrue(leastLoadedNode.hasNodeAvailableOrConnectionReady());
      }
  
      @Test
      public void testLeastLoadedNodeProvideDisconnectedNodesPrioritizedByLastConnectionTimestamp() {
          int nodeNumber = 3;
          NetworkClient client = createNetworkClientWithMultipleNodes(0, connectionSetupTimeoutMsTest, nodeNumber);
  
          Set<Node> providedNodeIds = new HashSet<>();
          for (int i = 0; i < nodeNumber * 10; i++) {
              Node node = client.leastLoadedNode(time.milliseconds()).node();
              assertNotNull(node, "Should provide a node");
              providedNodeIds.add(node);
              client.ready(node, time.milliseconds());
              client.disconnect(node.idString());
              time.sleep(connectionSetupTimeoutMsTest + 1);
              client.poll(0, time.milliseconds());
              // Define a round as nodeNumber of nodes have been provided
              // In each round every node should be provided exactly once
              if ((i + 1) % nodeNumber == 0) {
                  assertEquals(nodeNumber, providedNodeIds.size(), "All the nodes should be provided");
                  providedNodeIds.clear();
              }
          }
      }
  
      @Test
      public void testAuthenticationFailureWithInFlightMetadataRequest() {
          int refreshBackoffMs = 50;
  
          MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith(2, Collections.emptyMap());
          Metadata metadata = new Metadata(refreshBackoffMs, refreshBackoffMs, 5000, new LogContext(), new ClusterResourceListeners());
          metadata.updateWithCurrentRequestVersion(metadataResponse, false, time.milliseconds());
  
          Cluster cluster = metadata.fetch();
          Node node1 = cluster.nodes().get(0);
          Node node2 = cluster.nodes().get(1);
  
          NetworkClient client = createNetworkClientWithNoVersionDiscovery(metadata);
  
          awaitReady(client, node1);
  
          metadata.requestUpdate(true);
          time.sleep(refreshBackoffMs);
  
          client.poll(0, time.milliseconds());
  
          Optional<Node> nodeWithPendingMetadataOpt = cluster.nodes().stream()
                  .filter(node -> client.hasInFlightRequests(node.idString()))
                  .findFirst();
          assertEquals(Optional.of(node1), nodeWithPendingMetadataOpt);
  
          assertFalse(client.ready(node2, time.milliseconds()));
          selector.serverAuthenticationFailed(node2.idString());
          client.poll(0, time.milliseconds());
          assertNotNull(client.authenticationException(node2));
  
          ByteBuffer requestBuffer = selector.completedSendBuffers().get(0).buffer();
          RequestHeader header = parseHeader(requestBuffer);
          assertEquals(ApiKeys.METADATA, header.apiKey());
  
          ByteBuffer responseBuffer = RequestTestUtils.serializeResponseWithHeader(metadataResponse, header.apiVersion(), header.correlationId());
          selector.delayedReceive(new DelayedReceive(node1.idString(), new NetworkReceive(node1.idString(), responseBuffer)));
  
          int initialUpdateVersion = metadata.updateVersion();
          client.poll(0, time.milliseconds());
          assertEquals(initialUpdateVersion + 1, metadata.updateVersion());
      }
  
      @Test
      public void testLeastLoadedNodeConsidersThrottledConnections() {
          client.ready(node, time.milliseconds());
          awaitReady(client, node);
          client.poll(1, time.milliseconds());
          assertTrue(client.isReady(node, time.milliseconds()), "The client should be ready");
  
          int correlationId = sendEmptyProduceRequest();
          client.poll(1, time.milliseconds());
  
          sendThrottledProduceResponse(correlationId, 100, PRODUCE.latestVersion());
          client.poll(1, time.milliseconds());
  
          // leastloadednode should return null since the node is throttled
          assertNull(client.leastLoadedNode(time.milliseconds()).node());
      }
  
      @Test
      public void testConnectionDelayWithNoExponentialBackoff() {
          long now = time.milliseconds();
          long delay = clientWithNoExponentialBackoff.connectionDelay(node, now);
  
          assertEquals(0, delay);
      }
  
      @Test
      public void testConnectionDelayConnectedWithNoExponentialBackoff() {
          awaitReady(clientWithNoExponentialBackoff, node);
  
          long now = time.milliseconds();
          long delay = clientWithNoExponentialBackoff.connectionDelay(node, now);
  
          assertEquals(Long.MAX_VALUE, delay);
      }
  
      @Test
      public void testConnectionDelayDisconnectedWithNoExponentialBackoff() {
          awaitReady(clientWithNoExponentialBackoff, node);
  
          selector.serverDisconnect(node.idString());
          clientWithNoExponentialBackoff.poll(defaultRequestTimeoutMs, time.milliseconds());
          long delay = clientWithNoExponentialBackoff.connectionDelay(node, time.milliseconds());
  
          assertEquals(reconnectBackoffMsTest, delay);
  
          // Sleep until there is no connection delay
          time.sleep(delay);
          assertEquals(0, clientWithNoExponentialBackoff.connectionDelay(node, time.milliseconds()));
  
          // Start connecting and disconnect before the connection is established
          client.ready(node, time.milliseconds());
          selector.serverDisconnect(node.idString());
          client.poll(defaultRequestTimeoutMs, time.milliseconds());
  
          // Second attempt should have the same behaviour as exponential backoff is disabled
          assertEquals(reconnectBackoffMsTest, delay);
      }
  
      @Test
      public void testConnectionDelay() {
          long now = time.milliseconds();
          long delay = client.connectionDelay(node, now);
  
          assertEquals(0, delay);
      }
  
      @Test
      public void testConnectionDelayConnected() {
          awaitReady(client, node);
  
          long now = time.milliseconds();
          long delay = client.connectionDelay(node, now);
  
          assertEquals(Long.MAX_VALUE, delay);
      }
  
      @Test
      public void testConnectionDelayDisconnected() {
          awaitReady(client, node);
  
          // First disconnection
          selector.serverDisconnect(node.idString());
          client.poll(defaultRequestTimeoutMs, time.milliseconds());
          long delay = client.connectionDelay(node, time.milliseconds());
          long expectedDelay = reconnectBackoffMsTest;
          double jitter = 0.3;
          assertEquals(expectedDelay, delay, expectedDelay * jitter);
  
          // Sleep until there is no connection delay
          time.sleep(delay);
          assertEquals(0, client.connectionDelay(node, time.milliseconds()));
  
          // Start connecting and disconnect before the connection is established
          client.ready(node, time.milliseconds());
          selector.serverDisconnect(node.idString());
          client.poll(defaultRequestTimeoutMs, time.milliseconds());
  
          // Second attempt should take twice as long with twice the jitter
          expectedDelay = Math.round(delay * 2);
          delay = client.connectionDelay(node, time.milliseconds());
          jitter = 0.6;
          assertEquals(expectedDelay, delay, expectedDelay * jitter);
      }
  
      @Test
      public void testDisconnectDuringUserMetadataRequest() {
          // this test ensures that the default metadata updater does not intercept a user-initiated
          // metadata request when the remote node disconnects with the request in-flight.
          awaitReady(client, node);
  
          MetadataRequest.Builder builder = new MetadataRequest.Builder(Collections.emptyList(), true);
          long now = time.milliseconds();
          ClientRequest request = client.newClientRequest(node.idString(), builder, now, true);
          client.send(request, now);
          client.poll(defaultRequestTimeoutMs, now);
          assertEquals(1, client.inFlightRequestCount(node.idString()));
          assertTrue(client.hasInFlightRequests(node.idString()));
          assertTrue(client.hasInFlightRequests());
  
          selector.close(node.idString());
          List<ClientResponse> responses = client.poll(defaultRequestTimeoutMs, time.milliseconds());
          assertEquals(1, responses.size());
          assertTrue(responses.iterator().next().wasDisconnected());
      }
  
      @Test
      public void testServerDisconnectAfterInternalApiVersionRequest() throws Exception {
          final long numIterations = 5;
          double reconnectBackoffMaxExp = Math.log(reconnectBackoffMaxMsTest / (double) Math.max(reconnectBackoffMsTest, 1))
              / Math.log(reconnectBackoffExpBase);
          for (int i = 0; i < numIterations; i++) {
              selector.clear();
              awaitInFlightApiVersionRequest();
              selector.serverDisconnect(node.idString());
  
              // The failed ApiVersion request should not be forwarded to upper layers
              List<ClientResponse> responses = client.poll(0, time.milliseconds());
              assertFalse(client.hasInFlightRequests(node.idString()));
              assertTrue(responses.isEmpty());
  
              long expectedBackoff = Math.round(Math.pow(reconnectBackoffExpBase, Math.min(i, reconnectBackoffMaxExp))
                  * reconnectBackoffMsTest);
              long delay = client.connectionDelay(node, time.milliseconds());
              assertEquals(expectedBackoff, delay, reconnectBackoffJitter * expectedBackoff);
              if (i == numIterations - 1) {
                  break;
              }
              time.sleep(delay + 1);
          }
      }
  
      @Test
      public void testClientDisconnectAfterInternalApiVersionRequest() throws Exception {
          awaitInFlightApiVersionRequest();
          client.disconnect(node.idString());
          assertFalse(client.hasInFlightRequests(node.idString()));
  
          // The failed ApiVersion request should not be forwarded to upper layers
          List<ClientResponse> responses = client.poll(0, time.milliseconds());
          assertTrue(responses.isEmpty());
      }
  
      @Test
      public void testDisconnectWithMultipleInFlights() {
          NetworkClient client = this.clientWithNoVersionDiscovery;
          awaitReady(client, node);
          assertTrue(client.isReady(node, time.milliseconds()),
              "Expected NetworkClient to be ready to send to node " + node.idString());
  
          MetadataRequest.Builder builder = new MetadataRequest.Builder(Collections.emptyList(), true);
          long now = time.milliseconds();
  
          final List<ClientResponse> callbackResponses = new ArrayList<>();
          RequestCompletionHandler callback = callbackResponses::add;
  
          ClientRequest request1 = client.newClientRequest(node.idString(), builder, now, true, defaultRequestTimeoutMs, callback);
          client.send(request1, now);
          client.poll(0, now);
  
          ClientRequest request2 = client.newClientRequest(node.idString(), builder, now, true, defaultRequestTimeoutMs, callback);
          client.send(request2, now);
          client.poll(0, now);
  
          assertNotEquals(request1.correlationId(), request2.correlationId());
  
          assertEquals(2, client.inFlightRequestCount());
          assertEquals(2, client.inFlightRequestCount(node.idString()));
  
          client.disconnect(node.idString());
  
          List<ClientResponse> responses = client.poll(0, time.milliseconds());
          assertEquals(2, responses.size());
          assertEquals(responses, callbackResponses);
          assertEquals(0, client.inFlightRequestCount());
          assertEquals(0, client.inFlightRequestCount(node.idString()));
  
          // Ensure that the responses are returned in the order they were sent
          ClientResponse response1 = responses.get(0);
          assertTrue(response1.wasDisconnected());
          assertEquals(request1.correlationId(), response1.requestHeader().correlationId());
  
          ClientResponse response2 = responses.get(1);
          assertTrue(response2.wasDisconnected());
          assertEquals(request2.correlationId(), response2.requestHeader().correlationId());
      }
  
      @Test
      public void testCallDisconnect() {
          awaitReady(client, node);
          assertTrue(client.isReady(node, time.milliseconds()),
              "Expected NetworkClient to be ready to send to node " + node.idString());
          assertFalse(client.connectionFailed(node),
              "Did not expect connection to node " + node.idString() + " to be failed");
          client.disconnect(node.idString());
          assertFalse(client.isReady(node, time.milliseconds()),
              "Expected node " + node.idString() + " to be disconnected.");
          assertTrue(client.connectionFailed(node),
              "Expected connection to node " + node.idString() + " to be failed after disconnect");
          assertFalse(client.canConnect(node, time.milliseconds()));
  
          // ensure disconnect does not reset backoff period if already disconnected
          time.sleep(reconnectBackoffMaxMsTest);
          assertTrue(client.canConnect(node, time.milliseconds()));
          client.disconnect(node.idString());
          assertTrue(client.canConnect(node, time.milliseconds()));
      }
  
      @Test
      public void testCorrelationId() {
          int count = 100;
          Set<Integer> ids = IntStream.range(0, count)
              .mapToObj(i -> client.nextCorrelationId())
              .collect(Collectors.toSet());
          assertEquals(count, ids.size());
          ids.forEach(id -> assertTrue(id < SaslClientAuthenticator.MIN_RESERVED_CORRELATION_ID));
      }
  
      @Test
      public void testReconnectAfterAddressChange() {
          AddressChangeHostResolver mockHostResolver = new AddressChangeHostResolver(
                  initialAddresses.toArray(new InetAddress[0]), newAddresses.toArray(new InetAddress[0]));
          AtomicInteger initialAddressConns = new AtomicInteger();
          AtomicInteger newAddressConns = new AtomicInteger();
          MockSelector selector = new MockSelector(this.time, inetSocketAddress -> {
              InetAddress inetAddress = inetSocketAddress.getAddress();
              if (initialAddresses.contains(inetAddress)) {
                  initialAddressConns.incrementAndGet();
              } else if (newAddresses.contains(inetAddress)) {
                  newAddressConns.incrementAndGet();
              }
              return (mockHostResolver.useNewAddresses() && newAddresses.contains(inetAddress)) ||
                     (!mockHostResolver.useNewAddresses() && initialAddresses.contains(inetAddress));
          });
  
          ClientTelemetrySender mockClientTelemetrySender = mock(ClientTelemetrySender.class);
          when(mockClientTelemetrySender.timeToNextUpdate(anyLong())).thenReturn(0L);
  
          NetworkClient client = new NetworkClient(metadataUpdater, null, selector, "mock", Integer.MAX_VALUE,
                  reconnectBackoffMsTest, reconnectBackoffMaxMsTest, 64 * 1024, 64 * 1024,
                  defaultRequestTimeoutMs, connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest,
                  time, false, new ApiVersions(), null, new LogContext(), mockHostResolver, mockClientTelemetrySender,
                  Long.MAX_VALUE, MetadataRecoveryStrategy.NONE);
  
          // Connect to one the initial addresses, then change the addresses and disconnect
          client.ready(node, time.milliseconds());
          time.sleep(connectionSetupTimeoutMaxMsTest);
          client.poll(0, time.milliseconds());
          assertTrue(client.isReady(node, time.milliseconds()));
          // First poll should try to update the node but couldn't because node remains in connecting state
          // i.e. connection handling is completed after telemetry update.
          assertNull(client.telemetryConnectedNode());
  
          client.poll(0, time.milliseconds());
          assertEquals(node, client.telemetryConnectedNode());
  
          mockHostResolver.changeAddresses();
          selector.serverDisconnect(node.idString());
          client.poll(0, time.milliseconds());
          assertFalse(client.isReady(node, time.milliseconds()));
          assertNull(client.telemetryConnectedNode());
  
          time.sleep(reconnectBackoffMaxMsTest);
          client.ready(node, time.milliseconds());
          time.sleep(connectionSetupTimeoutMaxMsTest);
          client.poll(0, time.milliseconds());
          assertTrue(client.isReady(node, time.milliseconds()));
          assertNull(client.telemetryConnectedNode());
  
          client.poll(0, time.milliseconds());
          assertEquals(node, client.telemetryConnectedNode());
  
          // We should have tried to connect to one initial address and one new address, and resolved DNS twice
          assertEquals(1, initialAddressConns.get());
          assertEquals(1, newAddressConns.get());
          assertEquals(2, mockHostResolver.resolutionCount());
          verify(mockClientTelemetrySender, times(5)).timeToNextUpdate(anyLong());
      }
  
      @Test
      public void testFailedConnectionToFirstAddress() {
          AddressChangeHostResolver mockHostResolver = new AddressChangeHostResolver(
                  initialAddresses.toArray(new InetAddress[0]), newAddresses.toArray(new InetAddress[0]));
          AtomicInteger initialAddressConns = new AtomicInteger();
          AtomicInteger newAddressConns = new AtomicInteger();
          MockSelector selector = new MockSelector(this.time, inetSocketAddress -> {
              InetAddress inetAddress = inetSocketAddress.getAddress();
              if (initialAddresses.contains(inetAddress)) {
                  initialAddressConns.incrementAndGet();
              } else if (newAddresses.contains(inetAddress)) {
                  newAddressConns.incrementAndGet();
              }
              // Refuse first connection attempt
              return initialAddressConns.get() > 1;
          });
  
          ClientTelemetrySender mockClientTelemetrySender = mock(ClientTelemetrySender.class);
          when(mockClientTelemetrySender.timeToNextUpdate(anyLong())).thenReturn(0L);
  
          NetworkClient client = new NetworkClient(metadataUpdater, null, selector, "mock", Integer.MAX_VALUE,
                  reconnectBackoffMsTest, reconnectBackoffMaxMsTest, 64 * 1024, 64 * 1024,
                  defaultRequestTimeoutMs, connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest,
                  time, false, new ApiVersions(), null, new LogContext(), mockHostResolver, mockClientTelemetrySender,
                  Long.MAX_VALUE, MetadataRecoveryStrategy.NONE);
  
          // First connection attempt should fail
          client.ready(node, time.milliseconds());
          time.sleep(connectionSetupTimeoutMaxMsTest);
          client.poll(0, time.milliseconds());
          assertFalse(client.isReady(node, time.milliseconds()));
          assertNull(client.telemetryConnectedNode());
  
          // Second connection attempt should succeed
          time.sleep(reconnectBackoffMaxMsTest);
          client.ready(node, time.milliseconds());
          time.sleep(connectionSetupTimeoutMaxMsTest);
          client.poll(0, time.milliseconds());
          assertTrue(client.isReady(node, time.milliseconds()));
          assertNull(client.telemetryConnectedNode());
  
          // Next client poll after handling connection setup should update telemetry node.
          client.poll(0, time.milliseconds());
          assertEquals(node, client.telemetryConnectedNode());
  
          // We should have tried to connect to two of the initial addresses, none of the new address, and should
          // only have resolved DNS once
          assertEquals(2, initialAddressConns.get());
          assertEquals(0, newAddressConns.get());
          assertEquals(1, mockHostResolver.resolutionCount());
          verify(mockClientTelemetrySender, times(3)).timeToNextUpdate(anyLong());
      }
  
      @Test
      public void testFailedConnectionToFirstAddressAfterReconnect() {
          AddressChangeHostResolver mockHostResolver = new AddressChangeHostResolver(
                  initialAddresses.toArray(new InetAddress[0]), newAddresses.toArray(new InetAddress[0]));
          AtomicInteger initialAddressConns = new AtomicInteger();
          AtomicInteger newAddressConns = new AtomicInteger();
          MockSelector selector = new MockSelector(this.time, inetSocketAddress -> {
              InetAddress inetAddress = inetSocketAddress.getAddress();
              if (initialAddresses.contains(inetAddress)) {
                  initialAddressConns.incrementAndGet();
              } else if (newAddresses.contains(inetAddress)) {
                  newAddressConns.incrementAndGet();
              }
              // Refuse first connection attempt to the new addresses
              return initialAddresses.contains(inetAddress) || newAddressConns.get() > 1;
          });
  
          ClientTelemetrySender mockClientTelemetrySender = mock(ClientTelemetrySender.class);
          when(mockClientTelemetrySender.timeToNextUpdate(anyLong())).thenReturn(0L);
  
          NetworkClient client = new NetworkClient(metadataUpdater, null, selector, "mock", Integer.MAX_VALUE,
                  reconnectBackoffMsTest, reconnectBackoffMaxMsTest, 64 * 1024, 64 * 1024,
                  defaultRequestTimeoutMs, connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest,
                  time, false, new ApiVersions(), null, new LogContext(), mockHostResolver, mockClientTelemetrySender,
                  Long.MAX_VALUE, MetadataRecoveryStrategy.NONE);
  
          // Connect to one the initial addresses, then change the addresses and disconnect
          client.ready(node, time.milliseconds());
          time.sleep(connectionSetupTimeoutMaxMsTest);
          client.poll(0, time.milliseconds());
          assertTrue(client.isReady(node, time.milliseconds()));
          assertNull(client.telemetryConnectedNode());
          // Next client poll after handling connection setup should update telemetry node.
          client.poll(0, time.milliseconds());
          assertEquals(node, client.telemetryConnectedNode());
  
          mockHostResolver.changeAddresses();
          selector.serverDisconnect(node.idString());
          client.poll(0, time.milliseconds());
          assertFalse(client.isReady(node, time.milliseconds()));
          assertNull(client.telemetryConnectedNode());
  
          // First connection attempt to new addresses should fail
          time.sleep(reconnectBackoffMaxMsTest);
          client.ready(node, time.milliseconds());
          time.sleep(connectionSetupTimeoutMaxMsTest);
          client.poll(0, time.milliseconds());
          assertFalse(client.isReady(node, time.milliseconds()));
          assertNull(client.telemetryConnectedNode());
  
          // Second connection attempt to new addresses should succeed
          time.sleep(reconnectBackoffMaxMsTest);
          client.ready(node, time.milliseconds());
          time.sleep(connectionSetupTimeoutMaxMsTest);
          client.poll(0, time.milliseconds());
          assertTrue(client.isReady(node, time.milliseconds()));
          assertNull(client.telemetryConnectedNode());
  
          // Next client poll after handling connection setup should update telemetry node.
          client.poll(0, time.milliseconds());
          assertEquals(node, client.telemetryConnectedNode());
  
          // We should have tried to connect to one of the initial addresses and two of the new addresses (the first one
          // failed), and resolved DNS twice, once for each set of addresses
          assertEquals(1, initialAddressConns.get());
          assertEquals(2, newAddressConns.get());
          assertEquals(2, mockHostResolver.resolutionCount());
          verify(mockClientTelemetrySender, times(6)).timeToNextUpdate(anyLong());
      }
  
      @Test
      public void testCloseConnectingNode() {
          Cluster cluster = TestUtils.clusterWith(2);
          Node node0 = cluster.nodeById(0);
          Node node1 = cluster.nodeById(1);
          client.ready(node0, time.milliseconds());
          selector.serverConnectionBlocked(node0.idString());
          client.poll(1, time.milliseconds());
          client.close(node0.idString());
  
          // Poll without any connections should return without exceptions
          client.poll(0, time.milliseconds());
          assertFalse(NetworkClientUtils.isReady(client, node0, time.milliseconds()));
          assertFalse(NetworkClientUtils.isReady(client, node1, time.milliseconds()));
  
          // Connection to new node should work
          client.ready(node1, time.milliseconds());
          ByteBuffer buffer = RequestTestUtils.serializeResponseWithHeader(defaultApiVersionsResponse(), ApiKeys.API_VERSIONS.latestVersion(), 0);
          selector.delayedReceive(new DelayedReceive(node1.idString(), new NetworkReceive(node1.idString(), buffer)));
          while (!client.ready(node1, time.milliseconds()))
              client.poll(1, time.milliseconds());
          assertTrue(client.isReady(node1, time.milliseconds()));
          selector.clear();
  
          // New connection to node closed earlier should work
          client.ready(node0, time.milliseconds());
          buffer = RequestTestUtils.serializeResponseWithHeader(defaultApiVersionsResponse(), ApiKeys.API_VERSIONS.latestVersion(), 1);
          selector.delayedReceive(new DelayedReceive(node0.idString(), new NetworkReceive(node0.idString(), buffer)));
          while (!client.ready(node0, time.milliseconds()))
              client.poll(1, time.milliseconds());
          assertTrue(client.isReady(node0, time.milliseconds()));
      }
  
      @Test
      public void testConnectionDoesNotRemainStuckInCheckingApiVersionsStateIfChannelNeverBecomesReady() {
          final Cluster cluster = TestUtils.clusterWith(1);
          final Node node = cluster.nodeById(0);
  
          // Channel is ready by default so we mark it as not ready.
          client.ready(node, time.milliseconds());
          selector.channelNotReady(node.idString());
  
          // Channel should not be ready.
          client.poll(0, time.milliseconds());
          assertFalse(NetworkClientUtils.isReady(client, node, time.milliseconds()));
  
          // Connection should time out if the channel does not become ready within
          // the connection setup timeout. This ensures that the client does not remain
          // stuck in the CHECKING_API_VERSIONS state.
          time.sleep((long) (connectionSetupTimeoutMsTest * 1.2) + 1);
          client.poll(0, time.milliseconds());
          assertTrue(client.connectionFailed(node));
      }
  
      @Test
      public void testTelemetryRequest() {
          ClientTelemetrySender mockClientTelemetrySender = mock(ClientTelemetrySender.class);
          when(mockClientTelemetrySender.timeToNextUpdate(anyLong())).thenReturn(0L);
  
          NetworkClient client = new NetworkClient(metadataUpdater, null, selector, "mock", Integer.MAX_VALUE,
              reconnectBackoffMsTest, reconnectBackoffMaxMsTest, 64 * 1024, 64 * 1024,
              defaultRequestTimeoutMs, connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest,
              time, true, new ApiVersions(), null, new LogContext(), new DefaultHostResolver(), mockClientTelemetrySender,
              Long.MAX_VALUE, MetadataRecoveryStrategy.NONE);
  
          // Send the ApiVersionsRequest
          client.ready(node, time.milliseconds());
          client.poll(0, time.milliseconds());
          assertNull(client.telemetryConnectedNode());
          assertTrue(client.hasInFlightRequests(node.idString()));
          delayedApiVersionsResponse(0, ApiKeys.API_VERSIONS.latestVersion(), TestUtils.defaultApiVersionsResponse(
              ApiMessageType.ListenerType.BROKER));
          // handle ApiVersionsResponse
          client.poll(0, time.milliseconds());
          // the ApiVersionsRequest is gone
          assertFalse(client.hasInFlightRequests(node.idString()));
          selector.clear();
  
          GetTelemetrySubscriptionsRequest.Builder getRequest = new GetTelemetrySubscriptionsRequest.Builder(
              new GetTelemetrySubscriptionsRequestData(), true);
          when(mockClientTelemetrySender.createRequest()).thenReturn(Optional.of(getRequest));
  
          GetTelemetrySubscriptionsResponse getResponse = new GetTelemetrySubscriptionsResponse(new GetTelemetrySubscriptionsResponseData());
          ByteBuffer buffer = RequestTestUtils.serializeResponseWithHeader(getResponse, ApiKeys.GET_TELEMETRY_SUBSCRIPTIONS.latestVersion(), 1);
          selector.completeReceive(new NetworkReceive(node.idString(), buffer));
  
          // Initiate poll to send GetTelemetrySubscriptions request
          client.poll(0, time.milliseconds());
          assertTrue(client.isReady(node, time.milliseconds()));
          assertEquals(node, client.telemetryConnectedNode());
          verify(mockClientTelemetrySender, times(1)).handleResponse(any(GetTelemetrySubscriptionsResponse.class));
          selector.clear();
  
          PushTelemetryRequest.Builder pushRequest = new PushTelemetryRequest.Builder(
              new PushTelemetryRequestData(), true);
          when(mockClientTelemetrySender.createRequest()).thenReturn(Optional.of(pushRequest));
  
          PushTelemetryResponse pushResponse = new PushTelemetryResponse(new PushTelemetryResponseData());
          ByteBuffer pushBuffer = RequestTestUtils.serializeResponseWithHeader(pushResponse, ApiKeys.PUSH_TELEMETRY.latestVersion(), 2);
          selector.completeReceive(new NetworkReceive(node.idString(), pushBuffer));
  
          // Initiate poll to send PushTelemetry request
          client.poll(0, time.milliseconds());
          assertTrue(client.isReady(node, time.milliseconds()));
          assertEquals(node, client.telemetryConnectedNode());
          verify(mockClientTelemetrySender, times(1)).handleResponse(any(PushTelemetryResponse.class));
          verify(mockClientTelemetrySender, times(4)).timeToNextUpdate(anyLong());
          verify(mockClientTelemetrySender, times(2)).createRequest();
      }
  
      private RequestHeader parseHeader(ByteBuffer buffer) {
          buffer.getInt(); // skip size
          return RequestHeader.parse(buffer.slice());
      }
  
      private void awaitInFlightApiVersionRequest() throws Exception {
          client.ready(node, time.milliseconds());
          TestUtils.waitForCondition(() -> {
              client.poll(0, time.milliseconds());
              return client.hasInFlightRequests(node.idString());
          }, 1000, "");
          assertFalse(client.isReady(node, time.milliseconds()));
      }
  
      private ApiVersionsResponse defaultApiVersionsResponse() {
          return TestUtils.defaultApiVersionsResponse(ApiMessageType.ListenerType.BROKER);
      }
  
      private static class TestCallbackHandler implements RequestCompletionHandler {
          public boolean executed = false;
          public ClientResponse response;
  
          public void onComplete(ClientResponse response) {
              this.executed = true;
              this.response = response;
          }
      }
  
      // ManualMetadataUpdater with ability to keep track of failures
      private static class TestMetadataUpdater extends ManualMetadataUpdater {
          KafkaException failure;
  
          public TestMetadataUpdater(List<Node> nodes) {
              super(nodes);
          }
  
          @Override
          public void handleServerDisconnect(long now, String destinationId, Optional<AuthenticationException> maybeAuthException) {
              maybeAuthException.ifPresent(exception ->
                  failure = exception
              );
              super.handleServerDisconnect(now, destinationId, maybeAuthException);
          }
  
          @Override
          public void handleFailedRequest(long now, Optional<KafkaException> maybeFatalException) {
              maybeFatalException.ifPresent(exception ->
                  failure = exception
              );
          }
  
          public KafkaException getAndClearFailure() {
              KafkaException failure = this.failure;
              this.failure = null;
              return failure;
          }
      }
  }
  