content: /*
   * Copyright 2021 Alibaba Group Holding Limited.
   *
   * Licensed 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.alibaba.graphscope.parallel.netty;
  
  import static com.alibaba.graphscope.parallel.netty.handler.NettyServerHandler.RESPONSE_BYTES;
  
  import static org.apache.giraph.conf.GiraphConstants.CLIENT_RECEIVE_BUFFER_SIZE;
  import static org.apache.giraph.conf.GiraphConstants.CLIENT_SEND_BUFFER_SIZE;
  import static org.apache.giraph.conf.GiraphConstants.MAX_CONN_TRY_ATTEMPTS;
  import static org.apache.giraph.utils.ByteUtils.SIZE_OF_BYTE;
  import static org.apache.giraph.utils.ByteUtils.SIZE_OF_INT;
  
  import com.alibaba.graphscope.parallel.netty.handler.NettyClientHandler;
  import com.alibaba.graphscope.parallel.netty.request.WritableRequest;
  import com.alibaba.graphscope.parallel.netty.request.serialization.WritableRequestEncoder;
  import com.alibaba.graphscope.parallel.utils.NetworkMap;
  import com.google.common.collect.MapMaker;
  
  import io.netty.bootstrap.Bootstrap;
  import io.netty.channel.Channel;
  import io.netty.channel.ChannelFuture;
  import io.netty.channel.ChannelHandlerContext;
  import io.netty.channel.ChannelInitializer;
  import io.netty.channel.ChannelOption;
  import io.netty.channel.ChannelPipeline;
  import io.netty.channel.EventLoopGroup;
  import io.netty.channel.nio.NioEventLoopGroup;
  import io.netty.channel.socket.SocketChannel;
  import io.netty.channel.socket.nio.NioSocketChannel;
  import io.netty.handler.codec.FixedLengthFrameDecoder;
  
  import org.apache.giraph.conf.GiraphConstants;
  import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;
  import org.apache.giraph.utils.ThreadUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.lang.Thread.UncaughtExceptionHandler;
  import java.net.InetSocketAddress;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Objects;
  import java.util.concurrent.TimeUnit;
  
  public class NettyClient {
  
      /**
       * 30 seconds to connect by default
       */
      public static final int MAX_CONNECTION_MILLISECONDS_DEFAULT = 30 * 1000;
  
      private static Logger logger = LoggerFactory.getLogger(NettyClient.class);
      /**
       * Send buffer size
       */
      private final int sendBufferSize;
      /**
       * Receive buffer size
       */
      private final int receiveBufferSize;
      /**
       * Warn if request size is bigger than the buffer size by this factor
       */
      private final float requestSizeWarningThreshold;
      /**
       * Maximum thread pool size
       */
      private final int maxPoolSize;
      /**
       * Key: worker id, value: ip:port
       */
      private final Map<Integer, InetSocketAddress> workerId2Address = new MapMaker().makeMap();
  
      private ImmutableClassesGiraphConfiguration conf;
      /**
       * Giraph using waiting connections to characterize this list. We don't do so since it is
       * unnecessary.
       */
      private Connection[] connections;
      /**
       * All connected channels. length = workerNum, [index] = null
       */
      private Channel[] channels;
  
      private NettyClientHandler[] handlers;
      //    private Map<Integer, LinkedList<ChannelFuture>> pendingRequests;
      /**
       * Map between dst frag id -> number of requests sent.
       */
      private Map<Integer, Integer> pendingRequests;
  
      private NetworkMap networkMap;
  
      private EventLoopGroup workGroup;
  
      private Bootstrap bootstrap;
  
      private int workerId;
  
      public NettyClient(
              ImmutableClassesGiraphConfiguration conf,
              NetworkMap networkMap,
              final UncaughtExceptionHandler exceptionHandler) {
          this.conf = conf;
          this.networkMap = networkMap;
          workerId = networkMap.getSelfWorkerId();
          /** Init constants */
          /** Number of threads for client to use, i.e. number of handlers*/
          maxPoolSize = GiraphConstants.NETTY_CLIENT_THREADS.get(conf);
          sendBufferSize = CLIENT_SEND_BUFFER_SIZE.get(conf) + SIZE_OF_INT + SIZE_OF_BYTE;
          receiveBufferSize = CLIENT_RECEIVE_BUFFER_SIZE.get(conf);
          requestSizeWarningThreshold = GiraphConstants.REQUEST_SIZE_WARNING_THRESHOLD.get(conf);
          /**
           * Don't connect to self. But as we send message using index, we still make it length=workerNum,
           * but channels[workerId] = null and connections[workerId] = null
           */
          connections = new Connection[networkMap.getWorkerNum()];
          channels = new Channel[networkMap.getWorkerNum()];
          handlers = new NettyClientHandler[networkMap.getWorkerNum()];
          pendingRequests = new HashMap<>();
          for (int i = 0; i < networkMap.getWorkerNum(); ++i) {
              //            pendingRequests.put(i, new LinkedList<ChannelFuture>());
              pendingRequests.put(i, 0);
          }
          /**
           * Start the client. connect to all address.
           */
          startClient(exceptionHandler);
      }
  
      void startClient(final UncaughtExceptionHandler exceptionHandler) {
          workGroup =
                  new NioEventLoopGroup(
                          1,
                          ThreadUtils.createThreadFactory(
                                  "netty-client-worker-" + networkMap.getSelfWorkerId() + "-%d",
                                  exceptionHandler));
  
          bootstrap = new Bootstrap();
          bootstrap
                  .option(ChannelOption.SO_KEEPALIVE, true)
                  .option(ChannelOption.CONNECT_TIMEOUT_MILLIS, MAX_CONNECTION_MILLISECONDS_DEFAULT)
                  .option(ChannelOption.TCP_NODELAY, true)
                  .option(ChannelOption.SO_SNDBUF, sendBufferSize)
                  .option(ChannelOption.SO_RCVBUF, receiveBufferSize)
                  .option(ChannelOption.ALLOCATOR, conf.getNettyAllocator())
                  .group(workGroup)
                  .channel(NioSocketChannel.class)
                  .handler(
                          new ChannelInitializer<SocketChannel>() {
                              @Override
                              protected void initChannel(SocketChannel ch) throws Exception {
                                  ChannelPipeline p = ch.pipeline();
                                  p.addLast(new FixedLengthFrameDecoder(RESPONSE_BYTES));
                                  p.addLast(new WritableRequestEncoder(conf));
                                  //                        p.addLast(new
                                  // WritableRequestDecoder(conf));
                                  //                        p.addLast()
                                  p.addLast(new NettyClientHandler(workerId));
                              }
  
                              @Override
                              public void channelUnregistered(ChannelHandlerContext ctx)
                                      throws Exception {
                                  super.channelUnregistered(ctx);
                                  logger.error(
                                          "NettyClient [{}]: Channel failed {}",
                                          workerId,
                                          ctx.channel());
                                  //
                                  // checkRequestsAfterChannelFailure(ctx.channel());
                              }
                          });
      }
  
      /**
       * Put connectToAllAddress out size of Constructor.
       */
      public void connectToAllAddress() {
          for (int dstWorkerId = 0; dstWorkerId < networkMap.getWorkerNum(); ++dstWorkerId) {
              if (dstWorkerId == networkMap.getSelfWorkerId()) {
                  connections[dstWorkerId] = null;
                  continue;
              }
              String hostName = networkMap.getHostNameForWorker(dstWorkerId);
              int port = networkMap.getPortForWorker(dstWorkerId);
              InetSocketAddress dstAddress = resolveAddress(hostName, port);
              // There are no duplicated connections in our settings.
              workerId2Address.put(dstWorkerId, dstAddress);
              logger.debug(
                      "NettyClient [{}]: Resolved address for worker: {}, dstAddr: {}",
                      workerId,
                      dstWorkerId,
                      dstAddress);
  
              ChannelFuture channelFuture = bootstrap.connect(dstAddress);
              connections[dstWorkerId] = new Connection(channelFuture, dstAddress, dstWorkerId);
          }
          waitAllConnections();
          logger.info("NettyClient [{}]: All connection established!", workerId);
      }
  
      /**
       * Wait for all connections established.
       */
      private void waitAllConnections() {
          int maxTries = MAX_CONN_TRY_ATTEMPTS.get(conf);
          int index = 0;
          int successCnt = 1;
          while (successCnt < connections.length) {
              if (Objects.nonNull(connections[index])) {
                  Connection connection = connections[index];
                  logger.info(
                          "NettyClient [{}]: try for connection to {} while success connection cnt:"
                                  + " {}",
                          workerId,
                          connection.dstWorkerId,
                          successCnt);
  
                  int failedCnt = 0;
  
                  Channel channel = null;
                  while (failedCnt < maxTries) {
                      // If this future fails, we need to update the future.
                      ChannelFuture future = connection.future;
                      boolean res = false;
                      try {
                          res = future.await(1, TimeUnit.SECONDS);
                      } catch (Exception e) {
                          throw new IllegalStateException("Future waiting exception");
                      }
                      if (res && future.isSuccess() && future.channel().isOpen()) {
                          logger.info(
                                  "NettyClient [{}]: connection: {} success for after {} times",
                                  workerId,
                                  connection,
                                  failedCnt);
                          channel = future.channel();
                          break;
                      } else {
                          logger.warn(
                                  "NettyClient [{}]: Failed connection to [{}], tries: {}/{},"
                                          + " success: {}, opened: {}, cause: {}",
                                  workerId,
                                  index,
                                  failedCnt,
                                  maxTries,
                                  future.isSuccess(),
                                  future.channel().isOpen(),
                                  future.cause());
                          try {
                              TimeUnit.SECONDS.sleep(1);
                              failedCnt += 1;
                              // When encounter failure, we update the future;
                              ChannelFuture newFuture = bootstrap.connect(connection.address);
                              connection.updateFuture(newFuture);
                          } catch (InterruptedException e) {
                              throw new IllegalStateException(
                                      "Interrupted when waiting sleep: " + connection.address);
                          }
                      }
                  }
                  if (Objects.isNull(channel)) {
                      logger.warn(
                              "NettyClient [{}]: Skip connection {} for next time. try others",
                              workerId,
                              connection);
                  } else {
                      successCnt += 1;
                      channels[index] = channel;
                      handlers[index] = (NettyClientHandler) channel.pipeline().last();
                  }
              } else {
                  logger.info(
                          "NettyClient [{}]: Connection to self is not needed [{}]", workerId, index);
              }
              index = (index + 1) % connections.length;
          }
          checkChannels();
      }
  
      private void checkChannels() {
          for (int i = 0; i < channels.length; ++i) {
              if (i != workerId && Objects.isNull(channels[i])) {
                  throw new IllegalStateException(
                          "Found invalid channel: " + channels[i] + " open:" + channels[i].isOpen());
              }
          }
      }
  
      /**
       * Verify whether an address is reachable.
       *
       * @param hostNameOrIp hostname or ip
       * @param port         dst port
       * @return resolved address
       */
      private InetSocketAddress resolveAddress(String hostNameOrIp, int port) {
          int resolveAttempts = 0;
          InetSocketAddress address = new InetSocketAddress(hostNameOrIp, port);
          while (address.isUnresolved() && resolveAttempts < conf.getInetAddressMaxResolveTime()) {
              ++resolveAttempts;
              logger.warn(
                      "resolveAddress: Failed to resolve "
                              + address
                              + " on attempt "
                              + resolveAttempts
                              + " of "
                              + conf.getInetAddressMaxResolveTime()
                              + " attempts, sleeping for 1 seconds");
              ThreadUtils.trySleep(1000);
              address = new InetSocketAddress(hostNameOrIp, address.getPort());
          }
          if (resolveAttempts >= conf.getInetAddressMaxResolveTime()) {
              throw new IllegalStateException(
                      "resolveAddress: Couldn't "
                              + "resolve "
                              + address
                              + " in "
                              + resolveAttempts
                              + " tries.");
          }
          return address;
      }
  
      public void sendMessage(int dstFragId, WritableRequest request) {
          if (dstFragId == workerId) {
              throw new IllegalStateException("Shouldn't reach here:" + dstFragId + ", " + workerId);
          }
          //        ChannelFuture requestFuture = channels[dstFragId].writeAndFlush(request);
          ChannelFuture requestFuture = channels[dstFragId].write(request);
          //        pendingRequests.get(dstFragId).offer(requestFuture);
          // Must already been initialized to 0 in presuperstep.
          pendingRequests.put(dstFragId, pendingRequests.get(dstFragId) + 1);
          logger.debug(
                  "NettyClient [{}]: send msg {} to [{}], corresponding pending request: {}",
                  workerId,
                  request,
                  dstFragId,
                  pendingRequests.get(dstFragId));
      }
  
      /**
       * Hide flush from message manager. From message manager view, all request send immediately.
       */
      private void flushMessages() {
          logger.debug("NettyClient [{}]: flushing messages", workerId);
          for (int i = 0; i < networkMap.getWorkerNum(); ++i) {
              if (i == workerId) {
                  continue;
              }
              channels[i].flush();
          }
          logger.debug("NettyClient [{}]: finish flushing messages", workerId);
      }
  
      public void postSuperStep() {
          for (int i = 0; i < networkMap.getWorkerNum(); ++i) {
              pendingRequests.put(i, 0);
              if (i != workerId) {
                  handlers[i].postSuperStep();
              }
          }
      }
  
      public void waitAllRequests() {
          flushMessages();
          logger.info("NettyClient [{}]: Message sending success, waiting for response", workerId);
          for (int i = 0; i < networkMap.getWorkerNum(); ++i) {
              if (i == workerId) {
                  continue;
              }
              NettyClientHandler handler = handlers[i];
              //            handler.waitForResponse(pendingRequests.get(i).size());
              handler.waitForResponse(pendingRequests.get(i));
              if (logger.isDebugEnabled()) {
                  logger.debug("NettyClient [{}]: finished waiting response from [{}]", workerId, i);
              }
          }
          logger.info("NettyClient [{}]: finish waiting sending all messages", workerId);
      }
  
      @Override
      public String toString() {
          String res = "NettyClient: [" + workerId + "] channels:";
          for (Channel channel : channels) {
              if (Objects.nonNull(channel)) {
                  res += channel.remoteAddress() + ",";
              }
          }
          return res;
      }
  
      public void close() {
          waitAllRequests();
          try {
              for (int i = 0; i < channels.length; ++i) {
                  if (Objects.nonNull(channels[i])) {
                      channels[i].close();
                  }
              }
          } catch (Exception e) {
              e.printStackTrace();
          }
          workGroup.shutdownGracefully();
          logger.info("NettyClient [{}]: Closing...", workerId);
      }
  
      /**
       * The abstraction of connection.
       */
      private static class Connection {
  
          /**
           * Address of the future
           */
          private final InetSocketAddress address;
          /**
           * grape worker id(MPI_COMM)
           */
          private final Integer dstWorkerId;
          /**
           * Future object
           */
          private ChannelFuture future;
  
          /**
           * Constructor.
           *
           * @param future   Immutable future
           * @param address  Immutable address
           * @param workerId Immutable taskId
           */
          Connection(ChannelFuture future, InetSocketAddress address, Integer workerId) {
              this.future = future;
              this.address = address;
              this.dstWorkerId = workerId;
          }
  
          /**
           * We can update the future but don't update address.
           *
           * @param future future obj.
           */
          public void updateFuture(ChannelFuture future) {
              this.future = future;
          }
  
          @Override
          public String toString() {
              return "(future="
                      + future
                      + ",address="
                      + address
                      + ",dstWorkerId="
                      + dstWorkerId
                      + ")";
          }
      }
  }
  