ComputeUtils { public static <R> Collection<R> affinityCallWithRetries(Ignite ignite, Collection<String> cacheNames, IgniteFunction<Integer, R> fun, int retries, int interval, DeployingContext deployingContext) { assert !cacheNames.isEmpty(); assert interval >= 0; String primaryCache = cacheNames.iterator().next(); Affinity<?> affinity = ignite.affinity(primaryCache); int partitions = affinity.partitions(); BitSet completionFlags = new BitSet(partitions); Collection<R> results = new ArrayList<>(); for (int t = 0; t <= retries; t++) { ClusterGroup clusterGrp = ignite.cluster().forDataNodes(primaryCache); Map<Integer, IgniteFuture<R>> futures = new HashMap<>(); for (int part = 0; part < partitions; part++) if (!completionFlags.get(part)) { final int currPart = part; futures.put( currPart, ignite.compute(clusterGrp).affinityCallAsync( cacheNames, currPart, new DeployableCallable<>(deployingContext, part, fun) ) ); } for (int part : futures.keySet()) try { R res = futures.get(part).get(); results.add(res); completionFlags.set(part); } catch (IgniteException ignore) { } if (completionFlags.cardinality() == partitions) return results; LockSupport.parkNanos(interval * 1_000_000); } throw new IllegalStateException(); }  static Collection<R> affinityCallWithRetries(Ignite ignite, Collection<String> cacheNames,
        IgniteFunction<Integer, R> fun, int retries, int interval, DeployingContext deployingContext); static Collection<R> affinityCallWithRetries(Ignite ignite, Collection<String> cacheNames,
        IgniteFunction<Integer, R> fun, int retries, DeployingContext deployingContext); static LearningEnvironment getLearningEnvironment(Ignite ignite,
        UUID datasetId,
        int part,
        LearningEnvironmentBuilder envBuilder); static D getData(
        Ignite ignite,
        String upstreamCacheName, IgniteBiPredicate<K, V> filter,
        UpstreamTransformerBuilder transformerBuilder,
        String datasetCacheName, UUID datasetId,
        PartitionDataBuilder<K, V, C, D> partDataBuilder,
        LearningEnvironment env,
        boolean isKeepBinary); static void removeData(Ignite ignite, UUID datasetId); static void removeLearningEnv(Ignite ignite, UUID datasetId); static void initContext(
        Ignite ignite,
        String upstreamCacheName,
        UpstreamTransformerBuilder transformerBuilder,
        IgniteBiPredicate<K, V> filter,
        String datasetCacheName,
        PartitionContextBuilder<K, V, C> ctxBuilder,
        LearningEnvironmentBuilder envBuilder,
        int retries,
        int interval,
        boolean isKeepBinary,
        DeployingContext deployingContext); static C getContext(Ignite ignite, String datasetCacheName, int part); static void saveContext(Ignite ignite, String datasetCacheName, int part, C ctx);  }
@Test public void testAffinityCallWithRetriesNegative() { ClusterNode node1 = grid(1).cluster().localNode(); ClusterNode node2 = grid(2).cluster().localNode(); String firstCacheName = "CACHE_1_" + UUID.randomUUID(); String secondCacheName = "CACHE_2_" + UUID.randomUUID(); CacheConfiguration<Integer, Integer> cacheConfiguration1 = new CacheConfiguration<>(); cacheConfiguration1.setName(firstCacheName); cacheConfiguration1.setAffinity(new TestAffinityFunction(node1)); IgniteCache<Integer, Integer> cache1 = ignite.createCache(cacheConfiguration1); CacheConfiguration<Integer, Integer> cacheConfiguration2 = new CacheConfiguration<>(); cacheConfiguration2.setName(secondCacheName); cacheConfiguration2.setAffinity(new TestAffinityFunction(node2)); IgniteCache<Integer, Integer> cache2 = ignite.createCache(cacheConfiguration2); try { try { ComputeUtils.affinityCallWithRetries( ignite, Arrays.asList(firstCacheName, secondCacheName), part -> part, 0, DeployingContext.unitialized() ); } catch (IllegalStateException expectedException) { return; } fail("Missing IllegalStateException"); } finally { cache1.destroy(); cache2.destroy(); } }
@Test public void testAffinityCallWithRetriesPositive() { ClusterNode node = grid(1).cluster().localNode(); String firstCacheName = "CACHE_1_" + UUID.randomUUID(); String secondCacheName = "CACHE_2_" + UUID.randomUUID(); CacheConfiguration<Integer, Integer> cacheConfiguration1 = new CacheConfiguration<>(); cacheConfiguration1.setName(firstCacheName); cacheConfiguration1.setAffinity(new TestAffinityFunction(node)); IgniteCache<Integer, Integer> cache1 = ignite.createCache(cacheConfiguration1); CacheConfiguration<Integer, Integer> cacheConfiguration2 = new CacheConfiguration<>(); cacheConfiguration2.setName(secondCacheName); cacheConfiguration2.setAffinity(new TestAffinityFunction(node)); IgniteCache<Integer, Integer> cache2 = ignite.createCache(cacheConfiguration2); try (IgniteAtomicLong cnt = ignite.atomicLong("COUNTER_" + UUID.randomUUID(), 0, true)) { ComputeUtils.affinityCallWithRetries(ignite, Arrays.asList(firstCacheName, secondCacheName), part -> { Ignite locIgnite = Ignition.localIgnite(); assertEquals(node, locIgnite.cluster().localNode()); cnt.incrementAndGet(); return part; }, 0, DeployingContext.unitialized()); assertEquals(1, cnt.get()); } finally { cache1.destroy(); cache2.destroy(); } }