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.celeborn.common.network.client;
  
  import java.io.Closeable;
  import java.io.IOException;
  import java.net.InetSocketAddress;
  import java.net.SocketAddress;
  import java.util.List;
  import java.util.Random;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicReference;
  import java.util.function.Supplier;
  
  import com.google.common.base.Preconditions;
  import com.google.common.base.Throwables;
  import com.google.common.collect.Lists;
  import io.netty.bootstrap.Bootstrap;
  import io.netty.buffer.ByteBufAllocator;
  import io.netty.channel.*;
  import io.netty.channel.socket.SocketChannel;
  import io.netty.handler.ssl.SslHandler;
  import io.netty.util.concurrent.Future;
  import io.netty.util.concurrent.GenericFutureListener;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.common.client.MasterNotLeaderException;
  import org.apache.celeborn.common.exception.CelebornIOException;
  import org.apache.celeborn.common.network.TransportContext;
  import org.apache.celeborn.common.network.sasl.registration.RegistrationClientBootstrap;
  import org.apache.celeborn.common.network.server.TransportChannelHandler;
  import org.apache.celeborn.common.network.util.*;
  import org.apache.celeborn.common.util.JavaUtils;
  import org.apache.celeborn.common.util.Utils;
  
  /**
   * Factory for creating {@link TransportClient}s by using createClient.
   *
   * <p>The factory maintains a connection pool to other hosts and should return the same
   * TransportClient for the same remote host. It also shares a single worker thread pool for all
   * TransportClients.
   *
   * <p>TransportClients will be reused whenever possible.
   */
  public class TransportClientFactory implements Closeable {
  
    /** A simple data structure to track the pool of clients between two peer nodes. */
    private static class ClientPool {
      TransportClient[] clients;
      Object[] locks;
  
      ClientPool(int size) {
        clients = new TransportClient[size];
        locks = new Object[size];
        for (int i = 0; i < size; i++) {
          locks[i] = new Object();
        }
      }
    }
  
    private static final Logger logger = LoggerFactory.getLogger(TransportClientFactory.class);
  
    private final TransportContext context;
    private final List<TransportClientBootstrap> clientBootstraps;
    private final ConcurrentHashMap<SocketAddress, ClientPool> connectionPool;
  
    /** Random number generator for picking connections between peers. */
    private final Random rand;
  
    private final int numConnectionsPerPeer;
  
    private final int connectTimeoutMs;
    private final int connectionTimeoutMs;
    private final int sslHandshakeTimeoutMs;
  
    private final int receiveBuf;
  
    private final int sendBuf;
    private final Class<? extends Channel> socketChannelClass;
    private EventLoopGroup workerGroup;
    protected ByteBufAllocator allocator;
    private final int maxClientConnectRetries;
    private final int maxClientConnectRetryWaitTimeMs;
  
    public TransportClientFactory(
        TransportContext context, List<TransportClientBootstrap> clientBootstraps) {
      this.context = Preconditions.checkNotNull(context);
      TransportConf conf = context.getConf();
      this.clientBootstraps = Lists.newArrayList(Preconditions.checkNotNull(clientBootstraps));
      this.connectionPool = JavaUtils.newConcurrentHashMap();
      this.numConnectionsPerPeer = conf.numConnectionsPerPeer();
      this.connectTimeoutMs = conf.connectTimeoutMs();
      this.connectionTimeoutMs = conf.connectionTimeoutMs();
      this.sslHandshakeTimeoutMs = conf.sslHandshakeTimeoutMs();
      this.receiveBuf = conf.receiveBuf();
      this.sendBuf = conf.sendBuf();
      this.rand = new Random();
  
      IOMode ioMode = IOMode.valueOf(conf.ioMode());
      this.socketChannelClass = NettyUtils.getClientChannelClass(ioMode);
      logger.info("Module {} mode {} threads {}", conf.getModuleName(), ioMode, conf.clientThreads());
      this.workerGroup =
          NettyUtils.createEventLoop(
              ioMode,
              conf.clientThreads(),
              conf.conflictAvoidChooserEnable(),
              conf.getModuleName() + "-client");
      // Always disable thread-local cache when creating pooled ByteBuf allocator for TransportClients
      // because the ByteBufs are allocated by the event loop thread, but released by the executor
      // thread rather than the event loop thread. Those thread-local caches actually delay the
      // recycling of buffers, leading to larger memory usage.
      this.allocator =
          NettyUtils.getByteBufAllocator(conf, context.getSource(), false, conf.clientThreads());
      this.maxClientConnectRetries = conf.maxIORetries();
      this.maxClientConnectRetryWaitTimeMs = conf.ioRetryWaitTimeMs();
    }
  
    /**
     * Create a {@link TransportClient} connecting to the given remote host / port.
     *
     * <p>We maintains an array of clients (size determined by
     * celeborn.$module.io.numConnectionsPerPeer) and randomly picks one to use. If no client was
     * previously created in the randomly selected spot, this function creates a new client and places
     * it there.
     *
     * <p>This blocks until a connection is successfully established and fully bootstrapped.
     *
     * <p>Concurrency: This method is safe to call from multiple threads.
     */
    public TransportClient createClient(String remoteHost, int remotePort, int partitionId)
        throws IOException, InterruptedException {
      return retryCreateClient(remoteHost, remotePort, partitionId, TransportFrameDecoder::new);
    }
  
    public TransportClient retryCreateClient(
        String remoteHost,
        int remotePort,
        int partitionId,
        Supplier<ChannelInboundHandlerAdapter> supplier)
        throws IOException, InterruptedException {
      int numTries = 0;
      while (numTries < maxClientConnectRetries) {
        try {
          return createClient(remoteHost, remotePort, partitionId, supplier.get());
        } catch (Exception e) {
          if (e instanceof InterruptedException) {
            Thread.currentThread().interrupt();
            throw e;
          }
          numTries++;
          logger.warn(
              "Retry create client, times {}/{} with error: {}",
              numTries,
              maxClientConnectRetries,
              e.getMessage(),
              e);
          if (numTries == maxClientConnectRetries) {
            throw e;
          }
  
          Thread.sleep(maxClientConnectRetryWaitTimeMs);
        }
      }
  
      return null;
    }
  
    public TransportClient createClient(
        String remoteHost, int remotePort, int partitionId, ChannelInboundHandlerAdapter decoder)
        throws IOException, InterruptedException {
      // Get connection from the connection pool first.
      // If it is not found or not active, create a new one.
      // Use unresolved address here to avoid DNS resolution each time we create a client.
      final InetSocketAddress unresolvedAddress =
          InetSocketAddress.createUnresolved(remoteHost, remotePort);
  
      // Create the ClientPool if we don't have it yet.
      ClientPool clientPool =
          connectionPool.computeIfAbsent(
              unresolvedAddress, key -> new ClientPool(numConnectionsPerPeer));
      int clientIndex =
          partitionId < 0 ? rand.nextInt(numConnectionsPerPeer) : partitionId % numConnectionsPerPeer;
      TransportClient cachedClient = clientPool.clients[clientIndex];
  
      if (cachedClient != null && cachedClient.isActive()) {
        // Make sure that the channel will not timeout by updating the last use time of the
        // handler. Then check that the client is still alive, in case it timed out before
        // this code was able to update things.
        TransportChannelHandler handler =
            cachedClient.getChannel().pipeline().get(TransportChannelHandler.class);
        if (handler != null) {
          synchronized (handler) {
            handler.getResponseHandler().updateTimeOfLastRequest();
          }
        }
  
        if (cachedClient.isActive()) {
          logger.debug(
              "Returning cached connection from {} to {}: {}",
              cachedClient.getChannel().localAddress(),
              cachedClient.getSocketAddress(),
              cachedClient);
          return cachedClient;
        }
      }
  
      // If we reach here, we don't have an existing connection open. Let's create a new one.
      // Multiple threads might race here to create new connections. Keep only one of them active.
      final long preResolveHost = System.nanoTime();
      final InetSocketAddress resolvedAddress = new InetSocketAddress(remoteHost, remotePort);
      final long hostResolveTimeMs =
          TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - preResolveHost);
      final String resolveMsg = resolvedAddress.isUnresolved() ? "failed" : "succeed";
      if (hostResolveTimeMs > 2000) {
        logger.warn(
            "DNS resolution {} for {} took {} ms", resolveMsg, resolvedAddress, hostResolveTimeMs);
      } else {
        logger.trace(
            "DNS resolution {} for {} took {} ms", resolveMsg, resolvedAddress, hostResolveTimeMs);
      }
  
      synchronized (clientPool.locks[clientIndex]) {
        cachedClient = clientPool.clients[clientIndex];
  
        if (cachedClient != null) {
          if (cachedClient.isActive()) {
            logger.debug(
                "Returning cached connection from {} to {}: {}",
                cachedClient.getChannel().localAddress(),
                resolvedAddress,
                cachedClient);
            return cachedClient;
          } else {
            logger.info("Found inactive connection to {}, creating a new one.", resolvedAddress);
          }
        }
        clientPool.clients[clientIndex] = internalCreateClient(resolvedAddress, decoder);
        return clientPool.clients[clientIndex];
      }
    }
  
    public TransportClient createClient(String remoteHost, int remotePort)
        throws IOException, InterruptedException {
      return createClient(remoteHost, remotePort, -1);
    }
  
    /**
     * Create a completely new {@link TransportClient} to the given remote host / port. This
     * connection is not pooled.
     *
     * <p>As with {@link #createClient(String, int)}, this method is blocking.
     */
    private TransportClient internalCreateClient(
        InetSocketAddress address, ChannelInboundHandlerAdapter decoder)
        throws IOException, InterruptedException {
      Bootstrap bootstrap = new Bootstrap();
      bootstrap
          .group(workerGroup)
          .channel(socketChannelClass)
          // Disable Nagle's Algorithm since we don't want packets to wait
          .option(ChannelOption.TCP_NODELAY, true)
          .option(ChannelOption.SO_KEEPALIVE, true)
          .option(ChannelOption.CONNECT_TIMEOUT_MILLIS, connectTimeoutMs)
          .option(ChannelOption.ALLOCATOR, allocator);
  
      if (receiveBuf > 0) {
        bootstrap.option(ChannelOption.SO_RCVBUF, receiveBuf);
      }
  
      if (sendBuf > 0) {
        bootstrap.option(ChannelOption.SO_SNDBUF, sendBuf);
      }
  
      final AtomicReference<TransportClient> clientRef = new AtomicReference<>();
  
      bootstrap.handler(
          new ChannelInitializer<SocketChannel>() {
            @Override
            public void initChannel(SocketChannel ch) {
              TransportChannelHandler clientHandler = context.initializePipeline(ch, decoder, true);
              clientRef.set(clientHandler.getClient());
            }
          });
  
      // Connect to the remote server
      long preConnect = System.nanoTime();
      ChannelFuture cf = bootstrap.connect(address);
      if (connectTimeoutMs <= 0) {
        cf.await();
        assert cf.isDone();
        if (cf.isCancelled()) {
          throw new IOException(String.format("Connecting to %s cancelled", address));
        } else if (!cf.isSuccess()) {
          throw new IOException(String.format("Failed to connect to %s", address), cf.cause());
        }
      } else if (!cf.await(connectTimeoutMs)) {
        closeChannel(cf);
        throw new CelebornIOException(
            String.format("Connecting to %s timed out (%s ms)", address, connectTimeoutMs));
      } else if (cf.cause() != null) {
        closeChannel(cf);
        throw new CelebornIOException(String.format("Failed to connect to %s", address), cf.cause());
      }
      if (context.sslEncryptionEnabled()) {
        final SslHandler sslHandler = cf.channel().pipeline().get(SslHandler.class);
        sslHandler.setHandshakeTimeoutMillis(sslHandshakeTimeoutMs);
        Future<Channel> future =
            sslHandler
                .handshakeFuture()
                .addListener(
                    new GenericFutureListener<Future<Channel>>() {
                      @Override
                      public void operationComplete(final Future<Channel> handshakeFuture) {
                        if (handshakeFuture.isSuccess()) {
                          logger.debug("successfully completed TLS handshake to {}", address);
                        } else {
                          logger.info(
                              "failed to complete TLS handshake to {}",
                              address,
                              handshakeFuture.cause());
                          closeChannel(cf);
                        }
                      }
                    });
        if (!future.await(connectionTimeoutMs)) {
          closeChannel(cf);
          throw new IOException(
              String.format("Failed to connect to %s within connection timeout", address));
        }
      }
  
      TransportClient client = clientRef.get();
      assert client != null : "Channel future completed successfully with null client";
  
      // Execute any client bootstraps synchronously before marking the Client as successful.
      long preBootstrap = System.nanoTime();
      logger.debug("Running bootstraps for {} ...", address);
      for (TransportClientBootstrap clientBootstrap : clientBootstraps) {
        try {
          clientBootstrap.doBootstrap(client);
        } catch (
            Exception e) { // catch non-RuntimeExceptions too as bootstrap may be written in Scala
          long bootstrapTime = System.nanoTime() - preBootstrap;
          if (clientBootstrap instanceof RegistrationClientBootstrap) {
            Exception processed = RegistrationClientBootstrap.processMasterNotLeaderException(e);
            String message =
                (processed instanceof MasterNotLeaderException)
                    ? String.format(
                        "Suggested leader is %s",
                        ((MasterNotLeaderException) processed).getSuggestedLeaderAddress())
                    : e.getMessage();
            logger.warn(
                "Attempted to register with a Master that is not the leader after {}: {}",
                Utils.nanoDurationToString(bootstrapTime),
                message);
          } else {
            logger.error(
                "Exception while bootstrapping client after {}",
                Utils.nanoDurationToString(bootstrapTime),
                e);
          }
          client.close();
          Throwables.throwIfUnchecked(e);
          throw new RuntimeException(e);
        }
      }
      long postBootstrap = System.nanoTime();
      logger.debug(
          "Successfully created connection to {} after {} ({} spent in bootstraps)",
          address,
          Utils.nanoDurationToString(postBootstrap - preConnect),
          Utils.nanoDurationToString(postBootstrap - preBootstrap));
  
      return client;
    }
  
    /** Close all connections in the connection pool, and shutdown the worker thread pool. */
    @Override
    public void close() {
      // Go through all clients and close them if they are active.
      for (ClientPool clientPool : connectionPool.values()) {
        for (int i = 0; i < clientPool.clients.length; i++) {
          TransportClient client = clientPool.clients[i];
          if (client != null) {
            clientPool.clients[i] = null;
            JavaUtils.closeQuietly(client);
          }
        }
      }
      connectionPool.clear();
  
      // SPARK-19147
      if (workerGroup != null && !workerGroup.isShuttingDown()) {
        workerGroup.shutdownGracefully();
      }
    }
  
    public TransportContext getContext() {
      return context;
    }
  
    private void closeChannel(ChannelFuture channelFuture) {
      try {
        channelFuture.channel().close();
      } catch (Exception e) {
        logger.warn("Failed to close channel", e);
      }
    }
  }
  