/*
 * 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)
    {
        FutureCallback<Object> callback = toFutureCallback(asyncResponse);
        Futures.addCallback(futureResponse, callback, 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);
            }
        };
    }
}

@@ -716,174 +716,8 @@
 rs;%0A
-import com.google.common.util.concurrent.FutureCallback;%0Aimport com.google.common.util.concurrent.Futures;%0Aimport com.google.common.util.concurrent.ListenableFuture;%0A
 impo
@@ -816,67 +816,58 @@
 s.co
-ntainer.TimeoutHandler;%0Aimport javax.ws.rs.core.Respons
+re.Response;%0A%0Aimport java.lang.ref.WeakReferenc
 e;%0A
-%0A
 impo
@@ -874,37 +874,48 @@
 rt java.
-lang.ref.WeakReferenc
+util.concurrent.CompletableFutur
 e;%0Aimpor
@@ -946,16 +946,54 @@
 xecutor;
+%0Aimport java.util.function.BiConsumer;
 %0A%0Aimport
@@ -1337,38 +1337,39 @@
 l WeakReference%3C
-Listen
+Complet
 ableFuture%3C?%3E%3E f
@@ -1447,38 +1447,39 @@
  asyncResponse, 
-Listen
+Complet
 ableFuture%3C?%3E fu
@@ -1910,35 +1910,16 @@
 ference%3C
-ListenableFuture%3C?%3E
 %3E(checkN
@@ -2062,22 +2062,23 @@
 sponse, 
-Listen
+Complet
 ableFutu
@@ -2148,127 +2148,76 @@
     
-F
+f
 uture
-Callback%3CObject%3E callback = toFutureCallback(asyncResponse);%0A        Futures.addCallback(futureResponse, callback
+Response.whenCompleteAsync(toCompletionCallback(asyncResponse)
 , ht
@@ -2689,22 +2689,16 @@
 timeout,
- final
  Respons
@@ -2866,22 +2866,16 @@
 timeout,
- final
  Supplie
@@ -2953,146 +2953,27 @@
 ler(
-new TimeoutHandler()%0A        %7B%0A            @Override%0A            public void handleTimeout(AsyncResponse asyncResponse)%0A            %7B%0A    
+asyncResponse -%3E %7B%0A
     
@@ -3037,20 +3037,16 @@
         
-    
 cancelFu
@@ -3053,30 +3053,16 @@
 ture();%0A
-            %7D%0A
         
@@ -3268,14 +3268,15 @@
     
-Listen
+Complet
 able
@@ -3697,34 +3697,45 @@
 %3CT%3E 
-FutureCallback%3CT%3E toFuture
+BiConsumer%3CT, Throwable%3E toCompletion
 Call
@@ -3739,22 +3739,16 @@
 allback(
-final 
 AsyncRes
@@ -3793,106 +3793,117 @@
 urn 
-new FutureCallback%3CT%3E()%0A 
+(value, throwable) -%3E %7B%0A
       
- %7B%0A
       
-      @Override%0A            public void onSuccess(T valu
+if (throwable != null) %7B%0A                asyncResponse.resume(throwabl
 e)
+;
 %0A   
@@ -3903,32 +3903,51 @@
 e);%0A            
+%7D%0A            else 
 %7B%0A              
@@ -4132,147 +4132,8 @@
    %7D
-%0A%0A            @Override%0A            public void onFailure(Throwable t)%0A            %7B%0A                asyncResponse.resume(t);%0A            %7D
 %0A   

Change AsyncResponseHandler to use CompletableFuture