content: /*
   * 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 com.google.errorprone.annotations.concurrent.GuardedBy;
  import io.airlift.drift.TException;
  import io.airlift.drift.protocol.TTransportException;
  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 java.io.IOException;
  
  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)
      {
          if (throwable instanceof IOException) {
              throwable = new TTransportException(throwable);
          }
          // exception in the future is expected to be a TException
          if (!(throwable instanceof Error) && !(throwable instanceof TException)) {
              throwable = new TException(throwable);
          }
          setException(throwable);
      }
  }
  