package com.palantir.atlasdb.keyvalue.cassandra;

import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

import org.apache.cassandra.thrift.ConsistencyLevel;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import com.palantir.atlasdb.cassandra.CassandraKeyValueServiceConfigManager;
import com.palantir.atlasdb.cassandra.ImmutableCassandraKeyValueServiceConfig;

public abstract class SchemaMutationLockTest {
    protected SchemaMutationLock schemaMutationLock;
    private final ExecutorService executorService = Executors.newFixedThreadPool(4);

    @Before
    public void setUp() {
        setUpWithCasSupportSetTo(true);
    }

    protected void setUpWithCasSupportSetTo(boolean supportsCas) {
        ImmutableCassandraKeyValueServiceConfig quickTimeoutConfig = CassandraTestSuite.CASSANDRA_KVS_CONFIG
                .withSchemaMutationTimeoutMillis(500);
        CassandraKeyValueServiceConfigManager simpleManager = CassandraKeyValueServiceConfigManager.createSimpleManager(quickTimeoutConfig);
        ConsistencyLevel writeConsistency = ConsistencyLevel.EACH_QUORUM;
        CassandraClientPool clientPool = new CassandraClientPool(simpleManager.getConfig());

        schemaMutationLock = new SchemaMutationLock(supportsCas, simpleManager, clientPool, writeConsistency);
    }

    @Test
    public void testLockAndUnlockWithoutContention() {
        long ourId = schemaMutationLock.waitForSchemaMutationLock();
        schemaMutationLock.schemaMutationUnlock(ourId);
    }

    @Test
    public void testOnlyOneLockCanBeLockedAtATime() throws InterruptedException, ExecutionException, TimeoutException {
        long firstLock = schemaMutationLock.waitForSchemaMutationLock();

        Future tryToAcquireSecondLock = async(() -> schemaMutationLock.waitForSchemaMutationLock());

        Thread.sleep(3 * 1000);
        assertThatFutureDidNotSucceedYet(tryToAcquireSecondLock);

        tryToAcquireSecondLock.cancel(true);
        schemaMutationLock.schemaMutationUnlock(firstLock);
    }

    @Test
    public void testUnlockIsSuccessful() throws InterruptedException, TimeoutException, ExecutionException {
        long id = schemaMutationLock.waitForSchemaMutationLock();
        Future future = async(() -> {
            long newId = schemaMutationLock.waitForSchemaMutationLock();
            schemaMutationLock.schemaMutationUnlock(newId);
        });
        Thread.sleep(100);
        Assert.assertFalse(future.isDone());
        schemaMutationLock.schemaMutationUnlock(id);
        future.get(3, TimeUnit.SECONDS);
    }

    protected Future async(Runnable callable) {
        return executorService.submit(callable);
    }

    private void assertThatFutureDidNotSucceedYet(Future future) throws InterruptedException {
        if (future.isDone()) {
            try {
                future.get();
                throw new AssertionError("Future task should have failed but finished successfully");
            } catch (ExecutionException e) {
                // if execution is done, we expect it to have failed
            }
        }
    }
}

@@ -207,24 +207,63 @@
 ent.Future;%0A
+import java.util.concurrent.Semaphore;%0A
 import java.
@@ -666,16 +666,94 @@
 kTest %7B%0A
+    public static final SchemaMutationLock.Action DO_NOTHING = () -%3E %7B%0A    %7D;%0A
     prot
@@ -1702,21 +1702,8 @@
     
- long ourId =
  sch
@@ -1722,304 +1722,164 @@
 ock.
-waitForSchemaMutationLock();%0A        schemaMutationLock.schemaMutationUnlock(ourId);%0A    %7D%0A%0A    @Test%0A    public void testOnlyOneLockCanBeLockedAtATime() throws InterruptedException, ExecutionException, TimeoutException %7B%0A        long firstLock = schemaMutationLock.waitForSchemaMutation
+runWithLock(() -%3E %7B%7D);%0A    %7D%0A%0A    @Test%0A    public void doesNotPerformAnActionIfTheLockIsAlreadyHeld() %7B%0A        schemaMutationLock.runWith
 Lock(
-);%0A%0A
+() -%3E %7B%0A    
     
@@ -1889,38 +1889,28 @@
  Future 
-tryToAcquireSecondLock
+getLockAgain
  = async
@@ -1939,39 +1939,39 @@
 ock.
-waitForSchemaMutationLock(
+runWithLock(DO_NOTHING
 ));%0A%0A
+    
     
@@ -1992,18 +1992,21 @@
 ep(3
- * 
+*
 1000);%0A
+%0A    
     
@@ -2046,136 +2046,32 @@
 Yet(
-tryToAcquireSecondLock);%0A%0A        tryToAcquireSecondLock.cancel(true);%0A        schemaMutationLock.schemaMutationUnlock(firstLock
+getLockAgain);%0A        %7D
 );%0A 

Change test to use new API