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.server;
  
  import java.io.IOException;
  import java.net.SocketAddress;
  import java.nio.ByteBuffer;
  
  import com.google.common.base.Throwables;
  import io.netty.channel.Channel;
  import io.netty.channel.ChannelFuture;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.common.network.buffer.NioManagedBuffer;
  import org.apache.celeborn.common.network.client.RpcResponseCallback;
  import org.apache.celeborn.common.network.client.TransportClient;
  import org.apache.celeborn.common.network.protocol.*;
  import org.apache.celeborn.common.network.util.NettyUtils;
  
  /**
   * A handler that processes requests from clients and writes chunk data back. Each handler is
   * attached to a single Netty channel, and keeps track of which streams have been fetched via this
   * channel, in order to clean them up if the channel is terminated (see #channelUnregistered).
   *
   * <p>The messages should have been processed by the pipeline setup by {@link TransportServer}.
   */
  public class TransportRequestHandler extends MessageHandler<RequestMessage> {
  
    private static final Logger logger = LoggerFactory.getLogger(TransportRequestHandler.class);
  
    /** The Netty channel that this handler is associated with. */
    private final Channel channel;
  
    /** Client on the same channel allowing us to talk back to the requester. */
    private final TransportClient reverseClient;
  
    /** Handles all RPC messages. */
    private final BaseMessageHandler msgHandler;
  
    public TransportRequestHandler(
        Channel channel, TransportClient reverseClient, BaseMessageHandler msgHandler) {
      this.channel = channel;
      this.reverseClient = reverseClient;
      this.msgHandler = msgHandler;
    }
  
    @Override
    public void exceptionCaught(Throwable cause) {
      msgHandler.exceptionCaught(cause, reverseClient);
    }
  
    @Override
    public void channelActive() {
      msgHandler.channelActive(reverseClient);
    }
  
    @Override
    public void channelInactive() {
      msgHandler.channelInactive(reverseClient);
    }
  
    @Override
    public void handle(RequestMessage request) {
      if (logger.isTraceEnabled()) {
        logger.trace("Received request {} from {}", request.getClass().getName(), reverseClient);
      }
      if (checkRegistered(request)) {
        if (request instanceof RpcRequest) {
          processRpcRequest((RpcRequest) request);
        } else if (request instanceof OneWayMessage) {
          processOneWayMessage((OneWayMessage) request);
        } else {
          processOtherMessages(request);
        }
      } else {
        if (request.body() != null) {
          request.body().release();
        }
      }
    }
  
    protected void processHeartbeat() {
      try {
        logger.trace("Process heartbeat from {}.", NettyUtils.getRemoteAddress(channel));
        respond(new Heartbeat());
      } catch (Exception e) {
        logger.error(
            "Error while invoking handler#respond() on heartbeat from {}.",
            NettyUtils.getRemoteAddress(channel),
            e);
      }
    }
  
    private void processRpcRequest(final RpcRequest req) {
      try {
        logger.trace("Process rpc request {}", req.requestId);
        msgHandler.receive(
            reverseClient,
            req,
            new RpcResponseCallback() {
              @Override
              public void onSuccess(ByteBuffer response) {
                respond(new RpcResponse(req.requestId, new NioManagedBuffer(response)));
              }
  
              @Override
              public void onFailure(Throwable e) {
                respond(new RpcFailure(req.requestId, Throwables.getStackTraceAsString(e)));
              }
            });
      } catch (Exception e) {
        logger.error(
            "Error while invoking handler#receive() on RPC id {} from {}",
            req.requestId,
            NettyUtils.getRemoteAddress(channel),
            e);
        respond(new RpcFailure(req.requestId, Throwables.getStackTraceAsString(e)));
      } finally {
        req.body().release();
      }
    }
  
    private void processOneWayMessage(OneWayMessage req) {
      try {
        logger.trace("Process one way request");
        msgHandler.receive(reverseClient, req);
      } catch (Exception e) {
        logger.error(
            "Error while invoking handler#receive() for one-way message from {}.",
            NettyUtils.getRemoteAddress(channel),
            e);
      } finally {
        req.body().release();
      }
    }
  
    private void processOtherMessages(RequestMessage req) {
      try {
        logger.trace("delegating to handler to process other request");
        msgHandler.receive(reverseClient, req);
      } catch (Exception e) {
        logger.error(
            "Error while invoking handler#receive() for other message from {}.",
            NettyUtils.getRemoteAddress(channel),
            e);
      } finally {
        if (req.body() != null) {
          req.body().release();
        }
      }
    }
  
    private boolean checkRegistered(RequestMessage req) {
      if (!msgHandler.checkRegistered()) {
        IOException e = new IOException("Worker Not Registered!");
        if (req instanceof RpcRequest) {
          respond(new RpcFailure(((RpcRequest) req).requestId, Throwables.getStackTraceAsString(e)));
        } else if (req instanceof ChunkFetchRequest) {
          respond(
              new ChunkFetchFailure(
                  ((ChunkFetchRequest) req).streamChunkSlice, Throwables.getStackTraceAsString(e)));
        } else if (req instanceof OneWayMessage) {
          logger.warn("Ignore OneWayMessage since worker is not registered!");
        }
        return false;
      }
      return true;
    }
  
    /**
     * Responds to a single message with some Encodable object. If a failure occurs while sending, it
     * will be logged and the channel closed.
     */
    private ChannelFuture respond(Encodable result) {
      SocketAddress remoteAddress = channel.remoteAddress();
      return channel
          .writeAndFlush(result)
          .addListener(
              future -> {
                if (future.isSuccess()) {
                  logger.trace("Sent result {} to client {}", result, remoteAddress);
                } else {
                  logger.warn(
                      String.format(
                          "Fail to sending result %s to %s; closing connection",
                          result, remoteAddress),
                      future.cause());
                }
              });
    }
  }
  