text:
package org.ehcache.transactions.xa.internal;

import org.ehcache.Cache;
import org.ehcache.config.EvictionAdvisor;
import org.ehcache.config.builders.ExpiryPolicyBuilder;
import org.ehcache.config.builders.ResourcePoolsBuilder;
import org.ehcache.config.units.EntryUnit;
import org.ehcache.config.units.MemoryUnit;
import org.ehcache.core.events.StoreEventDispatcher;
import org.ehcache.core.spi.ServiceLocator;
import org.ehcache.core.spi.service.CacheManagerProviderService;
import org.ehcache.core.spi.service.DiskResourceService;
import org.ehcache.core.spi.store.Store;
import org.ehcache.core.statistics.DefaultStatisticsService;
import org.ehcache.core.store.StoreConfigurationImpl;
import org.ehcache.expiry.ExpiryPolicy;
import org.ehcache.impl.config.copy.DefaultCopyProviderConfiguration;
import org.ehcache.core.events.NullStoreEventDispatcher;
import org.ehcache.impl.internal.sizeof.NoopSizeOfEngine;
import org.ehcache.impl.internal.spi.copy.DefaultCopyProvider;
import org.ehcache.impl.internal.store.heap.OnHeapStore;
import org.ehcache.impl.internal.store.offheap.MemorySizeParser;
import org.ehcache.impl.internal.store.offheap.OffHeapStore;
import org.ehcache.impl.internal.store.offheap.OffHeapStoreLifecycleHelper;
import org.ehcache.impl.internal.store.tiering.TieredStore;
import org.ehcache.internal.TestTimeSource;
import org.ehcache.spi.service.ServiceProvider;
import org.ehcache.spi.copy.Copier;
import org.ehcache.spi.copy.CopyProvider;
import org.ehcache.spi.serialization.Serializer;
import org.ehcache.spi.service.Service;
import org.ehcache.spi.service.ServiceConfiguration;
import org.ehcache.transactions.xa.XACacheException;
import org.ehcache.transactions.xa.configuration.XAStoreConfiguration;
import org.ehcache.transactions.xa.internal.journal.Journal;
import org.ehcache.transactions.xa.internal.journal.TransientJournal;
import org.ehcache.transactions.xa.internal.txmgr.NullXAResourceRegistry;
import org.ehcache.transactions.xa.txmgr.TransactionManagerWrapper;
import org.ehcache.transactions.xa.txmgr.provider.TransactionManagerProvider;
import org.ehcache.transactions.xa.utils.JavaSerializer;
import org.ehcache.transactions.xa.utils.TestXid;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TestName;
import org.mockito.Answers;

import java.time.Duration;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Supplier;

import javax.transaction.HeuristicMixedException;
import javax.transaction.HeuristicRollbackException;
import javax.transaction.RollbackException;
import javax.transaction.Status;
import javax.transaction.Synchronization;
import javax.transaction.SystemException;
import javax.transaction.Transaction;
import javax.transaction.TransactionManager;
import javax.transaction.xa.XAException;
import javax.transaction.xa.XAResource;

import static java.time.Duration.ofSeconds;
import static java.util.Collections.emptySet;
import static org.ehcache.config.builders.ExpiryPolicyBuilder.timeToLiveExpiration;
import static org.ehcache.core.spi.ServiceLocator.dependencySet;
import static org.hamcrest.Matchers.containsString;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.nullValue;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.fail;
import static org.mockito.Mockito.mock;

/**
 * Tests for {@link XAStore} and {@link org.ehcache.transactions.xa.internal.XAStore.Provider}.
 */
public class XAStoreTest {

  private static final Supplier<Boolean> SUPPLY_TRUE = () -> true;
  private static final Supplier<Boolean> SUPPLY_FALSE = () -> false;

  @Rule
  public TestName testName = new TestName();

  @SuppressWarnings("unchecked")
  private final Class<SoftLock<String>> valueClass = (Class) SoftLock.class;
  private final TestTransactionManager testTransactionManager = new TestTransactionManager();
  private TransactionManagerWrapper transactionManagerWrapper;
  private OnHeapStore<Long, SoftLock<String>> onHeapStore;
  private Journal<Long> journal;
  private TestTimeSource testTimeSource;
  private ClassLoader classLoader;
  private Serializer<Long> keySerializer;
  private Serializer<SoftLock<String>> valueSerializer;
  private StoreEventDispatcher<Long, SoftLock<String>> eventDispatcher;
  private final ExpiryPolicy<Object, Object> expiry = timeToLiveExpiration(ofSeconds(1));
  private Copier<Long> keyCopier;
  private Copier<SoftLock<String>> valueCopier;

  @Before
  public void setUp() {
    transactionManagerWrapper = new TransactionManagerWrapper(testTransactionManager, new NullXAResourceRegistry());
    classLoader = ClassLoader.getSystemClassLoader();
    keySerializer = new JavaSerializer<>(classLoader);
    valueSerializer = new JavaSerializer<>(classLoader);
    CopyProvider copyProvider = new DefaultCopyProvider(new DefaultCopyProviderConfiguration());
    keyCopier = copyProvider.createKeyCopier(Long.class, keySerializer);
    valueCopier = copyProvider.createValueCopier(valueClass, valueSerializer);
    Store.Configuration<Long, SoftLock<String>> onHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass,
      null, classLoader, ExpiryPolicyBuilder.noExpiration(), ResourcePoolsBuilder.newResourcePoolsBuilder()
      .heap(10, EntryUnit.ENTRIES)
      .build(),
      0, keySerializer, valueSerializer);
    testTimeSource = new TestTimeSource();
    eventDispatcher = NullStoreEventDispatcher.nullStoreEventDispatcher();
    onHeapStore = new OnHeapStore<>(onHeapConfig, testTimeSource, keyCopier, valueCopier, new NoopSizeOfEngine(), eventDispatcher, new DefaultStatisticsService());
    journal = new TransientJournal<>();
  }

  @Test
  public void testXAStoreProviderFailsToRankWhenNoTMProviderConfigured() throws Exception {
    XAStore.Provider provider = new XAStore.Provider();
      provider.start(new ServiceProvider<Service>() {
        @Override
        public <U extends Service> U getService(Class<U> serviceType) {
          return null;
        }

        @Override
        public <U extends Service> Collection<U> getServicesOfType(Class<U> serviceType) {
          return emptySet();
        }
      });
    try {
      provider.wrapperStoreRank(Collections.singleton(mock(XAStoreConfiguration.class)));
      fail("Expected exception");
    } catch (IllegalStateException e) {
      assertThat(e.getMessage(), containsString("TransactionManagerProvider"));
    }
  }

  @Test
  public void testSimpleGetPutRemove() throws Exception {
    XAStore<Long, String> xaStore = getXAStore(onHeapStore);

    testTransactionManager.begin();
    {
      assertThat(xaStore.remove(1L), equalTo(false));
      assertThat(xaStore.get(1L), is(nullValue()));
      assertThat(xaStore.put(1L, "1"), equalTo(Store.PutStatus.PUT));
      assertThat(xaStore.put(1L, "one"), equalTo(Store.PutStatus.PUT));
      assertThat(xaStore.get(1L).get(), equalTo("one"));
    }
    testTransactionManager.rollback();

    assertMapping(xaStore, 1L, null);

    testTransactionManager.begin();
    {
      assertThat(xaStore.get(1L), is(nullValue()));
      assertThat(xaStore.put(1L, "1"), equalTo(Store.PutStatus.PUT));
      assertThat(xaStore.put(1L, "one"), equalTo(Store.PutStatus.PUT));
      assertThat(xaStore.get(1L).get(), equalTo("one"));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "one");

    testTransactionManager.begin();
    {
      assertThat(xaStore.remove(1L), equalTo(true));
      assertThat(xaStore.remove(1L), equalTo(false));
      assertThat(xaStore.get(1L), is(nullValue()));
      assertThat(xaStore.put(1L, "1"), equalTo(Store.PutStatus.PUT));
    }
    testTransactionManager.rollback();

    assertMapping(xaStore, 1L, "one");

    testTransactionManager.begin();
    {
      assertThat(xaStore.put(1L, "un"), equalTo(Store.PutStatus.PUT));
      assertThat(xaStore.remove(1L), equalTo(true));
      assertThat(xaStore.remove(1L), equalTo(false));
      assertThat(xaStore.get(1L), is(nullValue()));
      assertThat(xaStore.put(1L, "un"), equalTo(Store.PutStatus.PUT));
      assertThat(xaStore.get(1L).get(), equalTo("un"));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "un");
  }

  @Test
  public void testConflictingGetPutRemove() throws Exception {
    final XAStore<Long, String> xaStore = getXAStore(onHeapStore);
    final AtomicReference<Throwable> exception = new AtomicReference<>();

    testTransactionManager.begin();
    {
      xaStore.put(1L, "one");
    }
    testTransactionManager.commit();

    testTransactionManager.begin();
    {
      assertThat(xaStore.put(1L, "un"), equalTo(Store.PutStatus.PUT));

      executeWhileIn2PC(exception, () -> {
        testTransactionManager.begin();
        assertThat(xaStore.put(1L, "uno"), equalTo(Store.PutStatus.NOOP));
        testTransactionManager.commit();
        return null;
      });

      assertThat(xaStore.put(1L, "eins"), equalTo(Store.PutStatus.PUT));
    }
    testTransactionManager.commit();
    assertThat(exception.get(), is(nullValue()));

    assertMapping(xaStore, 1L, null);

    testTransactionManager.begin();
    {
      xaStore.put(1L, "one");
    }
    testTransactionManager.commit();

    testTransactionManager.begin();
    {
      assertThat(xaStore.put(1L, "un"), equalTo(Store.PutStatus.PUT));

      executeWhileIn2PC(exception, () -> {
        testTransactionManager.begin();

        assertThat(xaStore.remove(1L), is(false));

        testTransactionManager.commit();
        return null;
      });

      assertThat(xaStore.put(1L, "een"), equalTo(Store.PutStatus.PUT));
    }
    testTransactionManager.commit();

    assertThat(exception.get(), is(nullValue()));
    assertMapping(xaStore, 1L, null);

    testTransactionManager.begin();
    {
      xaStore.put(1L, "one");
    }
    testTransactionManager.commit();

    testTransactionManager.begin();
    {
      assertThat(xaStore.put(1L, "un"), equalTo(Store.PutStatus.PUT));

      executeWhileIn2PC(exception, () -> {
        testTransactionManager.begin();

        assertThat(xaStore.get(1L), is(nullValue()));

        testTransactionManager.commit();
        return null;
      });

      assertThat(xaStore.put(1L, "yksi"), equalTo(Store.PutStatus.PUT));
    }
    testTransactionManager.commit();

    assertThat(exception.get(), is(nullValue()));
    assertMapping(xaStore, 1L, null);
  }

  private void executeWhileIn2PC(AtomicReference<Throwable> exception, Callable<?> callable) {
    testTransactionManager.getCurrentTransaction().registerTwoPcListener(() -> {
      try {
        Thread t = new Thread(() -> {
          try {
            // this runs while the committing TX is in-doubt
            callable.call();
          } catch (Throwable t1) {
            exception.set(t1);
          }
        });
        t.start();
        t.join();
      } catch (Throwable e) {
        exception.set(e);
      }
    });
  }

  @Test
  public void testIterate() throws Exception {
    XAStore<Long, String> xaStore = getXAStore(onHeapStore);

    testTransactionManager.begin();
    {
      xaStore.put(1L, "one");
      xaStore.put(2L, "two");
      xaStore.put(3L, "three");
    }
    testTransactionManager.commit();

    testTransactionManager.begin();
    {
      xaStore.put(0L, "zero");
      xaStore.put(1L, "un");
      xaStore.put(2L, "two");
      xaStore.remove(3L);

      Map<Long, String> iterated = new HashMap<>();
      Store.Iterator<Cache.Entry<Long, Store.ValueHolder<String>>> iterator = xaStore.iterator();
      while (iterator.hasNext()) {
        Cache.Entry<Long, Store.ValueHolder<String>> next = iterator.next();
        iterated.put(next.getKey(), next.getValue().get());
      }
      assertThat(iterated.size(), is(3));
      assertThat(iterated.get(0L), equalTo("zero"));
      assertThat(iterated.get(1L), equalTo("un"));
      assertThat(iterated.get(2L), equalTo("two"));
    }
    testTransactionManager.commit();

    testTransactionManager.begin();
    {
      Map<Long, String> iterated = new HashMap<>();
      Store.Iterator<Cache.Entry<Long, Store.ValueHolder<String>>> iterator = xaStore.iterator();
      while (iterator.hasNext()) {
        Cache.Entry<Long, Store.ValueHolder<String>> next = iterator.next();
        iterated.put(next.getKey(), next.getValue().get());
      }
      assertThat(iterated.size(), is(3));
      assertThat(iterated.get(0L), equalTo("zero"));
      assertThat(iterated.get(1L), equalTo("un"));
      assertThat(iterated.get(2L), equalTo("two"));
    }
    testTransactionManager.commit();

    Store.Iterator<Cache.Entry<Long, Store.ValueHolder<String>>> iterator;
    testTransactionManager.begin();
    {
      iterator = xaStore.iterator();
      iterator.next();
    }
    testTransactionManager.commit();

    // cannot use iterator outside of tx context
    try {
      iterator.hasNext();
      fail();
    } catch (XACacheException e) {
      // expected
    }
    try {
      iterator.next();
      fail();
    } catch (XACacheException e) {
      // expected
    }

    // cannot use iterator outside of original tx context
    testTransactionManager.begin();
    {
      try {
        iterator.hasNext();
        fail();
      } catch (IllegalStateException e) {
        // expected
      }
      try {
        iterator.next();
        fail();
      } catch (IllegalStateException e) {
        // expected
      }
    }
    testTransactionManager.commit();
  }

  @Test
  public void testPutIfAbsent() throws Exception {
    final XAStore<Long, String> xaStore = getXAStore(onHeapStore);
    final AtomicReference<Throwable> exception = new AtomicReference<>();

    testTransactionManager.begin();
    {
      assertThat(xaStore.putIfAbsent(1L, "one", b -> {}), is(nullValue()));
      assertThat(xaStore.get(1L).get(), equalTo("one"));
      assertThat(xaStore.putIfAbsent(1L, "un", b -> {}).get(), equalTo("one"));
      assertThat(xaStore.get(1L).get(), equalTo("one"));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "one");

    testTransactionManager.begin();
    {
      assertThat(xaStore.putIfAbsent(1L, "un", b -> {}).get(), equalTo("one"));
      assertThat(xaStore.get(1L).get(), equalTo("one"));
      assertThat(xaStore.remove(1L), equalTo(true));
      assertThat(xaStore.putIfAbsent(1L, "uno", b -> {}), is(nullValue()));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "uno");

    testTransactionManager.begin();
    {
      xaStore.put(1L, "eins");
      executeWhileIn2PC(exception, () -> {
        testTransactionManager.begin();

        assertThat(xaStore.putIfAbsent(1L, "un", b -> {}), is(nullValue()));

        testTransactionManager.commit();
        return null;
      });
    }
    testTransactionManager.commit();
    assertThat(exception.get(), is(nullValue()));

    assertMapping(xaStore, 1L, null);
  }

  @Test
  public void testRemove2Args() throws Exception {
    final XAStore<Long, String> xaStore = getXAStore(onHeapStore);
    final AtomicReference<Throwable> exception = new AtomicReference<>();

    testTransactionManager.begin();
    {
      assertThat(xaStore.remove(1L, "one"), equalTo(Store.RemoveStatus.KEY_MISSING));
      assertThat(xaStore.put(1L, "one"), equalTo(Store.PutStatus.PUT));
      assertThat(xaStore.remove(1L, "un"), equalTo(Store.RemoveStatus.KEY_PRESENT));
      assertThat(xaStore.remove(1L, "one"), equalTo(Store.RemoveStatus.REMOVED));
      assertThat(xaStore.remove(1L, "eins"), equalTo(Store.RemoveStatus.KEY_MISSING));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, null);

    testTransactionManager.begin();
    {
      assertThat(xaStore.put(1L, "one"), equalTo(Store.PutStatus.PUT));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "one");

    testTransactionManager.begin();
    {
      assertThat(xaStore.remove(1L, "een"), equalTo(Store.RemoveStatus.KEY_PRESENT));
      assertThat(xaStore.remove(1L, "one"), equalTo(Store.RemoveStatus.REMOVED));
      assertThat(xaStore.remove(1L, "eins"), equalTo(Store.RemoveStatus.KEY_MISSING));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, null);

    testTransactionManager.begin();
    {
      xaStore.put(1L, "eins");
      executeWhileIn2PC(exception, () -> {
        testTransactionManager.begin();

        assertThat(xaStore.remove(1L, "un"), equalTo(Store.RemoveStatus.KEY_MISSING));

        testTransactionManager.commit();
        return null;
      });
    }
    testTransactionManager.commit();
    assertThat(exception.get(), is(nullValue()));

    assertMapping(xaStore, 1L, null);

    testTransactionManager.begin();
    {
      assertThat(xaStore.put(1L, "one"), equalTo(Store.PutStatus.PUT));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "one");

    testTransactionManager.begin();
    {
      xaStore.put(1L, "eins");
      executeWhileIn2PC(exception, () -> {
        testTransactionManager.begin();

        assertThat(xaStore.remove(1L, "un"), equalTo(Store.RemoveStatus.KEY_MISSING));

        testTransactionManager.commit();
        return null;
      });
    }
    testTransactionManager.commit();
    assertThat(exception.get(), is(nullValue()));

    assertMapping(xaStore, 1L, null);
  }

  @Test
  public void testReplace2Args() throws Exception {
    final XAStore<Long, String> xaStore = getXAStore(onHeapStore);
    final AtomicReference<Throwable> exception = new AtomicReference<>();

    testTransactionManager.begin();
    {
      assertThat(xaStore.replace(1L, "one"), is(nullValue()));
      assertThat(xaStore.put(1L, "one"), equalTo(Store.PutStatus.PUT));
      assertThat(xaStore.replace(1L, "un").get(), equalTo("one"));
      assertThat(xaStore.replace(1L, "uno").get(), equalTo("un"));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "uno");

    testTransactionManager.begin();
    {
      assertThat(xaStore.replace(1L, "een").get(), equalTo("uno"));
      assertThat(xaStore.replace(1L, "eins").get(), equalTo("een"));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "eins");

    testTransactionManager.begin();
    {
      assertThat(xaStore.remove(1L), is(true));
      assertThat(xaStore.replace(1L, "yksi"), is(nullValue()));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, null);

    testTransactionManager.begin();
    {
      xaStore.put(1L, "eins");
      executeWhileIn2PC(exception, () -> {
        testTransactionManager.begin();

        assertThat(xaStore.replace(1L, "un"), is(nullValue()));

        testTransactionManager.commit();
        return null;
      });
    }
    testTransactionManager.commit();
    assertThat(exception.get(), is(nullValue()));

    assertMapping(xaStore, 1L, null);

    testTransactionManager.begin();
    {
      assertThat(xaStore.put(1L, "one"), is(Store.PutStatus.PUT));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "one");

    testTransactionManager.begin();
    {
      xaStore.put(1L, "eins");
      executeWhileIn2PC(exception, () -> {
        testTransactionManager.begin();

        assertThat(xaStore.replace(1L, "un"), is(nullValue()));

        testTransactionManager.commit();
        return null;
      });
    }
    testTransactionManager.commit();
    assertThat(exception.get(), is(nullValue()));

    assertMapping(xaStore, 1L, null);
  }

  @Test
  public void testReplace3Args() throws Exception {
    final XAStore<Long, String> xaStore = getXAStore(onHeapStore);
    final AtomicReference<Throwable> exception = new AtomicReference<>();

    testTransactionManager.begin();
    {
      assertThat(xaStore.replace(1L, "one", "un"), equalTo(Store.ReplaceStatus.MISS_NOT_PRESENT));
      assertThat(xaStore.put(1L, "one"), equalTo(Store.PutStatus.PUT));
      assertThat(xaStore.replace(1L, "eins", "un"), equalTo(Store.ReplaceStatus.MISS_PRESENT));
      assertThat(xaStore.replace(1L, "one", "un"), equalTo(Store.ReplaceStatus.HIT));
      assertThat(xaStore.get(1L).get(), equalTo("un"));
      assertThat(xaStore.replace(1L, "eins", "een"), equalTo(Store.ReplaceStatus.MISS_PRESENT));
      assertThat(xaStore.replace(1L, "un", "uno"), equalTo(Store.ReplaceStatus.HIT));
      assertThat(xaStore.get(1L).get(), equalTo("uno"));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "uno");

    testTransactionManager.begin();
    {
      assertThat(xaStore.replace(1L, "one", "uno"), equalTo(Store.ReplaceStatus.MISS_PRESENT));
      assertThat(xaStore.replace(1L, "uno", "un"), equalTo(Store.ReplaceStatus.HIT));
      assertThat(xaStore.get(1L).get(), equalTo("un"));
      assertThat(xaStore.remove(1L), equalTo(true));
      assertThat(xaStore.replace(1L, "un", "eins"), equalTo(Store.ReplaceStatus.MISS_NOT_PRESENT));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, null);

    testTransactionManager.begin();
    {
      xaStore.put(1L, "eins");
      executeWhileIn2PC(exception, () -> {
        testTransactionManager.begin();

        assertThat(xaStore.replace(1L, "eins", "one"), is(Store.ReplaceStatus.MISS_NOT_PRESENT));

        testTransactionManager.commit();
        return null;
      });
    }
    testTransactionManager.commit();
    assertThat(exception.get(), is(nullValue()));

    assertMapping(xaStore, 1L, null);

    testTransactionManager.begin();
    {
      assertThat(xaStore.put(1L, "one"), is(Store.PutStatus.PUT));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "one");

    testTransactionManager.begin();
    {
      xaStore.put(1L, "eins");
      executeWhileIn2PC(exception, () -> {
        testTransactionManager.begin();

        assertThat(xaStore.replace(1L, "one", "un"), is(Store.ReplaceStatus.MISS_NOT_PRESENT));

        testTransactionManager.commit();
        return null;
      });
    }
    testTransactionManager.commit();
    assertThat(exception.get(), is(nullValue()));

    assertMapping(xaStore, 1L, null);
  }

  @Test
  public void testGetAndCompute() throws Exception {
    Store.Configuration<Long, SoftLock<String>> offHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass,
            null, classLoader, ExpiryPolicyBuilder.noExpiration(), ResourcePoolsBuilder.newResourcePoolsBuilder()
            .offheap(10, MemoryUnit.MB)
            .build(),
            0, keySerializer, valueSerializer);
    OffHeapStore<Long, SoftLock<String>> offHeapStore = new OffHeapStore<>(offHeapConfig, testTimeSource, eventDispatcher, MemorySizeParser
            .parse("10M"), new DefaultStatisticsService());
    OffHeapStoreLifecycleHelper.init(offHeapStore);
    TieredStore<Long, SoftLock<String>> tieredStore = new TieredStore<>(onHeapStore, offHeapStore);

    XAStore<Long, String> xaStore = getXAStore(tieredStore);

    testTransactionManager.begin();
    {
      Store.ValueHolder<String> computed1 = xaStore.getAndCompute(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, is(nullValue()));
        return "one";
      });
      assertThat(computed1, is(nullValue()));
      assertThat(xaStore.get(1L).get(), equalTo("one"));
      Store.ValueHolder<String> computed2 = xaStore.getAndCompute(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, equalTo("one"));
        return "un";
      });
      assertThat(computed2.get(), equalTo("one"));
      assertThat(xaStore.get(1L).get(), equalTo("un"));
      Store.ValueHolder<String> computed3 = xaStore.getAndCompute(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, equalTo("un"));
        return null;
      });
      assertThat(computed3.get(), equalTo("un"));
      assertThat(xaStore.get(1L), is(nullValue()));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, null);

    testTransactionManager.begin();
    {
      Store.ValueHolder<String> computed1 = xaStore.getAndCompute(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, is(nullValue()));
        return "one";
      });
      assertThat(computed1, is(nullValue()));
      assertThat(xaStore.get(1L).get(), equalTo("one"));
      Store.ValueHolder<String> computed2 = xaStore.getAndCompute(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, equalTo("one"));
        return null;
      });
      assertThat(computed2.get(), equalTo("one"));
      assertThat(xaStore.get(1L), is(nullValue()));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, null);

    testTransactionManager.begin();
    {
      Store.ValueHolder<String> computed1 = xaStore.getAndCompute(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, is(nullValue()));
        return "one";
      });
      assertThat(computed1, is(nullValue()));
      assertThat(xaStore.get(1L).get(), equalTo("one"));
      Store.ValueHolder<String> computed2 = xaStore.getAndCompute(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, equalTo("one"));
        return null;
      });
      assertThat(computed2.get(), equalTo("one"));
      assertThat(xaStore.get(1L), is(nullValue()));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, null);

    testTransactionManager.begin();
    {
      Store.ValueHolder<String> computed1 = xaStore.getAndCompute(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, is(nullValue()));
        return "one";
      });
      assertThat(computed1, is(nullValue()));
      assertThat(xaStore.get(1L).get(), equalTo("one"));
      Store.ValueHolder<String> computed2 = xaStore.getAndCompute(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, equalTo("one"));
        return "un";
      });
      assertThat(computed2.get(), equalTo("one"));
      assertThat(xaStore.get(1L).get(), equalTo("un"));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "un");

    testTransactionManager.begin();
    {
      Store.ValueHolder<String> computed = xaStore.getAndCompute(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, equalTo("un"));
        return "eins";
      });
      assertThat(computed.get(), equalTo("un"));
      assertThat(xaStore.get(1L).get(), equalTo("eins"));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "eins");

    testTransactionManager.begin();
    {
      Store.ValueHolder<String> computed = xaStore.getAndCompute(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, equalTo("eins"));
        return null;
      });
      assertThat(computed.get(), equalTo("eins"));
      assertThat(xaStore.get(1L), is(nullValue()));
    }
    testTransactionManager.rollback();

    assertMapping(xaStore, 1L, "eins");

    testTransactionManager.begin();
    {
      Store.ValueHolder<String> computed1 = xaStore.getAndCompute(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, equalTo("eins"));
        return null;
      });
      assertThat(computed1.get(), equalTo("eins"));
      assertThat(xaStore.get(1L), is(nullValue()));
      Store.ValueHolder<String> computed2 = xaStore.getAndCompute(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, is(nullValue()));
        return null;
      });
      assertThat(computed2, is(nullValue()));
      assertThat(xaStore.get(1L), is(nullValue()));
      Store.ValueHolder<String> computed3 = xaStore.getAndCompute(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, is(nullValue()));
        return "uno";
      });
      assertThat(computed3, is(nullValue()));
      assertThat(xaStore.get(1L).get(), equalTo("uno"));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "uno");

    testTransactionManager.begin();
    {
      xaStore.remove(1L);
    }
    testTransactionManager.commit();

    testTransactionManager.begin();
    {
      assertThat(xaStore.containsKey(1L), is(false));
      xaStore.put(1L, "uno");
      assertThat(xaStore.containsKey(1L), is(true));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "uno");

    testTransactionManager.begin();
    {
      assertThat(xaStore.containsKey(1L), is(true));
      xaStore.remove(1L);
      assertThat(xaStore.containsKey(1L), is(false));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, null);

    OffHeapStoreLifecycleHelper.close(offHeapStore);
  }

  @Test
  public void testCompute() throws Exception {
    Store.Configuration<Long, SoftLock<String>> offHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass,
      null, classLoader, ExpiryPolicyBuilder.noExpiration(), ResourcePoolsBuilder.newResourcePoolsBuilder()
      .offheap(10, MemoryUnit.MB)
      .build(),
      0, keySerializer, valueSerializer);
    OffHeapStore<Long, SoftLock<String>> offHeapStore = new OffHeapStore<>(offHeapConfig, testTimeSource, eventDispatcher, MemorySizeParser
      .parse("10M"), new DefaultStatisticsService());
    OffHeapStoreLifecycleHelper.init(offHeapStore);
    TieredStore<Long, SoftLock<String>> tieredStore = new TieredStore<>(onHeapStore, offHeapStore);

    XAStore<Long, String> xaStore = getXAStore(tieredStore);

    testTransactionManager.begin();
    {
      Store.ValueHolder<String> computed1 = xaStore.computeAndGet(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, is(nullValue()));
        return "one";
      }, SUPPLY_TRUE, SUPPLY_FALSE);
      assertThat(computed1.get(), equalTo("one"));
      Store.ValueHolder<String> computed2 = xaStore.computeAndGet(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, equalTo("one"));
        return "un";
      }, SUPPLY_TRUE, SUPPLY_FALSE);
      assertThat(computed2.get(), equalTo("un"));
      Store.ValueHolder<String> computed3 = xaStore.computeAndGet(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, equalTo("un"));
        return null;
      }, SUPPLY_TRUE, SUPPLY_FALSE);
      assertThat(computed3, is(nullValue()));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, null);

    testTransactionManager.begin();
    {
      Store.ValueHolder<String> computed1 = xaStore.computeAndGet(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, is(nullValue()));
        return "one";
      }, SUPPLY_FALSE, SUPPLY_FALSE);
      assertThat(computed1.get(), equalTo("one"));
      Store.ValueHolder<String> computed2 = xaStore.computeAndGet(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, equalTo("one"));
        return null;
      }, SUPPLY_FALSE, SUPPLY_FALSE);
      assertThat(computed2, is(nullValue()));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, null);

    testTransactionManager.begin();
    {
      Store.ValueHolder<String> computed1 = xaStore.computeAndGet(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, is(nullValue()));
        return "one";
      }, SUPPLY_TRUE, SUPPLY_FALSE);
      assertThat(computed1.get(), equalTo("one"));
      Store.ValueHolder<String> computed2 = xaStore.computeAndGet(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, equalTo("one"));
        return null;
      }, SUPPLY_TRUE, SUPPLY_FALSE);
      assertThat(computed2, is(nullValue()));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, null);

    testTransactionManager.begin();
    {
      Store.ValueHolder<String> computed1 = xaStore.computeAndGet(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, is(nullValue()));
        return "one";
      }, SUPPLY_TRUE, SUPPLY_FALSE);
      assertThat(computed1.get(), equalTo("one"));
      Store.ValueHolder<String> computed2 = xaStore.computeAndGet(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, equalTo("one"));
        return "un";
      }, SUPPLY_TRUE, SUPPLY_FALSE);
      assertThat(computed2.get(), equalTo("un"));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "un");

    testTransactionManager.begin();
    {
      Store.ValueHolder<String> computed = xaStore.computeAndGet(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, equalTo("un"));
        return "eins";
      }, SUPPLY_TRUE, SUPPLY_FALSE);
      assertThat(computed.get(), equalTo("eins"));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "eins");

    testTransactionManager.begin();
    {
      Store.ValueHolder<String> computed = xaStore.computeAndGet(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, equalTo("eins"));
        return null;
      }, SUPPLY_TRUE, SUPPLY_FALSE);
      assertThat(computed, is(nullValue()));
    }
    testTransactionManager.rollback();

    assertMapping(xaStore, 1L, "eins");

    testTransactionManager.begin();
    {
      Store.ValueHolder<String> computed1 = xaStore.computeAndGet(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, equalTo("eins"));
        return null;
      }, SUPPLY_TRUE, SUPPLY_FALSE);
      assertThat(computed1, is(nullValue()));
      Store.ValueHolder<String> computed2 = xaStore.computeAndGet(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, is(nullValue()));
        return null;
      }, SUPPLY_TRUE, SUPPLY_FALSE);
      assertThat(computed2, is(nullValue()));
      Store.ValueHolder<String> computed3 = xaStore.computeAndGet(1L, (aLong, s) -> {
        assertThat(aLong, is(1L));
        assertThat(s, is(nullValue()));
        return "uno";
      }, SUPPLY_TRUE, SUPPLY_FALSE);
      assertThat(computed3.get(), equalTo("uno"));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "uno");

    testTransactionManager.begin();
    {
      xaStore.remove(1L);
    }
    testTransactionManager.commit();

    testTransactionManager.begin();
    {
      assertThat(xaStore.containsKey(1L), is(false));
      xaStore.put(1L, "uno");
      assertThat(xaStore.containsKey(1L), is(true));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "uno");

    testTransactionManager.begin();
    {
      assertThat(xaStore.containsKey(1L), is(true));
      xaStore.remove(1L);
      assertThat(xaStore.containsKey(1L), is(false));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, null);

    OffHeapStoreLifecycleHelper.close(offHeapStore);
  }

  @Test
  public void testComputeIfAbsent() throws Exception {
    Store.Configuration<Long, SoftLock<String>> offHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass, null,
      classLoader, ExpiryPolicyBuilder.noExpiration(), ResourcePoolsBuilder.newResourcePoolsBuilder()
      .offheap(10, MemoryUnit.MB)
      .build(),
      0, keySerializer, valueSerializer);
    OffHeapStore<Long, SoftLock<String>> offHeapStore = new OffHeapStore<>(offHeapConfig, testTimeSource, eventDispatcher, MemorySizeParser
      .parse("10M"), new DefaultStatisticsService());
    OffHeapStoreLifecycleHelper.init(offHeapStore);
    TieredStore<Long, SoftLock<String>> tieredStore = new TieredStore<>(onHeapStore, offHeapStore);

    XAStore<Long, String> xaStore = getXAStore(tieredStore);

    testTransactionManager.begin();
    {
      Store.ValueHolder<String> computed1 = xaStore.computeIfAbsent(1L, aLong -> {
        assertThat(aLong, is(1L));
        return "one";
      });
      assertThat(computed1.get(), equalTo("one"));
      Store.ValueHolder<String> computed2 = xaStore.computeIfAbsent(1L, aLong -> {
        fail("should not be absent");
        throw new AssertionError();
      });
      assertThat(computed2.get(), equalTo("one"));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "one");

    testTransactionManager.begin();
    {
      Store.ValueHolder<String> computed1 = xaStore.computeIfAbsent(1L, aLong -> {
        fail("should not be absent");
        throw new AssertionError();
      });
      assertThat(computed1.get(), equalTo("one"));

      xaStore.remove(1L);

      Store.ValueHolder<String> computed2 = xaStore.computeIfAbsent(1L, aLong -> {
        assertThat(aLong, is(1L));
        return "un";
      });
      assertThat(computed2.get(), equalTo("un"));
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "un");

    OffHeapStoreLifecycleHelper.close(offHeapStore);
  }

  @Test
  public void testExpiry() throws Exception {
    Store.Configuration<Long, SoftLock<String>> onHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass,
      null, classLoader, expiry, ResourcePoolsBuilder.newResourcePoolsBuilder().heap(10, EntryUnit.ENTRIES).build(),
      0, keySerializer, valueSerializer);
    onHeapStore = new OnHeapStore<>(onHeapConfig, testTimeSource, keyCopier, valueCopier, new NoopSizeOfEngine(), eventDispatcher, new DefaultStatisticsService());
    Store.Configuration<Long, SoftLock<String>> offHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass, null,
      classLoader, expiry, ResourcePoolsBuilder.newResourcePoolsBuilder().offheap(10, MemoryUnit.MB).build(),
      0, keySerializer, valueSerializer);
    OffHeapStore<Long, SoftLock<String>> offHeapStore = new OffHeapStore<>(offHeapConfig, testTimeSource, eventDispatcher, MemorySizeParser
      .parse("10M"), new DefaultStatisticsService());
    OffHeapStoreLifecycleHelper.init(offHeapStore);
    TieredStore<Long, SoftLock<String>> tieredStore = new TieredStore<>(onHeapStore, offHeapStore);

    XAStore<Long, String> xaStore = getXAStore(tieredStore);

    testTransactionManager.begin();
    {
      xaStore.put(1L, "one");
    }
    testTransactionManager.commit();

    assertMapping(xaStore, 1L, "one");

    testTimeSource.advanceTime(2000);

    assertMapping(xaStore, 1L, null);

    OffHeapStoreLifecycleHelper.close(offHeapStore);
  }

  @Test
  public void testExpiryCreateException() throws Exception {
    ExpiryPolicy<Object, Object> expiry = new ExpiryPolicy<Object, Object>() {

      @Override
      public Duration getExpiryForCreation(Object key, Object value) {
        throw new RuntimeException();
      }

      @Override
      public Duration getExpiryForAccess(Object key, Supplier<?> value) {
        throw new AssertionError();
      }

      @Override
      public Duration getExpiryForUpdate(Object key, Supplier<?> oldValue, Object newValue) {
        throw new AssertionError();
      }
    };
    Store.Configuration<Long, SoftLock<String>> onHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass, null,
      classLoader, expiry, ResourcePoolsBuilder.newResourcePoolsBuilder().heap(10, EntryUnit.ENTRIES).build(),
      0, keySerializer, valueSerializer);
    OnHeapStore<Long, SoftLock<String>> onHeapStore = new OnHeapStore<>(onHeapConfig, testTimeSource, keyCopier, valueCopier, new NoopSizeOfEngine(), eventDispatcher, new DefaultStatisticsService());
    Store.Configuration<Long, SoftLock<String>> offHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass, null,
      classLoader, expiry, ResourcePoolsBuilder.newResourcePoolsBuilder().offheap(10, MemoryUnit.MB).build(),
      0, keySerializer, valueSerializer);
    OffHeapStore<Long, SoftLock<String>> offHeapStore = new OffHeapStore<>(offHeapConfig, testTimeSource, eventDispatcher, MemorySizeParser
      .parse("10M"), new DefaultStatisticsService());
    OffHeapStoreLifecycleHelper.init(offHeapStore);
    TieredStore<Long, SoftLock<String>> tieredStore = new TieredStore<>(onHeapStore, offHeapStore);

    XAStore<Long, String> xaStore = getXAStore(tieredStore);

    testTransactionManager.begin();
    xaStore.put(1L, "one");
    testTransactionManager.commit();
    assertMapping(xaStore, 1L, null);
  }

  @Test
  public void testExpiryAccessException() throws Exception {
    String uniqueXAResourceId = "testExpiryAccessException";
    ExpiryPolicy<Object, Object> expiry = new ExpiryPolicy<Object, Object>() {

      @Override
      public Duration getExpiryForCreation(Object key, Object value) {
        return ExpiryPolicy.INFINITE;
      }

      @Override
      public Duration getExpiryForAccess(Object key, Supplier<?> value) {
        if (testTimeSource.getTimeMillis() > 0) {
          throw new RuntimeException();
        }
        return ExpiryPolicy.INFINITE;
      }

      @Override
      public Duration getExpiryForUpdate(Object key, Supplier<?> oldValue, Object newValue) {
        return ExpiryPolicy.INFINITE;
      }
    };
    Store.Configuration<Long, SoftLock<String>> onHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass, null,
      classLoader, expiry, ResourcePoolsBuilder.newResourcePoolsBuilder().heap(10, EntryUnit.ENTRIES).build(),
      0, keySerializer, valueSerializer);
    OnHeapStore<Long, SoftLock<String>> onHeapStore = new OnHeapStore<>(onHeapConfig, testTimeSource, keyCopier, valueCopier, new NoopSizeOfEngine(), eventDispatcher, new DefaultStatisticsService());
    Store.Configuration<Long, SoftLock<String>> offHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass, null,
      classLoader, expiry, ResourcePoolsBuilder.newResourcePoolsBuilder().offheap(10, MemoryUnit.MB).build(),
      0, keySerializer, valueSerializer);
    OffHeapStore<Long, SoftLock<String>> offHeapStore = new OffHeapStore<>(offHeapConfig, testTimeSource, eventDispatcher, MemorySizeParser
      .parse("10M"), new DefaultStatisticsService());
    OffHeapStoreLifecycleHelper.init(offHeapStore);
    TieredStore<Long, SoftLock<String>> tieredStore = new TieredStore<>(onHeapStore, offHeapStore);

    XAStore<Long, String> xaStore = getXAStore(tieredStore);

    testTransactionManager.begin();
    xaStore.put(1L, "one");
    testTransactionManager.commit();

    testTimeSource.advanceTime(1000);
    testTransactionManager.begin();
    assertThat(xaStore.get(1L).get(), is("one"));
    testTransactionManager.commit();

    testTransactionManager.begin();
    assertThat(xaStore.get(1L), nullValue());
    testTransactionManager.commit();
  }

  @Test
  public void testExpiryUpdateException() throws Exception{
    ExpiryPolicy<Object, Object> expiry = new ExpiryPolicy<Object, Object>() {

      @Override
      public Duration getExpiryForCreation(Object key, Object value) {
        return ExpiryPolicy.INFINITE;
      }

      @Override
      public Duration getExpiryForAccess(Object key, Supplier<?> value) {
        return ExpiryPolicy.INFINITE;
      }

      @Override
      public Duration getExpiryForUpdate(Object key, Supplier<?> oldValue, Object newValue) {
        if (testTimeSource.getTimeMillis() > 0) {
          throw new RuntimeException();
        }
        return ExpiryPolicy.INFINITE;
      }
    };
    Store.Configuration<Long, SoftLock<String>> onHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass, null,
      classLoader, expiry, ResourcePoolsBuilder.newResourcePoolsBuilder().heap(10, EntryUnit.ENTRIES).build(),
      0, keySerializer, valueSerializer);
    OnHeapStore<Long, SoftLock<String>> onHeapStore = new OnHeapStore<>(onHeapConfig, testTimeSource, keyCopier, valueCopier, new NoopSizeOfEngine(), eventDispatcher, new DefaultStatisticsService());
    Store.Configuration<Long, SoftLock<String>> offHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass, null,
      classLoader, expiry, ResourcePoolsBuilder.newResourcePoolsBuilder().offheap(10, MemoryUnit.MB).build(),
      0, keySerializer, valueSerializer);
    OffHeapStore<Long, SoftLock<String>> offHeapStore = new OffHeapStore<>(offHeapConfig, testTimeSource, eventDispatcher, MemorySizeParser
      .parse("10M"), new DefaultStatisticsService());
    OffHeapStoreLifecycleHelper.init(offHeapStore);
    TieredStore<Long, SoftLock<String>> tieredStore = new TieredStore<>(onHeapStore, offHeapStore);

    XAStore<Long, String> xaStore = getXAStore(tieredStore);

    testTransactionManager.begin();
    xaStore.put(1L, "one");
    xaStore.get(1L);
    testTransactionManager.commit();

    testTimeSource.advanceTime(1000);
    testTransactionManager.begin();
    xaStore.put(1L, "two");
    testTransactionManager.commit();
    assertMapping(xaStore, 1L, null);
  }

  @Test
  public void testBulkCompute() throws Exception {
    ExpiryPolicy<Object, Object> expiry = ExpiryPolicyBuilder.timeToLiveExpiration(Duration.ofSeconds(1));
    Store.Configuration<Long, SoftLock<String>> onHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass, null,
      classLoader, expiry, ResourcePoolsBuilder.newResourcePoolsBuilder().heap(10, EntryUnit.ENTRIES).build(),
      0, keySerializer, valueSerializer);
    OnHeapStore<Long, SoftLock<String>> onHeapStore = new OnHeapStore<>(onHeapConfig, testTimeSource, keyCopier, valueCopier, new NoopSizeOfEngine(), eventDispatcher, new DefaultStatisticsService());
    Store.Configuration<Long, SoftLock<String>> offHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass, null,
      classLoader, expiry, ResourcePoolsBuilder.newResourcePoolsBuilder().offheap(10, MemoryUnit.MB).build(),
      0, keySerializer, valueSerializer);
    OffHeapStore<Long, SoftLock<String>> offHeapStore = new OffHeapStore<>(offHeapConfig, testTimeSource, eventDispatcher, MemorySizeParser
      .parse("10M"), new DefaultStatisticsService());
    OffHeapStoreLifecycleHelper.init(offHeapStore);
    TieredStore<Long, SoftLock<String>> tieredStore = new TieredStore<>(onHeapStore, offHeapStore);

    XAStore<Long, String> xaStore = getXAStore(tieredStore);

    testTransactionManager.begin();
    {
      Map<Long, Store.ValueHolder<String>> computedMap = xaStore.bulkCompute(asSet(1L, 2L, 3L), entries -> {
        Map<Long, String> result = new HashMap<>();
        for (Map.Entry<? extends Long, ? extends String> entry : entries) {
          Long key = entry.getKey();
          String value = entry.getValue();
          assertThat(value, is(nullValue()));
          result.put(key, "stuff#" + key);
        }
        return result.entrySet();
      });

      assertThat(computedMap.size(), is(3));
      assertThat(computedMap.get(1L).get(), equalTo("stuff#1"));
      assertThat(computedMap.get(2L).get(), equalTo("stuff#2"));
      assertThat(computedMap.get(3L).get(), equalTo("stuff#3"));

      computedMap = xaStore.bulkCompute(asSet(0L, 1L, 3L), entries -> {
        Map<Long, String> result = new HashMap<>();
        for (Map.Entry<? extends Long, ? extends String> entry : entries) {
          Long key = entry.getKey();
          String value = entry.getValue();

          switch (key.intValue()) {
            case 0:
              assertThat(value, is(nullValue()));
              break;
            case 1:
            case 3:
              assertThat(value, equalTo("stuff#" + key));
              break;
          }

          if (key != 3L) {
            result.put(key, "otherStuff#" + key);
          } else {
            result.put(key, null);
          }
        }
        return result.entrySet();
      });

      assertThat(computedMap.size(), is(3));
      assertThat(computedMap.get(0L).get(), equalTo("otherStuff#0"));
      assertThat(computedMap.get(1L).get(), equalTo("otherStuff#1"));
      assertThat(computedMap.get(3L), is(nullValue()));
    }
    testTransactionManager.commit();

    assertSize(xaStore, 3);
    assertMapping(xaStore, 0L, "otherStuff#0");
    assertMapping(xaStore, 1L, "otherStuff#1");
    assertMapping(xaStore, 2L, "stuff#2");

    OffHeapStoreLifecycleHelper.close(offHeapStore);
  }

  @Test
  public void testBulkComputeIfAbsent() throws Exception {
    ExpiryPolicy<Object, Object> expiry = ExpiryPolicyBuilder.timeToLiveExpiration(Duration.ofSeconds(1));
    Store.Configuration<Long, SoftLock<String>> onHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass, null,
      classLoader, expiry, ResourcePoolsBuilder.newResourcePoolsBuilder().heap(10, EntryUnit.ENTRIES).build(),
      0, keySerializer, valueSerializer);
    OnHeapStore<Long, SoftLock<String>> onHeapStore = new OnHeapStore<>(onHeapConfig, testTimeSource, keyCopier, valueCopier, new NoopSizeOfEngine(), eventDispatcher, new DefaultStatisticsService());
    Store.Configuration<Long, SoftLock<String>> offHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass, null,
      classLoader, expiry, ResourcePoolsBuilder.newResourcePoolsBuilder().offheap(10, MemoryUnit.MB).build(),
      0, keySerializer, valueSerializer);
    OffHeapStore<Long, SoftLock<String>> offHeapStore = new OffHeapStore<>(offHeapConfig, testTimeSource, eventDispatcher, MemorySizeParser
      .parse("10M"), new DefaultStatisticsService());
    OffHeapStoreLifecycleHelper.init(offHeapStore);
    TieredStore<Long, SoftLock<String>> tieredStore = new TieredStore<>(onHeapStore, offHeapStore);

    XAStore<Long, String> xaStore = getXAStore(tieredStore);

    testTransactionManager.begin();
    {
      Map<Long, Store.ValueHolder<String>> computedMap = xaStore.bulkComputeIfAbsent(asSet(1L, 2L, 3L), keys -> {
        Map<Long, String> result = new HashMap<>();
        for (Long key : keys) {
          result.put(key, "stuff#" + key);
        }
        return result.entrySet();
      });

      assertThat(computedMap.size(), is(3));
      assertThat(computedMap.get(1L).get(), equalTo("stuff#1"));
      assertThat(computedMap.get(2L).get(), equalTo("stuff#2"));
      assertThat(computedMap.get(3L).get(), equalTo("stuff#3"));

      computedMap = xaStore.bulkComputeIfAbsent(asSet(0L, 1L, 3L), keys -> {
        Map<Long, String> result = new HashMap<>();
        for (Long key : keys) {
          switch (key.intValue()) {
            case 0:
              result.put(key, "otherStuff#" + key);
              break;
            case 1:
            case 3:
              fail("key " + key + " should not be absent");
              break;
          }
        }
        return result.entrySet();
      });

      assertThat(computedMap.size(), is(3));
      assertThat(computedMap.get(0L).get(), equalTo("otherStuff#0"));
      assertThat(computedMap.get(1L).get(), equalTo("stuff#1"));
      assertThat(computedMap.get(3L).get(), equalTo("stuff#3"));
    }
    testTransactionManager.commit();

    assertSize(xaStore, 4);
    assertMapping(xaStore, 0L, "otherStuff#0");
    assertMapping(xaStore, 1L, "stuff#1");
    assertMapping(xaStore, 2L, "stuff#2");
    assertMapping(xaStore, 3L, "stuff#3");

    OffHeapStoreLifecycleHelper.close(offHeapStore);
  }

  @Test
  public void testCustomEvictionAdvisor() throws Exception {
    final AtomicBoolean invoked = new AtomicBoolean();

    EvictionAdvisor<Long, SoftLock<String>> evictionAdvisor = (key, value) -> {
      invoked.set(true);
      return false;
    };
    Store.Configuration<Long, SoftLock<String>> onHeapConfig = new StoreConfigurationImpl<>(Long.class, valueClass,
      evictionAdvisor, classLoader, ExpiryPolicyBuilder.noExpiration(), ResourcePoolsBuilder.newResourcePoolsBuilder()
      .heap(10, EntryUnit.ENTRIES)
      .build(),
      0, keySerializer, valueSerializer);
    OnHeapStore<Long, SoftLock<String>> onHeapStore = new OnHeapStore<>(onHeapConfig, testTimeSource, keyCopier, valueCopier, new NoopSizeOfEngine(), eventDispatcher, new DefaultStatisticsService());

    final XAStore<Long, String> xaStore = getXAStore(onHeapStore);

    testTransactionManager.begin();
    {
      xaStore.put(1L, "1");
    }
    testTransactionManager.rollback();
    assertThat(invoked.get(), is(false));

    testTransactionManager.begin();
    {
      xaStore.put(1L, "1");
    }
    testTransactionManager.commit();
    assertThat(invoked.get(), is(true));
  }

  @Test
  public void testRank() throws Exception {
    XAStore.Provider provider = new XAStore.Provider();
    XAStoreConfiguration configuration = new XAStoreConfiguration("testXAResourceId");
    ServiceLocator serviceLocator = dependencySet()
      .with(provider)
      .with(Store.Provider.class)
      .with(mock(DiskResourceService.class))
      .with(mock(CacheManagerProviderService.class, Answers.RETURNS_DEEP_STUBS))
      .with(mock(TransactionManagerProvider.class)).build();

    serviceLocator.startAllServices();

    Set<ServiceConfiguration<?, ?>> xaStoreConfigs = Collections.singleton(configuration);
    assertThat(provider.wrapperStoreRank(xaStoreConfigs), is(1));

    Set<ServiceConfiguration<?, ?>> emptyConfigs = emptySet();
    assertThat(provider.wrapperStoreRank(emptyConfigs), is(0));

  }

  private Set<Long> asSet(Long... longs) {
    return new HashSet<>(Arrays.asList(longs));
  }

  private void assertMapping(XAStore<Long, String> xaStore, long key, String value) throws Exception {
    testTransactionManager.begin();

    Store.ValueHolder<String> valueHolder = xaStore.get(key);
    if (value != null) {
      assertThat(valueHolder.get(), equalTo(value));
    } else {
      assertThat(valueHolder, is(nullValue()));
    }

    testTransactionManager.commit();
  }

  private void assertSize(XAStore<Long, String> xaStore, int expectedSize) throws Exception {
    testTransactionManager.begin();

    int counter = 0;
    Store.Iterator<Cache.Entry<Long, Store.ValueHolder<String>>> iterator = xaStore.iterator();
    while (iterator.hasNext()) {
      iterator.next();
      counter++;
    }
    assertThat(counter, is(expectedSize));

    testTransactionManager.commit();
  }

  private XAStore<Long, String> getXAStore(Store<Long, SoftLock<String>> store) {
    return new XAStore<>(Long.class, String.class, store, transactionManagerWrapper, testTimeSource, journal, testName.getMethodName(), new DefaultStatisticsService());
  }

  static class TestTransactionManager implements TransactionManager {

    volatile TestTransaction currentTransaction;
    final AtomicLong gtridGenerator = new AtomicLong();

    public TestTransaction getCurrentTransaction() {
      return currentTransaction;
    }

    @Override
    public void begin() {
      currentTransaction = new TestTransaction(gtridGenerator.incrementAndGet());
    }

    @Override
    public void commit() throws RollbackException, HeuristicMixedException, HeuristicRollbackException, SecurityException, IllegalStateException, SystemException {
      currentTransaction.commit();
      currentTransaction = null;
    }

    @Override
    public int getStatus() {
      return 0;
    }

    @Override
    public Transaction getTransaction() {
      return currentTransaction;
    }

    @Override
    public void resume(Transaction tobj) {

    }

    @Override
    public void rollback() throws IllegalStateException, SecurityException, SystemException {
      currentTransaction.rollback();
      currentTransaction = null;
    }

    @Override
    public void setRollbackOnly() {

    }

    @Override
    public void setTransactionTimeout(int seconds) {

    }

    @Override
    public Transaction suspend() {
      return null;
    }
  }

  static class TestTransaction implements Transaction {

    final long gtrid;
    final Map<XAResource, TestXid> xids = new IdentityHashMap<>();
    final AtomicLong bqualGenerator = new AtomicLong();
    final List<Synchronization> synchronizations = new CopyOnWriteArrayList<>();
    final List<TwoPcListener> twoPcListeners = new CopyOnWriteArrayList<>();

    public TestTransaction(long gtrid) {
      this.gtrid = gtrid;
    }

    @Override
    public void commit() throws RollbackException, HeuristicMixedException, HeuristicRollbackException, SecurityException, IllegalStateException, SystemException {
      try {
        Set<Map.Entry<XAResource, TestXid>> entries = xids.entrySet();

        // delist
        for (Map.Entry<XAResource, TestXid> entry : entries) {
          try {
            entry.getKey().end(entry.getValue(), XAResource.TMSUCCESS);
          } catch (XAException e) {
            throw (SystemException) new SystemException(XAException.XAER_RMERR).initCause(e);
          }
        }

        fireBeforeCompletion();

        Set<XAResource> preparedResources = new HashSet<>();

        // prepare
        for (Map.Entry<XAResource, TestXid> entry : entries) {
          try {
            int prepareStatus = entry.getKey().prepare(entry.getValue());
            if (prepareStatus != XAResource.XA_RDONLY) {
              preparedResources.add(entry.getKey());
            }
          } catch (XAException e) {
            throw (SystemException) new SystemException(XAException.XAER_RMERR).initCause(e);
          }
        }

        fireInMiddleOf2PC();

        // commit
        for (Map.Entry<XAResource, TestXid> entry : entries) {
          try {
            if (preparedResources.contains(entry.getKey())) {
              entry.getKey().commit(entry.getValue(), false);
            }
          } catch (XAException e) {
            throw (SystemException) new SystemException(XAException.XAER_RMERR).initCause(e);
          }
        }
      } finally {
        fireAfterCompletion(Status.STATUS_COMMITTED);
      }
    }

    @Override
    public boolean delistResource(XAResource xaRes, int flag) {
      return true;
    }

    @Override
    public boolean enlistResource(XAResource xaRes) throws IllegalStateException, SystemException {
      TestXid testXid = xids.get(xaRes);
      if (testXid == null) {
        testXid = new TestXid(gtrid, bqualGenerator.incrementAndGet());
        xids.put(xaRes, testXid);
      }

      try {
        xaRes.start(testXid, XAResource.TMNOFLAGS);
      } catch (XAException e) {
        throw (SystemException) new SystemException(XAException.XAER_RMERR).initCause(e);
      }
      return true;
    }

    @Override
    public int getStatus() {
      return 0;
    }

    public void registerTwoPcListener(TwoPcListener listener) {
      twoPcListeners.add(listener);
    }

    @Override
    public void registerSynchronization(Synchronization sync) {
      synchronizations.add(sync);
    }

    @Override
    public void rollback() throws IllegalStateException, SystemException {
      try {
        Set<Map.Entry<XAResource, TestXid>> entries = xids.entrySet();

        // delist
        for (Map.Entry<XAResource, TestXid> entry : entries) {
          try {
            entry.getKey().end(entry.getValue(), XAResource.TMSUCCESS);
          } catch (XAException e) {
            throw (SystemException) new SystemException(XAException.XAER_RMERR).initCause(e);
          }
        }

        // rollback
        for (Map.Entry<XAResource, TestXid> entry : entries) {
          try {
            entry.getKey().rollback(entry.getValue());
          } catch (XAException e) {
            throw (SystemException) new SystemException(XAException.XAER_RMERR).initCause(e);
          }
        }
      } finally {
        fireAfterCompletion(Status.STATUS_ROLLEDBACK);
      }
    }

    private void fireBeforeCompletion() {
      for (Synchronization synchronization : synchronizations) {
        synchronization.beforeCompletion();
      }
    }

    private void fireAfterCompletion(int status) {
      for (Synchronization synchronization : synchronizations) {
        synchronization.afterCompletion(status);
      }
    }

    private void fireInMiddleOf2PC() {
      for (TwoPcListener twoPcListener : twoPcListeners) {
        twoPcListener.inMiddleOf2PC();
      }
    }

    @Override
    public void setRollbackOnly() {

    }
  }

  interface TwoPcListener {
    void inMiddleOf2PC();
  }

}


text:
package org.apache.flink.runtime.taskexecutor.slot;

import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.runtime.clusterframework.types.AllocationID;
import org.apache.flink.runtime.clusterframework.types.ResourceID;
import org.apache.flink.runtime.clusterframework.types.ResourceProfile;
import org.apache.flink.runtime.clusterframework.types.SlotID;
import org.apache.flink.runtime.concurrent.ComponentMainThreadExecutorServiceAdapter;
import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
import org.apache.flink.runtime.taskexecutor.SlotReport;
import org.apache.flink.runtime.taskexecutor.SlotStatus;
import org.apache.flink.util.TestLogger;
import org.apache.flink.util.function.TriFunctionWithException;

import org.apache.flink.shaded.guava18.com.google.common.collect.Sets;

import org.junit.Test;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.is;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.fail;

/** Tests for the {@link TaskSlotTable}. */
public class TaskSlotTableImplTest extends TestLogger {
    private static final Time SLOT_TIMEOUT = Time.seconds(100L);

    /** Tests that one can can mark allocated slots as active. */
    @Test
    public void testTryMarkSlotActive() throws Exception {
        final TaskSlotTableImpl<?> taskSlotTable = createTaskSlotTableAndStart(3);

        try {
            final JobID jobId1 = new JobID();
            final AllocationID allocationId1 = new AllocationID();
            taskSlotTable.allocateSlot(0, jobId1, allocationId1, SLOT_TIMEOUT);
            final AllocationID allocationId2 = new AllocationID();
            taskSlotTable.allocateSlot(1, jobId1, allocationId2, SLOT_TIMEOUT);
            final AllocationID allocationId3 = new AllocationID();
            final JobID jobId2 = new JobID();
            taskSlotTable.allocateSlot(2, jobId2, allocationId3, SLOT_TIMEOUT);

            taskSlotTable.markSlotActive(allocationId1);

            assertThat(taskSlotTable.isAllocated(0, jobId1, allocationId1), is(true));
            assertThat(taskSlotTable.isAllocated(1, jobId1, allocationId2), is(true));
            assertThat(taskSlotTable.isAllocated(2, jobId2, allocationId3), is(true));

            assertThat(
                    taskSlotTable.getActiveTaskSlotAllocationIdsPerJob(jobId1),
                    is(equalTo(Sets.newHashSet(allocationId1))));

            assertThat(taskSlotTable.tryMarkSlotActive(jobId1, allocationId1), is(true));
            assertThat(taskSlotTable.tryMarkSlotActive(jobId1, allocationId2), is(true));
            assertThat(taskSlotTable.tryMarkSlotActive(jobId1, allocationId3), is(false));

            assertThat(
                    taskSlotTable.getActiveTaskSlotAllocationIdsPerJob(jobId1),
                    is(equalTo(new HashSet<>(Arrays.asList(allocationId2, allocationId1)))));
        } finally {
            taskSlotTable.close();
            assertThat(taskSlotTable.isClosed(), is(true));
        }
    }

    /** Tests {@link TaskSlotTableImpl#getActiveTaskSlotAllocationIds()}. */
    @Test
    public void testRetrievingAllActiveSlots() throws Exception {
        try (final TaskSlotTableImpl<?> taskSlotTable = createTaskSlotTableAndStart(3)) {
            final JobID jobId1 = new JobID();
            final AllocationID allocationId1 = new AllocationID();
            taskSlotTable.allocateSlot(0, jobId1, allocationId1, SLOT_TIMEOUT);
            final AllocationID allocationId2 = new AllocationID();
            taskSlotTable.allocateSlot(1, jobId1, allocationId2, SLOT_TIMEOUT);
            final AllocationID allocationId3 = new AllocationID();
            final JobID jobId2 = new JobID();
            taskSlotTable.allocateSlot(2, jobId2, allocationId3, SLOT_TIMEOUT);

            taskSlotTable.markSlotActive(allocationId1);
            taskSlotTable.markSlotActive(allocationId3);

            assertThat(
                    taskSlotTable.getActiveTaskSlotAllocationIds(),
                    is(Sets.newHashSet(allocationId1, allocationId3)));
        }
    }

    /**
     * Tests that inconsistent static slot allocation with the same AllocationID to a different slot
     * is rejected.
     */
    @Test
    public void testInconsistentStaticSlotAllocation() throws Exception {
        try (final TaskSlotTable<TaskSlotPayload> taskSlotTable = createTaskSlotTableAndStart(2)) {
            final JobID jobId = new JobID();
            final AllocationID allocationId1 = new AllocationID();
            final AllocationID allocationId2 = new AllocationID();

            assertThat(taskSlotTable.allocateSlot(0, jobId, allocationId1, SLOT_TIMEOUT), is(true));
            assertThat(
                    taskSlotTable.allocateSlot(1, jobId, allocationId1, SLOT_TIMEOUT), is(false));
            assertThat(
                    taskSlotTable.allocateSlot(0, jobId, allocationId2, SLOT_TIMEOUT), is(false));

            assertThat(taskSlotTable.isAllocated(0, jobId, allocationId1), is(true));
            assertThat(taskSlotTable.isSlotFree(1), is(true));

            Iterator<TaskSlot<TaskSlotPayload>> allocatedSlots =
                    taskSlotTable.getAllocatedSlots(jobId);
            assertThat(allocatedSlots.next().getIndex(), is(0));
            assertThat(allocatedSlots.hasNext(), is(false));
        }
    }

    /**
     * Tests that inconsistent dynamic slot allocation with the same AllocationID to a different
     * slot is rejected.
     */
    @Test
    public void testInconsistentDynamicSlotAllocation() throws Exception {
        try (final TaskSlotTable<TaskSlotPayload> taskSlotTable = createTaskSlotTableAndStart(1)) {
            final JobID jobId1 = new JobID();
            final JobID jobId2 = new JobID();
            final AllocationID allocationId = new AllocationID();

            assertThat(
                    taskSlotTable.allocateSlot(-1, jobId1, allocationId, SLOT_TIMEOUT), is(true));
            assertThat(
                    taskSlotTable.allocateSlot(-1, jobId2, allocationId, SLOT_TIMEOUT), is(false));

            assertThat(taskSlotTable.isAllocated(1, jobId1, allocationId), is(true));

            Iterator<TaskSlot<TaskSlotPayload>> allocatedSlots =
                    taskSlotTable.getAllocatedSlots(jobId1);
            assertThat(allocatedSlots.next().getAllocationId(), is(allocationId));
            assertThat(allocatedSlots.hasNext(), is(false));
        }
    }

    @Test
    public void testDuplicateStaticSlotAllocation() throws Exception {
        try (final TaskSlotTable<TaskSlotPayload> taskSlotTable = createTaskSlotTableAndStart(2)) {
            final JobID jobId = new JobID();
            final AllocationID allocationId = new AllocationID();

            assertThat(taskSlotTable.allocateSlot(0, jobId, allocationId, SLOT_TIMEOUT), is(true));
            assertThat(taskSlotTable.allocateSlot(0, jobId, allocationId, SLOT_TIMEOUT), is(true));

            assertThat(taskSlotTable.isAllocated(0, jobId, allocationId), is(true));
            assertThat(taskSlotTable.isSlotFree(1), is(true));

            Iterator<TaskSlot<TaskSlotPayload>> allocatedSlots =
                    taskSlotTable.getAllocatedSlots(jobId);
            assertThat(allocatedSlots.next().getIndex(), is(0));
            assertThat(allocatedSlots.hasNext(), is(false));
        }
    }

    @Test
    public void testDuplicateDynamicSlotAllocation() throws Exception {
        try (final TaskSlotTable<TaskSlotPayload> taskSlotTable = createTaskSlotTableAndStart(1)) {
            final JobID jobId = new JobID();
            final AllocationID allocationId = new AllocationID();

            assertThat(taskSlotTable.allocateSlot(-1, jobId, allocationId, SLOT_TIMEOUT), is(true));
            Iterator<TaskSlot<TaskSlotPayload>> allocatedSlots =
                    taskSlotTable.getAllocatedSlots(jobId);
            TaskSlot<TaskSlotPayload> taskSlot1 = allocatedSlots.next();

            assertThat(taskSlotTable.allocateSlot(-1, jobId, allocationId, SLOT_TIMEOUT), is(true));
            allocatedSlots = taskSlotTable.getAllocatedSlots(jobId);
            TaskSlot<TaskSlotPayload> taskSlot2 = allocatedSlots.next();

            assertThat(taskSlotTable.isAllocated(1, jobId, allocationId), is(true));
            assertEquals(taskSlot1, taskSlot2);
            assertThat(allocatedSlots.hasNext(), is(false));
        }
    }

    @Test
    public void testFreeSlot() throws Exception {
        try (final TaskSlotTable<TaskSlotPayload> taskSlotTable = createTaskSlotTableAndStart(2)) {
            final JobID jobId = new JobID();
            final AllocationID allocationId1 = new AllocationID();
            final AllocationID allocationId2 = new AllocationID();

            assertThat(taskSlotTable.allocateSlot(0, jobId, allocationId1, SLOT_TIMEOUT), is(true));
            assertThat(taskSlotTable.allocateSlot(1, jobId, allocationId2, SLOT_TIMEOUT), is(true));

            assertThat(taskSlotTable.freeSlot(allocationId2), is(1));

            Iterator<TaskSlot<TaskSlotPayload>> allocatedSlots =
                    taskSlotTable.getAllocatedSlots(jobId);
            assertThat(allocatedSlots.next().getIndex(), is(0));
            assertThat(allocatedSlots.hasNext(), is(false));
            assertThat(taskSlotTable.isAllocated(1, jobId, allocationId1), is(false));
            assertThat(taskSlotTable.isAllocated(1, jobId, allocationId2), is(false));
            assertThat(taskSlotTable.isSlotFree(1), is(true));
        }
    }

    @Test
    public void testSlotAllocationWithDynamicSlotId() throws Exception {
        try (final TaskSlotTable<TaskSlotPayload> taskSlotTable = createTaskSlotTableAndStart(2)) {
            final JobID jobId = new JobID();
            final AllocationID allocationId = new AllocationID();

            assertThat(taskSlotTable.allocateSlot(-1, jobId, allocationId, SLOT_TIMEOUT), is(true));

            Iterator<TaskSlot<TaskSlotPayload>> allocatedSlots =
                    taskSlotTable.getAllocatedSlots(jobId);
            assertThat(allocatedSlots.next().getIndex(), is(2));
            assertThat(allocatedSlots.hasNext(), is(false));
            assertThat(taskSlotTable.isAllocated(2, jobId, allocationId), is(true));
        }
    }

    @Test
    public void testSlotAllocationWithResourceProfile() throws Exception {
        try (final TaskSlotTable<TaskSlotPayload> taskSlotTable = createTaskSlotTableAndStart(2)) {
            final JobID jobId = new JobID();
            final AllocationID allocationId = new AllocationID();
            final ResourceProfile resourceProfile =
                    TaskSlotUtils.DEFAULT_RESOURCE_PROFILE.merge(
                            ResourceProfile.newBuilder().setCpuCores(0.1).build());

            assertThat(
                    taskSlotTable.allocateSlot(
                            -1, jobId, allocationId, resourceProfile, SLOT_TIMEOUT),
                    is(true));

            Iterator<TaskSlot<TaskSlotPayload>> allocatedSlots =
                    taskSlotTable.getAllocatedSlots(jobId);
            TaskSlot<TaskSlotPayload> allocatedSlot = allocatedSlots.next();
            assertThat(allocatedSlot.getIndex(), is(2));
            assertThat(allocatedSlot.getResourceProfile(), is(resourceProfile));
            assertThat(allocatedSlots.hasNext(), is(false));
        }
    }

    @Test
    public void testSlotAllocationWithResourceProfileFailure() throws Exception {
        try (final TaskSlotTable<TaskSlotPayload> taskSlotTable = createTaskSlotTableAndStart(2)) {
            final JobID jobId = new JobID();
            final AllocationID allocationId = new AllocationID();
            ResourceProfile resourceProfile = TaskSlotUtils.DEFAULT_RESOURCE_PROFILE;
            resourceProfile = resourceProfile.merge(resourceProfile).merge(resourceProfile);

            assertThat(
                    taskSlotTable.allocateSlot(
                            -1, jobId, allocationId, resourceProfile, SLOT_TIMEOUT),
                    is(false));

            Iterator<TaskSlot<TaskSlotPayload>> allocatedSlots =
                    taskSlotTable.getAllocatedSlots(jobId);
            assertThat(allocatedSlots.hasNext(), is(false));
        }
    }

    @Test
    public void testGenerateSlotReport() throws Exception {
        try (final TaskSlotTable<TaskSlotPayload> taskSlotTable = createTaskSlotTableAndStart(3)) {
            final JobID jobId = new JobID();
            final AllocationID allocationId1 = new AllocationID();
            final AllocationID allocationId2 = new AllocationID();
            final AllocationID allocationId3 = new AllocationID();

            assertThat(
                    taskSlotTable.allocateSlot(0, jobId, allocationId1, SLOT_TIMEOUT),
                    is(true)); // index 0
            assertThat(
                    taskSlotTable.allocateSlot(-1, jobId, allocationId2, SLOT_TIMEOUT),
                    is(true)); // index 3
            assertThat(
                    taskSlotTable.allocateSlot(-1, jobId, allocationId3, SLOT_TIMEOUT),
                    is(true)); // index 4

            assertThat(taskSlotTable.freeSlot(allocationId2), is(3));

            ResourceID resourceId = ResourceID.generate();
            SlotReport slotReport = taskSlotTable.createSlotReport(resourceId);
            List<SlotStatus> slotStatuses = new ArrayList<>();
            slotReport.iterator().forEachRemaining(slotStatuses::add);

            assertThat(slotStatuses.size(), is(4));
            assertThat(
                    slotStatuses,
                    containsInAnyOrder(
                            is(
                                    new SlotStatus(
                                            new SlotID(resourceId, 0),
                                            TaskSlotUtils.DEFAULT_RESOURCE_PROFILE,
                                            jobId,
                                            allocationId1)),
                            is(
                                    new SlotStatus(
                                            new SlotID(resourceId, 1),
                                            TaskSlotUtils.DEFAULT_RESOURCE_PROFILE,
                                            null,
                                            null)),
                            is(
                                    new SlotStatus(
                                            new SlotID(resourceId, 2),
                                            TaskSlotUtils.DEFAULT_RESOURCE_PROFILE,
                                            null,
                                            null)),
                            is(
                                    new SlotStatus(
                                            new SlotID(resourceId, 4),
                                            TaskSlotUtils.DEFAULT_RESOURCE_PROFILE,
                                            jobId,
                                            allocationId3))));
        }
    }

    @Test
    public void testAllocateSlot() throws Exception {
        final JobID jobId = new JobID();
        final AllocationID allocationId = new AllocationID();
        try (final TaskSlotTable<TaskSlotPayload> taskSlotTable =
                createTaskSlotTableWithAllocatedSlot(
                        jobId, allocationId, new TestingSlotActionsBuilder().build())) {
            Iterator<TaskSlot<TaskSlotPayload>> allocatedSlots =
                    taskSlotTable.getAllocatedSlots(jobId);
            TaskSlot<TaskSlotPayload> nextSlot = allocatedSlots.next();
            assertThat(nextSlot.getIndex(), is(0));
            assertThat(nextSlot.getAllocationId(), is(allocationId));
            assertThat(nextSlot.getJobId(), is(jobId));
            assertThat(allocatedSlots.hasNext(), is(false));
        }
    }

    @Test
    public void testAddTask() throws Exception {
        final JobID jobId = new JobID();
        final ExecutionAttemptID executionAttemptId = new ExecutionAttemptID();
        final AllocationID allocationId = new AllocationID();
        TaskSlotPayload task =
                new TestingTaskSlotPayload(jobId, executionAttemptId, allocationId).terminate();
        try (final TaskSlotTable<TaskSlotPayload> taskSlotTable =
                createTaskSlotTableWithStartedTask(task)) {
            Iterator<TaskSlotPayload> tasks = taskSlotTable.getTasks(jobId);
            TaskSlotPayload nextTask = tasks.next();
            assertThat(nextTask.getExecutionId(), is(executionAttemptId));
            assertThat(nextTask.getAllocationId(), is(allocationId));
            assertThat(tasks.hasNext(), is(false));
        }
    }

    @Test(timeout = 10000)
    public void testRemoveTaskCallsFreeSlotAction() throws Exception {
        final JobID jobId = new JobID();
        final ExecutionAttemptID executionAttemptId = new ExecutionAttemptID();
        final AllocationID allocationId = new AllocationID();
        CompletableFuture<AllocationID> freeSlotFuture = new CompletableFuture<>();
        SlotActions slotActions =
                new TestingSlotActions(freeSlotFuture::complete, (aid, uid) -> {});
        TaskSlotPayload task =
                new TestingTaskSlotPayload(jobId, executionAttemptId, allocationId).terminate();
        try (final TaskSlotTable<TaskSlotPayload> taskSlotTable =
                createTaskSlotTableWithStartedTask(task, slotActions)) {
            // we have to initiate closing of the slot externally
            // to enable that the last remaining finished task does the final slot freeing
            taskSlotTable.freeSlot(allocationId);
            taskSlotTable.removeTask(executionAttemptId);
            assertThat(freeSlotFuture.get(), is(allocationId));
        }
    }

    @Test(timeout = 10000)
    public void testFreeSlotInterruptsSubmittedTask() throws Exception {
        TestingTaskSlotPayload task = new TestingTaskSlotPayload();
        try (final TaskSlotTable<TaskSlotPayload> taskSlotTable =
                createTaskSlotTableWithStartedTask(task)) {
            assertThat(taskSlotTable.freeSlot(task.getAllocationId()), is(-1));
            task.waitForFailure();
            task.terminate();
        }
    }

    @Test(timeout = 10000)
    public void testTableIsClosedOnlyWhenAllTasksTerminated() throws Exception {
        TestingTaskSlotPayload task = new TestingTaskSlotPayload();
        final TaskSlotTable<TaskSlotPayload> taskSlotTable =
                createTaskSlotTableWithStartedTask(task);
        assertThat(taskSlotTable.freeSlot(task.getAllocationId()), is(-1));
        CompletableFuture<Void> closingFuture = taskSlotTable.closeAsync();
        assertThat(closingFuture.isDone(), is(false));
        task.terminate();
        closingFuture.get();
    }

    @Test
    public void testAllocatedSlotTimeout() throws Exception {
        final CompletableFuture<AllocationID> timeoutFuture = new CompletableFuture<>();
        final TestingSlotActions testingSlotActions =
                new TestingSlotActionsBuilder()
                        .setTimeoutSlotConsumer(
                                (allocationID, uuid) -> timeoutFuture.complete(allocationID))
                        .build();

        try (final TaskSlotTableImpl<TaskSlotPayload> taskSlotTable =
                createTaskSlotTableAndStart(1, testingSlotActions)) {
            final AllocationID allocationId = new AllocationID();
            assertThat(
                    taskSlotTable.allocateSlot(0, new JobID(), allocationId, Time.milliseconds(1L)),
                    is(true));
            assertThat(timeoutFuture.join(), is(allocationId));
        }
    }

    @Test
    public void testMarkSlotActiveDeactivatesSlotTimeout() throws Exception {
        runDeactivateSlotTimeoutTest(
                (taskSlotTable, jobId, allocationId) -> taskSlotTable.markSlotActive(allocationId));
    }

    @Test
    public void testTryMarkSlotActiveDeactivatesSlotTimeout() throws Exception {
        runDeactivateSlotTimeoutTest(TaskSlotTable::tryMarkSlotActive);
    }

    private void runDeactivateSlotTimeoutTest(
            TriFunctionWithException<
                            TaskSlotTable<TaskSlotPayload>,
                            JobID,
                            AllocationID,
                            Boolean,
                            SlotNotFoundException>
                    taskSlotTableAction)
            throws Exception {
        final CompletableFuture<AllocationID> timeoutFuture = new CompletableFuture<>();
        final TestingSlotActions testingSlotActions =
                new TestingSlotActionsBuilder()
                        .setTimeoutSlotConsumer(
                                (allocationID, uuid) -> timeoutFuture.complete(allocationID))
                        .build();

        try (final TaskSlotTableImpl<TaskSlotPayload> taskSlotTable =
                createTaskSlotTableAndStart(1, testingSlotActions)) {
            final AllocationID allocationId = new AllocationID();
            final long timeout = 50L;
            final JobID jobId = new JobID();
            assertThat(
                    taskSlotTable.allocateSlot(0, jobId, allocationId, Time.milliseconds(timeout)),
                    is(true));
            assertThat(taskSlotTableAction.apply(taskSlotTable, jobId, allocationId), is(true));

            try {
                timeoutFuture.get(timeout, TimeUnit.MILLISECONDS);
                fail("The slot timeout should have been deactivated.");
            } catch (TimeoutException expected) {
            }
        }
    }

    private static TaskSlotTable<TaskSlotPayload> createTaskSlotTableWithStartedTask(
            final TaskSlotPayload task) throws SlotNotFoundException, SlotNotActiveException {
        return createTaskSlotTableWithStartedTask(task, new TestingSlotActionsBuilder().build());
    }

    private static TaskSlotTable<TaskSlotPayload> createTaskSlotTableWithStartedTask(
            final TaskSlotPayload task, final SlotActions slotActions)
            throws SlotNotFoundException, SlotNotActiveException {
        final TaskSlotTable<TaskSlotPayload> taskSlotTable =
                createTaskSlotTableWithAllocatedSlot(
                        task.getJobID(), task.getAllocationId(), slotActions);
        taskSlotTable.markSlotActive(task.getAllocationId());
        taskSlotTable.addTask(task);
        return taskSlotTable;
    }

    private static TaskSlotTable<TaskSlotPayload> createTaskSlotTableWithAllocatedSlot(
            final JobID jobId, final AllocationID allocationId, final SlotActions slotActions) {
        final TaskSlotTable<TaskSlotPayload> taskSlotTable =
                createTaskSlotTableAndStart(1, slotActions);
        assertThat(taskSlotTable.allocateSlot(0, jobId, allocationId, SLOT_TIMEOUT), is(true));
        return taskSlotTable;
    }

    private static TaskSlotTableImpl<TaskSlotPayload> createTaskSlotTableAndStart(
            final int numberOfSlots) {
        return createTaskSlotTableAndStart(numberOfSlots, new TestingSlotActionsBuilder().build());
    }

    private static TaskSlotTableImpl<TaskSlotPayload> createTaskSlotTableAndStart(
            final int numberOfSlots, final SlotActions slotActions) {
        final TaskSlotTableImpl<TaskSlotPayload> taskSlotTable =
                TaskSlotUtils.createTaskSlotTable(numberOfSlots);
        taskSlotTable.start(slotActions, ComponentMainThreadExecutorServiceAdapter.forMainThread());
        return taskSlotTable;
    }
}

redpajama_set_name: RedPajamaGithub


text:
/*
 * GeoIP C library binding for nodejs
 *
 * Licensed under the GNU LGPL 2.1 license
 */

#include "region.h"
#include "global.h"

Persistent<FunctionTemplate> geoip::Region::constructor_template;

void geoip::Region::Init(Handle<Object> target)
{
  HandleScope scope;

  Local<FunctionTemplate> t = FunctionTemplate::New(New);
  constructor_template = Persistent<FunctionTemplate>::New(t);
  constructor_template->InstanceTemplate()->SetInternalFieldCount(1);
  constructor_template->SetClassName(String::NewSymbol("geoip"));

  NODE_SET_PROTOTYPE_METHOD(constructor_template, "lookup", lookup);
  NODE_SET_PROTOTYPE_METHOD(constructor_template, "lookupSync", lookupSync);
  NODE_SET_PROTOTYPE_METHOD(constructor_template, "close", close);
  target->Set(String::NewSymbol("Region"), constructor_template->GetFunction());
}

/*
   geoip::Region::Region() :
   db_edition(0)
   {
   }

   geoip::Region::~Region()
   {
   } */

Handle<Value> geoip::Region::New(const Arguments& args)
{
  HandleScope scope;
  Region * r = new Region();

  String::Utf8Value file_str(args[0]->ToString());
  const char * file_cstr = ToCString(file_str);
  // Local<String> host_str = args[0]->ToString();
  // char host_cstr[host_str->Length()];
  // host_str->WriteAscii(host_cstr);
  bool cache_on = args[1]->ToBoolean()->Value();

  r->db = GeoIP_open(file_cstr, cache_on?GEOIP_MEMORY_CACHE:GEOIP_STANDARD);

  if (r->db != NULL) {
    // Successfully opened the file, return 1 (true)
    r->db_edition = GeoIP_database_edition(r->db);
    if (r->db_edition == GEOIP_REGION_EDITION_REV0 ||
        r->db_edition == GEOIP_REGION_EDITION_REV1) {
      r->Wrap(args.This());
      return scope.Close(args.This());
    } else {
      GeoIP_delete(r->db);	// free()'s the gi reference & closes its fd
      r->db = NULL;
      printf("edition is %d", r->db_edition);
      return scope.Close(ThrowException(String::New("Error: Not valid region database")));
    }
  } else {
    return scope.Close(ThrowException(String::New("Error: Cao not open database")));
  }
}

Handle<Value> geoip::Region::lookupSync(const Arguments &args) {
  HandleScope scope;

  Local<String> host_str = args[0]->ToString();
  Local<Object> data = Object::New();
  char host_cstr[host_str->Length()];
  host_str->WriteAscii(host_cstr);
  Region* r = ObjectWrap::Unwrap<Region>(args.This());

  uint32_t ipnum = _GeoIP_lookupaddress(host_cstr);

  if (ipnum <= 0) {
    return scope.Close(Null());
  }

  GeoIPRegion *region = GeoIP_region_by_ipnum(r->db, ipnum);

  if (region != NULL) {
    data->Set(String::NewSymbol("country_code"), String::New(region->country_code));
    data->Set(String::NewSymbol("region"), String::New(region->region));
    return scope.Close(data);
  } else {
    return scope.Close(ThrowException(String::New("Error: Can not find match data")));
  }
}

Handle<Value> geoip::Region::lookup(const Arguments& args)
{
  HandleScope scope;

  REQ_FUN_ARG(1, cb);

  Region * r = ObjectWrap::Unwrap<Region>(args.This());
  Local<String> host_str = args[0]->ToString();

  region_baton_t *baton = new region_baton_t();

  baton->r = r;
  host_str->WriteAscii(baton->host_cstr);
  baton->increment_by = 2;
  baton->sleep_for = 1;
  baton->cb = Persistent<Function>::New(cb);

  r->Ref();

  eio_custom(EIO_Region, EIO_PRI_DEFAULT, EIO_AfterRegion, baton);
  ev_ref(EV_DEFAULT_UC);

  return scope.Close(Undefined());
}

int geoip::Region::EIO_Region(eio_req *req)
{
  region_baton_t *baton = static_cast<region_baton_t *>(req->data);

  sleep(baton->sleep_for);

  uint32_t ipnum = _GeoIP_lookupaddress(baton->host_cstr);
  if (ipnum <= 0) {
    baton->region = NULL;
  } else {
    baton->region = GeoIP_region_by_ipnum(baton->r->db, ipnum);
  }

  return 0;
}

int geoip::Region::EIO_AfterRegion(eio_req *req)
{
  HandleScope scope;

  region_baton_t * baton = static_cast<region_baton_t *>(req->data);
  ev_unref(EV_DEFAULT_UC);
  baton->r->Unref();

  Handle<Value> argv[2];
  if (baton->region == NULL) {
    argv[0] = Exception::Error(String::New("Data not found"));
    argv[1] = Null();
  } else {
    Local<Object> data = Object::New();
    data->Set(String::NewSymbol("country_code"), String::New(baton->region->country_code));
    data->Set(String::NewSymbol("region"), String::New(baton->region->region));

    argv[0] = Null();
    argv[1] = data;
  }

  TryCatch try_catch;

  baton->cb->Call(Context::GetCurrent()->Global(), 2, argv);

  if (try_catch.HasCaught()) {
    FatalException(try_catch);
  }

  baton->cb.Dispose();

  delete baton;
  return 0;
}

Handle<Value> geoip::Region::close(const Arguments &args) {
  Region * r = ObjectWrap::Unwrap<Region>(args.This());
  GeoIP_delete(r->db);	// free()'s the gi reference & closes its fd
  r->db = NULL;
  HandleScope scope;	// Stick this down here since it seems to segfault when on top?
}