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.server.util;
  
  import org.apache.kafka.clients.ClientRequest;
  import org.apache.kafka.clients.ClientResponse;
  import org.apache.kafka.clients.KafkaClient;
  import org.apache.kafka.clients.RequestCompletionHandler;
  import org.apache.kafka.common.Node;
  import org.apache.kafka.common.errors.AuthenticationException;
  import org.apache.kafka.common.errors.DisconnectException;
  import org.apache.kafka.common.internals.FatalExitError;
  import org.apache.kafka.common.utils.Time;
  import org.apache.kafka.common.utils.Utils;
  
  import java.util.ArrayDeque;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.Set;
  
  /**
   * An inter-broker send thread that utilizes a non-blocking network client.
   */
  public abstract class InterBrokerSendThread extends ShutdownableThread {
  
      protected final KafkaClient networkClient;
  
      private final int requestTimeoutMs;
      private final Time time;
      private final UnsentRequests unsentRequests;
  
      protected InterBrokerSendThread(
          String name,
          KafkaClient networkClient,
          int requestTimeoutMs,
          Time time
      ) {
          this(name, networkClient, requestTimeoutMs, time, true);
      }
  
      protected InterBrokerSendThread(
          String name,
          KafkaClient networkClient,
          int requestTimeoutMs,
          Time time,
          boolean isInterruptible
      ) {
          super(name, isInterruptible);
          this.networkClient = networkClient;
          this.requestTimeoutMs = requestTimeoutMs;
          this.time = time;
          this.unsentRequests = new UnsentRequests();
      }
  
      public abstract Collection<RequestAndCompletionHandler> generateRequests();
  
      public boolean hasUnsentRequests() {
          return unsentRequests.iterator().hasNext();
      }
  
      @Override
      public void shutdown() throws InterruptedException {
          initiateShutdown();
          networkClient.initiateClose();
          awaitShutdown();
          Utils.closeQuietly(networkClient, "InterBrokerSendThread network client");
      }
  
      private void drainGeneratedRequests() {
          generateRequests().forEach(request ->
              unsentRequests.put(
                  request.destination,
                  networkClient.newClientRequest(
                      request.destination.idString(),
                      request.request,
                      request.creationTimeMs,
                      true,
                      requestTimeoutMs,
                      request.handler
                  )
              )
          );
      }
  
      protected void pollOnce(long maxTimeoutMs) {
          try {
              drainGeneratedRequests();
              long now = time.milliseconds();
              final long timeout = sendRequests(now, maxTimeoutMs);
              networkClient.poll(timeout, now);
              now = time.milliseconds();
              checkDisconnects(now);
              failExpiredRequests(now);
              unsentRequests.clean();
          } catch (FatalExitError fee) {
              throw fee;
          } catch (Throwable t) {
              if (t instanceof DisconnectException && !networkClient.active()) {
                  // DisconnectException is expected when NetworkClient#initiateClose is called
                  return;
              }
              if (t instanceof InterruptedException && !isRunning()) {
                  // InterruptedException is expected when shutting down. Throw the error to ShutdownableThread to handle
                  throw t;
              }
              log.error("unhandled exception caught in InterBrokerSendThread", t);
              // rethrow any unhandled exceptions as FatalExitError so the JVM will be terminated
              // as we will be in an unknown state with potentially some requests dropped and not
              // being able to make progress. Known and expected Errors should have been appropriately
              // dealt with already.
              throw new FatalExitError();
          }
      }
  
      @Override
      public void doWork() {
          pollOnce(Long.MAX_VALUE);
      }
  
      private long sendRequests(long now, long maxTimeoutMs) {
          long pollTimeout = maxTimeoutMs;
          for (Node node : unsentRequests.nodes()) {
              final Iterator<ClientRequest> requestIterator = unsentRequests.requestIterator(node);
              while (requestIterator.hasNext()) {
                  final ClientRequest request = requestIterator.next();
                  if (networkClient.ready(node, now)) {
                      networkClient.send(request, now);
                      requestIterator.remove();
                  } else {
                      pollTimeout = Math.min(pollTimeout, networkClient.connectionDelay(node, now));
                  }
              }
          }
          return pollTimeout;
      }
  
      private void checkDisconnects(long now) {
          // any disconnects affecting requests that have already been transmitted will be handled
          // by NetworkClient, so we just need to check whether connections for any of the unsent
          // requests have been disconnected; if they have, then we complete the corresponding future
          // and set the disconnect flag in the ClientResponse
          final Iterator<Map.Entry<Node, ArrayDeque<ClientRequest>>> iterator = unsentRequests.iterator();
          while (iterator.hasNext()) {
              final Map.Entry<Node, ArrayDeque<ClientRequest>> entry = iterator.next();
              final Node node = entry.getKey();
              final ArrayDeque<ClientRequest> requests = entry.getValue();
              if (!requests.isEmpty() && networkClient.connectionFailed(node)) {
                  iterator.remove();
                  for (ClientRequest request : requests) {
                      final AuthenticationException authenticationException = networkClient.authenticationException(node);
                      if (authenticationException != null) {
                          log.error("Failed to send the following request due to authentication error: {}", request);
                      }
                      completeWithDisconnect(request, now, authenticationException);
                  }
              }
          }
      }
  
      private void failExpiredRequests(long now) {
          // clear all expired unsent requests
          final Collection<ClientRequest> timedOutRequests = unsentRequests.removeAllTimedOut(now);
          for (ClientRequest request : timedOutRequests) {
              log.debug("Failed to send the following request after {} ms: {}", request.requestTimeoutMs(), request);
              completeWithDisconnect(request, now, null);
          }
      }
  
      private static void completeWithDisconnect(
          ClientRequest request,
          long now,
          AuthenticationException authenticationException
      ) {
          final RequestCompletionHandler handler = request.callback();
          handler.onComplete(
              new ClientResponse(
                  request.makeHeader(request.requestBuilder().latestAllowedVersion()),
                  handler,
                  request.destination(),
                  now /* createdTimeMs */,
                  now /* receivedTimeMs */,
                  true /* disconnected */,
                  null /* versionMismatch */,
                  authenticationException,
                  null
              )
          );
      }
  
      public void wakeup() {
          networkClient.wakeup();
      }
  
      private static final class UnsentRequests {
  
          private final Map<Node, ArrayDeque<ClientRequest>> unsent = new HashMap<>();
  
          void put(Node node, ClientRequest request) {
              ArrayDeque<ClientRequest> requests = unsent.computeIfAbsent(node, n -> new ArrayDeque<>());
              requests.add(request);
          }
  
          Collection<ClientRequest> removeAllTimedOut(long now) {
              final List<ClientRequest> expiredRequests = new ArrayList<>();
              for (ArrayDeque<ClientRequest> requests : unsent.values()) {
                  final Iterator<ClientRequest> requestIterator = requests.iterator();
                  boolean foundExpiredRequest = false;
                  while (requestIterator.hasNext() && !foundExpiredRequest) {
                      final ClientRequest request = requestIterator.next();
                      final long elapsedMs = Math.max(0, now - request.createdTimeMs());
                      if (elapsedMs > request.requestTimeoutMs()) {
                          expiredRequests.add(request);
                          requestIterator.remove();
                          foundExpiredRequest = true;
                      }
                  }
              }
              return expiredRequests;
          }
  
          void clean() {
              unsent.values().removeIf(ArrayDeque::isEmpty);
          }
  
          Iterator<Entry<Node, ArrayDeque<ClientRequest>>> iterator() {
              return unsent.entrySet().iterator();
          }
  
          Iterator<ClientRequest> requestIterator(Node node) {
              ArrayDeque<ClientRequest> requests = unsent.get(node);
              return (requests == null) ? Collections.emptyIterator() : requests.iterator();
          }
  
          Set<Node> nodes() {
              return unsent.keySet();
          }
      }
  }
  