NettyResponseFuture.java

/*
 *    Copyright (c) 2014-2024 AsyncHttpClient Project. All rights reserved.
 *
 *    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 org.asynchttpclient.netty;

import io.netty.channel.Channel;
import org.asynchttpclient.AsyncHandler;
import org.asynchttpclient.ListenableFuture;
import org.asynchttpclient.Realm;
import org.asynchttpclient.Request;
import org.asynchttpclient.channel.ChannelPoolPartitioning;
import org.asynchttpclient.netty.channel.ChannelState;
import org.asynchttpclient.netty.channel.Channels;
import org.asynchttpclient.netty.channel.ConnectionSemaphore;
import org.asynchttpclient.netty.request.NettyRequest;
import org.asynchttpclient.netty.timeout.TimeoutsHolder;
import org.asynchttpclient.proxy.ProxyServer;
import org.asynchttpclient.uri.Uri;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;

import static org.asynchttpclient.util.DateUtils.unpreciseMillisTime;

/**
 * A {@link Future} that can be used to track when an asynchronous HTTP request
 * has been fully processed.
 *
 * @param <V> the result type
 */
public final class NettyResponseFuture<V> implements ListenableFuture<V> {

    private static final Logger LOGGER = LoggerFactory.getLogger(NettyResponseFuture.class);

    @SuppressWarnings("rawtypes")
    private static final AtomicIntegerFieldUpdater<NettyResponseFuture> REDIRECT_COUNT_UPDATER = AtomicIntegerFieldUpdater
            .newUpdater(NettyResponseFuture.class, "redirectCount");
    @SuppressWarnings("rawtypes")
    private static final AtomicIntegerFieldUpdater<NettyResponseFuture> CURRENT_RETRY_UPDATER = AtomicIntegerFieldUpdater
            .newUpdater(NettyResponseFuture.class, "currentRetry");
    @SuppressWarnings("rawtypes")
    private static final AtomicIntegerFieldUpdater<NettyResponseFuture> IS_DONE_FIELD = AtomicIntegerFieldUpdater
            .newUpdater(NettyResponseFuture.class, "isDone");
    @SuppressWarnings("rawtypes")
    private static final AtomicIntegerFieldUpdater<NettyResponseFuture> IS_CANCELLED_FIELD = AtomicIntegerFieldUpdater
            .newUpdater(NettyResponseFuture.class, "isCancelled");
    @SuppressWarnings("rawtypes")
    private static final AtomicIntegerFieldUpdater<NettyResponseFuture> IN_AUTH_FIELD = AtomicIntegerFieldUpdater
            .newUpdater(NettyResponseFuture.class, "inAuth");
    @SuppressWarnings("rawtypes")
    private static final AtomicIntegerFieldUpdater<NettyResponseFuture> IN_PROXY_AUTH_FIELD = AtomicIntegerFieldUpdater
            .newUpdater(NettyResponseFuture.class, "inProxyAuth");
    @SuppressWarnings("rawtypes")
    private static final AtomicIntegerFieldUpdater<NettyResponseFuture> CONTENT_PROCESSED_FIELD = AtomicIntegerFieldUpdater
            .newUpdater(NettyResponseFuture.class, "contentProcessed");
    @SuppressWarnings("rawtypes")
    private static final AtomicIntegerFieldUpdater<NettyResponseFuture> ON_THROWABLE_CALLED_FIELD = AtomicIntegerFieldUpdater
            .newUpdater(NettyResponseFuture.class, "onThrowableCalled");
    @SuppressWarnings("rawtypes")
    private static final AtomicReferenceFieldUpdater<NettyResponseFuture, TimeoutsHolder> TIMEOUTS_HOLDER_FIELD = AtomicReferenceFieldUpdater
            .newUpdater(NettyResponseFuture.class, TimeoutsHolder.class, "timeoutsHolder");
    @SuppressWarnings("rawtypes")
    private static final AtomicReferenceFieldUpdater<NettyResponseFuture, Object> PARTITION_KEY_LOCK_FIELD = AtomicReferenceFieldUpdater
            .newUpdater(NettyResponseFuture.class, Object.class, "partitionKeyLock");

    private final long start = unpreciseMillisTime();
    private final ChannelPoolPartitioning connectionPoolPartitioning;
    private final ConnectionSemaphore connectionSemaphore;
    private final ProxyServer proxyServer;
    private final int maxRetry;
    private final CompletableFuture<V> future = new CompletableFuture<>();
    public Throwable pendingException;
    // state mutated from outside the event loop
    // TODO check if they are indeed mutated outside the event loop
    private volatile int isDone;
    private volatile int isCancelled;
    private volatile int inAuth;
    private volatile int inProxyAuth;
    @SuppressWarnings("unused")
    private volatile int contentProcessed;
    @SuppressWarnings("unused")
    private volatile int onThrowableCalled;
    @SuppressWarnings("unused")
    private volatile TimeoutsHolder timeoutsHolder;
    // partition key, when != null used to release lock in ChannelManager
    private volatile Object partitionKeyLock;
    // volatile where we need CAS ops
    private volatile int redirectCount;
    private volatile int currentRetry;
    // volatile where we don't need CAS ops
    private volatile long touch = unpreciseMillisTime();
    private volatile ChannelState channelState = ChannelState.NEW;
    // state mutated only inside the event loop
    private Channel channel;
    private boolean keepAlive = true;
    private Request targetRequest;
    private Request currentRequest;
    private NettyRequest nettyRequest;
    private AsyncHandler<V> asyncHandler;
    private boolean streamAlreadyConsumed;
    private boolean reuseChannel;
    private boolean headersAlreadyWrittenOnContinue;
    private boolean dontWriteBodyBecauseExpectContinue;
    private boolean allowConnect;
    private Realm realm;
    private Realm proxyRealm;

    public NettyResponseFuture(Request originalRequest,
                               AsyncHandler<V> asyncHandler,
                               NettyRequest nettyRequest,
                               int maxRetry,
                               ChannelPoolPartitioning connectionPoolPartitioning,
                               ConnectionSemaphore connectionSemaphore,
                               ProxyServer proxyServer) {

        this.asyncHandler = asyncHandler;
        targetRequest = currentRequest = originalRequest;
        this.nettyRequest = nettyRequest;
        this.connectionPoolPartitioning = connectionPoolPartitioning;
        this.connectionSemaphore = connectionSemaphore;
        this.proxyServer = proxyServer;
        this.maxRetry = maxRetry;
    }

    private void releasePartitionKeyLock() {
        if (connectionSemaphore == null) {
            return;
        }

        Object partitionKey = takePartitionKeyLock();
        if (partitionKey != null) {
            connectionSemaphore.releaseChannelLock(partitionKey);
        }
    }

    // Take partition key lock object,
    // but do not release channel lock.
    public Object takePartitionKeyLock() {
        // shortcut, much faster than getAndSet
        if (partitionKeyLock == null) {
            return null;
        }

        return PARTITION_KEY_LOCK_FIELD.getAndSet(this, null);
    }

    // java.util.concurrent.Future

    @Override
    public boolean isDone() {
        return isDone != 0 || isCancelled();
    }

    @Override
    public boolean isCancelled() {
        return isCancelled != 0;
    }

    @Override
    public boolean cancel(boolean force) {
        releasePartitionKeyLock();
        cancelTimeouts();

        if (IS_CANCELLED_FIELD.getAndSet(this, 1) != 0) {
            return false;
        }

        // cancel could happen before channel was attached
        if (channel != null) {
            Channels.setDiscard(channel);
            Channels.silentlyCloseChannel(channel);
        }

        if (ON_THROWABLE_CALLED_FIELD.getAndSet(this, 1) == 0) {
            try {
                asyncHandler.onThrowable(new CancellationException());
            } catch (Throwable t) {
                LOGGER.warn("cancel", t);
            }
        }

        future.cancel(false);
        return true;
    }

    @Override
    public V get() throws InterruptedException, ExecutionException {
        return future.get();
    }

    @Override
    public V get(long l, TimeUnit tu) throws InterruptedException, TimeoutException, ExecutionException {
        return future.get(l, tu);
    }

    private void loadContent() throws ExecutionException {
        if (future.isDone()) {
            try {
                future.get();
            } catch (InterruptedException e) {
                throw new RuntimeException("unreachable", e);
            }
        }

        // No more retry
        CURRENT_RETRY_UPDATER.set(this, maxRetry);
        if (CONTENT_PROCESSED_FIELD.getAndSet(this, 1) == 0) {
            try {
                future.complete(asyncHandler.onCompleted());
            } catch (Throwable ex) {
                if (ON_THROWABLE_CALLED_FIELD.getAndSet(this, 1) == 0) {
                    try {
                        try {
                            asyncHandler.onThrowable(ex);
                        } catch (Throwable t) {
                            LOGGER.debug("asyncHandler.onThrowable", t);
                        }
                    } finally {
                        cancelTimeouts();
                    }
                }
                future.completeExceptionally(ex);
            }
        }
        future.getNow(null);
    }

    // org.asynchttpclient.ListenableFuture

    private boolean terminateAndExit() {
        releasePartitionKeyLock();
        cancelTimeouts();
        channel = null;
        reuseChannel = false;
        return IS_DONE_FIELD.getAndSet(this, 1) != 0 || isCancelled != 0;
    }

    @Override
    public void done() {

        if (terminateAndExit()) {
            return;
        }

        try {
            loadContent();
        } catch (ExecutionException ignored) {

        } catch (RuntimeException t) {
            future.completeExceptionally(t);
        } catch (Throwable t) {
            future.completeExceptionally(t);
            throw t;
        }
    }

    @Override
    public void abort(final Throwable t) {

        if (terminateAndExit()) {
            return;
        }

        future.completeExceptionally(t);

        if (ON_THROWABLE_CALLED_FIELD.compareAndSet(this, 0, 1)) {
            try {
                asyncHandler.onThrowable(t);
            } catch (Throwable te) {
                LOGGER.debug("asyncHandler.onThrowable", te);
            }
        }
    }

    @Override
    public void touch() {
        touch = unpreciseMillisTime();
    }

    @Override
    public ListenableFuture<V> addListener(Runnable listener, Executor exec) {
        if (exec == null) {
            exec = Runnable::run;
        }
        future.whenCompleteAsync((r, v) -> listener.run(), exec);
        return this;
    }

    @Override
    public CompletableFuture<V> toCompletableFuture() {
        return future;
    }

    // INTERNAL

    public Uri getUri() {
        return targetRequest.getUri();
    }

    public ProxyServer getProxyServer() {
        return proxyServer;
    }

    public void cancelTimeouts() {
        TimeoutsHolder ref = TIMEOUTS_HOLDER_FIELD.getAndSet(this, null);
        if (ref != null) {
            ref.cancel();
        }
    }

    public Request getTargetRequest() {
        return targetRequest;
    }

    public void setTargetRequest(Request targetRequest) {
        this.targetRequest = targetRequest;
    }

    public Request getCurrentRequest() {
        return currentRequest;
    }

    public void setCurrentRequest(Request currentRequest) {
        this.currentRequest = currentRequest;
    }

    public NettyRequest getNettyRequest() {
        return nettyRequest;
    }

    public void setNettyRequest(NettyRequest nettyRequest) {
        this.nettyRequest = nettyRequest;
    }

    public AsyncHandler<V> getAsyncHandler() {
        return asyncHandler;
    }

    public void setAsyncHandler(AsyncHandler<V> asyncHandler) {
        this.asyncHandler = asyncHandler;
    }

    public boolean isKeepAlive() {
        return keepAlive;
    }

    public void setKeepAlive(final boolean keepAlive) {
        this.keepAlive = keepAlive;
    }

    public int incrementAndGetCurrentRedirectCount() {
        return REDIRECT_COUNT_UPDATER.incrementAndGet(this);
    }

    public TimeoutsHolder getTimeoutsHolder() {
        return TIMEOUTS_HOLDER_FIELD.get(this);
    }

    public void setTimeoutsHolder(TimeoutsHolder timeoutsHolder) {
        TimeoutsHolder ref = TIMEOUTS_HOLDER_FIELD.getAndSet(this, timeoutsHolder);
        if (ref != null) {
            ref.cancel();
        }
    }

    public boolean isInAuth() {
        return inAuth != 0;
    }

    public void setInAuth(boolean inAuth) {
        this.inAuth = inAuth ? 1 : 0;
    }

    public boolean isAndSetInAuth(boolean set) {
        return IN_AUTH_FIELD.getAndSet(this, set ? 1 : 0) != 0;
    }

    public boolean isInProxyAuth() {
        return inProxyAuth != 0;
    }

    public void setInProxyAuth(boolean inProxyAuth) {
        this.inProxyAuth = inProxyAuth ? 1 : 0;
    }

    public boolean isAndSetInProxyAuth(boolean inProxyAuth) {
        return IN_PROXY_AUTH_FIELD.getAndSet(this, inProxyAuth ? 1 : 0) != 0;
    }

    public ChannelState getChannelState() {
        return channelState;
    }

    public void setChannelState(ChannelState channelState) {
        this.channelState = channelState;
    }

    public boolean isStreamConsumed() {
        return streamAlreadyConsumed;
    }

    public void setStreamConsumed(boolean streamConsumed) {
        streamAlreadyConsumed = streamConsumed;
    }

    public long getLastTouch() {
        return touch;
    }

    public boolean isHeadersAlreadyWrittenOnContinue() {
        return headersAlreadyWrittenOnContinue;
    }

    public void setHeadersAlreadyWrittenOnContinue(boolean headersAlreadyWrittenOnContinue) {
        this.headersAlreadyWrittenOnContinue = headersAlreadyWrittenOnContinue;
    }

    public boolean isDontWriteBodyBecauseExpectContinue() {
        return dontWriteBodyBecauseExpectContinue;
    }

    public void setDontWriteBodyBecauseExpectContinue(boolean dontWriteBodyBecauseExpectContinue) {
        this.dontWriteBodyBecauseExpectContinue = dontWriteBodyBecauseExpectContinue;
    }

    public boolean isConnectAllowed() {
        return allowConnect;
    }

    public void setConnectAllowed(boolean allowConnect) {
        this.allowConnect = allowConnect;
    }

    public void attachChannel(Channel channel, boolean reuseChannel) {

        // future could have been cancelled first
        if (isDone()) {
            Channels.silentlyCloseChannel(channel);
        }

        this.channel = channel;
        this.reuseChannel = reuseChannel;
    }

    public Channel channel() {
        return channel;
    }

    public boolean isReuseChannel() {
        return reuseChannel;
    }

    public void setReuseChannel(boolean reuseChannel) {
        this.reuseChannel = reuseChannel;
    }

    public boolean incrementRetryAndCheck() {
        return maxRetry > 0 && CURRENT_RETRY_UPDATER.incrementAndGet(this) <= maxRetry;
    }

    /**
     * Return true if the {@link Future} can be recovered. There is some scenario
     * where a connection can be closed by an unexpected IOException, and in some
     * situation we can recover from that exception.
     *
     * @return true if that {@link Future} cannot be recovered.
     */
    public boolean isReplayPossible() {
        return !isDone() && !(Channels.isChannelActive(channel) && !"https".equalsIgnoreCase(getUri().getScheme()))
                && inAuth == 0 && inProxyAuth == 0;
    }

    public long getStart() {
        return start;
    }

    public Object getPartitionKey() {
        return connectionPoolPartitioning.getPartitionKey(targetRequest.getUri(), targetRequest.getVirtualHost(),
                proxyServer);
    }

    public void acquirePartitionLockLazily() throws IOException {
        if (connectionSemaphore == null || partitionKeyLock != null) {
            return;
        }

        Object partitionKey = getPartitionKey();
        connectionSemaphore.acquireChannelLock(partitionKey);
        Object prevKey = PARTITION_KEY_LOCK_FIELD.getAndSet(this, partitionKey);
        if (prevKey != null) {
            // self-check

            connectionSemaphore.releaseChannelLock(prevKey);
            releasePartitionKeyLock();

            throw new IllegalStateException("Trying to acquire partition lock concurrently. Please report.");
        }

        if (isDone()) {
            // may be cancelled while we acquired a lock
            releasePartitionKeyLock();
        }
    }

    public Realm getRealm() {
        return realm;
    }

    public void setRealm(Realm realm) {
        this.realm = realm;
    }

    public Realm getProxyRealm() {
        return proxyRealm;
    }

    public void setProxyRealm(Realm proxyRealm) {
        this.proxyRealm = proxyRealm;
    }

    @Override
    public String toString() {
        return "NettyResponseFuture{" + //
                "currentRetry=" + currentRetry + //
                ",\n\tisDone=" + isDone + //
                ",\n\tisCancelled=" + isCancelled + //
                ",\n\tasyncHandler=" + asyncHandler + //
                ",\n\tnettyRequest=" + nettyRequest + //
                ",\n\tfuture=" + future + //
                ",\n\turi=" + getUri() + //
                ",\n\tkeepAlive=" + keepAlive + //
                ",\n\tredirectCount=" + redirectCount + //
                ",\n\ttimeoutsHolder=" + TIMEOUTS_HOLDER_FIELD.get(this) + //
                ",\n\tinAuth=" + inAuth + //
                ",\n\ttouch=" + touch + //
                '}';
    }
}