com.github.hrpc.rpc.Client.java Source code

Java tutorial

Introduction

Here is the source code for com.github.hrpc.rpc.Client.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 com.github.hrpc.rpc;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.FilterInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InterruptedIOException;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.net.SocketTimeoutException;
import java.net.UnknownHostException;
import java.util.Arrays;
import java.util.Hashtable;
import java.util.Iterator;
import java.util.Map.Entry;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

import javax.net.SocketFactory;

import com.github.hrpc.io.DataOutputBuffer;
import com.github.hrpc.io.Writable;
import com.github.hrpc.net.ConnectTimeoutException;
import com.github.hrpc.net.NetUtils;
import com.github.hrpc.util.*;
import com.google.common.base.Preconditions;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import com.google.protobuf.CodedOutputStream;
import com.github.hrpc.rpc.protobuf.RpcHeaderProtos.RpcRequestHeaderProto;
import com.github.hrpc.rpc.protobuf.RpcHeaderProtos.RpcResponseHeaderProto;
import com.github.hrpc.rpc.protobuf.RpcHeaderProtos.RpcRequestHeaderProto.OperationProto;
import com.github.hrpc.rpc.protobuf.IpcConnectionContextProtos.IpcConnectionContextProto;
import com.github.hrpc.rpc.protobuf.RpcHeaderProtos.RpcResponseHeaderProto.RpcStatusProto;
import com.github.hrpc.rpc.protobuf.RpcHeaderProtos.RpcResponseHeaderProto.RpcErrorCodeProto;
import com.github.hrpc.rpc.RPC.RpcKind;
import com.github.hrpc.rpc.ProtobufRpcEngine.RpcRequestMessageWrapper;
import com.github.hrpc.rpc.RetryPolicy.RetryAction;

/** A client for an IPC service.  IPC calls take a single {@link Writable} as a
 * parameter, and return a {@link Writable} as their value.  A service runs on
 * a port and is defined by a parameter class and a value class.
 *
 * @see Server
 */
public class Client {

    public static final Log LOG = LogFactory.getLog(Client.class);

    /** A counter for generating call IDs. */
    private static final AtomicInteger callIdCounter = new AtomicInteger();

    private static final ThreadLocal<Integer> callId = new ThreadLocal<Integer>();
    private static final ThreadLocal<Integer> retryCount = new ThreadLocal<Integer>();

    /** Set call id and retry count for the next call. */
    public static void setCallIdAndRetryCount(int cid, int rc) {
        Preconditions.checkArgument(cid != RpcConstants.INVALID_CALL_ID);
        Preconditions.checkState(callId.get() == null);
        Preconditions.checkArgument(rc != RpcConstants.INVALID_RETRY_COUNT);

        callId.set(cid);
        retryCount.set(rc);
    }

    private Hashtable<ConnectionId, Connection> connections = new Hashtable<ConnectionId, Connection>();

    private Class<? extends Writable> valueClass; // class of call values
    private AtomicBoolean running = new AtomicBoolean(true); // if client runs
    final private Option conf;

    private SocketFactory socketFactory; // how to create sockets
    private int refCount = 1;

    private final int connectionTimeout;

    private final boolean fallbackAllowed;
    private final byte[] clientId;

    final static int CONNECTION_CONTEXT_CALL_ID = -3;

    /**
     * Executor on which IPC calls' parameters are sent. Deferring
     * the sending of parameters to a separate thread isolates them
     * from thread interruptions in the calling code.
     */
    private static final ExecutorService SEND_PARAMS_EXECUTOR = Executors.newCachedThreadPool(
            new ThreadFactoryBuilder().setDaemon(true).setNameFormat("IPC Parameter Sending Thread #%d").build());

    /**
     * set the ping interval value in configuration
     *
     * @param conf Configuration
     * @param pingInterval the ping interval
     */
    final public static void setPingInterval(Option conf, int pingInterval) {
        conf.setInt(CommonConfigurationKeys.IPC_PING_INTERVAL_KEY, pingInterval);
    }

    /**
     * Get the ping interval from configuration;
     * If not set in the configuration, return the default value.
     *
     * @param conf Configuration
     * @return the ping interval
     */
    final static int getPingInterval(Option conf) {
        return conf.getInt(CommonConfigurationKeys.IPC_PING_INTERVAL_KEY,
                CommonConfigurationKeys.IPC_PING_INTERVAL_DEFAULT);
    }

    /**
     * The time after which a RPC will timeout.
     * If ping is not enabled (via ipc.client.ping), then the timeout value is the
     * same as the pingInterval.
     * If ping is enabled, then there is no timeout value.
     *
     * @param conf Configuration
     * @return the timeout period in milliseconds. -1 if no timeout value is set
     */
    final public static int getTimeout(Option conf) {
        if (!conf.getBoolean(CommonConfigurationKeys.IPC_CLIENT_PING_KEY, true)) {
            return getPingInterval(conf);
        }
        return -1;
    }

    /**
     * set the connection timeout value in configuration
     *
     * @param conf Configuration
     * @param timeout the socket connect timeout value
     */
    public static final void setConnectTimeout(Option conf, int timeout) {
        conf.setInt(CommonConfigurationKeys.IPC_CLIENT_CONNECT_TIMEOUT_KEY, timeout);
    }

    /**
     * Increment this client's reference count
     *
     */
    synchronized void incCount() {
        refCount++;
    }

    /**
     * Decrement this client's reference count
     *
     */
    synchronized void decCount() {
        refCount--;
    }

    /**
     * Return if this client has no reference
     *
     * @return true if this client has no reference; false otherwise
     */
    synchronized boolean isZeroReference() {
        return refCount == 0;
    }

    /** Check the rpc response header. */
    void checkResponse(RpcResponseHeaderProto header) throws IOException {
        if (header == null) {
            throw new IOException("Response is null.");
        }
        if (header.hasClientId()) {
            // check client IDs
            final byte[] id = header.getClientId().toByteArray();
            if (!Arrays.equals(id, RpcConstants.DUMMY_CLIENT_ID)) {
                if (!Arrays.equals(id, clientId)) {
                    throw new IOException("Client IDs not matched: local ID="
                            + StringUtils.byteToHexString(clientId) + ", ID in reponse="
                            + StringUtils.byteToHexString(header.getClientId().toByteArray()));
                }
            }
        }
    }

    Call createCall(RPC.RpcKind rpcKind, Writable rpcRequest) {
        return new Call(rpcKind, rpcRequest);
    }

    /**
     * Class that represents an RPC call
     */
    static class Call {
        final int id; // call id
        final int retry; // retry count
        final Writable rpcRequest; // the serialized rpc request
        Writable rpcResponse; // null if rpc has error
        IOException error; // exception, null if success
        final RPC.RpcKind rpcKind; // Rpc EngineKind
        boolean done; // true when call is done

        private Call(RPC.RpcKind rpcKind, Writable param) {
            this.rpcKind = rpcKind;
            this.rpcRequest = param;

            final Integer id = callId.get();
            if (id == null) {
                this.id = nextCallId();
            } else {
                callId.set(null);
                this.id = id;
            }

            final Integer rc = retryCount.get();
            if (rc == null) {
                this.retry = 0;
            } else {
                this.retry = rc;
            }
        }

        /** Indicate when the call is complete and the
         * value or error are available.  Notifies by default.  */
        protected synchronized void callComplete() {
            this.done = true;
            notify(); // notify caller
        }

        /** Set the exception when there is an error.
         * Notify the caller the call is done.
         *
         * @param error exception thrown by the call; either local or remote
         */
        public synchronized void setException(IOException error) {
            this.error = error;
            callComplete();
        }

        /** Set the return value when there is no error.
         * Notify the caller the call is done.
         *
         * @param rpcResponse return value of the rpc call.
         */
        public synchronized void setRpcResponse(Writable rpcResponse) {
            this.rpcResponse = rpcResponse;
            callComplete();
        }

        public synchronized Writable getRpcResponse() {
            return rpcResponse;
        }
    }

    /** Thread that reads responses and notifies callers.  Each connection owns a
     * socket connected to a remote address.  Calls are multiplexed through this
     * socket: responses may be delivered out of order. */
    private class Connection extends Thread {
        private InetSocketAddress server; // server ip:port
        private final ConnectionId remoteId; // connection id
        private int serviceClass;

        private Socket socket = null; // connected socket
        private DataInputStream in;
        private DataOutputStream out;
        private int rpcTimeout;
        private int maxIdleTime; //connections will be culled if it was idle for
        //maxIdleTime msecs
        private final RetryPolicy connectionRetryPolicy;
        private int maxRetriesOnSocketTimeouts;
        private boolean tcpNoDelay; // if T then disable Nagle's Algorithm
        private boolean doPing; //do we need to send ping message
        private int pingInterval; // how often sends ping to the server in msecs
        private ByteArrayOutputStream pingRequest; // ping message

        // currently active calls
        private Hashtable<Integer, Call> calls = new Hashtable<Integer, Call>();
        private AtomicLong lastActivity = new AtomicLong();// last I/O activity time
        private AtomicBoolean shouldCloseConnection = new AtomicBoolean(); // indicate if the connection is closed
        private IOException closeException; // close reason

        private final Object sendRpcRequestLock = new Object();

        public Connection(ConnectionId remoteId, int serviceClass) throws IOException {
            this.remoteId = remoteId;
            this.server = remoteId.getAddress();
            if (server.isUnresolved()) {
                throw NetUtils.wrapException(server.getHostName(), server.getPort(), null, 0,
                        new UnknownHostException());
            }
            this.rpcTimeout = remoteId.getRpcTimeout();
            this.maxIdleTime = remoteId.getMaxIdleTime();
            this.connectionRetryPolicy = remoteId.connectionRetryPolicy;
            this.maxRetriesOnSocketTimeouts = remoteId.getMaxRetriesOnSocketTimeouts();
            this.tcpNoDelay = remoteId.getTcpNoDelay();
            this.doPing = remoteId.getDoPing();
            if (doPing) {
                // construct a RPC header with the callId as the ping callId
                pingRequest = new ByteArrayOutputStream();
                RpcRequestHeaderProto pingHeader = ProtoUtil.makeRpcRequestHeader(RpcKind.RPC_PROTOCOL_BUFFER,
                        OperationProto.RPC_FINAL_PACKET, RpcConstants.PING_CALL_ID,
                        RpcConstants.INVALID_RETRY_COUNT, clientId);
                pingHeader.writeDelimitedTo(pingRequest);
            }
            this.pingInterval = remoteId.getPingInterval();
            this.serviceClass = serviceClass;
            if (LOG.isDebugEnabled()) {
                LOG.debug("The ping interval is " + this.pingInterval + " ms.");
            }

            this.setName("IPC Client (" + socketFactory.hashCode() + ") connection to " + server.toString());
            this.setDaemon(true);
        }

        /** Update lastActivity with the current time. */
        private void touch() {
            lastActivity.set(Time.now());
        }

        /**
         * Add a call to this connection's call queue and notify
         * a listener; synchronized.
         * Returns false if called during shutdown.
         * @param call to add
         * @return true if the call was added.
         */
        private synchronized boolean addCall(Call call) {
            if (shouldCloseConnection.get())
                return false;
            calls.put(call.id, call);
            notify();
            return true;
        }

        /** This class sends a ping to the remote side when timeout on
         * reading. If no failure is detected, it retries until at least
         * a byte is read.
         */
        private class PingInputStream extends FilterInputStream {
            /* constructor */
            protected PingInputStream(InputStream in) {
                super(in);
            }

            /* Process timeout exception
             * if the connection is not going to be closed or
             * is not configured to have a RPC timeout, send a ping.
             * (if rpcTimeout is not set to be 0, then RPC should timeout.
             * otherwise, throw the timeout exception.
             */
            private void handleTimeout(SocketTimeoutException e) throws IOException {
                if (shouldCloseConnection.get() || !running.get() || rpcTimeout > 0) {
                    throw e;
                } else {
                    sendPing();
                }
            }

            /** Read a byte from the stream.
             * Send a ping if timeout on read. Retries if no failure is detected
             * until a byte is read.
             * @throws IOException for any IO problem other than socket timeout
             */
            @Override
            public int read() throws IOException {
                do {
                    try {
                        return super.read();
                    } catch (SocketTimeoutException e) {
                        handleTimeout(e);
                    }
                } while (true);
            }

            /** Read bytes into a buffer starting from offset <code>off</code>
             * Send a ping if timeout on read. Retries if no failure is detected
             * until a byte is read.
             *
             * @return the total number of bytes read; -1 if the connection is closed.
             */
            @Override
            public int read(byte[] buf, int off, int len) throws IOException {
                do {
                    try {
                        return super.read(buf, off, len);
                    } catch (SocketTimeoutException e) {
                        handleTimeout(e);
                    }
                } while (true);
            }
        }

        /**
         * Update the server address if the address corresponding to the host
         * name has changed.
         *
         * @return true if an addr change was detected.
         * @throws IOException when the hostname cannot be resolved.
         */
        private synchronized boolean updateAddress() throws IOException {
            // Do a fresh lookup with the old host name.
            InetSocketAddress currentAddr = new InetSocketAddress(server.getHostName(), server.getPort());

            if (!server.equals(currentAddr)) {
                LOG.warn("Address change detected. Old: " + server.toString() + " New: " + currentAddr.toString());
                server = currentAddr;
                return true;
            }
            return false;
        }

        private synchronized void setupConnection() throws IOException {
            short ioFailures = 0;
            short timeoutFailures = 0;
            while (true) {
                try {
                    this.socket = socketFactory.createSocket();
                    this.socket.setTcpNoDelay(tcpNoDelay);

                    NetUtils.connect(this.socket, server, connectionTimeout);
                    if (rpcTimeout > 0) {
                        pingInterval = rpcTimeout; // rpcTimeout overwrites pingInterval
                    }
                    this.socket.setSoTimeout(pingInterval);
                    return;
                } catch (ConnectTimeoutException toe) {
                    /* Check for an address change and update the local reference.
                     * Reset the failure counter if the address was changed
                     */
                    if (updateAddress()) {
                        timeoutFailures = ioFailures = 0;
                    }
                    handleConnectionTimeout(timeoutFailures++, maxRetriesOnSocketTimeouts, toe);
                } catch (IOException ie) {
                    if (updateAddress()) {
                        timeoutFailures = ioFailures = 0;
                    }
                    handleConnectionFailure(ioFailures++, ie);
                }
            }
        }

        /** Connect to the server and set up the I/O streams. It then sends
         * a header to the server and starts
         * the connection thread that waits for responses.
         */
        private synchronized void setupIOstreams() {
            if (socket != null || shouldCloseConnection.get()) {
                return;
            }
            try {
                if (LOG.isDebugEnabled()) {
                    LOG.debug("Connecting to " + server);
                }
                while (true) {
                    setupConnection();
                    InputStream inStream = NetUtils.getInputStream(socket);
                    OutputStream outStream = NetUtils.getOutputStream(socket);
                    writeConnectionHeader(outStream);

                    if (doPing) {
                        inStream = new PingInputStream(inStream);
                    }
                    this.in = new DataInputStream(new BufferedInputStream(inStream));

                    // SASL may have already buffered the stream
                    if (!(outStream instanceof BufferedOutputStream)) {
                        outStream = new BufferedOutputStream(outStream);
                    }
                    this.out = new DataOutputStream(outStream);

                    writeConnectionContext(remoteId);

                    // update last activity time
                    touch();

                    // start the receiver thread after the socket connection has been set
                    // up
                    start();
                    return;
                }
            } catch (Throwable t) {
                if (t instanceof IOException) {
                    markClosed((IOException) t);
                } else {
                    markClosed(new IOException("Couldn't set up IO streams", t));
                }
                close();
            }
        }

        private void closeConnection() {
            if (socket == null) {
                return;
            }
            // close the current connection
            try {
                socket.close();
            } catch (IOException e) {
                LOG.warn("Not able to close a socket", e);
            }
            // set socket to null so that the next call to setupIOstreams
            // can start the process of connect all over again.
            socket = null;
        }

        /* Handle connection failures due to timeout on connect
         *
         * If the current number of retries is equal to the max number of retries,
         * stop retrying and throw the exception; Otherwise backoff 1 second and
         * try connecting again.
         *
         * This Method is only called from inside setupIOstreams(), which is
         * synchronized. Hence the sleep is synchronized; the locks will be retained.
         *
         * @param curRetries current number of retries
         * @param maxRetries max number of retries allowed
         * @param ioe failure reason
         * @throws IOException if max number of retries is reached
         */
        private void handleConnectionTimeout(int curRetries, int maxRetries, IOException ioe) throws IOException {

            closeConnection();

            // throw the exception if the maximum number of retries is reached
            if (curRetries >= maxRetries) {
                throw ioe;
            }
            LOG.info("Retrying connect to server: " + server + ". Already tried " + curRetries
                    + " time(s); maxRetries=" + maxRetries);
        }

        private void handleConnectionFailure(int curRetries, IOException ioe) throws IOException {
            closeConnection();

            final RetryAction action;
            try {
                action = connectionRetryPolicy.shouldRetry(ioe, curRetries, 0, true);
            } catch (Exception e) {
                throw e instanceof IOException ? (IOException) e : new IOException(e);
            }
            if (action.action == RetryAction.RetryDecision.FAIL) {
                if (action.reason != null) {
                    LOG.warn("Failed to connect to server: " + server + ": " + action.reason, ioe);
                }
                throw ioe;
            }

            try {
                Thread.sleep(action.delayMillis);
            } catch (InterruptedException e) {
                throw (IOException) new InterruptedIOException(
                        "Interrupted: action=" + action + ", retry policy=" + connectionRetryPolicy).initCause(e);
            }
            LOG.info("Retrying connect to server: " + server + ". Already tried " + curRetries
                    + " time(s); retry policy is " + connectionRetryPolicy);
        }

        /**
         * Write the connection header - this is sent when connection is established
         * +----------------------------------+
         * |  "hrpc" 4 bytes                  |
         * +----------------------------------+
         * |  Version (1 byte)                |
         * +----------------------------------+
         * |  Service Class (1 byte)          |
         * +----------------------------------+
         */
        private void writeConnectionHeader(OutputStream outStream) throws IOException {
            DataOutputStream out = new DataOutputStream(new BufferedOutputStream(outStream));
            // Write out the header, version and authentication method
            out.write(RpcConstants.HEADER.array());
            out.write(RpcConstants.CURRENT_VERSION);
            out.write(serviceClass);
            out.flush();
        }

        /* Write the connection context header for each connection
         * Out is not synchronized because only the first thread does this.
         */
        private void writeConnectionContext(ConnectionId remoteId) throws IOException {
            // Write out the ConnectionHeader
            IpcConnectionContextProto message = ProtoUtil
                    .makeIpcConnectionContext(RPC.getProtocolName(remoteId.getProtocol()));
            RpcRequestHeaderProto connectionContextHeader = ProtoUtil.makeRpcRequestHeader(
                    RpcKind.RPC_PROTOCOL_BUFFER, OperationProto.RPC_FINAL_PACKET, CONNECTION_CONTEXT_CALL_ID,
                    RpcConstants.INVALID_RETRY_COUNT, clientId);
            RpcRequestMessageWrapper request = new RpcRequestMessageWrapper(connectionContextHeader, message);

            // Write out the packet length
            out.writeInt(request.getLength());
            request.write(out);
        }

        /* wait till someone signals us to start reading RPC response or
         * it is idle too long, it is marked as to be closed,
         * or the client is marked as not running.
         *
         * Return true if it is time to read a response; false otherwise.
         */
        private synchronized boolean waitForWork() {
            if (calls.isEmpty() && !shouldCloseConnection.get() && running.get()) {
                long timeout = maxIdleTime - (Time.now() - lastActivity.get());
                if (timeout > 0) {
                    try {
                        wait(timeout);
                    } catch (InterruptedException e) {
                    }
                }
            }

            if (!calls.isEmpty() && !shouldCloseConnection.get() && running.get()) {
                return true;
            } else if (shouldCloseConnection.get()) {
                return false;
            } else if (calls.isEmpty()) { // idle connection closed or stopped
                markClosed(null);
                return false;
            } else { // get stopped but there are still pending requests
                markClosed((IOException) new IOException().initCause(new InterruptedException()));
                return false;
            }
        }

        public InetSocketAddress getRemoteAddress() {
            return server;
        }

        /* Send a ping to the server if the time elapsed
         * since last I/O activity is equal to or greater than the ping interval
         */
        private synchronized void sendPing() throws IOException {
            long curTime = Time.now();
            if (curTime - lastActivity.get() >= pingInterval) {
                lastActivity.set(curTime);
                synchronized (out) {
                    out.writeInt(pingRequest.size());
                    pingRequest.writeTo(out);
                    out.flush();
                }
            }
        }

        @Override
        public void run() {
            if (LOG.isDebugEnabled())
                LOG.debug(getName() + ": starting, having connections " + connections.size());

            try {
                while (waitForWork()) {//wait here for work - read or close connection
                    receiveRpcResponse();
                }
            } catch (Throwable t) {
                // This truly is unexpected, since we catch IOException in receiveResponse
                // -- this is only to be really sure that we don't leave a client hanging
                // forever.
                LOG.warn("Unexpected error reading responses on connection " + this, t);
                markClosed(new IOException("Error reading responses", t));
            }

            close();

            if (LOG.isDebugEnabled())
                LOG.debug(getName() + ": stopped, remaining connections " + connections.size());
        }

        /** Initiates a rpc call by sending the rpc request to the remote server.
         * Note: this is not called from the Connection thread, but by other
         * threads.
         * @param call - the rpc request
         */
        public void sendRpcRequest(final Call call) throws InterruptedException, IOException {
            if (shouldCloseConnection.get()) {
                return;
            }

            // Serialize the call to be sent. This is done from the actual
            // caller thread, rather than the SEND_PARAMS_EXECUTOR thread,
            // so that if the serialization throws an error, it is reported
            // properly. This also parallelizes the serialization.
            //
            // Format of a call on the wire:
            // 0) Length of rest below (1 + 2)
            // 1) RpcRequestHeader  - is serialized Delimited hence contains length
            // 2) RpcRequest
            //
            // Items '1' and '2' are prepared here.
            final DataOutputBuffer d = new DataOutputBuffer();
            RpcRequestHeaderProto header = ProtoUtil.makeRpcRequestHeader(call.rpcKind,
                    OperationProto.RPC_FINAL_PACKET, call.id, call.retry, clientId);
            header.writeDelimitedTo(d);
            call.rpcRequest.write(d);

            synchronized (sendRpcRequestLock) {
                Future<?> senderFuture = SEND_PARAMS_EXECUTOR.submit(new Runnable() {
                    @Override
                    public void run() {
                        try {
                            synchronized (Connection.this.out) {
                                if (shouldCloseConnection.get()) {
                                    return;
                                }

                                if (LOG.isDebugEnabled())
                                    LOG.debug(getName() + " sending #" + call.id);

                                byte[] data = d.getData();
                                int totalLength = d.getLength();
                                out.writeInt(totalLength); // Total Length
                                out.write(data, 0, totalLength);// RpcRequestHeader + RpcRequest
                                out.flush();
                            }
                        } catch (IOException e) {
                            // exception at this point would leave the connection in an
                            // unrecoverable state (eg half a call left on the wire).
                            // So, close the connection, killing any outstanding calls
                            markClosed(e);
                        } finally {
                            //the buffer is just an in-memory buffer, but it is still polite to
                            // close early
                            IOUtils.closeStream(d);
                        }
                    }
                });

                try {
                    senderFuture.get();
                } catch (ExecutionException e) {
                    Throwable cause = e.getCause();

                    // cause should only be a RuntimeException as the Runnable above
                    // catches IOException
                    if (cause instanceof RuntimeException) {
                        throw (RuntimeException) cause;
                    } else {
                        throw new RuntimeException("unexpected checked exception", cause);
                    }
                }
            }
        }

        /* Receive a response.
         * Because only one receiver, so no synchronization on in.
         */
        private void receiveRpcResponse() {
            if (shouldCloseConnection.get()) {
                return;
            }
            touch();

            try {
                int totalLen = in.readInt();
                RpcResponseHeaderProto header = RpcResponseHeaderProto.parseDelimitedFrom(in);
                checkResponse(header);

                int headerLen = header.getSerializedSize();
                headerLen += CodedOutputStream.computeRawVarint32Size(headerLen);

                int callId = header.getCallId();
                if (LOG.isDebugEnabled())
                    LOG.debug(getName() + " got value #" + callId);

                Call call = calls.get(callId);
                RpcStatusProto status = header.getStatus();
                if (status == RpcStatusProto.SUCCESS) {
                    Writable value = ReflectionUtils.newInstance(valueClass);
                    value.readFields(in); // read value
                    calls.remove(callId);
                    call.setRpcResponse(value);

                    // verify that length was correct
                    // only for ProtobufEngine where len can be verified easily
                    if (call.getRpcResponse() instanceof ProtobufRpcEngine.RpcWrapper) {
                        ProtobufRpcEngine.RpcWrapper resWrapper = (ProtobufRpcEngine.RpcWrapper) call
                                .getRpcResponse();
                        if (totalLen != headerLen + resWrapper.getLength()) {
                            throw new RpcClientException("RPC response length mismatch on rpc success");
                        }
                    }
                } else { // Rpc Request failed
                    // Verify that length was correct
                    if (totalLen != headerLen) {
                        throw new RpcClientException("RPC response length mismatch on rpc error");
                    }

                    final String exceptionClassName = header.hasExceptionClassName()
                            ? header.getExceptionClassName()
                            : "ServerDidNotSetExceptionClassName";
                    final String errorMsg = header.hasErrorMsg() ? header.getErrorMsg() : "ServerDidNotSetErrorMsg";
                    final RpcErrorCodeProto erCode = (header.hasErrorDetail() ? header.getErrorDetail() : null);
                    if (erCode == null) {
                        LOG.warn("Detailed error code not set by server on rpc error");
                    }
                    RemoteException re = ((erCode == null) ? new RemoteException(exceptionClassName, errorMsg)
                            : new RemoteException(exceptionClassName, errorMsg, erCode));
                    if (status == RpcStatusProto.ERROR) {
                        calls.remove(callId);
                        call.setException(re);
                    } else if (status == RpcStatusProto.FATAL) {
                        // Close the connection
                        markClosed(re);
                    }
                }
            } catch (IOException e) {
                markClosed(e);
            }
        }

        private synchronized void markClosed(IOException e) {
            if (shouldCloseConnection.compareAndSet(false, true)) {
                closeException = e;
                notifyAll();
            }
        }

        /** Close the connection. */
        private synchronized void close() {
            if (!shouldCloseConnection.get()) {
                LOG.error("The connection is not in the closed state");
                return;
            }

            // release the resources
            // first thing to do;take the connection out of the connection list
            synchronized (connections) {
                if (connections.get(remoteId) == this) {
                    connections.remove(remoteId);
                }
            }

            // close the streams and therefore the socket
            IOUtils.closeStream(out);
            IOUtils.closeStream(in);

            // clean up all calls
            if (closeException == null) {
                if (!calls.isEmpty()) {
                    LOG.warn("A connection is closed for no cause and calls are not empty");

                    // clean up calls anyway
                    closeException = new IOException("Unexpected closed connection");
                    cleanupCalls();
                }
            } else {
                // log the info
                if (LOG.isDebugEnabled()) {
                    LOG.debug("closing ipc connection to " + server + ": " + closeException.getMessage(),
                            closeException);
                }

                // cleanup calls
                cleanupCalls();
            }
            if (LOG.isDebugEnabled())
                LOG.debug(getName() + ": closed");
        }

        /* Cleanup all calls and mark them as done */
        private void cleanupCalls() {
            Iterator<Entry<Integer, Call>> itor = calls.entrySet().iterator();
            while (itor.hasNext()) {
                Call c = itor.next().getValue();
                itor.remove();
                c.setException(closeException); // local exception
            }
        }
    }

    /** Construct an IPC client whose values are of the given {@link Writable}
     * class. */
    public Client(Class<? extends Writable> valueClass, Option conf, SocketFactory factory) {
        this.valueClass = valueClass;
        this.conf = conf;
        this.socketFactory = factory;
        this.connectionTimeout = conf.getInt(CommonConfigurationKeys.IPC_CLIENT_CONNECT_TIMEOUT_KEY,
                CommonConfigurationKeys.IPC_CLIENT_CONNECT_TIMEOUT_DEFAULT);
        this.fallbackAllowed = conf.getBoolean(
                CommonConfigurationKeys.IPC_CLIENT_FALLBACK_TO_SIMPLE_AUTH_ALLOWED_KEY,
                CommonConfigurationKeys.IPC_CLIENT_FALLBACK_TO_SIMPLE_AUTH_ALLOWED_DEFAULT);
        this.clientId = ClientId.getClientId();
    }

    /**
     * Construct an IPC client with the default SocketFactory
     * @param valueClass
     * @param conf
     */
    public Client(Class<? extends Writable> valueClass, Option conf) {
        this(valueClass, conf, NetUtils.getDefaultSocketFactory());
    }

    /** Return the socket factory of this client
     *
     * @return this client's socket factory
     */
    SocketFactory getSocketFactory() {
        return socketFactory;
    }

    /** Stop all threads related to this client.  No further calls may be made
     * using this client. */
    public void stop() {
        if (LOG.isDebugEnabled()) {
            LOG.debug("Stopping client");
        }

        if (!running.compareAndSet(true, false)) {
            return;
        }

        // wake up all connections
        synchronized (connections) {
            for (Connection conn : connections.values()) {
                conn.interrupt();
            }
        }

        // wait until all connections are closed
        while (!connections.isEmpty()) {
            try {
                Thread.sleep(100);
            } catch (InterruptedException e) {
            }
        }
    }

    /**
     * Same as {@link #call(RPC.RpcKind, Writable, ConnectionId)}
     *  for RPC_BUILTIN
     */
    public Writable call(Writable param, InetSocketAddress address) throws IOException {
        return call(RPC.RpcKind.RPC_BUILTIN, param, address);

    }

    /** Make a call, passing <code>param</code>, to the IPC server running at
     * <code>address</code>, returning the value.  Throws exceptions if there are
     * network problems or if the remote code threw an exception.
     * @deprecated Use {@link #call(RPC.RpcKind, Writable,
     *  ConnectionId)} instead
     */
    @Deprecated
    public Writable call(RPC.RpcKind rpcKind, Writable param, InetSocketAddress addr) throws IOException {
        ConnectionId remoteId = ConnectionId.getConnectionId(addr, null, 0, conf);
        return call(rpcKind, param, remoteId);
    }

    /** Make a call, passing <code>param</code>, to the IPC server running at
     * <code>address</code> which is servicing the <code>protocol</code> protocol,
     * with the <code>ticket</code> credentials and <code>rpcTimeout</code> as
     * timeout, returning the value.
     * Throws exceptions if there are network problems or if the remote code
     * threw an exception.
     * @deprecated Use {@link #call(RPC.RpcKind, Writable,
     *  ConnectionId)} instead
     */
    @Deprecated
    public Writable call(RPC.RpcKind rpcKind, Writable param, InetSocketAddress addr, Class<?> protocol,
            int rpcTimeout) throws IOException {
        ConnectionId remoteId = ConnectionId.getConnectionId(addr, protocol, rpcTimeout, conf);
        return call(rpcKind, param, remoteId);
    }

    /**
     * Same as {@link #call(RPC.RpcKind, Writable, InetSocketAddress,
     * Class, int, Option)}
     * except that rpcKind is writable.
     */
    public Writable call(Writable param, InetSocketAddress addr, Class<?> protocol, int rpcTimeout, Option conf)
            throws IOException {
        ConnectionId remoteId = ConnectionId.getConnectionId(addr, protocol, rpcTimeout, conf);
        return call(RPC.RpcKind.RPC_BUILTIN, param, remoteId);
    }

    /**
     * Same as {@link #call(Writable, InetSocketAddress,
     * Class, int, Option)}
     * except that specifying serviceClass.
     */
    public Writable call(Writable param, InetSocketAddress addr, Class<?> protocol, int rpcTimeout,
            int serviceClass, Option conf) throws IOException {
        ConnectionId remoteId = ConnectionId.getConnectionId(addr, protocol, rpcTimeout, conf);
        return call(RPC.RpcKind.RPC_BUILTIN, param, remoteId, serviceClass);
    }

    /**
     * Make a call, passing <code>param</code>, to the IPC server running at
     * <code>address</code> which is servicing the <code>protocol</code> protocol,
     * with the <code>ticket</code> credentials, <code>rpcTimeout</code> as
     * timeout and <code>conf</code> as conf for this connection, returning the
     * value. Throws exceptions if there are network problems or if the remote
     * code threw an exception.
     */
    public Writable call(RPC.RpcKind rpcKind, Writable param, InetSocketAddress addr, Class<?> protocol,
            int rpcTimeout, Option conf) throws IOException {
        ConnectionId remoteId = ConnectionId.getConnectionId(addr, protocol, rpcTimeout, conf);
        return call(rpcKind, param, remoteId);
    }

    /**
     * Same as {link {@link #call(RPC.RpcKind, Writable, ConnectionId)}
     * except the rpcKind is RPC_BUILTIN
     */
    public Writable call(Writable param, ConnectionId remoteId) throws IOException {
        return call(RPC.RpcKind.RPC_BUILTIN, param, remoteId);
    }

    /**
     * Make a call, passing <code>rpcRequest</code>, to the IPC server defined by
     * <code>remoteId</code>, returning the rpc respond.
     *
     * @param rpcKind
     * @param rpcRequest -  contains serialized method and method parameters
     * @param remoteId - the target rpc server
     * @returns the rpc response
     * Throws exceptions if there are network problems or if the remote code
     * threw an exception.
     */
    public Writable call(RPC.RpcKind rpcKind, Writable rpcRequest, ConnectionId remoteId) throws IOException {
        return call(rpcKind, rpcRequest, remoteId, RPC.RPC_SERVICE_CLASS_DEFAULT);
    }

    /**
     * Make a call, passing <code>rpcRequest</code>, to the IPC server defined by
     * <code>remoteId</code>, returning the rpc respond.
     *
     * @param rpcKind
     * @param rpcRequest -  contains serialized method and method parameters
     * @param remoteId - the target rpc server
     * @param serviceClass - service class for RPC
     * @returns the rpc response
     * Throws exceptions if there are network problems or if the remote code
     * threw an exception.
     */
    public Writable call(RPC.RpcKind rpcKind, Writable rpcRequest, ConnectionId remoteId, int serviceClass)
            throws IOException {
        final Call call = createCall(rpcKind, rpcRequest);
        Connection connection = getConnection(remoteId, call, serviceClass);
        try {
            connection.sendRpcRequest(call); // send the rpc request
        } catch (RejectedExecutionException e) {
            throw new IOException("connection has been closed", e);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            LOG.warn("interrupted waiting to send rpc request to server", e);
            throw new IOException(e);
        }

        boolean interrupted = false;
        synchronized (call) {
            while (!call.done) {
                try {
                    call.wait(); // wait for the result
                } catch (InterruptedException ie) {
                    // save the fact that we were interrupted
                    interrupted = true;
                }
            }

            if (interrupted) {
                // set the interrupt flag now that we are done waiting
                Thread.currentThread().interrupt();
            }

            if (call.error != null) {
                if (call.error instanceof RemoteException) {
                    call.error.fillInStackTrace();
                    throw call.error;
                } else { // local exception
                    InetSocketAddress address = connection.getRemoteAddress();
                    throw NetUtils.wrapException(address.getHostName(), address.getPort(), NetUtils.getHostname(),
                            0, call.error);
                }
            } else {
                return call.getRpcResponse();
            }
        }
    }

    // for unit testing only
    Set<ConnectionId> getConnectionIds() {
        synchronized (connections) {
            return connections.keySet();
        }
    }

    /** Get a connection from the pool, or create a new one and add it to the
     * pool.  Connections to a given ConnectionId are reused. */
    private Connection getConnection(ConnectionId remoteId, Call call, int serviceClass) throws IOException {
        if (!running.get()) {
            // the client is stopped
            throw new IOException("The client is stopped");
        }
        Connection connection;
        /* we could avoid this allocation for each RPC by having a  
         * connectionsId object and with set() method. We need to manage the
         * refs for keys in HashMap properly. For now its ok.
         */
        do {
            synchronized (connections) {
                connection = connections.get(remoteId);
                if (connection == null) {
                    connection = new Connection(remoteId, serviceClass);
                    connections.put(remoteId, connection);
                }
            }
        } while (!connection.addCall(call));

        //we don't invoke the method below inside "synchronized (connections)"
        //block above. The reason for that is if the server happens to be slow,
        //it will take longer to establish a connection and that will slow the
        //entire system down.
        connection.setupIOstreams();
        return connection;
    }

    /**
     * This class holds the address. The client connections
     * to servers are uniquely identified by <remoteAddress, protocol, ticket>
     */
    public static class ConnectionId {
        InetSocketAddress address;
        final Class<?> protocol;
        private static final int PRIME = 16777619;
        private final int rpcTimeout;
        private final int maxIdleTime; //connections will be culled if it was idle for
        //maxIdleTime msecs
        private final RetryPolicy connectionRetryPolicy;
        // the max. no. of retries for socket connections on time out exceptions
        private final int maxRetriesOnSocketTimeouts;
        private final boolean tcpNoDelay; // if T then disable Nagle's Algorithm
        private final boolean doPing; //do we need to send ping message
        private final int pingInterval; // how often sends ping to the server in msecs
        private String saslQop; // here for testing

        ConnectionId(InetSocketAddress address, Class<?> protocol, int rpcTimeout, int maxIdleTime,
                RetryPolicy connectionRetryPolicy, int maxRetriesOnSocketTimeouts, boolean tcpNoDelay,
                boolean doPing, int pingInterval) {
            this.protocol = protocol;
            this.address = address;
            this.rpcTimeout = rpcTimeout;
            this.maxIdleTime = maxIdleTime;
            this.connectionRetryPolicy = connectionRetryPolicy;
            this.maxRetriesOnSocketTimeouts = maxRetriesOnSocketTimeouts;
            this.tcpNoDelay = tcpNoDelay;
            this.doPing = doPing;
            this.pingInterval = pingInterval;
        }

        InetSocketAddress getAddress() {
            return address;
        }

        Class<?> getProtocol() {
            return protocol;
        }

        private int getRpcTimeout() {
            return rpcTimeout;
        }

        int getMaxIdleTime() {
            return maxIdleTime;
        }

        /** max connection retries on socket time outs */
        public int getMaxRetriesOnSocketTimeouts() {
            return maxRetriesOnSocketTimeouts;
        }

        boolean getTcpNoDelay() {
            return tcpNoDelay;
        }

        boolean getDoPing() {
            return doPing;
        }

        int getPingInterval() {
            return pingInterval;
        }

        static ConnectionId getConnectionId(InetSocketAddress addr, Class<?> protocol, int rpcTimeout, Option conf)
                throws IOException {
            return getConnectionId(addr, protocol, rpcTimeout, null, conf);
        }

        /**
         * Returns a ConnectionId object.
         * @param addr Remote address for the connection.
         * @param protocol Protocol for RPC.
         * @param rpcTimeout timeout
         * @param conf Configuration object
         * @return A ConnectionId instance
         * @throws IOException
         */
        static ConnectionId getConnectionId(InetSocketAddress addr, Class<?> protocol, int rpcTimeout,
                RetryPolicy connectionRetryPolicy, Option conf) throws IOException {

            if (connectionRetryPolicy == null) {
                final int max = conf.getInt(CommonConfigurationKeysPublic.IPC_CLIENT_CONNECT_MAX_RETRIES_KEY,
                        CommonConfigurationKeysPublic.IPC_CLIENT_CONNECT_MAX_RETRIES_DEFAULT);
                connectionRetryPolicy = RetryPolicies.retryUpToMaximumCountWithFixedSleep(max, 1, TimeUnit.SECONDS);
            }

            boolean doPing = conf.getBoolean(CommonConfigurationKeys.IPC_CLIENT_PING_KEY, true);
            return new ConnectionId(addr, protocol, rpcTimeout,
                    conf.getInt(CommonConfigurationKeysPublic.IPC_CLIENT_CONNECTION_MAXIDLETIME_KEY,
                            CommonConfigurationKeysPublic.IPC_CLIENT_CONNECTION_MAXIDLETIME_DEFAULT),
                    connectionRetryPolicy,
                    conf.getInt(CommonConfigurationKeysPublic.IPC_CLIENT_CONNECT_MAX_RETRIES_ON_SOCKET_TIMEOUTS_KEY,
                            CommonConfigurationKeysPublic.IPC_CLIENT_CONNECT_MAX_RETRIES_ON_SOCKET_TIMEOUTS_DEFAULT),
                    conf.getBoolean(CommonConfigurationKeysPublic.IPC_CLIENT_TCPNODELAY_KEY,
                            CommonConfigurationKeysPublic.IPC_CLIENT_TCPNODELAY_DEFAULT),
                    doPing, (doPing ? Client.getPingInterval(conf) : 0));
        }

        static boolean isEqual(Object a, Object b) {
            return a == null ? b == null : a.equals(b);
        }

        @Override
        public boolean equals(Object obj) {
            if (obj == this) {
                return true;
            }
            if (obj instanceof ConnectionId) {
                ConnectionId that = (ConnectionId) obj;
                return isEqual(this.address, that.address) && this.doPing == that.doPing
                        && this.maxIdleTime == that.maxIdleTime
                        && isEqual(this.connectionRetryPolicy, that.connectionRetryPolicy)
                        && this.pingInterval == that.pingInterval && isEqual(this.protocol, that.protocol)
                        && this.rpcTimeout == that.rpcTimeout && this.tcpNoDelay == that.tcpNoDelay;
            }
            return false;
        }

        @Override
        public int hashCode() {
            int result = connectionRetryPolicy.hashCode();
            result = PRIME * result + ((address == null) ? 0 : address.hashCode());
            result = PRIME * result + (doPing ? 1231 : 1237);
            result = PRIME * result + maxIdleTime;
            result = PRIME * result + pingInterval;
            result = PRIME * result + ((protocol == null) ? 0 : protocol.hashCode());
            result = PRIME * result + rpcTimeout;
            result = PRIME * result + (tcpNoDelay ? 1231 : 1237);
            return result;
        }

        @Override
        public String toString() {
            return address.toString();
        }
    }

    /**
     * Returns the next valid sequential call ID by incrementing an atomic counter
     * and masking off the sign bit.  Valid call IDs are non-negative integers in
     * the range [ 0, 2^31 - 1 ].  Negative numbers are reserved for special
     * purposes.  The values can overflow back to 0 and be reused.  Note that prior
     * versions of the client did not mask off the sign bit, so a server may still
     * see a negative call ID if it receives connections from an old client.
     *
     * @return next call ID
     */
    public static int nextCallId() {
        return callIdCounter.getAndIncrement() & 0x7FFFFFFF;
    }
}