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.admin.internals;
  
  import org.apache.kafka.clients.CommonClientConfigs;
  import org.apache.kafka.common.Node;
  import org.apache.kafka.common.errors.DisconnectException;
  import org.apache.kafka.common.errors.UnsupportedVersionException;
  import org.apache.kafka.common.requests.AbstractRequest;
  import org.apache.kafka.common.requests.AbstractResponse;
  import org.apache.kafka.common.requests.FindCoordinatorRequest.NoBatchedFindCoordinatorsException;
  import org.apache.kafka.common.requests.OffsetFetchRequest.NoBatchedOffsetFetchRequestException;
  import org.apache.kafka.common.utils.ExponentialBackoff;
  import org.apache.kafka.common.utils.LogContext;
  
  import org.slf4j.Logger;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Objects;
  import java.util.Optional;
  import java.util.OptionalInt;
  import java.util.Set;
  import java.util.function.BiFunction;
  import java.util.function.Function;
  import java.util.stream.Collectors;
  
  /**
   * The `KafkaAdminClient`'s internal `Call` primitive is not a good fit for multi-stage
   * request workflows such as we see with the group coordinator APIs or any request which
   * needs to be sent to a partition leader. Typically these APIs have two concrete stages:
   *
   * 1. Lookup: Find the broker that can fulfill the request (e.g. partition leader or group
   *            coordinator)
   * 2. Fulfillment: Send the request to the broker found in the first step
   *
   * This is complicated by the fact that `Admin` APIs are typically batched, which
   * means the Lookup stage may result in a set of brokers. For example, take a `ListOffsets`
   * request for a set of topic partitions. In the Lookup stage, we will find the partition
   * leaders for this set of partitions; in the Fulfillment stage, we will group together
   * partition according to the IDs of the discovered leaders.
   *
   * Additionally, the flow between these two stages is bi-directional. We may find after
   * sending a `ListOffsets` request to an expected leader that there was a leader change.
   * This would result in a topic partition being sent back to the Lookup stage.
   *
   * Managing this complexity by chaining together `Call` implementations is challenging
   * and messy, so instead we use this class to do the bookkeeping. It handles both the
   * batching aspect as well as the transitions between the Lookup and Fulfillment stages.
   *
   * Note that the interpretation of the `retries` configuration becomes ambiguous
   * for this kind of pipeline. We could treat it as an overall limit on the number
   * of requests that can be sent, but that is not very useful because each pipeline
   * has a minimum number of requests that need to be sent in order to satisfy the request.
   * Instead, we treat this number of retries independently at each stage so that each
   * stage has at least one opportunity to complete. So if a user sets `retries=1`, then
   * the full pipeline can still complete as long as there are no request failures.
   *
   * @param <K> The key type, which is also the granularity of the request routing (e.g.
   *            this could be `TopicPartition` in the case of requests intended for a partition
   *            leader or the `GroupId` in the case of consumer group requests intended for
   *            the group coordinator)
   * @param <V> The fulfillment type for each key (e.g. this could be consumer group state
   *            when the key type is a consumer `GroupId`)
   */
  public class AdminApiDriver<K, V> {
      private final Logger log;
      private final ExponentialBackoff retryBackoff;
      private final long deadlineMs;
      private final AdminApiHandler<K, V> handler;
      private final AdminApiFuture<K, V> future;
  
      private final BiMultimap<ApiRequestScope, K> lookupMap = new BiMultimap<>();
      private final BiMultimap<FulfillmentScope, K> fulfillmentMap = new BiMultimap<>();
      private final Map<ApiRequestScope, RequestState> requestStates = new HashMap<>();
  
      public AdminApiDriver(
          AdminApiHandler<K, V> handler,
          AdminApiFuture<K, V> future,
          long deadlineMs,
          long retryBackoffMs,
          long retryBackoffMaxMs,
          LogContext logContext
      ) {
          this.handler = handler;
          this.future = future;
          this.deadlineMs = deadlineMs;
          this.retryBackoff = new ExponentialBackoff(
              retryBackoffMs,
              CommonClientConfigs.RETRY_BACKOFF_EXP_BASE,
              retryBackoffMaxMs,
              CommonClientConfigs.RETRY_BACKOFF_JITTER);
          this.log = logContext.logger(AdminApiDriver.class);
  
          // For any lookup keys for which we do not have cached information, we will need to look up
          // metadata. For all cached keys, they can proceed straight to the fulfillment map.
          // Note that the cache is only used on the initial calls, and any errors that result
          // in additional lookups use the full set of lookup keys.
          retryLookup(future.uncachedLookupKeys());
          future.cachedKeyBrokerIdMapping().forEach((key, brokerId) -> fulfillmentMap.put(new FulfillmentScope(brokerId), key));
      }
  
      /**
       * Associate a key with a brokerId. This is called after a response in the Lookup
       * stage reveals the mapping (e.g. when the `FindCoordinator` tells us the group
       * coordinator for a specific consumer group).
       */
      private void map(K key, Integer brokerId) {
          lookupMap.remove(key);
          fulfillmentMap.put(new FulfillmentScope(brokerId), key);
      }
  
      /**
       * Disassociate a key from the currently mapped brokerId. This will send the key
       * back to the Lookup stage, which will allow us to attempt lookup again.
       */
      private void unmap(K key) {
          fulfillmentMap.remove(key);
  
          ApiRequestScope lookupScope = handler.lookupStrategy().lookupScope(key);
          OptionalInt destinationBrokerId = lookupScope.destinationBrokerId();
  
          if (destinationBrokerId.isPresent()) {
              fulfillmentMap.put(new FulfillmentScope(destinationBrokerId.getAsInt()), key);
          } else {
              lookupMap.put(handler.lookupStrategy().lookupScope(key), key);
          }
      }
  
      private void clear(Collection<K> keys) {
          keys.forEach(key -> {
              lookupMap.remove(key);
              fulfillmentMap.remove(key);
          });
      }
  
      OptionalInt keyToBrokerId(K key) {
          Optional<FulfillmentScope> scope = fulfillmentMap.getKey(key);
          return scope
              .map(fulfillmentScope -> OptionalInt.of(fulfillmentScope.destinationBrokerId))
              .orElseGet(OptionalInt::empty);
      }
  
      /**
       * Complete the future associated with the given key exceptionally. After is called,
       * the key will be taken out of both the Lookup and Fulfillment stages so that request
       * are not retried.
       */
      private void completeExceptionally(Map<K, Throwable> errors) {
          if (!errors.isEmpty()) {
              future.completeExceptionally(errors);
              clear(errors.keySet());
          }
      }
  
      private void completeLookupExceptionally(Map<K, Throwable> errors) {
          if (!errors.isEmpty()) {
              future.completeLookupExceptionally(errors);
              clear(errors.keySet());
          }
      }
  
      private void retryLookup(Collection<K> keys) {
          keys.forEach(this::unmap);
      }
  
      /**
       * Complete the future associated with the given key. After this is called, all keys will
       * be taken out of both the Lookup and Fulfillment stages so that request are not retried.
       */
      private void complete(Map<K, V> values) {
          if (!values.isEmpty()) {
              future.complete(values);
              clear(values.keySet());
          }
      }
  
      private void completeLookup(Map<K, Integer> brokerIdMapping) {
          if (!brokerIdMapping.isEmpty()) {
              future.completeLookup(brokerIdMapping);
              brokerIdMapping.forEach(this::map);
          }
      }
  
      /**
       * Check whether any requests need to be sent. This should be called immediately
       * after the driver is constructed and then again after each request returns
       * (i.e. after {@link #onFailure(long, RequestSpec, Throwable)} or
       * {@link #onResponse(long, RequestSpec, AbstractResponse, Node)}).
       *
       * @return A list of requests that need to be sent
       */
      public List<RequestSpec<K>> poll() {
          List<RequestSpec<K>> requests = new ArrayList<>();
          collectLookupRequests(requests);
          collectFulfillmentRequests(requests);
          return requests;
      }
  
      /**
       * Callback that is invoked when a `Call` returns a response successfully.
       */
      public void onResponse(
          long currentTimeMs,
          RequestSpec<K> spec,
          AbstractResponse response,
          Node node
      ) {
          clearInflightRequest(currentTimeMs, spec);
  
          if (spec.scope instanceof FulfillmentScope) {
              AdminApiHandler.ApiResult<K, V> result = handler.handleResponse(
                  node,
                  spec.keys,
                  response
              );
              complete(result.completedKeys);
              completeExceptionally(result.failedKeys);
              retryLookup(result.unmappedKeys);
          } else {
              AdminApiLookupStrategy.LookupResult<K> result = handler.lookupStrategy().handleResponse(
                  spec.keys,
                  response
              );
  
              result.completedKeys.forEach(lookupMap::remove);
              completeLookup(result.mappedKeys);
              completeLookupExceptionally(result.failedKeys);
          }
      }
  
      /**
       * Callback that is invoked when a `Call` is failed.
       */
      public void onFailure(
          long currentTimeMs,
          RequestSpec<K> spec,
          Throwable t
      ) {
          clearInflightRequest(currentTimeMs, spec);
          if (t instanceof DisconnectException) {
              log.debug("Node disconnected before response could be received for request {}. " +
                  "Will attempt retry", spec.request);
  
              // After a disconnect, we want the driver to attempt to lookup the key
              // again. This gives us a chance to find a new coordinator or partition
              // leader for example.
              Set<K> keysToUnmap = spec.keys.stream()
                  .filter(future.lookupKeys()::contains)
                  .collect(Collectors.toSet());
              retryLookup(keysToUnmap);
  
          } else if (t instanceof NoBatchedFindCoordinatorsException || t instanceof NoBatchedOffsetFetchRequestException) {
              ((CoordinatorStrategy) handler.lookupStrategy()).disableBatch();
              Set<K> keysToUnmap = spec.keys.stream()
                  .filter(future.lookupKeys()::contains)
                  .collect(Collectors.toSet());
              retryLookup(keysToUnmap);
          } else if (t instanceof UnsupportedVersionException) {
              if (spec.scope instanceof FulfillmentScope) {
                  int brokerId = ((FulfillmentScope) spec.scope).destinationBrokerId;
                  Map<K, Throwable> unrecoverableFailures =
                      handler.handleUnsupportedVersionException(
                          brokerId,
                          (UnsupportedVersionException) t,
                          spec.keys);
                  completeExceptionally(unrecoverableFailures);
              } else {
                  Map<K, Throwable> unrecoverableLookupFailures =
                      handler.lookupStrategy().handleUnsupportedVersionException(
                          (UnsupportedVersionException) t,
                          spec.keys);
                  completeLookupExceptionally(unrecoverableLookupFailures);
                  Set<K> keysToUnmap = spec.keys.stream()
                      .filter(k -> !unrecoverableLookupFailures.containsKey(k))
                      .collect(Collectors.toSet());
                  retryLookup(keysToUnmap);
              }
          } else {
              Map<K, Throwable> errors = spec.keys.stream().collect(Collectors.toMap(
                  Function.identity(),
                  key -> t
              ));
              if (spec.scope instanceof FulfillmentScope) {
                  completeExceptionally(errors);
              } else {
                  completeLookupExceptionally(errors);
              }
          }
      }
  
      private void clearInflightRequest(long currentTimeMs, RequestSpec<K> spec) {
          RequestState requestState = requestStates.get(spec.scope);
          if (requestState != null) {
              // Only apply backoff if it's not a retry of a lookup request
              if (spec.scope instanceof FulfillmentScope) {
                  requestState.clearInflightAndBackoff(currentTimeMs);
              } else {
                  requestState.clearInflight(currentTimeMs);
              }
          }
      }
  
      private <T extends ApiRequestScope> void collectRequests(
          List<RequestSpec<K>> requests,
          BiMultimap<T, K> multimap,
          BiFunction<Set<K>, T, Collection<AdminApiHandler.RequestAndKeys<K>>> buildRequest
      ) {
          for (Map.Entry<T, Set<K>> entry : multimap.entrySet()) {
              T scope = entry.getKey();
  
              Set<K> keys = entry.getValue();
              if (keys.isEmpty()) {
                  continue;
              }
  
              RequestState requestState = requestStates.computeIfAbsent(scope, c -> new RequestState());
              if (requestState.hasInflight()) {
                  continue;
              }
  
              // Copy the keys to avoid exposing the underlying mutable set
              Set<K> copyKeys = Set.copyOf(keys);
  
              Collection<AdminApiHandler.RequestAndKeys<K>> newRequests = buildRequest.apply(copyKeys, scope);
              if (newRequests.isEmpty()) {
                  return;
              }
  
              // Only process the first request; all the remaining requests will be targeted at the same broker
              // and we don't want to issue more than one fulfillment request per broker at a time
              AdminApiHandler.RequestAndKeys<K> newRequest = newRequests.iterator().next();
              RequestSpec<K> spec = new RequestSpec<>(
                  handler.apiName() + "(api=" + newRequest.request.apiKey() + ")",
                  scope,
                  newRequest.keys,
                  newRequest.request,
                  requestState.nextAllowedRetryMs,
                  deadlineMs,
                  requestState.tries
              );
  
              requestState.setInflight(spec);
              requests.add(spec);
          }
      }
  
      private void collectLookupRequests(List<RequestSpec<K>> requests) {
          collectRequests(
              requests,
              lookupMap,
              (keys, scope) -> Collections.singletonList(new AdminApiHandler.RequestAndKeys<>(handler.lookupStrategy().buildRequest(keys), keys))
          );
      }
  
      private void collectFulfillmentRequests(List<RequestSpec<K>> requests) {
          collectRequests(
              requests,
              fulfillmentMap,
              (keys, scope) -> handler.buildRequest(scope.destinationBrokerId, keys)
          );
      }
  
      /**
       * This is a helper class which helps us to map requests that need to be sent
       * to the internal `Call` implementation that is used internally in
       * {@link org.apache.kafka.clients.admin.KafkaAdminClient}.
       */
      public static class RequestSpec<K> {
          public final String name;
          public final ApiRequestScope scope;
          public final Set<K> keys;
          public final AbstractRequest.Builder<?> request;
          public final long nextAllowedTryMs;
          public final long deadlineMs;
          public final int tries;
  
          public RequestSpec(
              String name,
              ApiRequestScope scope,
              Set<K> keys,
              AbstractRequest.Builder<?> request,
              long nextAllowedTryMs,
              long deadlineMs,
              int tries
          ) {
              this.name = name;
              this.scope = scope;
              this.keys = keys;
              this.request = request;
              this.nextAllowedTryMs = nextAllowedTryMs;
              this.deadlineMs = deadlineMs;
              this.tries = tries;
          }
  
          @Override
          public String toString() {
              return "RequestSpec(" +
                  "name=" + name +
                  ", scope=" + scope +
                  ", keys=" + keys +
                  ", request=" + request +
                  ", nextAllowedTryMs=" + nextAllowedTryMs +
                  ", deadlineMs=" + deadlineMs +
                  ", tries=" + tries +
                  ')';
          }
      }
  
      /**
       * Helper class used to track the request state within each request scope.
       * This class enforces a maximum number of inflight request and keeps track
       * of backoff/retry state.
       */
      private class RequestState {
          private Optional<RequestSpec<K>> inflightRequest = Optional.empty();
          private int tries = 0;
          private long nextAllowedRetryMs = 0;
  
          boolean hasInflight() {
              return inflightRequest.isPresent();
          }
  
          public void clearInflight(long currentTimeMs) {
              this.inflightRequest = Optional.empty();
              this.nextAllowedRetryMs = currentTimeMs;
          }
  
          public void clearInflightAndBackoff(long currentTimeMs) {
              clearInflight(currentTimeMs + retryBackoff.backoff(tries >= 1 ? tries - 1 : 0));
          }
  
          public void setInflight(RequestSpec<K> spec) {
              this.inflightRequest = Optional.of(spec);
              this.tries++;
          }
      }
  
      /**
       * Completion of the Lookup stage results in a destination broker to send the
       * fulfillment request to. Each destination broker in the Fulfillment stage
       * gets its own request scope.
       */
      private static class FulfillmentScope implements ApiRequestScope {
          public final int destinationBrokerId;
  
          private FulfillmentScope(int destinationBrokerId) {
              this.destinationBrokerId = destinationBrokerId;
          }
  
          @Override
          public OptionalInt destinationBrokerId() {
              return OptionalInt.of(destinationBrokerId);
          }
  
          @Override
          public boolean equals(Object o) {
              if (this == o) return true;
              if (o == null || getClass() != o.getClass()) return false;
              FulfillmentScope that = (FulfillmentScope) o;
              return destinationBrokerId == that.destinationBrokerId;
          }
  
          @Override
          public int hashCode() {
              return Objects.hash(destinationBrokerId);
          }
      }
  
      /**
       * Helper class which maintains a bi-directional mapping from a key to a set of values.
       * Each value can map to one and only one key, but many values can be associated with
       * a single key.
       *
       * @param <K> The key type
       * @param <V> The value type
       */
      private static class BiMultimap<K, V> {
          private final Map<V, K> reverseMap = new HashMap<>();
          private final Map<K, Set<V>> map = new HashMap<>();
  
          void put(K key, V value) {
              remove(value);
              reverseMap.put(value, key);
              map.computeIfAbsent(key, k -> new HashSet<>()).add(value);
          }
  
          void remove(V value) {
              K key = reverseMap.remove(value);
              if (key != null) {
                  Set<V> set = map.get(key);
                  if (set != null) {
                      set.remove(value);
                      if (set.isEmpty()) {
                          map.remove(key);
                      }
                  }
              }
          }
  
          Optional<K> getKey(V value) {
              return Optional.ofNullable(reverseMap.get(value));
          }
  
          Set<Map.Entry<K, Set<V>>> entrySet() {
              return map.entrySet();
          }
      }
  
  }
  