/*
 * 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.http.server;

import com.google.common.annotations.Beta;
import com.google.common.base.Supplier;
import com.google.common.base.Suppliers;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import io.airlift.units.Duration;

import javax.ws.rs.container.AsyncResponse;
import javax.ws.rs.container.TimeoutHandler;
import javax.ws.rs.core.Response;
import java.lang.ref.WeakReference;
import java.util.concurrent.Executor;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static javax.ws.rs.core.Response.status;

@Beta
public class AsyncResponseHandler
{
    private final AsyncResponse asyncResponse;
    private final WeakReference<ListenableFuture<?>> futureResponseReference;

    private AsyncResponseHandler(AsyncResponse asyncResponse, ListenableFuture<?> futureResponse)
    {
        this.asyncResponse = checkNotNull(asyncResponse, "asyncResponse is null");
        // the jaxrs implementation can hold on to the async timeout 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.futureResponseReference = new WeakReference<ListenableFuture<?>>(checkNotNull(futureResponse, "futureResponse is null"));
    }

    public static AsyncResponseHandler bindAsyncResponse(AsyncResponse asyncResponse, ListenableFuture<?> futureResponse, Executor httpResponseExecutor)
    {
        Futures.addCallback(futureResponse, toFutureCallback(asyncResponse), httpResponseExecutor);
        return new AsyncResponseHandler(asyncResponse, futureResponse);
    }

    public AsyncResponseHandler withTimeout(Duration timeout)
    {
        return withTimeout(timeout,
                status(Response.Status.SERVICE_UNAVAILABLE)
                        .entity("Timed out after waiting for " + timeout.convertToMostSuccinctTimeUnit())
                        .build());
    }

    public AsyncResponseHandler withTimeout(Duration timeout, final Response timeoutResponse)
    {
        return withTimeout(timeout, Suppliers.ofInstance(timeoutResponse));
    }

    public AsyncResponseHandler withTimeout(Duration timeout, final Supplier<Response> timeoutResponse)
    {
        asyncResponse.setTimeoutHandler(new TimeoutHandler()
        {
            @Override
            public void handleTimeout(AsyncResponse asyncResponse)
            {
                asyncResponse.resume(timeoutResponse.get());
                cancelFuture();
            }
        });
        asyncResponse.setTimeout(timeout.toMillis(), MILLISECONDS);
        return this;
    }

    private void cancelFuture()
    {
        // Cancel the original future if it still exists
        ListenableFuture<?> futureResponse = futureResponseReference.get();
        if (futureResponse != null) {
            try {
                // Do not interrupt the future if it is running. Jersey uses
                // the calling thread to write the response to the wire.
                futureResponse.cancel(false);
            }
            catch (Exception ignored) {
            }
        }
    }

    private static <T> FutureCallback<T> toFutureCallback(final AsyncResponse asyncResponse)
    {
        return new FutureCallback<T>()
        {
            @Override
            public void onSuccess(T value)
            {
                checkArgument(!(value instanceof Response.ResponseBuilder), "Value is a ResponseBuilder. Did you forget to call build?");
                asyncResponse.resume(value);
            }

            @Override
            public void onFailure(Throwable t)
            {
                asyncResponse.resume(t);
            }
        };
    }
}

@@ -2295,21 +2295,16 @@
 ture
-s.add
 Callback
 (fut
@@ -2303,33 +2303,80 @@
 back
-(futureResponse, to
+%3CObject%3E callback = toFutureCallback(asyncResponse);%0A        
 Future
+s.add
 Call
@@ -2380,29 +2380,30 @@
 allback(
-async
+future
 Response
 ), httpR
@@ -2394,17 +2394,26 @@
 Response
-)
+, callback
 , httpRe

Fix type inference compile error on Java 8