/*
 * 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 com.facebook.presto.execution;

import com.google.common.collect.ImmutableSet;
import com.google.common.util.concurrent.AbstractFuture;
import com.google.common.util.concurrent.ListenableFuture;

import javax.annotation.Nullable;
import javax.annotation.concurrent.GuardedBy;
import javax.annotation.concurrent.ThreadSafe;

import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.CancellationException;
import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicBoolean;

import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
import static java.util.Objects.requireNonNull;

@ThreadSafe
public class FutureStateChange<T>
{
    // Use a separate future for each listener so canceled listeners can be removed
    @GuardedBy("listeners")
    private final Set<StateTrackingFuture<T>> listeners = new HashSet<>();

    public ListenableFuture<T> createNewListener()
    {
        StateTrackingFuture<T> listener = StateTrackingFuture.create();
        synchronized (listeners) {
            listeners.add(listener);
        }

        // remove the listener when the future completes
        listener.addListener(
                () -> {
                    synchronized (listeners) {
                        listeners.remove(listener);
                    }
                },
                directExecutor());

        return listener;
    }

    public void complete(T newState)
    {
        fireStateChange(newState, directExecutor());
    }

    public void complete(T newState, Executor executor)
    {
        fireStateChange(newState, executor);
    }

    private void fireStateChange(T newState, Executor executor)
    {
        requireNonNull(executor, "executor is null");
        Set<StateTrackingFuture<T>> futures;
        synchronized (listeners) {
            futures = ImmutableSet.copyOf(listeners);
            listeners.clear();
        }

        for (StateTrackingFuture<T> future : futures) {
            executor.execute(() -> future.set(newState));
        }
    }

    /**
     * Future cancellation can be expensive, because when a future is canceled
     * a CancellationException is created and thrown under the hood. To get rid
     * of the cancellation overhead this future implementation just cancels the
     * future with the same exception instance.
     */
    private static final class StateTrackingFuture<V>
            extends AbstractFuture<V>
    {
        private static final CancellationException EXCEPTION = new CancellationException("Future is canceled");
        private AtomicBoolean canceled = new AtomicBoolean();

        public static <V> StateTrackingFuture<V> create()
        {
            return new StateTrackingFuture<V>();
        }

        @Override
        public boolean cancel(boolean mayInterruptIfRunning)
        {
            if (canceled.compareAndSet(false, true)) {
                setException(EXCEPTION);
                return true;
            }
            return false;
        }

        @Override
        protected boolean set(@Nullable V value)
        {
            return super.set(value);
        }

        @Override
        public boolean isCancelled()
        {
            return canceled.get();
        }
    }
}

@@ -683,24 +683,26 @@
 current.
-Abstract
+Listenable
 Future;%0A
@@ -734,38 +734,36 @@
 util.concurrent.
-Listen
+Sett
 ableFuture;%0A%0Aimp
@@ -763,42 +763,8 @@
 e;%0A%0A
-import javax.annotation.Nullable;%0A
 impo
@@ -933,117 +933,16 @@
 ent.
-CancellationException;%0Aimport java.util.concurrent.Executor;%0Aimport java.util.concurrent.atomic.AtomicBoolean
+Executor
 ;%0A%0Ai
@@ -1246,36 +1246,31 @@
  final Set%3CS
-tateTracking
+ettable
 Future%3CT%3E%3E l
@@ -1364,28 +1364,23 @@
        S
-tateTracking
+ettable
 Future%3CT
@@ -1393,28 +1393,23 @@
 ener = S
-tateTracking
+ettable
 Future.c
@@ -2173,28 +2173,23 @@
    Set%3CS
-tateTracking
+ettable
 Future%3CT
@@ -2345,28 +2345,23 @@
   for (S
-tateTracking
+ettable
 Future%3CT
@@ -2459,1218 +2459,7 @@
    %7D
-%0A%0A    /**%0A     * Future cancellation can be expensive, because when a future is canceled%0A     * a CancellationException is created and thrown under the hood. To get rid%0A     * of the cancellation overhead this future implementation just cancels the%0A     * future with the same exception instance.%0A     */%0A    private static final class StateTrackingFuture%3CV%3E%0A            extends AbstractFuture%3CV%3E%0A    %7B%0A        private static final CancellationException EXCEPTION = new CancellationException(%22Future is canceled%22);%0A        private AtomicBoolean canceled = new AtomicBoolean();%0A%0A        public static %3CV%3E StateTrackingFuture%3CV%3E create()%0A        %7B%0A            return new StateTrackingFuture%3CV%3E();%0A        %7D%0A%0A        @Override%0A        public boolean cancel(boolean mayInterruptIfRunning)%0A        %7B%0A            if (canceled.compareAndSet(false, true)) %7B%0A                setException(EXCEPTION);%0A                return true;%0A            %7D%0A            return false;%0A        %7D%0A%0A        @Override%0A        protected boolean set(@Nullable V value)%0A        %7B%0A            return super.set(value);%0A        %7D%0A%0A        @Override%0A        public boolean isCancelled()%0A        %7B%0A            return canceled.get();%0A        %7D%0A    %7D
 %0A%7D%0A

Remove performance optimization for Future cancellation