/*
 * Copyright (C) 2013 Facebook, Inc.
 *
 * 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 io.airlift.drift.transport.netty.client;

import com.google.common.util.concurrent.AbstractFuture;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import io.airlift.drift.TException;
import io.airlift.drift.transport.client.ConnectionFailedException;
import io.airlift.drift.transport.client.InvokeRequest;
import io.airlift.drift.transport.netty.client.ConnectionManager.ConnectionParameters;
import io.airlift.drift.transport.netty.client.ThriftClientHandler.ThriftRequest;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.util.concurrent.Future;

import javax.annotation.concurrent.GuardedBy;

import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
import static java.util.Objects.requireNonNull;

class InvocationResponseFuture
        extends AbstractFuture<Object>
{
    private final InvokeRequest request;
    private final ConnectionParameters connectionParameters;
    private final ConnectionManager connectionManager;

    @GuardedBy("this")
    private Future<Channel> connectionFuture;

    @GuardedBy("this")
    private ThriftRequest invocationFuture;

    static InvocationResponseFuture createInvocationResponseFuture(InvokeRequest request, ConnectionParameters connectionParameters, ConnectionManager connectionManager)
    {
        InvocationResponseFuture future = new InvocationResponseFuture(request, connectionParameters, connectionManager);
        // invocation can not be started from constructor, because it may start threads that can call back into the unpublished object
        future.tryConnect();
        return future;
    }

    private InvocationResponseFuture(InvokeRequest request, ConnectionParameters connectionParameters, ConnectionManager connectionManager)
    {
        this.request = requireNonNull(request, "request is null");
        this.connectionParameters = requireNonNull(connectionParameters, "connectionConfig is null");
        this.connectionManager = requireNonNull(connectionManager, "connectionManager is null");

        // if this invocation is canceled, cancel the tasks
        super.addListener(() -> {
            if (super.isCancelled()) {
                onCancel(wasInterrupted());
            }
        }, directExecutor());
    }

    private synchronized void tryConnect()
    {
        try {
            connectionFuture = connectionManager.getConnection(connectionParameters, request.getAddress().getHostAndPort());
            connectionFuture.addListener(channelFuture -> {
                try {
                    if (channelFuture.isSuccess()) {
                        // Netty future listener generic type declaration requires a cast when used with a lambda
                        tryInvocation((Channel) channelFuture.getNow());
                    }
                    else {
                        fatalError(new ConnectionFailedException(request.getAddress(), channelFuture.cause()));
                    }
                }
                catch (Throwable t) {
                    fatalError(t);
                }
            });
        }
        catch (Throwable t) {
            fatalError(t);
        }
    }

    private synchronized void tryInvocation(Channel channel)
    {
        // is request already canceled
        if (isCancelled()) {
            connectionManager.returnConnection(channel);
            return;
        }

        try {
            invocationFuture = new ThriftRequest(request.getMethod(), request.getParameters(), request.getHeaders());
            Futures.addCallback(invocationFuture, new FutureCallback<Object>()
                    {
                        @Override
                        public void onSuccess(Object result)
                        {
                            try {
                                connectionManager.returnConnection(channel);
                                set(result);
                            }
                            catch (Throwable t) {
                                fatalError(t);
                            }
                        }

                        @Override
                        public void onFailure(Throwable t)
                        {
                            try {
                                connectionManager.returnConnection(channel);
                            }
                            finally {
                                fatalError(t);
                            }
                        }
                    },
                    directExecutor());

            ChannelFuture sendFuture = channel.writeAndFlush(invocationFuture);
            sendFuture.addListener(channelFuture -> {
                try {
                    if (!channelFuture.isSuccess()) {
                        fatalError(channelFuture.cause());
                    }
                }
                catch (Throwable t) {
                    fatalError(t);
                }
            });
        }
        catch (Throwable t) {
            try {
                connectionManager.returnConnection(channel);
            }
            finally {
                fatalError(t);
            }
        }
    }

    private synchronized void onCancel(boolean wasInterrupted)
    {
        if (connectionFuture != null) {
            connectionFuture.cancel(wasInterrupted);
        }
        if (invocationFuture != null) {
            invocationFuture.cancel(wasInterrupted);
        }
    }

    private void fatalError(Throwable throwable)
    {
        // exception in the future is expected to be a TException
        if (!(throwable instanceof Error) && !(throwable instanceof TException)) {
            throwable = new TException(throwable);
        }
        setException(throwable);
    }
}

@@ -837,32 +837,86 @@
 ift.TException;%0A
+import io.airlift.drift.protocol.TTransportException;%0A
 import io.airlif
@@ -1352,16 +1352,45 @@
 dedBy;%0A%0A
+import java.io.IOException;%0A%0A
 import s
@@ -6254,32 +6254,150 @@
 hrowable)%0A    %7B%0A
+        if (throwable instanceof IOException) %7B%0A            throwable = new TTransportException(throwable);%0A        %7D%0A
         // excep

Allow retry of idempotent methods for any network error