/*
 * 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.scheduler.group;

import com.facebook.presto.execution.Lifespan;
import com.facebook.presto.execution.scheduler.BucketNodeMap;
import com.facebook.presto.execution.scheduler.SourceScheduler;
import com.facebook.presto.metadata.InternalNode;
import com.facebook.presto.spi.connector.ConnectorPartitionHandle;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.SettableFuture;
import it.unimi.dsi.fastutil.ints.IntArrayList;
import it.unimi.dsi.fastutil.ints.IntListIterator;

import javax.annotation.concurrent.GuardedBy;

import java.util.ArrayList;
import java.util.List;
import java.util.OptionalInt;
import java.util.stream.IntStream;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.base.Verify.verify;
import static java.util.Collections.unmodifiableList;
import static java.util.Objects.requireNonNull;

/**
 * See {@link LifespanScheduler} about thread safety
 */
public class DynamicLifespanScheduler
        implements LifespanScheduler
{
    private final BucketNodeMap bucketNodeMap;
    private final List<InternalNode> allNodes;
    private final List<ConnectorPartitionHandle> partitionHandles;
    private final OptionalInt concurrentLifespansPerTask;

    private final IntListIterator driverGroups;

    // initialScheduled does not need to be guarded because this object
    // is safely published after its mutation.
    private boolean initialScheduled;
    // Write to newDriverGroupReady field is guarded. Read of the reference
    // is either guarded, or is guaranteed to happen in the same thread as the write.
    private SettableFuture<?> newDriverGroupReady = SettableFuture.create();

    @GuardedBy("this")
    private final List<Lifespan> recentlyCompletelyExecutedDriverGroups = new ArrayList<>();
    @GuardedBy("this")
    private int totalLifespanExecutionFinished;

    public DynamicLifespanScheduler(BucketNodeMap bucketNodeMap, List<InternalNode> allNodes, List<ConnectorPartitionHandle> partitionHandles, OptionalInt concurrentLifespansPerTask)
    {
        this.bucketNodeMap = requireNonNull(bucketNodeMap, "bucketNodeMap is null");
        this.allNodes = requireNonNull(allNodes, "allNodes is null");
        this.partitionHandles = unmodifiableList(new ArrayList<>(
                requireNonNull(partitionHandles, "partitionHandles is null")));

        this.concurrentLifespansPerTask = requireNonNull(concurrentLifespansPerTask, "concurrentLifespansPerTask is null");
        concurrentLifespansPerTask.ifPresent(lifespansPerTask -> checkArgument(lifespansPerTask >= 1, "concurrentLifespansPerTask must be great or equal to 1 if present"));

        int bucketCount = partitionHandles.size();
        verify(bucketCount > 0);
        this.driverGroups = new IntArrayList(IntStream.range(0, bucketCount).toArray()).iterator();
    }

    @Override
    public void scheduleInitial(SourceScheduler scheduler)
    {
        checkState(!initialScheduled);
        initialScheduled = true;

        int driverGroupsScheduledPerTask = 0;
        while (driverGroups.hasNext()) {
            for (int i = 0; i < allNodes.size() && driverGroups.hasNext(); i++) {
                int driverGroupId = driverGroups.nextInt();
                checkState(!bucketNodeMap.getAssignedNode(driverGroupId).isPresent());
                bucketNodeMap.assignBucketToNode(driverGroupId, allNodes.get(i));
                scheduler.startLifespan(Lifespan.driverGroup(driverGroupId), partitionHandles.get(driverGroupId));
            }

            driverGroupsScheduledPerTask++;
            if (concurrentLifespansPerTask.isPresent() && driverGroupsScheduledPerTask == concurrentLifespansPerTask.getAsInt()) {
                break;
            }
        }
    }

    @Override
    public void onLifespanExecutionFinished(Iterable<Lifespan> newlyCompletelyExecutedDriverGroups)
    {
        checkState(initialScheduled);

        SettableFuture<?> newDriverGroupReady;
        synchronized (this) {
            for (Lifespan newlyCompletelyExecutedDriverGroup : newlyCompletelyExecutedDriverGroups) {
                checkArgument(!newlyCompletelyExecutedDriverGroup.isTaskWide());
                recentlyCompletelyExecutedDriverGroups.add(newlyCompletelyExecutedDriverGroup);
                totalLifespanExecutionFinished++;
            }
            newDriverGroupReady = this.newDriverGroupReady;
        }
        newDriverGroupReady.set(null);
    }

    @Override
    public SettableFuture schedule(SourceScheduler scheduler)
    {
        // Return a new future even if newDriverGroupReady has not finished.
        // Returning the same SettableFuture instance could lead to ListenableFuture retaining too many listener objects.

        checkState(initialScheduled);

        List<Lifespan> recentlyCompletedDriverGroups;
        synchronized (this) {
            recentlyCompletedDriverGroups = ImmutableList.copyOf(this.recentlyCompletelyExecutedDriverGroups);
            this.recentlyCompletelyExecutedDriverGroups.clear();
            newDriverGroupReady = SettableFuture.create();
        }

        for (Lifespan driverGroup : recentlyCompletedDriverGroups) {
            if (!driverGroups.hasNext()) {
                break;
            }
            int driverGroupId = driverGroups.nextInt();

            InternalNode nodeForCompletedDriverGroup = bucketNodeMap.getAssignedNode(driverGroup.getId()).orElseThrow(IllegalStateException::new);
            bucketNodeMap.assignBucketToNode(driverGroupId, nodeForCompletedDriverGroup);
            scheduler.startLifespan(Lifespan.driverGroup(driverGroupId), partitionHandles.get(driverGroupId));
        }

        return newDriverGroupReady;
    }

    @Override
    public synchronized boolean allLifespanExecutionFinished()
    {
        return totalLifespanExecutionFinished == partitionHandles.size();
    }
}

@@ -1776,24 +1776,28 @@
 alNode%3E 
-allNodes
+nodeByTaskId
 ;%0A    pr
@@ -2634,24 +2634,28 @@
 alNode%3E 
-allNodes
+nodeByTaskId
 , List%3CC
@@ -2841,24 +2841,28 @@
    this.
-allNodes
+nodeByTaskId
  = requi
@@ -2875,27 +2875,35 @@
 ull(
-allNodes, %22allNodes
+nodeByTaskId, %22nodeByTaskId
  is 
@@ -3817,24 +3817,28 @@
  0; i %3C 
-allNodes
+nodeByTaskId
 .size() 
@@ -4086,16 +4086,20 @@
 Id, 
-allNodes
+nodeByTaskId
 .get

Rename allNodes to NodeByTaskId