/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.flink.runtime.jobmaster.slotpool;

import org.apache.flink.api.common.time.Deadline;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.runtime.clusterframework.types.AllocationID;
import org.apache.flink.runtime.clusterframework.types.ResourceProfile;
import org.apache.flink.runtime.executiongraph.TestingComponentMainThreadExecutor;
import org.apache.flink.runtime.executiongraph.utils.SimpleAckingTaskManagerGateway;
import org.apache.flink.runtime.jobmanager.slots.TaskManagerGateway;
import org.apache.flink.runtime.jobmaster.SlotRequestId;
import org.apache.flink.runtime.resourcemanager.SlotRequest;
import org.apache.flink.runtime.resourcemanager.utils.TestingResourceManagerGateway;
import org.apache.flink.runtime.taskexecutor.slot.SlotOffer;
import org.apache.flink.runtime.taskmanager.LocalTaskManagerLocation;
import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
import org.apache.flink.runtime.testutils.CommonTestUtils;
import org.apache.flink.runtime.testutils.TestingUtils;
import org.apache.flink.util.ExceptionUtils;
import org.apache.flink.util.TestLogger;

import org.junit.ClassRule;
import org.junit.Test;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeoutException;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

/** Tests for the {@link DeclarativeSlotPoolBridge} interactions. */
public class SlotPoolInteractionsTest extends TestLogger {

    private static final Time fastTimeout = Time.milliseconds(1L);

    @ClassRule
    public static final TestingComponentMainThreadExecutor.Resource EXECUTOR_RESOURCE =
            new TestingComponentMainThreadExecutor.Resource(10L);

    private final TestingComponentMainThreadExecutor testMainThreadExecutor =
            EXECUTOR_RESOURCE.getComponentMainThreadTestExecutor();

    // ------------------------------------------------------------------------
    //  tests
    // ------------------------------------------------------------------------

    @Test
    public void testSlotAllocationNoResourceManager() throws Exception {

        try (SlotPool pool = createAndSetUpSlotPoolWithoutResourceManager()) {

            final CompletableFuture<PhysicalSlot> future =
                    testMainThreadExecutor.execute(
                            () ->
                                    pool.requestNewAllocatedSlot(
                                            new SlotRequestId(),
                                            ResourceProfile.UNKNOWN,
                                            fastTimeout));

            try {
                future.get();
                fail("We expected an ExecutionException.");
            } catch (ExecutionException e) {
                assertTrue(ExceptionUtils.stripExecutionException(e) instanceof TimeoutException);
            }
        }
    }

    @Test
    public void testCancelSlotAllocationWithoutResourceManager() throws Exception {

        try (DeclarativeSlotPoolBridge pool = createAndSetUpSlotPoolWithoutResourceManager()) {

            final CompletableFuture<PhysicalSlot> future =
                    testMainThreadExecutor.execute(
                            () ->
                                    pool.requestNewAllocatedSlot(
                                            new SlotRequestId(),
                                            ResourceProfile.UNKNOWN,
                                            fastTimeout));

            try {
                future.get();
                fail("We expected a TimeoutException.");
            } catch (ExecutionException e) {
                assertTrue(ExceptionUtils.stripExecutionException(e) instanceof TimeoutException);
            }

            CommonTestUtils.waitUntilCondition(
                    () -> pool.getNumPendingRequests() == 0,
                    Deadline.fromNow(TestingUtils.TESTING_DURATION));
        }
    }

    /** Tests that a slot allocation times out wrt to the specified time out. */
    @Test
    public void testSlotAllocationTimeout() throws Exception {

        try (DeclarativeSlotPoolBridge pool = createAndSetUpSlotPool()) {

            final CompletableFuture<PhysicalSlot> future =
                    testMainThreadExecutor.execute(
                            () ->
                                    pool.requestNewAllocatedSlot(
                                            new SlotRequestId(),
                                            ResourceProfile.UNKNOWN,
                                            fastTimeout));

            try {
                future.get();
                fail("We expected a TimeoutException.");
            } catch (ExecutionException e) {
                assertTrue(ExceptionUtils.stripExecutionException(e) instanceof TimeoutException);
            }

            CommonTestUtils.waitUntilCondition(
                    () -> pool.getNumPendingRequests() == 0,
                    Deadline.fromNow(TestingUtils.TESTING_DURATION));
        }
    }

    /** Tests that extra slots are kept by the {@link SlotPoolImpl}. */
    @Test
    public void testExtraSlotsAreKept() throws Exception {
        final CompletableFuture<AllocationID> allocationIdFuture = new CompletableFuture<>();

        TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
        resourceManagerGateway.setRequestSlotConsumer(
                (SlotRequest slotRequest) ->
                        allocationIdFuture.complete(slotRequest.getAllocationId()));

        try (TestingSlotPoolImpl pool =
                new SlotPoolBuilder(testMainThreadExecutor.getMainThreadExecutor())
                        .setResourceManagerGateway(resourceManagerGateway)
                        .build()) {

            final CompletableFuture<SlotRequestId> slotRequestTimeoutFuture =
                    new CompletableFuture<>();
            pool.setTimeoutPendingSlotRequestConsumer(slotRequestTimeoutFuture::complete);

            final CompletableFuture<PhysicalSlot> future =
                    testMainThreadExecutor.execute(
                            () ->
                                    pool.requestNewAllocatedSlot(
                                            new SlotRequestId(),
                                            ResourceProfile.UNKNOWN,
                                            fastTimeout));

            try {
                future.get();
                fail("We expected a TimeoutException.");
            } catch (ExecutionException e) {
                assertTrue(ExceptionUtils.stripExecutionException(e) instanceof TimeoutException);
            }

            // wait until we have timed out the slot request
            slotRequestTimeoutFuture.get();

            assertEquals(0L, pool.getNumberOfPendingRequests());

            AllocationID allocationId = allocationIdFuture.get();
            final SlotOffer slotOffer = new SlotOffer(allocationId, 0, ResourceProfile.ANY);
            final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
            final TaskManagerGateway taskManagerGateway = new SimpleAckingTaskManagerGateway();

            testMainThreadExecutor.execute(
                    () -> pool.registerTaskManager(taskManagerLocation.getResourceID()));

            assertTrue(
                    testMainThreadExecutor.execute(
                            () ->
                                    pool.offerSlot(
                                            taskManagerLocation, taskManagerGateway, slotOffer)));

            assertTrue(pool.containsAvailableSlot(allocationId));
        }
    }

    private DeclarativeSlotPoolBridge createAndSetUpSlotPool() throws Exception {
        return new DeclarativeSlotPoolBridgeBuilder(testMainThreadExecutor.getMainThreadExecutor())
                .build();
    }

    private DeclarativeSlotPoolBridge createAndSetUpSlotPoolWithoutResourceManager()
            throws Exception {
        return new DeclarativeSlotPoolBridgeBuilder(testMainThreadExecutor.getMainThreadExecutor())
                .setResourceManagerGateway(null)
                .build();
    }
}

@@ -954,77 +954,8 @@
 me;%0A
-import org.apache.flink.runtime.clusterframework.types.AllocationID;%0A
 impo
@@ -1141,527 +1141,31 @@
 ime.
-executiongraph.utils.SimpleAckingTaskManagerGateway;%0Aimport org.apache.flink.runtime.jobmanager.slots.TaskManagerGateway;%0Aimport org.apache.flink.runtime.jobmaster.SlotRequestId;%0Aimport org.apache.flink.runtime.resourcemanager.SlotRequest;%0Aimport org.apache.flink.runtime.resourcemanager.utils.TestingResourceManagerGateway;%0Aimport org.apache.flink.runtime.taskexecutor.slot.SlotOffer;%0Aimport org.apache.flink.runtime.taskmanager.LocalTaskManagerLocation;%0Aimport org.apache.flink.runtime.taskmanager.TaskManagerLocation
+jobmaster.SlotRequestId
 ;%0Aim
@@ -1562,53 +1562,8 @@
 n;%0A%0A
-import static org.junit.Assert.assertEquals;%0A
 impo
@@ -5336,2652 +5336,8 @@
  %7D%0A%0A
-    /** Tests that extra slots are kept by the %7B@link SlotPoolImpl%7D. */%0A    @Test%0A    public void testExtraSlotsAreKept() throws Exception %7B%0A        final CompletableFuture%3CAllocationID%3E allocationIdFuture = new CompletableFuture%3C%3E();%0A%0A        TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();%0A        resourceManagerGateway.setRequestSlotConsumer(%0A                (SlotRequest slotRequest) -%3E%0A                        allocationIdFuture.complete(slotRequest.getAllocationId()));%0A%0A        try (TestingSlotPoolImpl pool =%0A                new SlotPoolBuilder(testMainThreadExecutor.getMainThreadExecutor())%0A                        .setResourceManagerGateway(resourceManagerGateway)%0A                        .build()) %7B%0A%0A            final CompletableFuture%3CSlotRequestId%3E slotRequestTimeoutFuture =%0A                    new CompletableFuture%3C%3E();%0A            pool.setTimeoutPendingSlotRequestConsumer(slotRequestTimeoutFuture::complete);%0A%0A            final CompletableFuture%3CPhysicalSlot%3E future =%0A                    testMainThreadExecutor.execute(%0A                            () -%3E%0A                                    pool.requestNewAllocatedSlot(%0A                                            new SlotRequestId(),%0A                                            ResourceProfile.UNKNOWN,%0A                                            fastTimeout));%0A%0A            try %7B%0A                future.get();%0A                fail(%22We expected a TimeoutException.%22);%0A            %7D catch (ExecutionException e) %7B%0A                assertTrue(ExceptionUtils.stripExecutionException(e) instanceof TimeoutException);%0A            %7D%0A%0A            // wait until we have timed out the slot request%0A            slotRequestTimeoutFuture.get();%0A%0A            assertEquals(0L, pool.getNumberOfPendingRequests());%0A%0A            AllocationID allocationId = allocationIdFuture.get();%0A            final SlotOffer slotOffer = new SlotOffer(allocationId, 0, ResourceProfile.ANY);%0A            final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();%0A            final TaskManagerGateway taskManagerGateway = new SimpleAckingTaskManagerGateway();%0A%0A            testMainThreadExecutor.execute(%0A                    () -%3E pool.registerTaskManager(taskManagerLocation.getResourceID()));%0A%0A            assertTrue(%0A                    testMainThreadExecutor.execute(%0A                            () -%3E%0A                                    pool.offerSlot(%0A                                            taskManagerLocation, taskManagerGateway, slotOffer)));%0A%0A            assertTrue(pool.containsAvailableSlot(allocationId));%0A        %7D%0A    %7D%0A%0A
     

Remove SlotPoolInteractionsTest.testExtraSlotsAreKept