package io.airlift.http.client.jetty;

import com.google.common.collect.ImmutableList;
import com.google.common.net.HostAndPort;
import com.google.common.primitives.Ints;
import io.airlift.http.client.BodyGenerator;
import io.airlift.http.client.FileBodyGenerator;
import io.airlift.http.client.HttpClientConfig;
import io.airlift.http.client.HttpRequestFilter;
import io.airlift.http.client.Request;
import io.airlift.http.client.RequestStats;
import io.airlift.http.client.ResponseHandler;
import io.airlift.http.client.StaticBodyGenerator;
import io.airlift.http.client.jetty.HttpClientLogger.RequestInfo;
import io.airlift.http.client.jetty.HttpClientLogger.ResponseInfo;
import io.airlift.http.client.spnego.KerberosConfig;
import io.airlift.http.client.spnego.SpnegoAuthentication;
import io.airlift.http.client.spnego.SpnegoAuthenticationProtocolHandler;
import io.airlift.http.client.spnego.SpnegoAuthenticationStore;
import io.airlift.security.pem.PemReader;
import io.airlift.units.Duration;
import org.eclipse.jetty.client.DuplexConnectionPool;
import org.eclipse.jetty.client.HttpClient;
import org.eclipse.jetty.client.HttpClientTransport;
import org.eclipse.jetty.client.HttpExchange;
import org.eclipse.jetty.client.HttpRequest;
import org.eclipse.jetty.client.PoolingHttpDestination;
import org.eclipse.jetty.client.Socks4Proxy;
import org.eclipse.jetty.client.WWWAuthenticationProtocolHandler;
import org.eclipse.jetty.client.api.AuthenticationStore;
import org.eclipse.jetty.client.api.ContentProvider;
import org.eclipse.jetty.client.api.Destination;
import org.eclipse.jetty.client.api.Response;
import org.eclipse.jetty.client.http.HttpClientTransportOverHTTP;
import org.eclipse.jetty.client.http.HttpConnectionOverHTTP;
import org.eclipse.jetty.client.util.BytesContentProvider;
import org.eclipse.jetty.client.util.InputStreamResponseListener;
import org.eclipse.jetty.client.util.PathContentProvider;
import org.eclipse.jetty.http2.client.HTTP2Client;
import org.eclipse.jetty.http2.client.http.HttpClientTransportOverHTTP2;
import org.eclipse.jetty.io.ByteBufferPool;
import org.eclipse.jetty.io.ConnectionStatistics;
import org.eclipse.jetty.io.MappedByteBufferPool;
import org.eclipse.jetty.util.HttpCookieStore;
import org.eclipse.jetty.util.component.LifeCycle;
import org.eclipse.jetty.util.ssl.SslContextFactory;
import org.eclipse.jetty.util.thread.QueuedThreadPool;
import org.eclipse.jetty.util.thread.ScheduledExecutorScheduler;
import org.eclipse.jetty.util.thread.Scheduler;
import org.eclipse.jetty.util.thread.Sweeper;
import org.weakref.jmx.Flatten;
import org.weakref.jmx.Managed;
import org.weakref.jmx.Nested;

import javax.annotation.PreDestroy;

import java.io.File;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.net.URI;
import java.nio.ByteBuffer;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.GeneralSecurityException;
import java.security.KeyStore;
import java.util.List;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.Optional;
import java.util.Queue;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;

import static com.google.common.base.MoreObjects.toStringHelper;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Throwables.throwIfUnchecked;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static io.airlift.http.client.jetty.AuthorizationPreservingHttpClient.setPreserveAuthorization;
import static java.lang.Math.max;
import static java.lang.Math.toIntExact;
import static java.lang.String.format;
import static java.util.Objects.requireNonNull;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static java.util.concurrent.TimeUnit.NANOSECONDS;

public class JettyHttpClient
        implements io.airlift.http.client.HttpClient
{
    static {
        JettyLogging.setup();
    }

    private static final String STATS_KEY = "airlift_stats";
    private static final long SWEEP_PERIOD_MILLIS = 5000;

    private static final AtomicLong NAME_COUNTER = new AtomicLong();

    private final HttpClient httpClient;
    private final long maxContentLength;
    private final long requestTimeoutMillis;
    private final long idleTimeoutMillis;
    private final boolean recordRequestComplete;
    private final boolean logEnabled;
    private final QueuedThreadPoolMBean queuedThreadPoolMBean;
    private final ConnectionStats connectionStats;
    private final RequestStats stats = new RequestStats();
    private final CachedDistribution queuedRequestsPerDestination;
    private final CachedDistribution activeConnectionsPerDestination;
    private final CachedDistribution idleConnectionsPerDestination;

    private final CachedDistribution currentQueuedTime;
    private final CachedDistribution currentRequestTime;
    private final CachedDistribution currentRequestSendTime;
    private final CachedDistribution currentResponseWaitTime;
    private final CachedDistribution currentResponseProcessTime;

    private final List<HttpRequestFilter> requestFilters;
    private final Exception creationLocation = new Exception();
    private final String name;

    private final HttpClientLogger requestLogger;
    private final JettyClientDiagnostics clientDiagnostics;

    public JettyHttpClient()
    {
        this(new HttpClientConfig());
    }

    public JettyHttpClient(HttpClientConfig config)
    {
        this(uniqueName(), config);
    }

    public JettyHttpClient(String name, HttpClientConfig config)
    {
        this(name, config, new KerberosConfig(), ImmutableList.of());
    }

    public JettyHttpClient(
            String name,
            HttpClientConfig config,
            KerberosConfig kerberosConfig,
            Iterable<? extends HttpRequestFilter> requestFilters)
    {
        this.name = requireNonNull(name, "name is null");

        requireNonNull(config, "config is null");
        requireNonNull(requestFilters, "requestFilters is null");

        maxContentLength = config.getMaxContentLength().toBytes();
        requestTimeoutMillis = config.getRequestTimeout().toMillis();
        idleTimeoutMillis = config.getIdleTimeout().toMillis();
        recordRequestComplete = config.getRecordRequestComplete();

        creationLocation.fillInStackTrace();

        SslContextFactory sslContextFactory = new SslContextFactory();
        sslContextFactory.setEndpointIdentificationAlgorithm("HTTPS");
        if (config.getKeyStorePath() != null) {
            Optional<KeyStore> pemKeyStore = tryLoadPemKeyStore(config);
            if (pemKeyStore.isPresent()) {
                sslContextFactory.setKeyStore(pemKeyStore.get());
                sslContextFactory.setKeyStorePassword("");
            }
            else {
                sslContextFactory.setKeyStorePath(config.getKeyStorePath());
                sslContextFactory.setKeyStorePassword(config.getKeyStorePassword());
            }
        }
        if (config.getTrustStorePath() != null) {
            Optional<KeyStore> pemTrustStore = tryLoadPemTrustStore(config);
            if (pemTrustStore.isPresent()) {
                sslContextFactory.setTrustStore(pemTrustStore.get());
                sslContextFactory.setTrustStorePassword("");
            }
            else {
                sslContextFactory.setTrustStorePath(config.getTrustStorePath());
                sslContextFactory.setTrustStorePassword(config.getTrustStorePassword());
            }
        }
        sslContextFactory.setSecureRandomAlgorithm(config.getSecureRandomAlgorithm());
        List<String> includedCipherSuites = config.getHttpsIncludedCipherSuites();
        List<String> excludedCipherSuites = config.getHttpsExcludedCipherSuites();
        sslContextFactory.setIncludeCipherSuites(includedCipherSuites.toArray(new String[0]));
        sslContextFactory.setExcludeCipherSuites(excludedCipherSuites.toArray(new String[0]));

        HttpClientTransport transport;
        if (config.isHttp2Enabled()) {
            HTTP2Client client = new HTTP2Client();
            client.setInitialSessionRecvWindow(toIntExact(config.getHttp2InitialSessionReceiveWindowSize().toBytes()));
            client.setInitialStreamRecvWindow(toIntExact(config.getHttp2InitialStreamReceiveWindowSize().toBytes()));
            client.setInputBufferSize(toIntExact(config.getHttp2InputBufferSize().toBytes()));
            client.setSelectors(config.getSelectorCount());
            transport = new HttpClientTransportOverHTTP2(client);
        }
        else {
            transport = new HttpClientTransportOverHTTP(config.getSelectorCount());
        }

        httpClient = new AuthorizationPreservingHttpClient(transport, sslContextFactory);

        // request and response buffer size
        httpClient.setRequestBufferSize(toIntExact(config.getRequestBufferSize().toBytes()));
        httpClient.setResponseBufferSize(toIntExact(config.getResponseBufferSize().toBytes()));

        // Kerberos authentication
        if (config.getAuthenticationEnabled()) {
            AuthenticationStore store = new SpnegoAuthenticationStore(new SpnegoAuthentication(
                    kerberosConfig.getKeytab(),
                    kerberosConfig.getConfig(),
                    kerberosConfig.getCredentialCache(),
                    config.getKerberosServicePrincipalPattern(),
                    config.getKerberosPrincipal(),
                    config.getKerberosRemoteServiceName(),
                    config.getKerberosNameType(),
                    kerberosConfig.isUseCanonicalHostname()));
            httpClient.setAuthenticationStore(store);
            httpClient.getProtocolHandlers().remove(WWWAuthenticationProtocolHandler.NAME);
            httpClient.getProtocolHandlers().put(new SpnegoAuthenticationProtocolHandler(httpClient));
        }

        httpClient.setMaxConnectionsPerDestination(config.getMaxConnectionsPerServer());
        httpClient.setMaxRequestsQueuedPerDestination(config.getMaxRequestsQueuedPerDestination());

        // disable cookies
        httpClient.setCookieStore(new HttpCookieStore.Empty());

        // remove default user agent
        httpClient.setUserAgentField(null);

        // timeouts
        httpClient.setIdleTimeout(idleTimeoutMillis);
        httpClient.setConnectTimeout(config.getConnectTimeout().toMillis());
        httpClient.setAddressResolutionTimeout(config.getConnectTimeout().toMillis());

        httpClient.setConnectBlocking(config.isConnectBlocking());

        HostAndPort socksProxy = config.getSocksProxy();
        if (socksProxy != null) {
            httpClient.getProxyConfiguration().getProxies().add(new Socks4Proxy(socksProxy.getHost(), socksProxy.getPortOrDefault(1080)));
        }

        httpClient.setByteBufferPool(new MappedByteBufferPool());
        QueuedThreadPool queuedThreadPool = createExecutor(name, config.getMinThreads(), config.getMaxThreads());
        httpClient.setExecutor(queuedThreadPool);
        // add queuedThreadPool as a managed bean to get its state in the client dumps
        httpClient.addBean(queuedThreadPool, true);
        httpClient.setScheduler(createScheduler(name, config.getTimeoutConcurrency(), config.getTimeoutThreads()));

        httpClient.setSocketAddressResolver(new JettyAsyncSocketAddressResolver(
                httpClient.getExecutor(),
                httpClient.getScheduler(),
                config.getConnectTimeout().toMillis()));

        // Jetty client connections can sometimes get stuck while closing which reduces
        // the available connections.  The Jetty Sweeper periodically scans the active
        // connection pool looking for connections in the closed state, and if a connection
        // is observed in the closed state multiple times, it logs, and destroys the connection.
        httpClient.addBean(new Sweeper(httpClient.getScheduler(), SWEEP_PERIOD_MILLIS), true);

        // track connection statistics
        ConnectionStatistics connectionStats = new ConnectionStatistics();
        httpClient.addBean(connectionStats);
        this.connectionStats = new ConnectionStats(connectionStats);

        // configure logging
        this.logEnabled = config.isLogEnabled();
        if (logEnabled) {
            String logFilePath = Paths.get(config.getLogPath(), format("%s-http-client.log", name)).toAbsolutePath().toString();
            requestLogger = new DefaultHttpClientLogger(
                    logFilePath,
                    config.getLogHistory(),
                    config.getLogQueueSize(),
                    config.getLogBufferSize(),
                    config.getLogFlushInterval(),
                    config.getLogMaxFileSize().toBytes(),
                    config.isLogCompressionEnabled());
        }
        else {
            requestLogger = new NoopLogger();
        }

        try {
            httpClient.start();

            // remove the GZIP encoding from the client
            // TODO: there should be a better way to to do this
            httpClient.getContentDecoderFactories().clear();
        }
        catch (Exception e) {
            if (e instanceof InterruptedException) {
                Thread.currentThread().interrupt();
            }
            throwIfUnchecked(e);
            throw new RuntimeException(e);
        }

        this.clientDiagnostics = new JettyClientDiagnostics();

        this.requestFilters = ImmutableList.copyOf(requestFilters);

        this.queuedThreadPoolMBean = new QueuedThreadPoolMBean((QueuedThreadPool) httpClient.getExecutor());

        this.activeConnectionsPerDestination = new ConnectionPoolDistribution(httpClient,
                (distribution, connectionPool) -> distribution.add(connectionPool.getActiveConnections().size()));

        this.idleConnectionsPerDestination = new ConnectionPoolDistribution(httpClient,
                (distribution, connectionPool) -> distribution.add(connectionPool.getIdleConnections().size()));

        this.queuedRequestsPerDestination = new DestinationDistribution(httpClient,
                (distribution, destination) -> distribution.add(destination.getHttpExchanges().size()));

        this.currentQueuedTime = new RequestDistribution(httpClient, (distribution, listener, now) -> {
            long started = listener.getRequestStarted();
            if (started == 0) {
                started = now;
            }
            distribution.add(NANOSECONDS.toMillis(started - listener.getCreated()));
        });

        this.currentRequestTime = new RequestDistribution(httpClient, (distribution, listener, now) -> {
            long started = listener.getRequestStarted();
            if (started == 0) {
                return;
            }
            long finished = listener.getResponseFinished();
            if (finished == 0) {
                finished = now;
            }
            distribution.add(NANOSECONDS.toMillis(finished - started));
        });

        this.currentRequestSendTime = new RequestDistribution(httpClient, (distribution, listener, now) -> {
            long started = listener.getRequestStarted();
            if (started == 0) {
                return;
            }
            long requestSent = listener.getRequestFinished();
            if (requestSent == 0) {
                requestSent = now;
            }
            distribution.add(NANOSECONDS.toMillis(requestSent - started));
        });

        this.currentResponseWaitTime = new RequestDistribution(httpClient, (distribution, listener, now) -> {
            long requestSent = listener.getRequestFinished();
            if (requestSent == 0) {
                return;
            }
            long responseStarted = listener.getResponseStarted();
            if (responseStarted == 0) {
                responseStarted = now;
            }
            distribution.add(NANOSECONDS.toMillis(responseStarted - requestSent));
        });

        this.currentResponseProcessTime = new RequestDistribution(httpClient, (distribution, listener, now) -> {
            long responseStarted = listener.getResponseStarted();
            if (responseStarted == 0) {
                return;
            }
            long finished = listener.getResponseFinished();
            if (finished == 0) {
                finished = now;
            }
            distribution.add(NANOSECONDS.toMillis(finished - responseStarted));
        });
    }

    private static Optional<KeyStore> tryLoadPemKeyStore(HttpClientConfig config)
    {
        File keyStoreFile = new File(config.getKeyStorePath());
        try {
            if (!PemReader.isPem(keyStoreFile)) {
                return Optional.empty();
            }
        }
        catch (IOException e) {
            throw new IllegalArgumentException("Error reading key store file: " + keyStoreFile, e);
        }

        try {
            return Optional.of(PemReader.loadKeyStore(keyStoreFile, keyStoreFile, Optional.ofNullable(config.getKeyStorePassword())));
        }
        catch (IOException | GeneralSecurityException e) {
            throw new IllegalArgumentException("Error loading PEM key store: " + keyStoreFile, e);
        }
    }

    private static Optional<KeyStore> tryLoadPemTrustStore(HttpClientConfig config)
    {
        File trustStoreFile = new File(config.getTrustStorePath());
        try {
            if (!PemReader.isPem(trustStoreFile)) {
                return Optional.empty();
            }
        }
        catch (IOException e) {
            throw new IllegalArgumentException("Error reading trust store file: " + trustStoreFile, e);
        }

        try {
            if (PemReader.readCertificateChain(trustStoreFile).isEmpty()) {
                throw new IllegalArgumentException("PEM trust store file does not contain any certificates: " + trustStoreFile);
            }
            return Optional.of(PemReader.loadTrustStore(trustStoreFile));
        }
        catch (IOException | GeneralSecurityException e) {
            throw new IllegalArgumentException("Error loading PEM trust store: " + trustStoreFile, e);
        }
    }

    private static QueuedThreadPool createExecutor(String name, int minThreads, int maxThreads)
    {
        try {
            QueuedThreadPool pool = new QueuedThreadPool(maxThreads, minThreads, 60000, null);
            pool.setName("http-client-" + name);
            pool.setDaemon(true);
            pool.start();
            pool.setStopTimeout(2000);
            pool.setDetailedDump(true);
            return pool;
        }
        catch (Exception e) {
            throwIfUnchecked(e);
            throw new RuntimeException(e);
        }
    }

    private static Scheduler createScheduler(String name, int timeoutConcurrency, int timeoutThreads)
    {
        Scheduler scheduler;
        String threadName = "http-client-" + name + "-scheduler";
        if ((timeoutConcurrency == 1) && (timeoutThreads == 1)) {
            scheduler = new ScheduledExecutorScheduler(threadName, true);
        }
        else {
            checkArgument(timeoutConcurrency >= 1, "timeoutConcurrency must be at least one");
            int threads = max(1, timeoutThreads / timeoutConcurrency);
            scheduler = new ConcurrentScheduler(timeoutConcurrency, threads, threadName);
        }

        try {
            scheduler.start();
        }
        catch (Exception e) {
            throwIfUnchecked(e);
            throw new RuntimeException(e);
        }

        return scheduler;
    }

    @Override
    public <T, E extends Exception> T execute(Request request, ResponseHandler<T, E> responseHandler)
            throws E
    {
        long requestStart = System.nanoTime();

        // apply filters
        request = applyRequestFilters(request);

        // create jetty request and response listener
        JettyRequestListener requestListener = new JettyRequestListener(request.getUri());
        HttpRequest jettyRequest = buildJettyRequest(request, requestListener);
        InputStreamResponseListener listener = new InputStreamResponseListener()
        {
            @Override
            public void onContent(Response response, ByteBuffer content)
            {
                // ignore empty blocks
                if (content.remaining() == 0) {
                    return;
                }
                super.onContent(response, content);
            }
        };

        long requestTimestamp = System.currentTimeMillis();
        RequestInfo requestInfo = RequestInfo.from(jettyRequest, requestTimestamp);
        if (logEnabled) {
            addLoggingListener(jettyRequest, requestTimestamp);
        }

        // fire the request
        jettyRequest.send(listener);

        // wait for response to begin
        Response response;
        try {
            response = listener.get(httpClient.getIdleTimeout(), MILLISECONDS);
        }
        catch (InterruptedException e) {
            stats.recordRequestFailed();
            requestLogger.log(requestInfo, ResponseInfo.failed(Optional.empty(), Optional.of(e)));
            jettyRequest.abort(e);
            Thread.currentThread().interrupt();
            return responseHandler.handleException(request, e);
        }
        catch (TimeoutException e) {
            stats.recordRequestFailed();
            requestLogger.log(requestInfo, ResponseInfo.failed(Optional.empty(), Optional.of(e)));
            jettyRequest.abort(e);
            return responseHandler.handleException(request, e);
        }
        catch (ExecutionException e) {
            stats.recordRequestFailed();
            requestLogger.log(requestInfo, ResponseInfo.failed(Optional.empty(), Optional.of(e)));
            Throwable cause = e.getCause();
            if (cause instanceof Exception) {
                return responseHandler.handleException(request, (Exception) cause);
            }
            else if ((cause instanceof NoClassDefFoundError) && cause.getMessage().endsWith("ALPNClientConnection")) {
                return responseHandler.handleException(request, new RuntimeException("HTTPS cannot be used when HTTP/2 is enabled", cause));
            }
            else {
                return responseHandler.handleException(request, new RuntimeException(cause));
            }
        }

        // process response
        long responseStart = System.nanoTime();

        JettyResponse jettyResponse = null;
        T value;
        try {
            jettyResponse = new JettyResponse(response, listener.getInputStream());
            value = responseHandler.handle(request, jettyResponse);
        }
        finally {
            if (recordRequestComplete) {
                recordRequestComplete(stats, request, requestStart, jettyResponse, responseStart);
            }
        }
        return value;
    }

    @Override
    public <T, E extends Exception> HttpResponseFuture<T> executeAsync(Request request, ResponseHandler<T, E> responseHandler)
    {
        requireNonNull(request, "request is null");
        requireNonNull(responseHandler, "responseHandler is null");

        request = applyRequestFilters(request);

        HttpRequest jettyRequest = buildJettyRequest(request, new JettyRequestListener(request.getUri()));

        JettyResponseFuture<T, E> future = new JettyResponseFuture<>(request, jettyRequest, responseHandler, stats, recordRequestComplete);

        BufferingResponseListener listener = new BufferingResponseListener(future, Ints.saturatedCast(maxContentLength));

        long requestTimestamp = System.currentTimeMillis();

        if (logEnabled) {
            addLoggingListener(jettyRequest, requestTimestamp);
        }

        try {
            jettyRequest.send(listener);
        }
        catch (RuntimeException e) {
            if (!(e instanceof RejectedExecutionException)) {
                e = new RejectedExecutionException(e);
            }
            // normally this is a rejected execution exception because the client has been closed
            future.failed(e);
            requestLogger.log(RequestInfo.from(jettyRequest, requestTimestamp), ResponseInfo.failed(Optional.empty(), Optional.of(e)));
        }
        return future;
    }

    private void addLoggingListener(HttpRequest jettyRequest, long requestTimestamp)
    {
        HttpClientLoggingListener loggingListener = new HttpClientLoggingListener(jettyRequest, requestTimestamp, requestLogger);
        jettyRequest.listener(loggingListener);
        jettyRequest.onResponseBegin(loggingListener);
        jettyRequest.onComplete(loggingListener);
    }

    private Request applyRequestFilters(Request request)
    {
        for (HttpRequestFilter requestFilter : requestFilters) {
            request = requestFilter.filterRequest(request);
        }
        return request;
    }

    private HttpRequest buildJettyRequest(Request finalRequest, JettyRequestListener listener)
    {
        HttpRequest jettyRequest = (HttpRequest) httpClient.newRequest(finalRequest.getUri());
        jettyRequest.onRequestBegin(request -> listener.onRequestBegin());
        jettyRequest.onRequestSuccess(request -> listener.onRequestEnd());
        jettyRequest.onResponseBegin(response -> listener.onResponseBegin());
        jettyRequest.onComplete(result -> listener.onFinish());
        jettyRequest.onComplete(result -> {
            if (result.isFailed() && result.getFailure() instanceof TimeoutException) {
                clientDiagnostics.logDiagnosticsInfo(httpClient);
            }
        });

        jettyRequest.attribute(STATS_KEY, listener);

        jettyRequest.method(finalRequest.getMethod());

        for (Entry<String, String> entry : finalRequest.getHeaders().entries()) {
            jettyRequest.header(entry.getKey(), entry.getValue());
        }

        BodyGenerator bodyGenerator = finalRequest.getBodyGenerator();
        if (bodyGenerator != null) {
            if (bodyGenerator instanceof StaticBodyGenerator) {
                StaticBodyGenerator staticBodyGenerator = (StaticBodyGenerator) bodyGenerator;
                jettyRequest.content(new BytesContentProvider(staticBodyGenerator.getBody()));
            }
            else if (bodyGenerator instanceof FileBodyGenerator) {
                Path path = ((FileBodyGenerator) bodyGenerator).getPath();
                jettyRequest.content(fileContentProvider(path));
            }
            else {
                jettyRequest.content(new BodyGeneratorContentProvider(bodyGenerator, httpClient.getExecutor()));
            }
        }

        jettyRequest.followRedirects(finalRequest.isFollowRedirects());

        setPreserveAuthorization(jettyRequest, finalRequest.isPreserveAuthorizationOnRedirect());

        // timeouts
        jettyRequest.timeout(requestTimeoutMillis, MILLISECONDS);
        jettyRequest.idleTimeout(idleTimeoutMillis, MILLISECONDS);

        return jettyRequest;
    }

    private static ContentProvider fileContentProvider(Path path)
    {
        try {
            PathContentProvider provider = new PathContentProvider(null, path);
            provider.setByteBufferPool(new ByteBufferPool()
            {
                @Override
                public ByteBuffer acquire(int size, boolean direct)
                {
                    return ByteBuffer.allocate(size);
                }

                @Override
                public void release(ByteBuffer buffer) {}
            });
            return provider;
        }
        catch (IOException e) {
            throw new UncheckedIOException(e);
        }
    }

    public List<HttpRequestFilter> getRequestFilters()
    {
        return requestFilters;
    }

    public long getRequestTimeoutMillis()
    {
        return requestTimeoutMillis;
    }

    @Override
    @Managed
    @Flatten
    public RequestStats getStats()
    {
        return stats;
    }

    @Override
    public long getMaxContentLength()
    {
        return maxContentLength;
    }

    @Managed
    @Nested
    public QueuedThreadPoolMBean getThreadPool()
    {
        return queuedThreadPoolMBean;
    }

    @Managed
    @Nested
    public ConnectionStats getConnectionStats()
    {
        return connectionStats;
    }

    @Managed
    @Nested
    public CachedDistribution getActiveConnectionsPerDestination()
    {
        return activeConnectionsPerDestination;
    }

    @Managed
    @Nested
    public CachedDistribution getIdleConnectionsPerDestination()
    {
        return idleConnectionsPerDestination;
    }

    @Managed
    @Nested
    public CachedDistribution getQueuedRequestsPerDestination()
    {
        return queuedRequestsPerDestination;
    }

    @Managed
    @Nested
    public CachedDistribution getCurrentQueuedTime()
    {
        return currentQueuedTime;
    }

    @Managed
    @Nested
    public CachedDistribution getCurrentRequestTime()
    {
        return currentRequestTime;
    }

    @Managed
    @Nested
    public CachedDistribution getCurrentRequestSendTime()
    {
        return currentRequestSendTime;
    }

    @Managed
    @Nested
    public CachedDistribution getCurrentResponseWaitTime()
    {
        return currentResponseWaitTime;
    }

    @Managed
    @Nested
    public CachedDistribution getCurrentResponseProcessTime()
    {
        return currentResponseProcessTime;
    }

    @Managed
    public String dump()
    {
        return httpClient.dump();
    }

    @Managed
    public void dumpStdErr()
    {
        httpClient.dumpStdErr();
    }

    @Managed
    public String dumpAllDestinations()
    {
        return format("%s\t%s\t%s\t%s\t%s\n", "URI", "queued", "request", "wait", "response") +
                httpClient.getDestinations().stream()
                        .map(JettyHttpClient::dumpDestination)
                        .collect(Collectors.joining("\n"));
    }

    @Managed
    public int getLoggerQueueSize()
    {
        return requestLogger.getQueueSize();
    }

    // todo this should be @Managed but operations with parameters are broken in jmx utils https://github.com/martint/jmxutils/issues/27
    @SuppressWarnings("UnusedDeclaration")
    public String dumpDestination(URI uri)
    {
        Destination destination = httpClient.getDestination(uri.getScheme(), uri.getHost(), uri.getPort());
        if (destination == null) {
            return null;
        }

        return dumpDestination(destination);
    }

    private static String dumpDestination(Destination destination)
    {
        long now = System.nanoTime();
        return getRequestListenersForDestination(destination).stream()
                .map(listener -> dumpRequest(now, listener))
                .sorted()
                .collect(Collectors.joining("\n"));
    }

    static List<JettyRequestListener> getRequestListenersForDestination(Destination destination)
    {
        return getRequestForDestination(destination).stream()
                .map(request -> request.getAttributes().get(STATS_KEY))
                .map(JettyRequestListener.class::cast)
                .filter(Objects::nonNull)
                .collect(Collectors.toList());
    }

    private static List<org.eclipse.jetty.client.api.Request> getRequestForDestination(Destination destination)
    {
        PoolingHttpDestination poolingHttpDestination = (PoolingHttpDestination) destination;
        Queue<HttpExchange> httpExchanges = poolingHttpDestination.getHttpExchanges();

        List<org.eclipse.jetty.client.api.Request> requests = httpExchanges.stream()
                .map(HttpExchange::getRequest)
                .collect(Collectors.toList());

        ((DuplexConnectionPool) poolingHttpDestination.getConnectionPool()).getActiveConnections().stream()
                .filter(HttpConnectionOverHTTP.class::isInstance)
                .map(HttpConnectionOverHTTP.class::cast)
                .map(connection -> connection.getHttpChannel().getHttpExchange())
                .filter(Objects::nonNull)
                .forEach(exchange -> requests.add(exchange.getRequest()));

        return requests.stream()
                .filter(Objects::nonNull)
                .collect(toImmutableList());
    }

    private static String dumpRequest(long now, JettyRequestListener listener)
    {
        long created = listener.getCreated();
        long requestStarted = listener.getRequestStarted();
        if (requestStarted == 0) {
            requestStarted = now;
        }
        long requestFinished = listener.getRequestFinished();
        if (requestFinished == 0) {
            requestFinished = now;
        }
        long responseStarted = listener.getResponseStarted();
        if (responseStarted == 0) {
            responseStarted = now;
        }
        long finished = listener.getResponseFinished();
        if (finished == 0) {
            finished = now;
        }
        return format("%s\t%.1f\t%.1f\t%.1f\t%.1f",
                listener.getUri(),
                nanosToMillis(requestStarted - created),
                nanosToMillis(requestFinished - requestStarted),
                nanosToMillis(responseStarted - requestFinished),
                nanosToMillis(finished - responseStarted));
    }

    private static double nanosToMillis(long nanos)
    {
        return new Duration(nanos, NANOSECONDS).getValue(MILLISECONDS);
    }

    @PreDestroy
    @Override
    public void close()
    {
        // client must be destroyed before the pools or
        // you will create a several second busy wait loop
        closeQuietly(httpClient);
        closeQuietly((LifeCycle) httpClient.getExecutor());
        closeQuietly(httpClient.getScheduler());
        requestLogger.close();
    }

    @Override
    public boolean isClosed()
    {
        return !httpClient.isRunning();
    }

    @Override
    public String toString()
    {
        return toStringHelper(this)
                .addValue(name)
                .toString();
    }

    @SuppressWarnings("UnusedDeclaration")
    public StackTraceElement[] getCreationLocation()
    {
        return creationLocation.getStackTrace();
    }

    private static void closeQuietly(LifeCycle service)
    {
        try {
            if (service != null) {
                service.stop();
            }
        }
        catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
        catch (Exception ignored) {
        }
    }

    private static String uniqueName()
    {
        return "anonymous" + NAME_COUNTER.incrementAndGet();
    }

    static void recordRequestComplete(RequestStats requestStats, Request request, long requestStart, JettyResponse response, long responseStart)
    {
        if (response == null) {
            return;
        }

        Duration responseProcessingTime = Duration.nanosSince(responseStart);
        Duration requestProcessingTime = new Duration(responseStart - requestStart, NANOSECONDS);

        requestStats.recordResponseReceived(request.getMethod(),
                response.getStatusCode(),
                response.getBytesRead(),
                response.getBytesRead(),
                requestProcessingTime,
                responseProcessingTime);
    }
}


/*
 * 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.stats.cardinality;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.Ordering;
import com.google.common.primitives.Ints;
import io.airlift.slice.BasicSliceInput;
import io.airlift.slice.DynamicSliceOutput;
import io.airlift.slice.SizeOf;
import io.airlift.slice.Slice;
import org.openjdk.jol.info.ClassLayout;

import javax.annotation.concurrent.NotThreadSafe;

import java.util.Arrays;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static io.airlift.slice.SizeOf.sizeOf;
import static io.airlift.stats.cardinality.Utils.computeIndex;
import static io.airlift.stats.cardinality.Utils.linearCounting;
import static io.airlift.stats.cardinality.Utils.numberOfBuckets;
import static io.airlift.stats.cardinality.Utils.numberOfLeadingZeros;
import static java.lang.Math.toIntExact;
import static java.util.Comparator.comparingInt;

@NotThreadSafe
final class SparseHll
        implements HllInstance
{
    private static final int SPARSE_INSTANCE_SIZE = ClassLayout.parseClass(SparseHll.class).instanceSize();

    // 6 bits to encode the number of zeros after the truncated hash
    // and be able to fit the encoded value in an integer
    private static final int VALUE_BITS = 6;
    private static final int VALUE_MASK = (1 << VALUE_BITS) - 1;
    private static final int EXTENDED_PREFIX_BITS = Integer.SIZE - VALUE_BITS;

    private final byte indexBitLength;
    private short numberOfEntries;
    private int[] entries;

    public SparseHll(int indexBitLength)
    {
        validatePrefixLength(indexBitLength);

        this.indexBitLength = (byte) indexBitLength;
        entries = new int[1];
    }

    public SparseHll(Slice serialized)
    {
        BasicSliceInput input = serialized.getInput();

        checkArgument(input.readByte() == Format.SPARSE_V2.getTag(), "invalid format tag");

        indexBitLength = input.readByte();
        validatePrefixLength(indexBitLength);

        numberOfEntries = input.readShort();

        entries = new int[numberOfEntries];
        for (int i = 0; i < numberOfEntries; i++) {
            entries[i] = input.readInt();
        }

        checkArgument(!input.isReadable(), "input is too big");
    }

    public static boolean canDeserialize(Slice serialized)
    {
        return serialized.getByte(0) == Format.SPARSE_V2.getTag();
    }

    public void insertHash(long hash)
    {
        // TODO: investigate whether accumulate, sort and merge results in better performance due to avoiding the shift+insert in every call

        int bucket = Utils.computeIndex(hash, EXTENDED_PREFIX_BITS);
        int position = searchBucket(bucket);

        // add entry if missing
        if (position < 0) {
            // ensure capacity
            if (numberOfEntries + 1 > entries.length) {
                entries = Arrays.copyOf(entries, entries.length + 10);
            }

            // shift right
            int insertionPoint = -(position + 1);
            if (insertionPoint < numberOfEntries) {
                System.arraycopy(entries, insertionPoint, entries, insertionPoint + 1, numberOfEntries - insertionPoint);
            }

            entries[insertionPoint] = encode(hash);
            numberOfEntries++;
        }
        else {
            int currentEntry = entries[position];
            int newValue = Utils.numberOfLeadingZeros(hash, EXTENDED_PREFIX_BITS);

            if (decodeBucketValue(currentEntry) < newValue) {
                entries[position] = encode(bucket, newValue);
            }
        }
    }

    private int encode(long hash)
    {
        return encode(computeIndex(hash, EXTENDED_PREFIX_BITS), numberOfLeadingZeros(hash, EXTENDED_PREFIX_BITS));
    }

    private static int encode(int bucketIndex, int value)
    {
        return (bucketIndex << VALUE_BITS) | value;
    }

    private static int decodeBucketIndex(int entry)
    {
        return decodeBucketIndex(EXTENDED_PREFIX_BITS, entry);
    }

    private static int decodeBucketValue(int entry)
    {
        return entry & VALUE_MASK;
    }

    private static int decodeBucketIndex(int indexBitLength, int entry)
    {
        return entry >>> (Integer.SIZE - indexBitLength);
    }

    public void mergeWith(SparseHll other)
    {
        entries = mergeEntries(other);
        numberOfEntries = (short) entries.length;
    }

    public DenseHll toDense()
    {
        DenseHll result = new DenseHll(indexBitLength);

        for (int i = 0; i < numberOfEntries; i++) {
            int entry = entries[i];

            // The leading EXTENDED_BITS_LENGTH are a proper subset of the original hash.
            // Since we're guaranteed that indexBitLength is <= EXTENDED_BITS_LENGTH,
            // the value stored in those bits corresponds to the bucket index in the dense HLL
            int bucket = decodeBucketIndex(indexBitLength, entry);

            // compute the number of zeros between indexBitLength and EXTENDED_BITS_LENGTH
            int zeros = Integer.numberOfLeadingZeros(entry << indexBitLength);

            // if zeros > EXTENDED_BITS_LENGTH - indexBits, it means all those bits were zeros,
            // so look at the entry value, which contains the number of leading 0 *after* EXTENDED_BITS_LENGTH
            int bits = EXTENDED_PREFIX_BITS - indexBitLength;
            if (zeros > bits) {
                zeros = bits + decodeBucketValue(entry);
            }

            result.insert(bucket, zeros + 1); // + 1 because HLL stores leading number of zeros + 1
        }

        return result;
    }

    @Override
    public long cardinality()
    {
        // Estimate the cardinality using linear counting over the theoretical 2^EXTENDED_BITS_LENGTH buckets available due
        // to the fact that we're recording the raw leading EXTENDED_BITS_LENGTH of the hash. This produces much better precision
        // while in the sparse regime.
        int totalBuckets = numberOfBuckets(EXTENDED_PREFIX_BITS);
        int zeroBuckets = totalBuckets - numberOfEntries;

        return Math.round(linearCounting(zeroBuckets, totalBuckets));
    }

    @Override
    public int estimatedInMemorySize()
    {
        return SPARSE_INSTANCE_SIZE + toIntExact(sizeOf(entries));
    }

    @Override
    public int getIndexBitLength()
    {
        return indexBitLength;
    }

    /**
     * Returns a index of the entry if found. Otherwise, it returns -(insertionPoint + 1)
     */
    private int searchBucket(int bucketIndex)
    {
        int low = 0;
        int high = numberOfEntries - 1;

        while (low <= high) {
            int middle = (low + high) >>> 1;

            int middleBucketIndex = decodeBucketIndex(entries[middle]);

            if (bucketIndex > middleBucketIndex) {
                low = middle + 1;
            }
            else if (bucketIndex < middleBucketIndex) {
                high = middle - 1;
            }
            else {
                return middle;
            }
        }

        return -(low + 1); // not found... return insertion point
    }

    private int[] mergeEntries(SparseHll other)
    {
        int[] result = new int[numberOfEntries + other.numberOfEntries];
        int leftIndex = 0;
        int rightIndex = 0;

        int index = 0;
        while (leftIndex < numberOfEntries && rightIndex < other.numberOfEntries) {
            int left = decodeBucketIndex(entries[leftIndex]);
            int right = decodeBucketIndex(other.entries[rightIndex]);

            if (left < right) {
                result[index++] = entries[leftIndex++];
            }
            else if (left > right) {
                result[index++] = other.entries[rightIndex++];
            }
            else {
                int value = Math.max(decodeBucketValue(entries[leftIndex]), decodeBucketValue(other.entries[rightIndex]));
                result[index++] = encode(left, value);
                leftIndex++;
                rightIndex++;
            }
        }

        while (leftIndex < numberOfEntries) {
            result[index++] = entries[leftIndex++];
        }

        while (rightIndex < other.numberOfEntries) {
            result[index++] = other.entries[rightIndex++];
        }

        return Arrays.copyOf(result, index);
    }

    public Slice serialize()
    {
        int size = SizeOf.SIZE_OF_BYTE + // format tag
                SizeOf.SIZE_OF_BYTE + // p
                SizeOf.SIZE_OF_SHORT + // number of entries
                SizeOf.SIZE_OF_INT * numberOfEntries;

        DynamicSliceOutput out = new DynamicSliceOutput(size)
                .appendByte(Format.SPARSE_V2.getTag())
                .appendByte(indexBitLength)
                .appendShort(numberOfEntries);

        for (int i = 0; i < numberOfEntries; i++) {
            out.appendInt(entries[i]);
        }

        return out.slice();
    }

    @Override
    public int estimatedSerializedSize()
    {
        return SizeOf.SIZE_OF_SHORT // type + version
                + SizeOf.SIZE_OF_BYTE  // p
                + SizeOf.SIZE_OF_SHORT // numberOfEntries
                + SizeOf.SIZE_OF_INT * numberOfEntries; // entries
    }

    private static void validatePrefixLength(int indexBitLength)
    {
        checkArgument(indexBitLength >= 1 && indexBitLength <= EXTENDED_PREFIX_BITS, "indexBitLength is out of range");
    }

    @VisibleForTesting
    public void verify()
    {
        checkState(numberOfEntries <= entries.length,
                "Expected number of hashes (%s) larger than array length (%s)",
                numberOfEntries, entries.length);

        checkState(Ordering.from(comparingInt(e -> decodeBucketIndex((Integer) e)))
                        .isOrdered(Ints.asList(Arrays.copyOf(entries, numberOfEntries))),
                "entries are not sorted");
    }
}


package io.airlift.stats;

import java.util.concurrent.TimeUnit;

import static com.google.common.base.Preconditions.checkArgument;

public final class ExponentialDecay
{
    private ExponentialDecay()
    {
    }

    public static double oneMinute()
    {
        // alpha for a target weight of 1/E at 1 minute
        return 1.0 / TimeUnit.MINUTES.toSeconds(1);
    }

    public static double fiveMinutes()
    {
        // alpha for a target weight of 1/E at 5 minutes
        return 1.0 / TimeUnit.MINUTES.toSeconds(5);
    }

    public static double fifteenMinutes()
    {
        // alpha for a target weight of 1/E at 15 minutes
        return 1.0 / TimeUnit.MINUTES.toSeconds(15);
    }

    public static double seconds(int seconds)
    {
        // alpha for a target weight of 1/E at the specified number of seconds
        return 1.0 / seconds;
    }

    /**
     * Compute the alpha decay factor such that the weight of an entry with age 'targetAgeInSeconds' is targetWeight'
     */
    public static double computeAlpha(double targetWeight, long targetAgeInSeconds)
    {
        checkArgument(targetAgeInSeconds > 0, "targetAgeInSeconds must be > 0");
        checkArgument(targetWeight > 0 && targetWeight < 1, "targetWeight must be in range (0, 1)");

        return -Math.log(targetWeight) / targetAgeInSeconds;
    }
}


package io.airlift.concurrent;

import com.google.common.util.concurrent.AsyncFunction;
import com.google.common.util.concurrent.FluentFuture;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture;
import io.airlift.units.Duration;

import javax.annotation.Nullable;

import java.lang.ref.WeakReference;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.Callable;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.Collectors;

import static com.google.common.base.MoreObjects.firstNonNull;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Throwables.propagateIfPossible;
import static com.google.common.base.Throwables.throwIfInstanceOf;
import static com.google.common.collect.Iterables.isEmpty;
import static com.google.common.util.concurrent.Futures.immediateFailedFuture;
import static com.google.common.util.concurrent.Futures.immediateFuture;
import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
import static java.util.Objects.requireNonNull;
import static java.util.concurrent.TimeUnit.MILLISECONDS;

public final class MoreFutures
{
    private MoreFutures() {}

    /**
     * Cancels the destination Future if the source Future is cancelled.
     */
    public static <X, Y> void propagateCancellation(ListenableFuture<? extends X> source, Future<? extends Y> destination, boolean mayInterruptIfRunning)
    {
        source.addListener(() -> {
            if (source.isCancelled()) {
                destination.cancel(mayInterruptIfRunning);
            }
        }, directExecutor());
    }

    /**
     * Mirrors all results of the source Future to the destination Future.
     *

     * This also propagates cancellations from the destination Future back to the source Future.
     */
    public static <T> void mirror(ListenableFuture<? extends T> source, SettableFuture<? super T> destination, boolean mayInterruptIfRunning)
    {
        FutureCallback<T> callback = new FutureCallback<T>()
        {
            @Override
            public void onSuccess(@Nullable T result)
            {
                destination.set(result);
            }

            @Override
            public void onFailure(Throwable t)
            {
                destination.setException(t);
            }
        };
        Futures.addCallback(source, callback, directExecutor());
        propagateCancellation(destination, source, mayInterruptIfRunning);
    }

    /**
     * Attempts to unwrap a throwable that has been wrapped in a {@link CompletionException}.
     */
    public static Throwable unwrapCompletionException(Throwable throwable)
    {
        if (throwable instanceof CompletionException) {
            return firstNonNull(throwable.getCause(), throwable);
        }
        return throwable;
    }

    /**
     * Returns a future that can not be completed or canceled.
     */
    @Deprecated
    public static <V> CompletableFuture<V> unmodifiableFuture(CompletableFuture<V> future)
    {
        return unmodifiableFuture(future, false);
    }

    /**
     * Returns a future that can not be completed or optionally canceled.
     */
    @Deprecated
    public static <V> CompletableFuture<V> unmodifiableFuture(CompletableFuture<V> future, boolean propagateCancel)
    {
        requireNonNull(future, "future is null");

        Function<Boolean, Boolean> onCancelFunction;
        if (propagateCancel) {
            onCancelFunction = future::cancel;
        }
        else {
            onCancelFunction = mayInterrupt -> false;
        }

        UnmodifiableCompletableFuture<V> unmodifiableFuture = new UnmodifiableCompletableFuture<>(onCancelFunction);
        future.whenComplete((value, exception) -> {
            if (exception != null) {
                unmodifiableFuture.internalCompleteExceptionally(exception);
            }
            else {
                unmodifiableFuture.internalComplete(value);
            }
        });
        return unmodifiableFuture;
    }

    /**
     * Returns a failed future containing the specified throwable.
     */
    @Deprecated
    public static <V> CompletableFuture<V> failedFuture(Throwable throwable)
    {
        requireNonNull(throwable, "throwable is null");
        CompletableFuture<V> future = new CompletableFuture<>();
        future.completeExceptionally(throwable);
        return future;
    }

    /**
     * Waits for the value from the future. If the future is failed, the exception
     * is thrown directly if unchecked or wrapped in a RuntimeException. If the
     * thread is interrupted, the thread interruption flag is set and the original
     * InterruptedException is wrapped in a RuntimeException and thrown.
     */
    public static <V> V getFutureValue(Future<V> future)
    {
        return getFutureValue(future, RuntimeException.class);
    }

    /**
     * Waits for the value from the future. If the future is failed, the exception
     * is thrown directly if it is an instance of the specified exception type or
     * unchecked, or it is wrapped in a RuntimeException. If the thread is
     * interrupted, the thread interruption flag is set and the original
     * InterruptedException is wrapped in a RuntimeException and thrown.
     */
    public static <V, E extends Exception> V getFutureValue(Future<V> future, Class<E> exceptionType)
            throws E
    {
        requireNonNull(future, "future is null");
        requireNonNull(exceptionType, "exceptionType is null");

        try {
            return future.get();
        }
        catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException("interrupted", e);
        }
        catch (ExecutionException e) {
            Throwable cause = e.getCause() == null ? e : e.getCause();
            propagateIfPossible(cause, exceptionType);
            throw new RuntimeException(cause);
        }
    }

    /**
     * Gets the current value of the future without waiting. If the future
     * value is null, an empty Optional is still returned, and in this case the caller
     * must check the future directly for the null value.
     */
    public static <T> Optional<T> tryGetFutureValue(Future<T> future)
    {
        requireNonNull(future, "future is null");
        if (!future.isDone()) {
            return Optional.empty();
        }

        return tryGetFutureValue(future, 0, MILLISECONDS);
    }

    /**
     * Waits for the the value from the future for the specified time.  If the future
     * value is null, an empty Optional is still returned, and in this case the caller
     * must check the future directly for the null value.  If the future is failed,
     * the exception is thrown directly if unchecked or wrapped in a RuntimeException.
     * If the thread is interrupted, the thread interruption flag is set and the original
     * InterruptedException is wrapped in a RuntimeException and thrown.
     */
    public static <V> Optional<V> tryGetFutureValue(Future<V> future, int timeout, TimeUnit timeUnit)
    {
        return tryGetFutureValue(future, timeout, timeUnit, RuntimeException.class);
    }

    /**
     * Waits for the the value from the future for the specified time.  If the future
     * value is null, an empty Optional is still returned, and in this case the caller
     * must check the future directly for the null value.  If the future is failed,
     * the exception is thrown directly if it is an instance of the specified exception
     * type or unchecked, or it is wrapped in a RuntimeException. If the thread is
     * interrupted, the thread interruption flag is set and the original
     * InterruptedException is wrapped in a RuntimeException and thrown.
     */
    public static <V, E extends Exception> Optional<V> tryGetFutureValue(Future<V> future, int timeout, TimeUnit timeUnit, Class<E> exceptionType)
            throws E
    {
        requireNonNull(future, "future is null");
        checkArgument(timeout >= 0, "timeout is negative");
        requireNonNull(timeUnit, "timeUnit is null");
        requireNonNull(exceptionType, "exceptionType is null");

        try {
            return Optional.ofNullable(future.get(timeout, timeUnit));
        }
        catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException("interrupted", e);
        }
        catch (ExecutionException e) {
            Throwable cause = e.getCause() == null ? e : e.getCause();
            propagateIfPossible(cause, exceptionType);
            throw new RuntimeException(cause);
        }
        catch (TimeoutException expected) {
            // expected
        }
        return Optional.empty();
    }

    /**
     * Returns the result of the input {@link Future}, which must have already completed.
     *

     * Similar to Futures{@link Futures#getDone(Future)}, but does not throw checked exceptions.
     */
    public static <T> T getDone(Future<T> future)
    {
        requireNonNull(future, "future is null");
        checkArgument(future.isDone(), "future not done yet");
        return getFutureValue(future);
    }

    /**
     * Checks that the completed future completed successfully.
     */
    public static void checkSuccess(Future<?> future, String errorMessage)
    {
        requireNonNull(future, "future is null");
        requireNonNull(errorMessage, "errorMessage is null");
        checkArgument(future.isDone(), "future not done yet");
        try {
            getFutureValue(future);
        }
        catch (RuntimeException e) {
            throw new IllegalArgumentException(errorMessage, e);
        }
    }

    /**
     * Creates a future that completes when the first future completes either normally
     * or exceptionally. Cancellation of the future propagates to the supplied futures.
     */
    public static <V> ListenableFuture<V> whenAnyComplete(Iterable<? extends ListenableFuture<? extends V>> futures)
    {
        requireNonNull(futures, "futures is null");
        checkArgument(!isEmpty(futures), "futures is empty");

        ExtendedSettableFuture<V> firstCompletedFuture = ExtendedSettableFuture.create();
        for (ListenableFuture<? extends V> future : futures) {
            firstCompletedFuture.setAsync(future);
        }
        return firstCompletedFuture;
    }

    /**
     * Creates a future that completes when the first future completes either normally
     * or exceptionally. All other futures are cancelled when one completes.
     * Cancellation of the returned future propagates to the supplied futures.
     *

     * It is critical for the performance of this function that
     * {@code guava.concurrent.generate_cancellation_cause} is false,
     * which is the default since Guava v20.
     */
    public static <V> ListenableFuture<V> whenAnyCompleteCancelOthers(Iterable<? extends ListenableFuture<? extends V>> futures)
    {
        requireNonNull(futures, "futures is null");
        checkArgument(!isEmpty(futures), "futures is empty");

        // wait for the first task to unblock and then cancel all futures to free up resources
        ListenableFuture<V> anyComplete = whenAnyComplete(futures);
        anyComplete.addListener(
                () -> {
                    for (ListenableFuture<?> future : futures) {
                        future.cancel(true);
                    }
                },
                directExecutor());
        return anyComplete;
    }

    /**
     * Creates a future that completes when the first future completes either normally
     * or exceptionally. Cancellation of the future does not propagate to the supplied
     * futures.
     */
    @Deprecated
    public static <V> CompletableFuture<V> firstCompletedFuture(Iterable<? extends CompletionStage<? extends V>> futures)
    {
        return firstCompletedFuture(futures, false);
    }

    /**
     * Creates a future that completes when the first future completes either normally
     * or exceptionally. Cancellation of the future will optionally propagate to the
     * supplied futures.
     */
    @Deprecated
    public static <V> CompletableFuture<V> firstCompletedFuture(Iterable<? extends CompletionStage<? extends V>> futures, boolean propagateCancel)
    {
        requireNonNull(futures, "futures is null");
        checkArgument(!isEmpty(futures), "futures is empty");

        CompletableFuture<V> future = new CompletableFuture<>();
        for (CompletionStage<? extends V> stage : futures) {
            stage.whenComplete((value, exception) -> {
                if (exception != null) {
                    future.completeExceptionally(exception);
                }
                else {
                    future.complete(value);
                }
            });
        }
        if (propagateCancel) {
            future.exceptionally(throwable -> {
                if (throwable instanceof CancellationException) {
                    for (CompletionStage<? extends V> sourceFuture : futures) {
                        if (sourceFuture instanceof Future) {
                            ((Future<?>) sourceFuture).cancel(true);
                        }
                    }
                }
                return null;
            });
        }
        return future;
    }

    /**
     * Returns an unmodifiable future that is completed when all of the given
     * futures complete. If any of the given futures complete exceptionally, then the
     * returned future also does so immediately, with a CompletionException holding this exception
     * as its cause. Otherwise, the results of the given futures are reflected in the
     * returned future as a list of results matching the input order. If no futures are
     * provided, returns a future completed with an empty list.
     */
    @Deprecated
    public static <V> CompletableFuture<List<V>> allAsList(List<CompletableFuture<? extends V>> futures)
    {
        CompletableFuture<Void> allDoneFuture = CompletableFuture.allOf(futures.toArray(new CompletableFuture[futures.size()]));

        // Eagerly propagate exceptions, rather than waiting for all the futures to complete first (default behavior)
        for (CompletableFuture<? extends V> future : futures) {
            future.whenComplete((v, throwable) -> {
                if (throwable != null) {
                    allDoneFuture.completeExceptionally(throwable);
                }
            });
        }

        return unmodifiableFuture(allDoneFuture.thenApply(v ->
                futures.stream()
                        .map(CompletableFuture::join)
                        .collect(Collectors.<V>toList())));
    }

    /**
     * Returns a new future that is completed when the supplied future completes or
     * when the timeout expires.  If the timeout occurs or the returned CompletableFuture
     * is canceled, the supplied future will be canceled.
     */
    public static <T> ListenableFuture<T> addTimeout(ListenableFuture<T> future, Callable<T> onTimeout, Duration timeout, ScheduledExecutorService executorService)
    {
        AsyncFunction<TimeoutException, T> timeoutHandler = timeoutException -> {
            try {
                return immediateFuture(onTimeout.call());
            }
            catch (Throwable throwable) {
                return immediateFailedFuture(throwable);
            }
        };
        return FluentFuture.from(future)
                .withTimeout(timeout.toMillis(), MILLISECONDS, executorService)
                .catchingAsync(TimeoutException.class, timeoutHandler, directExecutor());
    }

    /**
     * Returns a new future that is completed when the supplied future completes or
     * when the timeout expires.  If the timeout occurs or the returned CompletableFuture
     * is canceled, the supplied future will be canceled.
     */
    @Deprecated
    public static <T> CompletableFuture<T> addTimeout(CompletableFuture<T> future, Callable<T> onTimeout, Duration timeout, ScheduledExecutorService executorService)
    {
        requireNonNull(future, "future is null");
        requireNonNull(onTimeout, "timeoutValue is null");
        requireNonNull(timeout, "timeout is null");
        requireNonNull(executorService, "executorService is null");

        // if the future is already complete, just return it
        if (future.isDone()) {
            return future;
        }

        // create an unmodifiable future that propagates cancel
        // down cast is safe because this is our code
        UnmodifiableCompletableFuture<T> futureWithTimeout = (UnmodifiableCompletableFuture<T>) unmodifiableFuture(future, true);

        // schedule a task to complete the future when the time expires
        ScheduledFuture<?> timeoutTaskFuture = executorService.schedule(new TimeoutFutureTask<>(futureWithTimeout, onTimeout, future), timeout.toMillis(), MILLISECONDS);

        // when future completes, cancel the timeout task
        future.whenCompleteAsync((value, exception) -> timeoutTaskFuture.cancel(false), executorService);

        return futureWithTimeout;
    }

    /**
     * Converts a ListenableFuture to a CompletableFuture. Cancellation of the
     * CompletableFuture will be propagated to the ListenableFuture.
     */
    public static <V> CompletableFuture<V> toCompletableFuture(ListenableFuture<V> listenableFuture)
    {
        requireNonNull(listenableFuture, "listenableFuture is null");

        CompletableFuture<V> future = new CompletableFuture<>();
        future.exceptionally(throwable -> {
            if (throwable instanceof CancellationException) {
                listenableFuture.cancel(true);
            }
            return null;
        });

        FutureCallback<V> callback = new FutureCallback<V>()
        {
            @Override
            public void onSuccess(V result)
            {
                future.complete(result);
            }

            @Override
            public void onFailure(Throwable t)
            {
                future.completeExceptionally(t);
            }
        };
        Futures.addCallback(listenableFuture, callback, directExecutor());
        return future;
    }

    /**
     * Converts a CompletableFuture to a ListenableFuture. Cancellation of the
     * ListenableFuture will be propagated to the CompletableFuture.
     */
    public static <V> ListenableFuture<V> toListenableFuture(CompletableFuture<V> completableFuture)
    {
        requireNonNull(completableFuture, "completableFuture is null");
        SettableFuture<V> future = SettableFuture.create();
        propagateCancellation(future, completableFuture, true);

        completableFuture.whenComplete((value, exception) -> {
            if (exception != null) {
                future.setException(exception);
            }
            else {
                future.set(value);
            }
        });
        return future;
    }

    /**
     * Invokes the callback if the future completes successfully. Note, this uses the direct
     * executor, so the callback should not be resource intensive.
     */
    public static <T> void addSuccessCallback(ListenableFuture<T> future, Consumer<T> successCallback)
    {
        addSuccessCallback(future, successCallback, directExecutor());
    }

    /**
     * Invokes the callback, using the specified executor, if the future completes successfully.
     */
    public static <T> void addSuccessCallback(ListenableFuture<T> future, Consumer<T> successCallback, Executor executor)
    {
        requireNonNull(future, "future is null");
        requireNonNull(successCallback, "successCallback is null");

        FutureCallback<T> callback = new FutureCallback<T>()
        {
            @Override
            public void onSuccess(@Nullable T result)
            {
                successCallback.accept(result);
            }

            @Override
            public void onFailure(Throwable t) {}
        };
        Futures.addCallback(future, callback, executor);
    }

    /**
     * Invokes the callback if the future completes successfully. Note, this uses the direct
     * executor, so the callback should not be resource intensive.
     */
    public static <T> void addSuccessCallback(ListenableFuture<T> future, Runnable successCallback)
    {
        addSuccessCallback(future, successCallback, directExecutor());
    }

    /**
     * Invokes the callback, using the specified executor, if the future completes successfully.
     */
    public static <T> void addSuccessCallback(ListenableFuture<T> future, Runnable successCallback, Executor executor)
    {
        requireNonNull(successCallback, "successCallback is null");

        addSuccessCallback(future, t -> successCallback.run(), executor);
    }

    /**
     * Invokes the callback if the future fails. Note, this uses the direct
     * executor, so the callback should not be resource intensive.
     */
    public static <T> void addExceptionCallback(ListenableFuture<T> future, Consumer<Throwable> exceptionCallback)
    {
        addExceptionCallback(future, exceptionCallback, directExecutor());
    }

    /**
     * Invokes the callback, using the specified executor, if the future fails.
     */
    public static <T> void addExceptionCallback(ListenableFuture<T> future, Consumer<Throwable> exceptionCallback, Executor executor)
    {
        requireNonNull(future, "future is null");
        requireNonNull(exceptionCallback, "exceptionCallback is null");

        FutureCallback<T> callback = new FutureCallback<T>()
        {
            @Override
            public void onSuccess(@Nullable T result) {}

            @Override
            public void onFailure(Throwable t)
            {
                exceptionCallback.accept(t);
            }
        };
        Futures.addCallback(future, callback, executor);
    }

    /**
     * Invokes the callback if the future fails. Note, this uses the direct
     * executor, so the callback should not be resource intensive.
     */
    public static <T> void addExceptionCallback(ListenableFuture<T> future, Runnable exceptionCallback)
    {
        addExceptionCallback(future, exceptionCallback, directExecutor());
    }

    /**
     * Invokes the callback, using the specified executor, if the future fails.
     */
    public static <T> void addExceptionCallback(ListenableFuture<T> future, Runnable exceptionCallback, Executor executor)
    {
        requireNonNull(exceptionCallback, "exceptionCallback is null");

        addExceptionCallback(future, t -> exceptionCallback.run(), executor);
    }

    private static class UnmodifiableCompletableFuture<V>
            extends CompletableFuture<V>
    {
        private final Function<Boolean, Boolean> onCancel;

        public UnmodifiableCompletableFuture(Function<Boolean, Boolean> onCancel)
        {
            this.onCancel = requireNonNull(onCancel, "onCancel is null");
        }

        void internalComplete(V value)
        {
            super.complete(value);
        }

        void internalCompleteExceptionally(Throwable ex)
        {
            super.completeExceptionally(ex);
        }

        @Override
        public boolean cancel(boolean mayInterruptIfRunning)
        {
            return onCancel.apply(mayInterruptIfRunning);
        }

        @Override
        public boolean complete(V value)
        {
            throw new UnsupportedOperationException();
        }

        @Override
        public boolean completeExceptionally(Throwable ex)
        {
            if (ex instanceof CancellationException) {
                return cancel(false);
            }
            throw new UnsupportedOperationException();
        }

        @Override
        public void obtrudeValue(V value)
        {
            throw new UnsupportedOperationException();
        }

        @Override
        public void obtrudeException(Throwable ex)
        {
            throw new UnsupportedOperationException();
        }
    }

    private static class TimeoutFutureTask<T>
            implements Runnable
    {
        private final UnmodifiableCompletableFuture<T> settableFuture;
        private final Callable<T> timeoutValue;
        private final WeakReference<CompletableFuture<T>> futureReference;

        public TimeoutFutureTask(UnmodifiableCompletableFuture<T> settableFuture, Callable<T> timeoutValue, CompletableFuture<T> future)
        {
            this.settableFuture = settableFuture;
            this.timeoutValue = timeoutValue;

            // the scheduled executor can hold on to the timeout task for a long time, and
            // the future can reference large expensive objects.  Since we are only interested
            // in canceling this future on a timeout, only hold a weak reference to the future
            this.futureReference = new WeakReference<>(future);
        }

        @Override
        public void run()
        {
            if (settableFuture.isDone()) {
                return;
            }

            // run the timeout task and set the result into the future
            try {
                T result = timeoutValue.call();
                settableFuture.internalComplete(result);
            }
            catch (Throwable t) {
                settableFuture.internalCompleteExceptionally(t);
                throwIfInstanceOf(t, RuntimeException.class);
            }

            // cancel the original future, if it still exists
            Future<T> future = futureReference.get();
            if (future != null) {
                future.cancel(true);
            }
        }
    }
}