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 com.datastax.oss.driver.internal.core.pool;
  
  import com.datastax.oss.driver.api.core.AsyncAutoCloseable;
  import com.datastax.oss.driver.api.core.CqlIdentifier;
  import com.datastax.oss.driver.api.core.InvalidKeyspaceException;
  import com.datastax.oss.driver.api.core.UnsupportedProtocolVersionException;
  import com.datastax.oss.driver.api.core.auth.AuthenticationException;
  import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
  import com.datastax.oss.driver.api.core.config.DriverConfig;
  import com.datastax.oss.driver.api.core.connection.ReconnectionPolicy;
  import com.datastax.oss.driver.api.core.loadbalancing.NodeDistance;
  import com.datastax.oss.driver.api.core.metadata.Node;
  import com.datastax.oss.driver.api.core.metrics.DefaultNodeMetric;
  import com.datastax.oss.driver.internal.core.channel.ChannelEvent;
  import com.datastax.oss.driver.internal.core.channel.ChannelFactory;
  import com.datastax.oss.driver.internal.core.channel.ClusterNameMismatchException;
  import com.datastax.oss.driver.internal.core.channel.DriverChannel;
  import com.datastax.oss.driver.internal.core.channel.DriverChannelOptions;
  import com.datastax.oss.driver.internal.core.config.ConfigChangeEvent;
  import com.datastax.oss.driver.internal.core.context.EventBus;
  import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
  import com.datastax.oss.driver.internal.core.metadata.DefaultNode;
  import com.datastax.oss.driver.internal.core.metadata.TopologyEvent;
  import com.datastax.oss.driver.internal.core.util.Loggers;
  import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures;
  import com.datastax.oss.driver.internal.core.util.concurrent.Reconnection;
  import com.datastax.oss.driver.internal.core.util.concurrent.RunOrSchedule;
  import com.datastax.oss.driver.internal.core.util.concurrent.UncaughtExceptions;
  import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting;
  import com.datastax.oss.driver.shaded.guava.common.collect.Sets;
  import edu.umd.cs.findbugs.annotations.NonNull;
  import io.netty.util.concurrent.EventExecutor;
  import io.netty.util.concurrent.Future;
  import io.netty.util.concurrent.GenericFutureListener;
  import java.util.ArrayList;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Set;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.CompletionStage;
  import java.util.concurrent.atomic.AtomicInteger;
  import net.jcip.annotations.ThreadSafe;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * The channel pool maintains a set of {@link DriverChannel} instances connected to a given node.
   *
   * <p>It allows clients to obtain a channel to execute their requests.
   *
   * <p>If one or more channels go down, a reconnection process starts in order to replace them; it
   * runs until the channel count is back to its intended target.
   */
  @ThreadSafe
  public class ChannelPool implements AsyncAutoCloseable {
    private static final Logger LOG = LoggerFactory.getLogger(ChannelPool.class);
  
    /**
     * Initializes a new pool.
     *
     * <p>The returned completion stage will complete when all the underlying channels have finished
     * their initialization. If one or more channels fail, a reconnection will be started immediately.
     * Note that this method succeeds even if all channels fail, so you might get a pool that has no
     * channels (i.e. {@link #next()} return {@code null}) and is reconnecting.
     */
    public static CompletionStage<ChannelPool> init(
        Node node,
        CqlIdentifier keyspaceName,
        NodeDistance distance,
        InternalDriverContext context,
        String sessionLogPrefix) {
      ChannelPool pool = new ChannelPool(node, keyspaceName, distance, context, sessionLogPrefix);
      return pool.connect();
    }
  
    // This is read concurrently, but only mutated on adminExecutor (by methods in SingleThreaded)
    @VisibleForTesting final ChannelSet channels = new ChannelSet();
  
    private final Node node;
    private final CqlIdentifier initialKeyspaceName;
    private final EventExecutor adminExecutor;
    private final String sessionLogPrefix;
    private final String logPrefix;
    private final SingleThreaded singleThreaded;
    private volatile boolean invalidKeyspace;
  
    private ChannelPool(
        Node node,
        CqlIdentifier keyspaceName,
        NodeDistance distance,
        InternalDriverContext context,
        String sessionLogPrefix) {
      this.node = node;
      this.initialKeyspaceName = keyspaceName;
      this.adminExecutor = context.getNettyOptions().adminEventExecutorGroup().next();
      this.sessionLogPrefix = sessionLogPrefix;
      this.logPrefix = sessionLogPrefix + "|" + node.getEndPoint();
      this.singleThreaded = new SingleThreaded(keyspaceName, distance, context);
    }
  
    private CompletionStage<ChannelPool> connect() {
      RunOrSchedule.on(adminExecutor, singleThreaded::connect);
      return singleThreaded.connectFuture;
    }
  
    public Node getNode() {
      return node;
    }
  
    /**
     * The keyspace with which the pool was initialized (therefore a constant, it's not updated if the
     * keyspace is switched later).
     */
    public CqlIdentifier getInitialKeyspaceName() {
      return initialKeyspaceName;
    }
  
    /**
     * Whether all channels failed due to an invalid keyspace. This is only used at initialization. We
     * don't make the decision to close the pool here yet, that's done at the session level.
     */
    public boolean isInvalidKeyspace() {
      return invalidKeyspace;
    }
  
    /**
     * @return the channel that has the most available stream ids. This is called on the direct
     *     request path, and we want to avoid complex check-then-act semantics; therefore this might
     *     race and return a channel that is already closed, or {@code null}. In those cases, it is up
     *     to the caller to fail fast and move to the next node.
     */
    public DriverChannel next() {
      return channels.next();
    }
  
    /** @return the number of active channels in the pool. */
    public int size() {
      return channels.size();
    }
  
    /** @return the number of available stream ids on all channels in the pool. */
    public int getAvailableIds() {
      return channels.getAvailableIds();
    }
  
    /**
     * @return the number of requests currently executing on all channels in this pool (including
     *     {@link #getOrphanedIds() orphaned ids}).
     */
    public int getInFlight() {
      return channels.getInFlight();
    }
  
    /**
     * @return the number of stream ids for requests in all channels in this pool that have either
     *     timed out or been cancelled, but for which we can't release the stream id because a request
     *     might still come from the server.
     */
    public int getOrphanedIds() {
      return channels.getOrphanedIds();
    }
  
    /**
     * Sets a new distance for the node this pool belongs to. This method returns immediately, the new
     * distance will be set asynchronously.
     *
     * @param newDistance the new distance to set.
     */
    public void resize(NodeDistance newDistance) {
      RunOrSchedule.on(adminExecutor, () -> singleThreaded.resize(newDistance));
    }
  
    /**
     * Changes the keyspace name on all the channels in this pool.
     *
     * <p>Note that this is not called directly by the user, but happens only on a SetKeyspace
     * response after a successful "USE ..." query, so the name should be valid. If the keyspace
     * switch fails on any channel, that channel is closed and a reconnection is started.
     */
    public CompletionStage<Void> setKeyspace(CqlIdentifier newKeyspaceName) {
      return RunOrSchedule.on(adminExecutor, () -> singleThreaded.setKeyspace(newKeyspaceName));
    }
  
    public void reconnectNow() {
      RunOrSchedule.on(adminExecutor, singleThreaded::reconnectNow);
    }
  
    @NonNull
    @Override
    public CompletionStage<Void> closeFuture() {
      return singleThreaded.closeFuture;
    }
  
    @NonNull
    @Override
    public CompletionStage<Void> closeAsync() {
      RunOrSchedule.on(adminExecutor, singleThreaded::close);
      return singleThreaded.closeFuture;
    }
  
    @NonNull
    @Override
    public CompletionStage<Void> forceCloseAsync() {
      RunOrSchedule.on(adminExecutor, singleThreaded::forceClose);
      return singleThreaded.closeFuture;
    }
  
    /** Holds all administration tasks, that are confined to the admin executor. */
    private class SingleThreaded {
  
      private final DriverConfig config;
      private final ChannelFactory channelFactory;
      private final EventBus eventBus;
      // The channels that are currently connecting
      private final List<CompletionStage<DriverChannel>> pendingChannels = new ArrayList<>();
      private final Set<DriverChannel> closingChannels = new HashSet<>();
      private final Reconnection reconnection;
      private final Object configListenerKey;
  
      private NodeDistance distance;
      private int wantedCount;
      private final CompletableFuture<ChannelPool> connectFuture = new CompletableFuture<>();
      private boolean isConnecting;
      private final CompletableFuture<Void> closeFuture = new CompletableFuture<>();
      private boolean isClosing;
      private CompletableFuture<Void> setKeyspaceFuture;
  
      private CqlIdentifier keyspaceName;
  
      private SingleThreaded(
          CqlIdentifier keyspaceName, NodeDistance distance, InternalDriverContext context) {
        this.keyspaceName = keyspaceName;
        this.config = context.getConfig();
        this.distance = distance;
        this.wantedCount = getConfiguredSize(distance);
        this.channelFactory = context.getChannelFactory();
        this.eventBus = context.getEventBus();
        ReconnectionPolicy reconnectionPolicy = context.getReconnectionPolicy();
        this.reconnection =
            new Reconnection(
                logPrefix,
                adminExecutor,
                () -> reconnectionPolicy.newNodeSchedule(node),
                this::addMissingChannels,
                () -> eventBus.fire(ChannelEvent.reconnectionStarted(node)),
                () -> eventBus.fire(ChannelEvent.reconnectionStopped(node)));
        this.configListenerKey =
            eventBus.register(
                ConfigChangeEvent.class, RunOrSchedule.on(adminExecutor, this::onConfigChanged));
      }
  
      private void connect() {
        assert adminExecutor.inEventLoop();
        if (isConnecting) {
          return;
        }
        isConnecting = true;
        CompletionStage<ChannelPool> initialChannels =
            addMissingChannels()
                .thenApply(
                    allConnected -> {
                      if (!allConnected) {
                        reconnection.start();
                      }
                      return ChannelPool.this;
                    });
        CompletableFutures.completeFrom(initialChannels, connectFuture);
      }
  
      private CompletionStage<Boolean> addMissingChannels() {
        assert adminExecutor.inEventLoop();
        // We always wait for all attempts to succeed or fail before scheduling a reconnection
        assert pendingChannels.isEmpty();
  
        int missing = wantedCount - channels.size();
        LOG.debug("[{}] Trying to create {} missing channels", logPrefix, missing);
        DriverChannelOptions options =
            DriverChannelOptions.builder()
                .withKeyspace(keyspaceName)
                .withOwnerLogPrefix(sessionLogPrefix)
                .build();
        for (int i = 0; i < missing; i++) {
          CompletionStage<DriverChannel> channelFuture = channelFactory.connect(node, options);
          pendingChannels.add(channelFuture);
        }
        return CompletableFutures.allDone(pendingChannels)
            .thenApplyAsync(this::onAllConnected, adminExecutor);
      }
  
      private boolean onAllConnected(@SuppressWarnings("unused") Void v) {
        assert adminExecutor.inEventLoop();
        Throwable fatalError = null;
        int invalidKeyspaceErrors = 0;
        for (CompletionStage<DriverChannel> pendingChannel : pendingChannels) {
          CompletableFuture<DriverChannel> future = pendingChannel.toCompletableFuture();
          assert future.isDone();
          if (future.isCompletedExceptionally()) {
            Throwable error = CompletableFutures.getFailed(future);
            ((DefaultNode) node)
                .getMetricUpdater()
                .incrementCounter(
                    error instanceof AuthenticationException
                        ? DefaultNodeMetric.AUTHENTICATION_ERRORS
                        : DefaultNodeMetric.CONNECTION_INIT_ERRORS,
                    null);
            if (error instanceof ClusterNameMismatchException
                || error instanceof UnsupportedProtocolVersionException) {
              // This will likely be thrown by all channels, but finish the loop cleanly
              fatalError = error;
            } else if (error instanceof AuthenticationException) {
              // Always warn because this is most likely something the operator needs to fix.
              // Keep going to reconnect if it can be fixed without bouncing the client.
              Loggers.warnWithException(LOG, "[{}] Authentication error", logPrefix, error);
            } else if (error instanceof InvalidKeyspaceException) {
              invalidKeyspaceErrors += 1;
            } else {
              if (config
                  .getDefaultProfile()
                  .getBoolean(DefaultDriverOption.CONNECTION_WARN_INIT_ERROR)) {
                Loggers.warnWithException(
                    LOG, "[{}]  Error while opening new channel", logPrefix, error);
              } else {
                LOG.debug("[{}]  Error while opening new channel", logPrefix, error);
              }
            }
          } else {
            DriverChannel channel = CompletableFutures.getCompleted(future);
            if (isClosing) {
              LOG.debug(
                  "[{}] New channel added ({}) but the pool was closed, closing it",
                  logPrefix,
                  channel);
              channel.forceClose();
            } else {
              LOG.debug("[{}] New channel added {}", logPrefix, channel);
              channels.add(channel);
              eventBus.fire(ChannelEvent.channelOpened(node));
              channel
                  .closeStartedFuture()
                  .addListener(
                      f ->
                          adminExecutor
                              .submit(() -> onChannelCloseStarted(channel))
                              .addListener(UncaughtExceptions::log));
              channel
                  .closeFuture()
                  .addListener(
                      f ->
                          adminExecutor
                              .submit(() -> onChannelClosed(channel))
                              .addListener(UncaughtExceptions::log));
            }
          }
        }
        // If all channels failed, assume the keyspace is wrong
        invalidKeyspace =
            invalidKeyspaceErrors > 0 && invalidKeyspaceErrors == pendingChannels.size();
  
        pendingChannels.clear();
  
        if (fatalError != null) {
          Loggers.warnWithException(
              LOG,
              "[{}] Fatal error while initializing pool, forcing the node down",
              logPrefix,
              fatalError);
          // Note: getBroadcastRpcAddress() can only be empty for the control node (and not for modern
          // C* versions anyway). If we already have a control connection open to that node, it's
          // impossible to get a protocol version or cluster name mismatch error while creating the
          // pool, so it's safe to ignore this case.
          node.getBroadcastRpcAddress()
              .ifPresent(address -> eventBus.fire(TopologyEvent.forceDown(address)));
          // Don't bother continuing, the pool will get shut down soon anyway
          return true;
        }
  
        shrinkIfTooManyChannels(); // Can happen if the pool was shrinked during the reconnection
  
        int currentCount = channels.size();
        LOG.debug(
            "[{}] Reconnection attempt complete, {}/{} channels",
            logPrefix,
            currentCount,
            wantedCount);
        // Stop reconnecting if we have the wanted count
        return currentCount >= wantedCount;
      }
  
      private void onChannelCloseStarted(DriverChannel channel) {
        assert adminExecutor.inEventLoop();
        if (!isClosing) {
          LOG.debug("[{}] Channel {} started graceful shutdown", logPrefix, channel);
          channels.remove(channel);
          closingChannels.add(channel);
          eventBus.fire(ChannelEvent.channelClosed(node));
          reconnection.start();
        }
      }
  
      private void onChannelClosed(DriverChannel channel) {
        assert adminExecutor.inEventLoop();
        if (!isClosing) {
          // Either it was closed abruptly and was still in the live set, or it was an orderly
          // shutdown and it had moved to the closing set.
          if (channels.remove(channel)) {
            LOG.debug("[{}] Lost channel {}", logPrefix, channel);
            eventBus.fire(ChannelEvent.channelClosed(node));
            reconnection.start();
          } else {
            LOG.debug("[{}] Channel {} completed graceful shutdown", logPrefix, channel);
            closingChannels.remove(channel);
          }
        }
      }
  
      private void resize(NodeDistance newDistance) {
        assert adminExecutor.inEventLoop();
        distance = newDistance;
        int newChannelCount = getConfiguredSize(newDistance);
        if (newChannelCount > wantedCount) {
          LOG.debug("[{}] Growing ({} => {} channels)", logPrefix, wantedCount, newChannelCount);
          wantedCount = newChannelCount;
          reconnection.start();
        } else if (newChannelCount < wantedCount) {
          LOG.debug("[{}] Shrinking ({} => {} channels)", logPrefix, wantedCount, newChannelCount);
          wantedCount = newChannelCount;
          if (!reconnection.isRunning()) {
            shrinkIfTooManyChannels();
          } // else it will be handled at the end of the reconnection attempt
        }
      }
  
      private void shrinkIfTooManyChannels() {
        assert adminExecutor.inEventLoop();
        int extraCount = channels.size() - wantedCount;
        if (extraCount > 0) {
          LOG.debug("[{}] Closing {} extra channels", logPrefix, extraCount);
          Set<DriverChannel> toRemove = Sets.newHashSetWithExpectedSize(extraCount);
          for (DriverChannel channel : channels) {
            toRemove.add(channel);
            if (--extraCount == 0) {
              break;
            }
          }
          for (DriverChannel channel : toRemove) {
            channels.remove(channel);
            channel.close();
            eventBus.fire(ChannelEvent.channelClosed(node));
          }
        }
      }
  
      private void onConfigChanged(@SuppressWarnings("unused") ConfigChangeEvent event) {
        assert adminExecutor.inEventLoop();
        // resize re-reads the pool size from the configuration and does nothing if it hasn't changed,
        // which is exactly what we want.
        resize(distance);
      }
  
      private CompletionStage<Void> setKeyspace(CqlIdentifier newKeyspaceName) {
        assert adminExecutor.inEventLoop();
        if (setKeyspaceFuture != null && !setKeyspaceFuture.isDone()) {
          return CompletableFutures.failedFuture(
              new IllegalStateException(
                  "Can't call setKeyspace while a keyspace switch is already in progress"));
        }
        keyspaceName = newKeyspaceName;
        setKeyspaceFuture = new CompletableFuture<>();
  
        // Switch the keyspace on all live channels.
        // We can read the size before iterating because mutations are confined to this thread:
        int toSwitch = channels.size();
        if (toSwitch == 0) {
          setKeyspaceFuture.complete(null);
        } else {
          AtomicInteger remaining = new AtomicInteger(toSwitch);
          for (DriverChannel channel : channels) {
            channel
                .setKeyspace(newKeyspaceName)
                .addListener(
                    f -> {
                      // Don't handle errors: if a channel fails to switch the keyspace, it closes
                      if (remaining.decrementAndGet() == 0) {
                        setKeyspaceFuture.complete(null);
                      }
                    });
          }
        }
  
        // pending channels were scheduled with the old keyspace name, ensure they eventually switch
        for (CompletionStage<DriverChannel> channelFuture : pendingChannels) {
          // errors are swallowed here, this is fine because a setkeyspace error will close the
          // channel, so it will eventually get reported
          channelFuture.thenAccept(channel -> channel.setKeyspace(newKeyspaceName));
        }
  
        return setKeyspaceFuture;
      }
  
      private void reconnectNow() {
        assert adminExecutor.inEventLoop();
        // Don't force because if the reconnection is stopped, it means either we have enough channels
        // or the pool is shutting down.
        reconnection.reconnectNow(false);
      }
  
      private void close() {
        assert adminExecutor.inEventLoop();
        if (isClosing) {
          return;
        }
        isClosing = true;
  
        // If an attempt was in progress right now, it might open new channels but they will be
        // handled in onAllConnected
        reconnection.stop();
  
        eventBus.unregister(configListenerKey, ConfigChangeEvent.class);
  
        // Close all channels, the pool future completes when all the channels futures have completed
        int toClose = closingChannels.size() + channels.size();
        if (toClose == 0) {
          closeFuture.complete(null);
        } else {
          AtomicInteger remaining = new AtomicInteger(toClose);
          GenericFutureListener<Future<? super Void>> channelCloseListener =
              f -> {
                if (!f.isSuccess()) {
                  Loggers.warnWithException(LOG, "[{}] Error closing channel", logPrefix, f.cause());
                }
                if (remaining.decrementAndGet() == 0) {
                  closeFuture.complete(null);
                }
              };
          for (DriverChannel channel : channels) {
            eventBus.fire(ChannelEvent.channelClosed(node));
            channel.close().addListener(channelCloseListener);
          }
          for (DriverChannel channel : closingChannels) {
            // don't fire the close event, onChannelCloseStarted() already did it
            channel.closeFuture().addListener(channelCloseListener);
          }
        }
      }
  
      private void forceClose() {
        assert adminExecutor.inEventLoop();
        if (!isClosing) {
          close();
        }
        for (DriverChannel channel : channels) {
          channel.forceClose();
        }
        for (DriverChannel channel : closingChannels) {
          channel.forceClose();
        }
      }
  
      private int getConfiguredSize(NodeDistance distance) {
        return config
            .getDefaultProfile()
            .getInt(
                (distance == NodeDistance.LOCAL)
                    ? DefaultDriverOption.CONNECTION_POOL_LOCAL_SIZE
                    : DefaultDriverOption.CONNECTION_POOL_REMOTE_SIZE);
      }
    }
  }
  