AffinityTest.java example

Explorer
infinispan-master
package org.infinispan.query.affinity;

import static java.util.stream.IntStream.rangeClosed;
import static org.infinispan.hibernate.search.spi.InfinispanIntegration.DEFAULT_INDEXESDATA_CACHENAME;
import static org.infinispan.hibernate.search.spi.InfinispanIntegration.DEFAULT_INDEXESMETADATA_CACHENAME;
import static org.infinispan.hibernate.search.spi.InfinispanIntegration.DEFAULT_LOCKING_CACHENAME;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertTrue;

import java.util.List;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

import org.apache.lucene.search.MatchAllDocsQuery;
import org.infinispan.AdvancedCache;
import org.infinispan.Cache;
import org.infinispan.container.DataContainer;
import org.infinispan.distribution.ch.ConsistentHash;
import org.infinispan.lucene.IndexScopedKey;
import org.infinispan.manager.EmbeddedCacheManager;
import org.infinispan.query.CacheQuery;
import org.infinispan.query.Search;
import org.infinispan.remoting.transport.Address;
import org.infinispan.test.fwk.CleanupAfterMethod;
import org.testng.annotations.Test;

/**
 * Test index affinity for the AffinityIndexManager
 * @author gustavonalle
 * @since 8.2
 */
@Test(groups = "functional", testName = "query.AffinityTest")
@CleanupAfterMethod
public class AffinityTest extends BaseAffinityTest {

   @Override
   protected void createCacheManagers() throws Throwable {
      createClusteredCaches(3, getDefaultCacheConfigBuilder());
   }

   public void testConcurrentWrites() throws InterruptedException {
      int numThreads = 2;
      AtomicInteger counter = new AtomicInteger(0);
      ExecutorService executorService = Executors.newFixedThreadPool(numThreads);
      List<Cache<String, Entity>> cacheList = caches();

      log.info("Starting threads");
      List<Future<?>> futures = rangeClosed(1, numThreads).boxed().map(tid -> executorService.submit(() ->
            rangeClosed(1, getNumEntries()).forEach(entry -> {
               int id = counter.incrementAndGet();
               pickCache().put(String.valueOf(id), new Entity(id));
            }))).collect(Collectors.toList());

      log.info("Waiting for threads");
      futures.forEach(f -> {
         try {
            f.get();
         } catch (InterruptedException | ExecutionException e) {
            log.error("Error while waiting for thread.", e);
         }
      });

      log.info("Checking cache size");
      assertEquals(pickCache().size(), numThreads * getNumEntries());
      cacheList.forEach(c -> {
         CacheQuery<Entity> q = Search.getSearchManager(c).getQuery(new MatchAllDocsQuery(), Entity.class);
         eventuallyEquals(numThreads * getNumEntries(), () -> q.list().size());
      });

   }

   public void shouldHaveIndexAffinity() throws Exception {
      populate(1, getNumEntries() / 2);
      checkAffinity();

      addNode();
      populate(getNumEntries() / 2 + 1, getNumEntries());
      checkAffinity();

      CacheQuery<Entity> q = Search.getSearchManager(pickCache()).getQuery(new MatchAllDocsQuery(), Entity.class);
      assertEquals(getNumEntries(), pickCache().size());
      assertEquals(getNumEntries(), q.list().size());

      addNode();
      checkAffinity();
      assertEquals(getNumEntries(), pickCache().size());

      populate(getNumEntries() + 1, getNumEntries() * 2);
      checkAffinity();
      assertEquals(getNumEntries() * 2, q.list().size());
   }


   void checkAffinity() {
      for (EmbeddedCacheManager clusterMember : cacheManagers) {
         checkAffinity(clusterMember.getCache(DEFAULT_INDEXESDATA_CACHENAME));
         checkAffinity(clusterMember.getCache(DEFAULT_INDEXESMETADATA_CACHENAME));
         checkAffinity(clusterMember.getCache(DEFAULT_LOCKING_CACHENAME));
      }
   }

   private void checkAffinity(Cache<IndexScopedKey, ?> indexCache) {
      AdvancedCache<IndexScopedKey, ?> advancedCache = indexCache.getAdvancedCache();
      DataContainer<IndexScopedKey, ?> dataContainer = advancedCache.getDataContainer();
      ConsistentHash consistentHash = advancedCache.getDistributionManager().getConsistentHash();
      Address address = advancedCache.getRpcManager().getAddress();
      Set<Integer> ownedSegments = consistentHash.getSegmentsForOwner(address);
      dataContainer.forEach(entry -> {
         int segmentAffinity = entry.getKey().getAffinitySegmentId();
         assertTrue(ownedSegments.contains(segmentAffinity));
      });
   }

}