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.consumer.internals;
  
  import org.apache.kafka.clients.ClientRequest;
  import org.apache.kafka.clients.ClientResponse;
  import org.apache.kafka.clients.KafkaClient;
  import org.apache.kafka.clients.Metadata;
  import org.apache.kafka.clients.NetworkClientUtils;
  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.errors.InterruptException;
  import org.apache.kafka.common.errors.TimeoutException;
  import org.apache.kafka.common.errors.WakeupException;
  import org.apache.kafka.common.requests.AbstractRequest;
  import org.apache.kafka.common.utils.LogContext;
  import org.apache.kafka.common.utils.Time;
  import org.apache.kafka.common.utils.Timer;
  
  import org.slf4j.Logger;
  
  import java.io.Closeable;
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.Iterator;
  import java.util.List;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.ConcurrentLinkedQueue;
  import java.util.concurrent.ConcurrentMap;
  import java.util.concurrent.atomic.AtomicBoolean;
  import java.util.concurrent.locks.ReentrantLock;
  
  /**
   * Higher level consumer access to the network layer with basic support for request futures. This class
   * is thread-safe, but provides no synchronization for response callbacks. This guarantees that no locks
   * are held when they are invoked.
   */
  public class ConsumerNetworkClient implements Closeable {
      private static final int MAX_POLL_TIMEOUT_MS = 5000;
  
      // the mutable state of this class is protected by the object's monitor (excluding the wakeup
      // flag and the request completion queue below).
      private final Logger log;
      private final KafkaClient client;
      private final UnsentRequests unsent = new UnsentRequests();
      private final Metadata metadata;
      private final Time time;
      private final long retryBackoffMs;
      private final int maxPollTimeoutMs;
      private final int requestTimeoutMs;
      private final AtomicBoolean wakeupDisabled = new AtomicBoolean();
  
      // We do not need high throughput, so use a fair lock to try to avoid starvation
      private final ReentrantLock lock = new ReentrantLock(true);
  
      // when requests complete, they are transferred to this queue prior to invocation. The purpose
      // is to avoid invoking them while holding this object's monitor which can open the door for deadlocks.
      private final ConcurrentLinkedQueue<RequestFutureCompletionHandler> pendingCompletion = new ConcurrentLinkedQueue<>();
  
      private final ConcurrentLinkedQueue<Node> pendingDisconnects = new ConcurrentLinkedQueue<>();
  
      // this flag allows the client to be safely woken up without waiting on the lock above. It is
      // atomic to avoid the need to acquire the lock above in order to enable it concurrently.
      private final AtomicBoolean wakeup = new AtomicBoolean(false);
  
      public ConsumerNetworkClient(LogContext logContext,
                                   KafkaClient client,
                                   Metadata metadata,
                                   Time time,
                                   long retryBackoffMs,
                                   int requestTimeoutMs,
                                   int maxPollTimeoutMs) {
          this.log = logContext.logger(ConsumerNetworkClient.class);
          this.client = client;
          this.metadata = metadata;
          this.time = time;
          this.retryBackoffMs = retryBackoffMs;
          this.maxPollTimeoutMs = Math.min(maxPollTimeoutMs, MAX_POLL_TIMEOUT_MS);
          this.requestTimeoutMs = requestTimeoutMs;
      }
  
      public int defaultRequestTimeoutMs() {
          return requestTimeoutMs;
      }
  
      /**
       * Send a request with the default timeout. See {@link #send(Node, AbstractRequest.Builder, int)}.
       */
      public RequestFuture<ClientResponse> send(Node node, AbstractRequest.Builder<?> requestBuilder) {
          return send(node, requestBuilder, requestTimeoutMs);
      }
  
      /**
       * Send a new request. Note that the request is not actually transmitted on the
       * network until one of the {@link #poll(Timer)} variants is invoked. At this
       * point the request will either be transmitted successfully or will fail.
       * Use the returned future to obtain the result of the send. Note that there is no
       * need to check for disconnects explicitly on the {@link ClientResponse} object;
       * instead, the future will be failed with a {@link DisconnectException}.
       *
       * @param node The destination of the request
       * @param requestBuilder A builder for the request payload
       * @param requestTimeoutMs Maximum time in milliseconds to await a response before disconnecting the socket and
       *                         cancelling the request. The request may be cancelled sooner if the socket disconnects
       *                         for any reason.
       * @return A future which indicates the result of the send.
       */
      public RequestFuture<ClientResponse> send(Node node,
                                                AbstractRequest.Builder<?> requestBuilder,
                                                int requestTimeoutMs) {
          long now = time.milliseconds();
          RequestFutureCompletionHandler completionHandler = new RequestFutureCompletionHandler();
          ClientRequest clientRequest = client.newClientRequest(node.idString(), requestBuilder, now, true,
              requestTimeoutMs, completionHandler);
          unsent.put(node, clientRequest);
  
          // wakeup the client in case it is blocking in poll so that we can send the queued request
          client.wakeup();
          return completionHandler.future;
      }
  
      public Node leastLoadedNode() {
          lock.lock();
          try {
              return client.leastLoadedNode(time.milliseconds()).node();
          } finally {
              lock.unlock();
          }
      }
  
      public boolean hasReadyNodes(long now) {
          lock.lock();
          try {
              return client.hasReadyNodes(now);
          } finally {
              lock.unlock();
          }
      }
  
      /**
       * Block waiting on the metadata refresh with a timeout.
       *
       * @return true if update succeeded, false otherwise.
       */
      public boolean awaitMetadataUpdate(Timer timer) {
          int version = this.metadata.requestUpdate(false);
          do {
              poll(timer);
          } while (this.metadata.updateVersion() == version && timer.notExpired());
          return this.metadata.updateVersion() > version;
      }
  
      /**
       * Ensure our metadata is fresh (if an update is expected, this will block
       * until it has completed).
       */
      boolean ensureFreshMetadata(Timer timer) {
          if (this.metadata.updateRequested() || this.metadata.timeToNextUpdate(timer.currentTimeMs()) == 0) {
              return awaitMetadataUpdate(timer);
          } else {
              // the metadata is already fresh
              return true;
          }
      }
  
      /**
       * Wakeup an active poll. This will cause the polling thread to throw an exception either
       * on the current poll if one is active, or the next poll.
       */
      public void wakeup() {
          // wakeup should be safe without holding the client lock since it simply delegates to
          // Selector's wakeup, which is thread-safe
          log.debug("Received user wakeup");
          this.wakeup.set(true);
          this.client.wakeup();
      }
  
      /**
       * Block indefinitely until the given request future has finished.
       * @param future The request future to await.
       * @throws WakeupException if {@link #wakeup()} is called from another thread
       * @throws InterruptException if the calling thread is interrupted
       */
      public void poll(RequestFuture<?> future) {
          while (!future.isDone())
              poll(time.timer(Long.MAX_VALUE), future);
      }
  
      /**
       * Block until the provided request future request has finished or the timeout has expired.
       * @param future The request future to wait for
       * @param timer Timer bounding how long this method can block
       * @return true if the future is done, false otherwise
       * @throws WakeupException if {@link #wakeup()} is called from another thread
       * @throws InterruptException if the calling thread is interrupted
       */
      public boolean poll(RequestFuture<?> future, Timer timer) {
          return poll(future, timer, false);
      }
  
      /**
       * Block until the provided request future request has finished or the timeout has expired.
       *
       * @param future The request future to wait for
       * @param timer Timer bounding how long this method can block
       * @param disableWakeup true if we should not check for wakeups, false otherwise
       *
       * @return true if the future is done, false otherwise
       * @throws WakeupException if {@link #wakeup()} is called from another thread and `disableWakeup` is false
       * @throws InterruptException if the calling thread is interrupted
       */
      public boolean poll(RequestFuture<?> future, Timer timer, boolean disableWakeup) {
          do {
              poll(timer, future, disableWakeup);
          } while (!future.isDone() && timer.notExpired());
          return future.isDone();
      }
  
      /**
       * Poll for any network IO.
       * @param timer Timer bounding how long this method can block
       * @throws WakeupException if {@link #wakeup()} is called from another thread
       * @throws InterruptException if the calling thread is interrupted
       */
      public void poll(Timer timer) {
          poll(timer, null);
      }
  
      /**
       * Poll for any network IO.
       * @param timer Timer bounding how long this method can block
       * @param pollCondition Nullable blocking condition
       */
      public void poll(Timer timer, PollCondition pollCondition) {
          poll(timer, pollCondition, false);
      }
  
      /**
       * Poll for any network IO.
       * @param timer Timer bounding how long this method can block
       * @param pollCondition Nullable blocking condition
       * @param disableWakeup If TRUE disable triggering wake-ups
       */
      public void poll(Timer timer, PollCondition pollCondition, boolean disableWakeup) {
          // there may be handlers which need to be invoked if we woke up the previous call to poll
          firePendingCompletedRequests();
  
          lock.lock();
          try {
              // Handle async disconnects prior to attempting any sends
              handlePendingDisconnects();
  
              // send all the requests we can send now
              long pollDelayMs = trySend(timer.currentTimeMs());
  
              // check whether the poll is still needed by the caller. Note that if the expected completion
              // condition becomes satisfied after the call to shouldBlock() (because of a fired completion
              // handler), the client will be woken up.
              if (pendingCompletion.isEmpty() && (pollCondition == null || pollCondition.shouldBlock())) {
                  // if there are no requests in flight, do not block longer than the retry backoff
                  long pollTimeout = Math.min(timer.remainingMs(), pollDelayMs);
                  if (client.inFlightRequestCount() == 0)
                      pollTimeout = Math.min(pollTimeout, retryBackoffMs);
                  client.poll(pollTimeout, timer.currentTimeMs());
              } else {
                  client.poll(0, timer.currentTimeMs());
              }
              timer.update();
  
              // handle any disconnects by failing the active requests. note that disconnects must
              // be checked immediately following poll since any subsequent call to client.ready()
              // will reset the disconnect status
              checkDisconnects(timer.currentTimeMs());
              if (!disableWakeup) {
                  // trigger wakeups after checking for disconnects so that the callbacks will be ready
                  // to be fired on the next call to poll()
                  maybeTriggerWakeup();
              }
              // throw InterruptException if this thread is interrupted
              maybeThrowInterruptException();
  
              // try again to send requests since buffer space may have been
              // cleared or a connect finished in the poll
              trySend(timer.currentTimeMs());
  
              // fail requests that couldn't be sent if they have expired
              failExpiredRequests(timer.currentTimeMs());
  
              // clean unsent requests collection to keep the map from growing indefinitely
              unsent.clean();
          } finally {
              lock.unlock();
          }
  
          // called without the lock to avoid deadlock potential if handlers need to acquire locks
          firePendingCompletedRequests();
  
          metadata.maybeThrowAnyException();
      }
  
      /**
       * Poll for network IO and return immediately. This will not trigger wakeups.
       */
      public void pollNoWakeup() {
          poll(time.timer(0), null, true);
      }
  
      /**
       * Poll for network IO in best-effort only trying to transmit the ready-to-send request
       * Do not check any pending requests or metadata errors so that no exception should ever
       * be thrown, also no wakeups be triggered and no interrupted exception either.
       */
      public void transmitSends() {
          Timer timer = time.timer(0);
  
          // do not try to handle any disconnects, prev request failures, metadata exception etc;
          // just try once and return immediately
          lock.lock();
          try {
              // send all the requests we can send now
              trySend(timer.currentTimeMs());
  
              client.poll(0, timer.currentTimeMs());
          } finally {
              lock.unlock();
          }
      }
  
      /**
       * Block until all pending requests from the given node have finished.
       * @param node The node to await requests from
       * @param timer Timer bounding how long this method can block
       * @return true If all requests finished, false if the timeout expired first
       */
      public boolean awaitPendingRequests(Node node, Timer timer) {
          while (hasPendingRequests(node) && timer.notExpired()) {
              poll(timer);
          }
          return !hasPendingRequests(node);
      }
  
      /**
       * Get the count of pending requests to the given node. This includes both request that
       * have been transmitted (i.e. in-flight requests) and those which are awaiting transmission.
       * @param node The node in question
       * @return The number of pending requests
       */
      public int pendingRequestCount(Node node) {
          lock.lock();
          try {
              return unsent.requestCount(node) + client.inFlightRequestCount(node.idString());
          } finally {
              lock.unlock();
          }
      }
  
      /**
       * Check whether there is pending request to the given node. This includes both request that
       * have been transmitted (i.e. in-flight requests) and those which are awaiting transmission.
       * @param node The node in question
       * @return A boolean indicating whether there is pending request
       */
      public boolean hasPendingRequests(Node node) {
          if (unsent.hasRequests(node))
              return true;
          lock.lock();
          try {
              return client.hasInFlightRequests(node.idString());
          } finally {
              lock.unlock();
          }
      }
  
      /**
       * Get the total count of pending requests from all nodes. This includes both requests that
       * have been transmitted (i.e. in-flight requests) and those which are awaiting transmission.
       * @return The total count of pending requests
       */
      public int pendingRequestCount() {
          lock.lock();
          try {
              return unsent.requestCount() + client.inFlightRequestCount();
          } finally {
              lock.unlock();
          }
      }
  
      /**
       * Check whether there is pending request. This includes both requests that
       * have been transmitted (i.e. in-flight requests) and those which are awaiting transmission.
       * @return A boolean indicating whether there is pending request
       */
      public boolean hasPendingRequests() {
          if (unsent.hasRequests())
              return true;
          lock.lock();
          try {
              return client.hasInFlightRequests();
          } finally {
              lock.unlock();
          }
      }
  
      private void firePendingCompletedRequests() {
          boolean completedRequestsFired = false;
          for (;;) {
              RequestFutureCompletionHandler completionHandler = pendingCompletion.poll();
              if (completionHandler == null)
                  break;
  
              completionHandler.fireCompletion();
              completedRequestsFired = true;
          }
  
          // wakeup the client in case it is blocking in poll for this future's completion
          if (completedRequestsFired)
              client.wakeup();
      }
  
      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
          for (Node node : unsent.nodes()) {
              if (client.connectionFailed(node)) {
                  // Remove entry before invoking request callback to avoid callbacks handling
                  // coordinator failures traversing the unsent list again.
                  Collection<ClientRequest> requests = unsent.remove(node);
                  for (ClientRequest request : requests) {
                      RequestFutureCompletionHandler handler = (RequestFutureCompletionHandler) request.callback();
                      AuthenticationException authenticationException = client.authenticationException(node);
                      handler.onComplete(new ClientResponse(request.makeHeader(request.requestBuilder().latestAllowedVersion()),
                              request.callback(), request.destination(), request.createdTimeMs(), now, true,
                              null, authenticationException, null));
                  }
              }
          }
      }
  
      private void handlePendingDisconnects() {
          lock.lock();
          try {
              while (true) {
                  Node node = pendingDisconnects.poll();
                  if (node == null)
                      break;
  
                  failUnsentRequests(node, DisconnectException.INSTANCE);
                  client.disconnect(node.idString());
              }
          } finally {
              lock.unlock();
          }
      }
  
      public void disconnectAsync(Node node) {
          pendingDisconnects.offer(node);
          client.wakeup();
      }
  
      private void failExpiredRequests(long now) {
          // clear all expired unsent requests and fail their corresponding futures
          Collection<ClientRequest> expiredRequests = unsent.removeExpiredRequests(now);
          for (ClientRequest request : expiredRequests) {
              RequestFutureCompletionHandler handler = (RequestFutureCompletionHandler) request.callback();
              handler.onFailure(new TimeoutException("Failed to send request after " + request.requestTimeoutMs() + " ms."));
          }
      }
  
      private void failUnsentRequests(Node node, RuntimeException e) {
          // clear unsent requests to node and fail their corresponding futures
          lock.lock();
          try {
              Collection<ClientRequest> unsentRequests = unsent.remove(node);
              for (ClientRequest unsentRequest : unsentRequests) {
                  RequestFutureCompletionHandler handler = (RequestFutureCompletionHandler) unsentRequest.callback();
                  handler.onFailure(e);
              }
          } finally {
              lock.unlock();
          }
      }
  
      // Visible for testing
      long trySend(long now) {
          long pollDelayMs = maxPollTimeoutMs;
  
          // send any requests that can be sent now
          for (Node node : unsent.nodes()) {
              Iterator<ClientRequest> iterator = unsent.requestIterator(node);
              if (iterator.hasNext())
                  pollDelayMs = Math.min(pollDelayMs, client.pollDelayMs(node, now));
  
              while (iterator.hasNext()) {
                  ClientRequest request = iterator.next();
                  if (client.ready(node, now)) {
                      client.send(request, now);
                      iterator.remove();
                  } else {
                      // try next node when current node is not ready
                      break;
                  }
              }
          }
          return pollDelayMs;
      }
  
      public void maybeTriggerWakeup() {
          if (!wakeupDisabled.get() && wakeup.get()) {
              log.debug("Raising WakeupException in response to user wakeup");
              wakeup.set(false);
              throw new WakeupException();
          }
      }
  
      private void maybeThrowInterruptException() {
          if (Thread.interrupted()) {
              throw new InterruptException(new InterruptedException());
          }
      }
  
      public void disableWakeups() {
          wakeupDisabled.set(true);
      }
  
      @Override
      public void close() throws IOException {
          lock.lock();
          try {
              client.close();
          } finally {
              lock.unlock();
          }
      }
  
  
      /**
       * Check if the code is disconnected and unavailable for immediate reconnection (i.e. if it is in
       * reconnect backoff window following the disconnect).
       */
      public boolean isUnavailable(Node node) {
          lock.lock();
          try {
              return NetworkClientUtils.isUnavailable(client, node, time);
          } finally {
              lock.unlock();
          }
      }
  
      /**
       * Check for an authentication error on a given node and raise the exception if there is one.
       */
      public void maybeThrowAuthFailure(Node node) {
          lock.lock();
          try {
              NetworkClientUtils.maybeThrowAuthFailure(client, node);
          } finally {
              lock.unlock();
          }
      }
  
      /**
       * Initiate a connection if currently possible. This is only really useful for resetting the failed
       * status of a socket. If there is an actual request to send, then {@link #send(Node, AbstractRequest.Builder)}
       * should be used.
       * @param node The node to connect to
       */
      public void tryConnect(Node node) {
          lock.lock();
          try {
              NetworkClientUtils.tryConnect(client, node, time);
          } finally {
              lock.unlock();
          }
      }
  
      private class RequestFutureCompletionHandler implements RequestCompletionHandler {
          private final RequestFuture<ClientResponse> future;
          private ClientResponse response;
          private RuntimeException e;
  
          private RequestFutureCompletionHandler() {
              this.future = new RequestFuture<>();
          }
  
          public void fireCompletion() {
              if (e != null) {
                  future.raise(e);
              } else if (response.authenticationException() != null) {
                  future.raise(response.authenticationException());
              } else if (response.wasDisconnected()) {
                  log.debug("Cancelled request with header {} due to node {} being disconnected",
                          response.requestHeader(), response.destination());
                  future.raise(DisconnectException.INSTANCE);
              } else if (response.versionMismatch() != null) {
                  future.raise(response.versionMismatch());
              } else {
                  future.complete(response);
              }
          }
  
          public void onFailure(RuntimeException e) {
              this.e = e;
              pendingCompletion.add(this);
          }
  
          @Override
          public void onComplete(ClientResponse response) {
              this.response = response;
              pendingCompletion.add(this);
          }
      }
  
      /**
       * When invoking poll from a multi-threaded environment, it is possible that the condition that
       * the caller is awaiting has already been satisfied prior to the invocation of poll. We therefore
       * introduce this interface to push the condition checking as close as possible to the invocation
       * of poll. In particular, the check will be done while holding the lock used to protect concurrent
       * access to {@link org.apache.kafka.clients.NetworkClient}, which means implementations must be
       * very careful about locking order if the callback must acquire additional locks.
       */
      public interface PollCondition {
          /**
           * Return whether the caller is still awaiting an IO event.
           * @return true if so, false otherwise.
           */
          boolean shouldBlock();
      }
  
      /*
       * A thread-safe helper class to hold requests per node that have not been sent yet
       */
      private static final class UnsentRequests {
          private final ConcurrentMap<Node, ConcurrentLinkedQueue<ClientRequest>> unsent;
  
          private UnsentRequests() {
              unsent = new ConcurrentHashMap<>();
          }
  
          public void put(Node node, ClientRequest request) {
              // the lock protects the put from a concurrent removal of the queue for the node
              synchronized (unsent) {
                  ConcurrentLinkedQueue<ClientRequest> requests = unsent.computeIfAbsent(node, key -> new ConcurrentLinkedQueue<>());
                  requests.add(request);
              }
          }
  
          public int requestCount(Node node) {
              ConcurrentLinkedQueue<ClientRequest> requests = unsent.get(node);
              return requests == null ? 0 : requests.size();
          }
  
          public int requestCount() {
              int total = 0;
              for (ConcurrentLinkedQueue<ClientRequest> requests : unsent.values())
                  total += requests.size();
              return total;
          }
  
          public boolean hasRequests(Node node) {
              ConcurrentLinkedQueue<ClientRequest> requests = unsent.get(node);
              return requests != null && !requests.isEmpty();
          }
  
          public boolean hasRequests() {
              for (ConcurrentLinkedQueue<ClientRequest> requests : unsent.values())
                  if (!requests.isEmpty())
                      return true;
              return false;
          }
  
          private Collection<ClientRequest> removeExpiredRequests(long now) {
              List<ClientRequest> expiredRequests = new ArrayList<>();
              for (ConcurrentLinkedQueue<ClientRequest> requests : unsent.values()) {
                  Iterator<ClientRequest> requestIterator = requests.iterator();
                  while (requestIterator.hasNext()) {
                      ClientRequest request = requestIterator.next();
                      long elapsedMs = Math.max(0, now - request.createdTimeMs());
                      if (elapsedMs > request.requestTimeoutMs()) {
                          expiredRequests.add(request);
                          requestIterator.remove();
                      } else
                          break;
                  }
              }
              return expiredRequests;
          }
  
          public void clean() {
              // the lock protects removal from a concurrent put which could otherwise mutate the
              // queue after it has been removed from the map
              synchronized (unsent) {
                  unsent.values().removeIf(ConcurrentLinkedQueue::isEmpty);
              }
          }
  
          public Collection<ClientRequest> remove(Node node) {
              // the lock protects removal from a concurrent put which could otherwise mutate the
              // queue after it has been removed from the map
              synchronized (unsent) {
                  ConcurrentLinkedQueue<ClientRequest> requests = unsent.remove(node);
                  return requests == null ? Collections.emptyList() : requests;
              }
          }
  
          public Iterator<ClientRequest> requestIterator(Node node) {
              ConcurrentLinkedQueue<ClientRequest> requests = unsent.get(node);
              return requests == null ? Collections.emptyIterator() : requests.iterator();
          }
  
          public Collection<Node> nodes() {
              return unsent.keySet();
          }
      }
  
  }
  