org.apache.hive.spark.client.rpc.Rpc.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.hive.spark.client.rpc.Rpc.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.hive.spark.client.rpc;

import java.io.Closeable;
import java.io.IOException;
import java.util.Collection;
import java.util.Collections;
import java.util.Map;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import javax.security.auth.callback.Callback;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.callback.NameCallback;
import javax.security.auth.callback.PasswordCallback;
import javax.security.sasl.RealmCallback;
import javax.security.sasl.Sasl;
import javax.security.sasl.SaslClient;
import javax.security.sasl.SaslException;

import com.google.common.annotations.VisibleForTesting;
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.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelFutureListener;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import io.netty.channel.ChannelOption;
import io.netty.channel.embedded.EmbeddedChannel;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.logging.LoggingHandler;
import io.netty.handler.logging.LogLevel;
import io.netty.util.concurrent.EventExecutorGroup;
import io.netty.util.concurrent.Future;
import io.netty.util.concurrent.GenericFutureListener;
import io.netty.util.concurrent.ImmediateEventExecutor;
import io.netty.util.concurrent.Promise;
import io.netty.util.concurrent.ScheduledFuture;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.hadoop.hive.common.classification.InterfaceAudience;

/**
 * Encapsulates the RPC functionality. Provides higher-level methods to talk to the remote
 * endpoint.
 */
@InterfaceAudience.Private
public class Rpc implements Closeable {

    private static final Logger LOG = LoggerFactory.getLogger(Rpc.class);

    static final String SASL_REALM = "rsc";
    static final String SASL_USER = "rsc";
    static final String SASL_PROTOCOL = "rsc";
    static final String SASL_AUTH_CONF = "auth-conf";

    /**
     * Creates an RPC client for a server running on the given remote host and port.
     *
     * @param config RPC configuration data.
     * @param eloop Event loop for managing the connection.
     * @param host Host name or IP address to connect to.
     * @param port Port where server is listening.
     * @param clientId The client ID that identifies the connection.
     * @param secret Secret for authenticating the client with the server.
     * @param dispatcher Dispatcher used to handle RPC calls.
     * @return A future that can be used to monitor the creation of the RPC object.
     */
    public static Promise<Rpc> createClient(Map<String, String> config, final NioEventLoopGroup eloop, String host,
            int port, final String clientId, final String secret, final RpcDispatcher dispatcher) throws Exception {
        final RpcConfiguration rpcConf = new RpcConfiguration(config);
        int connectTimeoutMs = (int) rpcConf.getConnectTimeoutMs();

        final ChannelFuture cf = new Bootstrap().group(eloop).handler(new ChannelInboundHandlerAdapter() {
        }).channel(NioSocketChannel.class).option(ChannelOption.SO_KEEPALIVE, true)
                .option(ChannelOption.CONNECT_TIMEOUT_MILLIS, connectTimeoutMs).connect(host, port);

        final Promise<Rpc> promise = eloop.next().newPromise();
        final AtomicReference<Rpc> rpc = new AtomicReference<Rpc>();

        // Set up a timeout to undo everything.
        final Runnable timeoutTask = new Runnable() {
            @Override
            public void run() {
                promise.setFailure(new TimeoutException("Timed out waiting for RPC server connection."));
            }
        };
        final ScheduledFuture<?> timeoutFuture = eloop.schedule(timeoutTask, rpcConf.getServerConnectTimeoutMs(),
                TimeUnit.MILLISECONDS);

        // The channel listener instantiates the Rpc instance when the connection is established,
        // and initiates the SASL handshake.
        cf.addListener(new ChannelFutureListener() {
            @Override
            public void operationComplete(ChannelFuture cf) throws Exception {
                if (cf.isSuccess()) {
                    SaslClientHandler saslHandler = new SaslClientHandler(rpcConf, clientId, promise, timeoutFuture,
                            secret, dispatcher);
                    Rpc rpc = createRpc(rpcConf, saslHandler, (SocketChannel) cf.channel(), eloop);
                    saslHandler.rpc = rpc;
                    saslHandler.sendHello(cf.channel());
                } else {
                    promise.setFailure(cf.cause());
                }
            }
        });

        // Handle cancellation of the promise.
        promise.addListener(new GenericFutureListener<Promise<Rpc>>() {
            @Override
            public void operationComplete(Promise<Rpc> p) {
                if (p.isCancelled()) {
                    cf.cancel(true);
                }
            }
        });

        return promise;
    }

    static Rpc createServer(SaslHandler saslHandler, RpcConfiguration config, SocketChannel channel,
            EventExecutorGroup egroup) throws IOException {
        return createRpc(config, saslHandler, channel, egroup);
    }

    private static Rpc createRpc(RpcConfiguration config, SaslHandler saslHandler, SocketChannel client,
            EventExecutorGroup egroup) throws IOException {
        LogLevel logLevel = LogLevel.TRACE;
        if (config.getRpcChannelLogLevel() != null) {
            try {
                logLevel = LogLevel.valueOf(config.getRpcChannelLogLevel());
            } catch (Exception e) {
                LOG.warn("Invalid log level {}, reverting to default.", config.getRpcChannelLogLevel());
            }
        }

        boolean logEnabled = false;
        switch (logLevel) {
        case DEBUG:
            logEnabled = LOG.isDebugEnabled();
            break;
        case ERROR:
            logEnabled = LOG.isErrorEnabled();
            break;
        case INFO:
            logEnabled = LOG.isInfoEnabled();
            break;
        case TRACE:
            logEnabled = LOG.isTraceEnabled();
            break;
        case WARN:
            logEnabled = LOG.isWarnEnabled();
            break;
        }

        if (logEnabled) {
            client.pipeline().addLast("logger", new LoggingHandler(Rpc.class, logLevel));
        }

        KryoMessageCodec kryo = new KryoMessageCodec(config.getMaxMessageSize(), MessageHeader.class,
                NullMessage.class, SaslMessage.class);
        saslHandler.setKryoMessageCodec(kryo);
        client.pipeline().addLast("codec", kryo).addLast("sasl", saslHandler);
        return new Rpc(config, client, egroup);
    }

    @VisibleForTesting
    static Rpc createEmbedded(RpcDispatcher dispatcher) {
        EmbeddedChannel c = new EmbeddedChannel(new LoggingHandler(Rpc.class),
                new KryoMessageCodec(0, MessageHeader.class, NullMessage.class), dispatcher);
        Rpc rpc = new Rpc(new RpcConfiguration(Collections.<String, String>emptyMap()), c,
                ImmediateEventExecutor.INSTANCE);
        rpc.dispatcher = dispatcher;
        return rpc;
    }

    private final RpcConfiguration config;
    private final AtomicBoolean rpcClosed;
    private final AtomicLong rpcId;
    private final Channel channel;
    private final Collection<Listener> listeners;
    private final EventExecutorGroup egroup;
    private final Object channelLock;
    private volatile RpcDispatcher dispatcher;

    private Rpc(RpcConfiguration config, Channel channel, EventExecutorGroup egroup) {
        Preconditions.checkArgument(channel != null);
        Preconditions.checkArgument(egroup != null);
        this.config = config;
        this.channel = channel;
        this.channelLock = new Object();
        this.dispatcher = null;
        this.egroup = egroup;
        this.listeners = Lists.newLinkedList();
        this.rpcClosed = new AtomicBoolean();
        this.rpcId = new AtomicLong();

        // Note: this does not work for embedded channels.
        channel.pipeline().addLast("monitor", new ChannelInboundHandlerAdapter() {
            @Override
            public void channelInactive(ChannelHandlerContext ctx) {
                close();
            }
        });
    }

    public void addListener(Listener l) {
        synchronized (listeners) {
            listeners.add(l);
        }
    }

    /**
     * Send an RPC call to the remote endpoint and returns a future that can be used to monitor the
     * operation.
     */
    public Future<Void> call(Object msg) {
        return call(msg, Void.class);
    }

    public boolean isActive() {
        return channel.isActive();
    }

    /**
     * Send an RPC call to the remote endpoint and returns a future that can be used to monitor the
     * operation.
     *
     * @param msg RPC call to send.
     * @param retType Type of expected reply.
     * @return A future used to monitor the operation.
     */
    public <T> Future<T> call(Object msg, Class<T> retType) {
        Preconditions.checkArgument(msg != null);
        Preconditions.checkState(channel.isActive(), "RPC channel is closed.");
        try {
            final long id = rpcId.getAndIncrement();
            final Promise<T> promise = createPromise();
            ChannelFutureListener listener = new ChannelFutureListener() {
                @Override
                public void operationComplete(ChannelFuture cf) {
                    if (!cf.isSuccess() && !promise.isDone()) {
                        LOG.warn("Failed to send RPC, closing connection.", cf.cause());
                        promise.setFailure(cf.cause());
                        dispatcher.discardRpc(id);
                        close();
                    }
                }
            };

            dispatcher.registerRpc(id, promise, msg.getClass().getName());
            synchronized (channelLock) {
                channel.write(new MessageHeader(id, Rpc.MessageType.CALL)).addListener(listener);
                channel.writeAndFlush(msg).addListener(listener);
            }
            return promise;
        } catch (Exception e) {
            throw Throwables.propagate(e);
        }
    }

    /**
     * Creates a promise backed by this RPC's event loop.
     */
    public <T> Promise<T> createPromise() {
        return egroup.next().newPromise();
    }

    @VisibleForTesting
    Channel getChannel() {
        return channel;
    }

    void setDispatcher(RpcDispatcher dispatcher) {
        Preconditions.checkNotNull(dispatcher);
        Preconditions.checkState(this.dispatcher == null);
        this.dispatcher = dispatcher;
        channel.pipeline().addLast("dispatcher", dispatcher);
    }

    @Override
    public void close() {
        if (!rpcClosed.compareAndSet(false, true)) {
            return;
        }
        try {
            channel.close().sync();
        } catch (InterruptedException ie) {
            Thread.interrupted();
        } finally {
            synchronized (listeners) {
                for (Listener l : listeners) {
                    try {
                        l.rpcClosed(this);
                    } catch (Exception e) {
                        LOG.warn("Error caught in Rpc.Listener invocation.", e);
                    }
                }
            }
        }
    }

    public interface Listener {

        void rpcClosed(Rpc rpc);

    }

    static enum MessageType {
        CALL, REPLY, ERROR;
    }

    static class MessageHeader {
        final long id;
        final MessageType type;

        MessageHeader() {
            this(-1, null);
        }

        MessageHeader(long id, MessageType type) {
            this.id = id;
            this.type = type;
        }

    }

    static class NullMessage {

    }

    static class SaslMessage {
        final String clientId;
        final byte[] payload;

        SaslMessage() {
            this(null, null);
        }

        SaslMessage(byte[] payload) {
            this(null, payload);
        }

        SaslMessage(String clientId, byte[] payload) {
            this.clientId = clientId;
            this.payload = payload;
        }

    }

    private static class SaslClientHandler extends SaslHandler implements CallbackHandler {

        private final SaslClient client;
        private final String clientId;
        private final String secret;
        private final RpcDispatcher dispatcher;
        private Promise<Rpc> promise;
        private ScheduledFuture<?> timeout;

        // Can't be set in constructor due to circular dependency.
        private Rpc rpc;

        SaslClientHandler(RpcConfiguration config, String clientId, Promise<Rpc> promise,
                ScheduledFuture<?> timeout, String secret, RpcDispatcher dispatcher) throws IOException {
            super(config);
            this.clientId = clientId;
            this.promise = promise;
            this.timeout = timeout;
            this.secret = secret;
            this.dispatcher = dispatcher;
            this.client = Sasl.createSaslClient(new String[] { config.getSaslMechanism() }, null, SASL_PROTOCOL,
                    SASL_REALM, config.getSaslOptions(), this);
        }

        @Override
        protected boolean isComplete() {
            return client.isComplete();
        }

        @Override
        protected String getNegotiatedProperty(String name) {
            return (String) client.getNegotiatedProperty(name);
        }

        @Override
        protected SaslMessage update(SaslMessage challenge) throws IOException {
            byte[] response = client.evaluateChallenge(challenge.payload);
            return response != null ? new SaslMessage(response) : null;
        }

        @Override
        public byte[] wrap(byte[] data, int offset, int len) throws IOException {
            return client.wrap(data, offset, len);
        }

        @Override
        public byte[] unwrap(byte[] data, int offset, int len) throws IOException {
            return client.unwrap(data, offset, len);
        }

        @Override
        public void dispose() throws IOException {
            if (!client.isComplete()) {
                onError(new SaslException("Client closed before SASL negotiation finished."));
            }
            client.dispose();
        }

        @Override
        protected void onComplete() throws Exception {
            timeout.cancel(true);
            rpc.setDispatcher(dispatcher);
            promise.setSuccess(rpc);
            timeout = null;
            promise = null;
        }

        @Override
        protected void onError(Throwable error) {
            timeout.cancel(true);
            if (!promise.isDone()) {
                promise.setFailure(error);
            }
        }

        @Override
        public void handle(Callback[] callbacks) {
            for (Callback cb : callbacks) {
                if (cb instanceof NameCallback) {
                    ((NameCallback) cb).setName(clientId);
                } else if (cb instanceof PasswordCallback) {
                    ((PasswordCallback) cb).setPassword(secret.toCharArray());
                } else if (cb instanceof RealmCallback) {
                    RealmCallback rb = (RealmCallback) cb;
                    rb.setText(rb.getDefaultText());
                }
            }
        }

        void sendHello(Channel c) throws Exception {
            byte[] hello = client.hasInitialResponse() ? client.evaluateChallenge(new byte[0]) : new byte[0];
            c.writeAndFlush(new SaslMessage(clientId, hello));
        }

    }

}