org.apache.spark.network.netty.NettyTransportRequestHandler.java Source code

Introduction

Here is the source code for org.apache.spark.network.netty.NettyTransportRequestHandler.java

Source

/*
 * 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.spark.network.netty;

import java.nio.ByteBuffer;

import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelFutureListener;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.spark.network.server.TransportRequestHandler;
import org.apache.spark.network.server.RpcHandler;
import org.apache.spark.network.server.StreamManager;
import org.apache.spark.network.buffer.ManagedBuffer;
import org.apache.spark.network.buffer.NioManagedBuffer;
import org.apache.spark.network.client.RpcResponseCallback;
import org.apache.spark.network.client.TransportClient;
import org.apache.spark.network.protocol.ChunkFetchRequest;
import org.apache.spark.network.protocol.ChunkFetchFailure;
import org.apache.spark.network.protocol.ChunkFetchSuccess;
import org.apache.spark.network.protocol.Encodable;
import org.apache.spark.network.protocol.OneWayMessage;
import org.apache.spark.network.protocol.RequestMessage;
import org.apache.spark.network.protocol.RpcFailure;
import org.apache.spark.network.protocol.RpcRequest;
import org.apache.spark.network.protocol.RpcResponse;
import org.apache.spark.network.protocol.StreamFailure;
import org.apache.spark.network.protocol.StreamRequest;
import org.apache.spark.network.protocol.StreamResponse;
import org.apache.spark.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).
 *
 * The messages should have been processed by the pipeline setup by {@link TransportServer}.
 */
public class NettyTransportRequestHandler extends TransportRequestHandler {

    private final Logger logger = LoggerFactory.getLogger(NettyTransportRequestHandler.class);
    /** The Netty channel that this handler is associated with. */
    private final Channel channel;

    public NettyTransportRequestHandler(Channel channel, TransportClient reverseClient, RpcHandler rpcHandler) {
        super(NettyUtils.getRemoteAddress(channel), reverseClient, rpcHandler, rpcHandler.getStreamManager());
        this.channel = channel;
    }

    @Override
    public void channelUnregistered() {
        if (streamManager != null) {
            try {
                streamManager.connectionChannelTerminated(channel);
            } catch (RuntimeException e) {
                logger.error("StreamManager connectionTerminated() callback failed.", e);
            }
        }
        rpcHandler.connectionTerminated(reverseClient);
    }

    public void processFetchRequest(final ChunkFetchRequest req) {
        final String client = NettyUtils.getRemoteAddress(channel);

        logger.trace("Received req from {} to fetch block {}", client, req.streamChunkId);

        ManagedBuffer buf;
        try {
            streamManager.checkAuthorization(reverseClient, req.streamChunkId.streamId);
            streamManager.registerChannel(channel, req.streamChunkId.streamId);
            buf = streamManager.getChunk(req.streamChunkId.streamId, req.streamChunkId.chunkIndex);
        } catch (Exception e) {
            logger.error(String.format("Error opening block %s for request from %s", req.streamChunkId, client), e);
            respond(new ChunkFetchFailure(req.streamChunkId, Throwables.getStackTraceAsString(e)));
            return;
        }

        respond(new ChunkFetchSuccess(req.streamChunkId, buf));
    }

    /**
     * Responds to a single message with some Encodable object. If a failure occurs while sending,
     * it will be logged and the channel closed.
     */
    public void respond(final Encodable result) {
        final String remoteAddress = channel.remoteAddress().toString();
        channel.writeAndFlush(result).addListener(new ChannelFutureListener() {
            @Override
            public void operationComplete(ChannelFuture future) throws Exception {
                if (future.isSuccess()) {
                    logger.trace(String.format("Sent result %s to client %s", result, remoteAddress));
                } else {
                    logger.error(String.format("Error sending result %s to %s; closing connection", result,
                            remoteAddress), future.cause());
                    channel.close();
                }
            }
        });
    }
}