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.network;
  
  import org.apache.kafka.common.Endpoint;
  import org.apache.kafka.common.internals.Plugin;
  import org.apache.kafka.common.utils.LogContext;
  import org.apache.kafka.server.authorizer.Authorizer;
  import org.apache.kafka.server.authorizer.AuthorizerServerInfo;
  
  import org.slf4j.Logger;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.TreeSet;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.CompletionStage;
  
  /**
   * Manages a set of per-endpoint futures.
   */
  public class EndpointReadyFutures {
      public static class Builder {
          private LogContext logContext = null;
          private final Map<Endpoint, List<EndpointCompletionStage>> endpointStages = new HashMap<>();
          private final List<EndpointCompletionStage> stages = new ArrayList<>();
  
          /**
           * Add a readiness future that will block all endpoints.
           *
           * @param name          The future name.
           * @param future        The future object.
           *
           * @return              This builder object.
           */
          public Builder addReadinessFuture(
              String name,
              CompletableFuture<?> future
          ) {
              stages.add(new EndpointCompletionStage(name, future));
              return this;
          }
  
          /**
           * Add readiness futures for individual endpoints.
           *
           * @param name          The future name.
           * @param newFutures    A map from endpoints to futures.
           *
           * @return              This builder object.
           */
          public Builder addReadinessFutures(
              String name,
              Map<Endpoint, ? extends CompletionStage<?>> newFutures
          ) {
              newFutures.forEach((endpoint, future) -> endpointStages.computeIfAbsent(endpoint, __ -> new ArrayList<>()).
                  add(new EndpointCompletionStage(name, future)));
              return this;
          }
  
          /**
           * Build the EndpointReadyFutures object.
           *
           * @param authorizerPlugin    The authorizer to use, if any. Will be started.
           * @param info                Server information to be passed to the authorizer.
           *
           * @return              The new futures object.
           */
          public EndpointReadyFutures build(
              Optional<Plugin<Authorizer>> authorizerPlugin,
              AuthorizerServerInfo info
          ) {
              if (authorizerPlugin.isPresent()) {
                  return build(authorizerPlugin.get().get().start(info), info);
              } else {
                  return build(Map.of(), info);
              }
          }
  
          EndpointReadyFutures build(
              Map<Endpoint, ? extends CompletionStage<?>> authorizerStartFutures,
              AuthorizerServerInfo info
          ) {
              if (logContext == null) logContext = new LogContext();
              Map<Endpoint, CompletionStage<?>> effectiveStartFutures =
                      new HashMap<>(authorizerStartFutures);
              for (Endpoint endpoint : info.endpoints()) {
                  if (!effectiveStartFutures.containsKey(endpoint)) {
                      CompletableFuture<Void> completedFuture = CompletableFuture.completedFuture(null);
                      effectiveStartFutures.put(endpoint, completedFuture);
                  }
              }
              if (info.endpoints().size() != effectiveStartFutures.size()) {
                  List<String> notInInfo = new ArrayList<>();
                  for (Endpoint endpoint : effectiveStartFutures.keySet()) {
                      if (!info.endpoints().contains(endpoint)) {
                          notInInfo.add(endpoint.listener());
                      }
                  }
                  throw new RuntimeException("Found authorizer futures that weren't included " +
                          "in AuthorizerServerInfo: " + notInInfo);
              }
              addReadinessFutures("authorizerStart", effectiveStartFutures);
              stages.forEach(stage -> {
                  Map<Endpoint, CompletionStage<?>> newReadinessFutures = new HashMap<>();
                  info.endpoints().forEach(endpoint -> newReadinessFutures.put(endpoint, stage.future));
                  addReadinessFutures(stage.name, newReadinessFutures);
              });
              return new EndpointReadyFutures(logContext,
                      endpointStages);
          }
      }
  
      static class EndpointCompletionStage {
          final String name;
          final CompletionStage<?> future;
  
          EndpointCompletionStage(String name, CompletionStage<?> future) {
              this.name = name;
              this.future = future;
          }
      }
  
      class EndpointReadyFuture {
          final String endpointName;
          final TreeSet<String> incomplete;
          final CompletableFuture<Void> future;
  
          EndpointReadyFuture(Endpoint endpoint, Collection<String> stageNames) {
              this.endpointName = endpoint.listener();
              this.incomplete = new TreeSet<>(stageNames);
              this.future = new CompletableFuture<>();
          }
  
          void completeStage(String stageName) {
              boolean done = false;
              synchronized (EndpointReadyFuture.this) {
                  if (incomplete.remove(stageName)) {
                      if (incomplete.isEmpty()) {
                          done = true;
                      } else {
                          log.info("{} completed for endpoint {}. Still waiting for {}.",
                                  stageName, endpointName, incomplete);
                      }
                  }
              }
              if (done) {
                  if (future.complete(null)) {
                      log.info("{} completed for endpoint {}. Endpoint is now READY.",
                              stageName, endpointName);
                  }
              }
          }
  
          void failStage(String what, Throwable exception) {
              if (future.completeExceptionally(exception)) {
                  synchronized (EndpointReadyFuture.this) {
                      incomplete.clear();
                  }
                  log.warn("Endpoint {} will never become ready because we encountered an {} exception",
                          endpointName, what, exception);
              }
          }
      }
  
      private final Logger log;
  
      private final Map<Endpoint, CompletableFuture<Void>> futures;
  
      private EndpointReadyFutures(
          LogContext logContext,
          Map<Endpoint, List<EndpointCompletionStage>> endpointStages
      ) {
          this.log = logContext.logger(EndpointReadyFutures.class);
          Map<Endpoint, CompletableFuture<Void>> newFutures = new HashMap<>();
          endpointStages.forEach((endpoint, stages) -> {
              List<String> stageNames = new ArrayList<>();
              stages.forEach(stage -> stageNames.add(stage.name));
              EndpointReadyFuture readyFuture = new EndpointReadyFuture(endpoint, stageNames);
              newFutures.put(endpoint, readyFuture.future);
              stages.forEach(stage -> stage.future.whenComplete((__, exception) -> {
                  if (exception != null) {
                      readyFuture.failStage(stage.name, exception);
                  } else {
                      readyFuture.completeStage(stage.name);
                  }
              }));
          });
          this.futures = newFutures;
      }
  
      public Map<Endpoint, CompletableFuture<Void>> futures() {
          return futures;
      }
  }
  