/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch 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.elasticsearch.index.shard;

import org.elasticsearch.action.support.PlainActionFuture;
import org.elasticsearch.common.inject.internal.Nullable;
import org.elasticsearch.common.lease.Releasable;
import org.elasticsearch.test.ESTestCase;
import org.elasticsearch.threadpool.TestThreadPool;
import org.elasticsearch.threadpool.ThreadPool;
import org.junit.After;
import org.junit.AfterClass;
import org.junit.Before;
import org.junit.BeforeClass;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;

import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.instanceOf;

public class IndexShardOperationsLockTests extends ESTestCase {

    private static ThreadPool threadPool;

    private IndexShardOperationsLock block;

    @BeforeClass
    public static void setupThreadPool() {
        threadPool = new TestThreadPool("IndexShardOperationsLockTests");
    }

    @AfterClass
    public static void shutdownThreadPool() {
        ThreadPool.terminate(threadPool, 30, TimeUnit.SECONDS);
        threadPool = null;
    }

    @Before
    public void createIndexShardOperationsLock() {
         block = new IndexShardOperationsLock(new ShardId("blubb", "id", 0), logger, threadPool);
    }

    @After
    public void checkNoInflightOperations() {
        assertThat(block.semaphore.availablePermits(), equalTo(Integer.MAX_VALUE));
        assertThat(block.getActiveOperationsCount(), equalTo(0));
    }

    public void testAllOperationsInvoked() throws InterruptedException, TimeoutException, ExecutionException {
        int numThreads = 10;

        List<PlainActionFuture<Releasable>> futures = new ArrayList<>();
        List<Thread> operationThreads = new ArrayList<>();
        CountDownLatch latch = new CountDownLatch(numThreads / 2);
        for (int i = 0; i < numThreads; i++) {
            PlainActionFuture<Releasable> future = new PlainActionFuture<Releasable>() {
                @Override
                public void onResponse(Releasable releasable) {
                    releasable.close();
                    super.onResponse(releasable);
                }
            };
            Thread thread = new Thread() {
                public void run() {
                    latch.countDown();
                    block.acquire(future, ThreadPool.Names.GENERIC, true);
                }
            };
            futures.add(future);
            operationThreads.add(thread);
        }

        threadPool.generic().execute(() -> {
            try {
                latch.await();
                blockAndWait().close();
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        });

        for (Thread thread : operationThreads) {
            thread.start();
        }

        for (PlainActionFuture<Releasable> future : futures) {
            assertNotNull(future.get(1, TimeUnit.MINUTES));
        }

        for (Thread thread : operationThreads) {
            thread.join();
        }
    }


    public void testOperationsInvokedImmediatelyIfNoBlock() throws ExecutionException, InterruptedException {
        PlainActionFuture<Releasable> future = new PlainActionFuture<>();
        block.acquire(future, ThreadPool.Names.GENERIC, true);
        assertTrue(future.isDone());
        future.get().close();
    }

    public void testOperationsIfClosed() throws ExecutionException, InterruptedException {
        PlainActionFuture<Releasable> future = new PlainActionFuture<>();
        block.close();
        block.acquire(future, ThreadPool.Names.GENERIC, true);
        ExecutionException exception = expectThrows(ExecutionException.class, future::get);
        assertThat(exception.getCause(), instanceOf(IndexShardClosedException.class));
    }

    public void testBlockIfClosed() throws ExecutionException, InterruptedException {
        block.close();
        expectThrows(IndexShardClosedException.class, () -> block.blockOperations(randomInt(10), TimeUnit.MINUTES,
            () -> { throw new IllegalArgumentException("fake error"); }));
    }

    public void testOperationsDelayedIfBlock() throws ExecutionException, InterruptedException, TimeoutException {
        PlainActionFuture<Releasable> future = new PlainActionFuture<>();
        try (Releasable releasable = blockAndWait()) {
            block.acquire(future, ThreadPool.Names.GENERIC, true);
            assertFalse(future.isDone());
        }
        future.get(1, TimeUnit.MINUTES).close();
    }

    protected Releasable blockAndWait() throws InterruptedException {
        CountDownLatch blockAcquired = new CountDownLatch(1);
        CountDownLatch releaseBlock = new CountDownLatch(1);
        CountDownLatch blockReleased = new CountDownLatch(1);
        boolean throwsException = randomBoolean();
        IndexShardClosedException exception = new IndexShardClosedException(new ShardId("blubb", "id", 0));
        threadPool.generic().execute(() -> {
                try {
                    block.blockOperations(1, TimeUnit.MINUTES, () -> {
                        try {
                            blockAcquired.countDown();
                            releaseBlock.await();
                            if (throwsException) {
                                throw exception;
                            }
                        } catch (InterruptedException e) {
                            throw new RuntimeException();
                        }
                    });
                } catch (Exception e) {
                    if (e != exception) {
                        throw new RuntimeException(e);
                    }
                } finally {
                    blockReleased.countDown();
                }
            });
        blockAcquired.await();
        return () -> {
            releaseBlock.countDown();
            try {
                blockReleased.await();
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        };
    }

    public void testActiveOperationsCount() throws ExecutionException, InterruptedException {
        PlainActionFuture<Releasable> future1 = new PlainActionFuture<>();
        block.acquire(future1, ThreadPool.Names.GENERIC, true);
        assertTrue(future1.isDone());
        assertThat(block.getActiveOperationsCount(), equalTo(1));

        PlainActionFuture<Releasable> future2 = new PlainActionFuture<>();
        block.acquire(future2, ThreadPool.Names.GENERIC, true);
        assertTrue(future2.isDone());
        assertThat(block.getActiveOperationsCount(), equalTo(2));

        future1.get().close();
        assertThat(block.getActiveOperationsCount(), equalTo(1));
        future1.get().close(); // check idempotence
        assertThat(block.getActiveOperationsCount(), equalTo(1));
        future2.get().close();
        assertThat(block.getActiveOperationsCount(), equalTo(0));

        try (Releasable releasable = blockAndWait()) {
            assertThat(block.getActiveOperationsCount(), equalTo(0));
        }

        PlainActionFuture<Releasable> future3 = new PlainActionFuture<>();
        block.acquire(future3, ThreadPool.Names.GENERIC, true);
        assertTrue(future3.isDone());
        assertThat(block.getActiveOperationsCount(), equalTo(1));
        future3.get().close();
        assertThat(block.getActiveOperationsCount(), equalTo(0));
    }
}

@@ -3441,32 +3441,94 @@
 ad);%0A        %7D%0A%0A
+        CountDownLatch blockFinished = new CountDownLatch(1);%0A
         threadPo
@@ -3637,32 +3637,75 @@
 Wait().close();%0A
+                blockFinished.countDown();%0A
             %7D ca
@@ -4112,24 +4112,56 @@
 );%0A        %7D
+%0A%0A        blockFinished.await();
 %0A    %7D%0A%0A%0A   

Make testAllOperationsInvoked properly clean up after itself