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;
  
  import java.io.Closeable;
  import java.util.Collections;
  import java.util.List;
  
  import javax.annotation.Nullable;
  
  import io.netty.channel.Channel;
  import io.netty.channel.ChannelDuplexHandler;
  import io.netty.channel.ChannelInboundHandlerAdapter;
  import io.netty.channel.ChannelPipeline;
  import io.netty.channel.socket.SocketChannel;
  import io.netty.handler.ssl.SslHandler;
  import io.netty.handler.stream.ChunkedWriteHandler;
  import io.netty.handler.timeout.IdleStateHandler;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.common.metrics.source.AbstractSource;
  import org.apache.celeborn.common.network.client.TransportClient;
  import org.apache.celeborn.common.network.client.TransportClientBootstrap;
  import org.apache.celeborn.common.network.client.TransportClientFactory;
  import org.apache.celeborn.common.network.client.TransportResponseHandler;
  import org.apache.celeborn.common.network.protocol.MessageEncoder;
  import org.apache.celeborn.common.network.protocol.SslMessageEncoder;
  import org.apache.celeborn.common.network.server.*;
  import org.apache.celeborn.common.network.ssl.SSLFactory;
  import org.apache.celeborn.common.network.util.FrameDecoder;
  import org.apache.celeborn.common.network.util.NettyLogger;
  import org.apache.celeborn.common.network.util.TransportConf;
  import org.apache.celeborn.common.network.util.TransportFrameDecoder;
  
  /**
   * Contains the context to create a {@link TransportServer}, {@link TransportClientFactory}, and to
   * setup Netty Channel pipelines with a {@link TransportChannelHandler}.
   *
   * <p>There are two communication protocols that the TransportClient provides, control-plane RPCs
   * and data-plane "chunk fetching". The handling of the RPCs is performed outside the scope of the
   * TransportContext (i.e., by a user-provided handler), and it is responsible for setting up streams
   * which can be streamed through the data plane in chunks using zero-copy IO.
   *
   * <p>The TransportServer and TransportClientFactory both create a TransportChannelHandler for each
   * channel. As each TransportChannelHandler contains a TransportClient, this enables server
   * processes to send messages back to the client on an existing channel.
   */
  public class TransportContext implements Closeable {
    private static final Logger logger = LoggerFactory.getLogger(TransportContext.class);
  
    private static final NettyLogger nettyLogger = new NettyLogger();
    private final TransportConf conf;
    private final BaseMessageHandler msgHandler;
    private final ChannelDuplexHandler channelsLimiter;
    private final boolean closeIdleConnections;
    // Non-null if SSL is enabled, null otherwise.
    @Nullable private final SSLFactory sslFactory;
    private final boolean enableHeartbeat;
    private final AbstractSource source;
  
    private static final MessageEncoder ENCODER = MessageEncoder.INSTANCE;
    private static final SslMessageEncoder SSL_ENCODER = SslMessageEncoder.INSTANCE;
  
    public TransportContext(
        TransportConf conf,
        BaseMessageHandler msgHandler,
        boolean closeIdleConnections,
        ChannelDuplexHandler channelsLimiter,
        boolean enableHeartbeat,
        AbstractSource source) {
      this.conf = conf;
      this.msgHandler = msgHandler;
      this.closeIdleConnections = closeIdleConnections;
      this.sslFactory = SSLFactory.createSslFactory(conf);
      this.channelsLimiter = channelsLimiter;
      this.enableHeartbeat = enableHeartbeat;
      this.source = source;
  
      if (null != this.sslFactory) {
        logger.info(
            "SSL factory created for module {}, has keys ? {}",
            conf.getModuleName(),
            this.sslFactory.hasKeyManagers());
      } else {
        logger.info("SSL not enabled for module = {}", conf.getModuleName());
      }
    }
  
    public TransportContext(
        TransportConf conf,
        BaseMessageHandler msgHandler,
        boolean closeIdleConnections,
        boolean enableHeartbeat,
        AbstractSource source) {
      this(conf, msgHandler, closeIdleConnections, null, enableHeartbeat, source);
    }
  
    public TransportContext(
        TransportConf conf,
        BaseMessageHandler msgHandler,
        boolean closeIdleConnections,
        boolean enableHeartbeat,
        AbstractSource source,
        boolean collectFetchChunkDetailMetrics) {
      this(conf, msgHandler, closeIdleConnections, null, enableHeartbeat, source);
      if (collectFetchChunkDetailMetrics) {
        ENCODER.setSource(source);
      }
    }
  
    public TransportContext(
        TransportConf conf, BaseMessageHandler msgHandler, boolean closeIdleConnections) {
      this(conf, msgHandler, closeIdleConnections, null, false, null);
    }
  
    public TransportContext(TransportConf conf, BaseMessageHandler msgHandler) {
      this(conf, msgHandler, false, false, null);
    }
  
    public TransportClientFactory createClientFactory(List<TransportClientBootstrap> bootstraps) {
      return new TransportClientFactory(this, bootstraps);
    }
  
    public TransportClientFactory createClientFactory() {
      return createClientFactory(Collections.emptyList());
    }
  
    /** Create a server which will attempt to bind to a specific host and port. */
    public TransportServer createServer(String host, int port) {
      return new TransportServer(this, host, port, Collections.emptyList());
    }
  
    public TransportServer createServer(
        String host, int port, List<TransportServerBootstrap> bootstraps) {
      return new TransportServer(this, host, port, bootstraps);
    }
  
    public TransportServer createServer(List<TransportServerBootstrap> bootstraps) {
      return createServer(null, 0, bootstraps);
    }
  
    public TransportServer createServer(int port) {
      return createServer(null, port, Collections.emptyList());
    }
  
    public TransportServer createServer(int port, List<TransportServerBootstrap> bootstraps) {
      return createServer(null, port, bootstraps);
    }
  
    /** For Suite only */
    public TransportServer createServer() {
      return createServer(null, 0, Collections.emptyList());
    }
  
    public boolean sslEncryptionEnabled() {
      return this.sslFactory != null;
    }
  
    public TransportChannelHandler initializePipeline(
        SocketChannel channel, ChannelInboundHandlerAdapter decoder, boolean isClient) {
      return initializePipeline(channel, decoder, msgHandler, isClient);
    }
  
    public TransportChannelHandler initializePipeline(
        SocketChannel channel, BaseMessageHandler resolvedMsgHandler, boolean isClient) {
      return initializePipeline(channel, new TransportFrameDecoder(), resolvedMsgHandler, isClient);
    }
  
    public TransportChannelHandler initializePipeline(
        SocketChannel channel,
        ChannelInboundHandlerAdapter decoder,
        BaseMessageHandler resolvedMsgHandler,
        boolean isClient) {
      try {
        ChannelPipeline pipeline = channel.pipeline();
        if (nettyLogger.getLoggingHandler() != null) {
          pipeline.addLast("loggingHandler", nettyLogger.getLoggingHandler());
        }
        if (sslEncryptionEnabled()) {
          if (!isClient && !sslFactory.hasKeyManagers()) {
            throw new IllegalStateException("Not a client connection and no keys configured");
          }
  
          SslHandler sslHandler;
          try {
            sslHandler = new SslHandler(sslFactory.createSSLEngine(isClient, channel.alloc()));
            sslHandler.setHandshakeTimeoutMillis(conf.sslHandshakeTimeoutMs());
          } catch (Exception e) {
            throw new IllegalStateException("Error creating Netty SslHandler", e);
          }
          pipeline.addFirst("NettySslEncryptionHandler", sslHandler);
          // Cannot use zero-copy with HTTPS, so we add in our ChunkedWriteHandler just before the
          // MessageEncoder
          pipeline.addLast("chunkedWriter", new ChunkedWriteHandler());
        }
  
        if (channelsLimiter != null) {
          pipeline.addLast("limiter", channelsLimiter);
        }
        TransportChannelHandler channelHandler = createChannelHandler(channel, resolvedMsgHandler);
        pipeline
            .addLast("encoder", sslEncryptionEnabled() ? SSL_ENCODER : ENCODER)
            .addLast(FrameDecoder.HANDLER_NAME, decoder)
            .addLast(
                "idleStateHandler",
                enableHeartbeat
                    ? new IdleStateHandler(conf.connectionTimeoutMs() / 1000, 0, 0)
                    : new IdleStateHandler(0, 0, conf.connectionTimeoutMs() / 1000))
            .addLast("handler", channelHandler);
        return channelHandler;
      } catch (RuntimeException e) {
        logger.error("Error while initializing Netty pipeline", e);
        throw e;
      }
    }
  
    private TransportChannelHandler createChannelHandler(
        Channel channel, BaseMessageHandler msgHandler) {
      TransportResponseHandler responseHandler = new TransportResponseHandler(conf, channel);
      TransportClient client = new TransportClient(channel, responseHandler);
      TransportRequestHandler requestHandler =
          new TransportRequestHandler(channel, client, msgHandler);
      return new TransportChannelHandler(
          client,
          responseHandler,
          requestHandler,
          conf.connectionTimeoutMs(),
          closeIdleConnections,
          enableHeartbeat,
          conf.channelHeartbeatInterval(),
          this);
    }
  
    public TransportConf getConf() {
      return conf;
    }
  
    public BaseMessageHandler getMsgHandler() {
      return msgHandler;
    }
  
    public AbstractSource getSource() {
      return source;
    }
  
    @Override
    public void close() {
      if (sslFactory != null) {
        sslFactory.destroy();
      }
    }
  }
  