org.infinispan.stream.impl.termop.primitive.ForEachIntOperation Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of infinispan-embedded
Infinispan Embedded All-in-One module
There is a newer version: 9.1.7.Final
package org.infinispan.stream.impl.termop.primitive;

import org.infinispan.Cache;
import org.infinispan.container.entries.CacheEntry;
import org.infinispan.container.entries.ImmortalCacheEntry;
import org.infinispan.factories.ComponentRegistry;
import org.infinispan.stream.impl.KeyTrackingTerminalOperation;
import org.infinispan.stream.impl.intops.IntermediateOperation;
import org.infinispan.stream.impl.termop.BaseTerminalOperation;
import org.infinispan.stream.impl.termop.object.NoMapIteratorOperation;
import org.infinispan.stream.CacheAware;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.IntConsumer;
import java.util.function.Supplier;
import java.util.stream.BaseStream;
import java.util.stream.IntStream;
import java.util.stream.Stream;

/**
 * Terminal rehash aware operation that handles for each where no flat map operations are defined on a
 * {@link IntStream}. Note this means it is an implied map intermediate operation.
 * @param  key type of the supplied stream
 */
public class ForEachIntOperation extends BaseTerminalOperation implements KeyTrackingTerminalOperation {
   private final int batchSize;
   private final IntConsumer consumer;

   public ForEachIntOperation(Iterable intermediateOperations,
           Supplier> supplier, int batchSize, IntConsumer consumer) {
      super(intermediateOperations, supplier);
      this.batchSize = batchSize;
      this.consumer = consumer;
   }

   @Override
   public boolean lostSegment(boolean stopIfLost) {
      // TODO: stop this early
      return true;
   }

   @Override
   public List performOperation(IntermediateCollector> response) {
      /**
       * This is for rehash only! {@link NoMapIteratorOperation} should always be used for non rehash
       */
      throw new UnsupportedOperationException();
   }

   @Override
   public Collection> performOperationRehashAware(
           IntermediateCollector>> response) {
      // We only support sequential streams for iterator rehash aware
      BaseStream stream = supplier.get().sequential();

      List> collectedValues = new ArrayList(batchSize);

      int[] list = new int[batchSize];
      AtomicInteger offset = new AtomicInteger();
      Object[] currentKey = new Object[1];
      stream = ((Stream>) stream).peek(e -> {
         if (offset.get() > 0) {
            collectedValues.add(new ImmortalCacheEntry(currentKey[0], currentKey[0]));
            if (collectedValues.size() >= batchSize) {
               for (int i = 0; i < offset.get(); ++i) {
                  consumer.accept(list[i]);
               }
               response.sendDataResonse(collectedValues);
               collectedValues.clear();
               offset.set(0);
            }
         }
         currentKey[0] = e.getKey();
      });
      for (IntermediateOperation intermediateOperation : intermediateOperations) {
         stream = intermediateOperation.perform(stream);
      }

      IntStream convertedStream = ((IntStream)stream);
      // We rely on the fact that iterator processes 1 entry at a time when sequential
      convertedStream.forEach(i -> list[offset.getAndIncrement()] = i);
      if (offset.get() > 0) {
         for (int i = 0; i < offset.get(); ++i) {
            consumer.accept(list[i]);
         }
         collectedValues.add(new ImmortalCacheEntry(currentKey[0], currentKey[0]));
      }
      return collectedValues;
   }

   public int getBatchSize() {
      return batchSize;
   }

   public IntConsumer getConsumer() {
      return consumer;
   }

   @Override
   public void handleInjection(ComponentRegistry registry) {
      super.handleInjection(registry);
      if (consumer instanceof CacheAware) {
         ((CacheAware) consumer).injectCache(registry.getComponent(Cache.class));
      }
   }
}