org.infinispan.stream.impl.termop.AbstractForEachOperation 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;

import org.infinispan.commons.util.ByRef;
import org.infinispan.container.entries.CacheEntry;
import org.infinispan.container.entries.ImmortalCacheEntry;
import org.infinispan.stream.impl.KeyTrackingTerminalOperation;
import org.infinispan.stream.impl.intops.IntermediateOperation;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.function.Supplier;
import java.util.stream.BaseStream;
import java.util.stream.Stream;

/**
 * This is a base operation class for the use of the for each terminal operator.  This class can be used for any
 * forEach configuration, however since it relies on generics it may not be as performant as a primitive based
 * for each operation.
 * This class assumes the stream is composed of {@link java.util.Map.Entry} instances where the key is typed the same
 * as defined K type.
 * @param  key type of underlying stream
 * @param  value type of transformed stream
 * @param  type of the transformed stream
 */
public abstract class AbstractForEachOperation> extends BaseTerminalOperation
        implements KeyTrackingTerminalOperation {
   private final int batchSize;

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

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

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

   protected abstract void handleList(List list);

   protected abstract void handleStreamForEach(S stream, List list);

   @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);

      List currentList = new ArrayList<>();
      ByRef currentKey = new ByRef<>(null);
      stream = ((Stream>) stream).peek(e -> {
         if (!currentList.isEmpty()) {
            collectedValues.add(new ImmortalCacheEntry(currentKey.get(), null));
            if (collectedValues.size() >= batchSize) {
               handleList(currentList);
               response.sendDataResonse(collectedValues);
               collectedValues.clear();
               currentList.clear();
            }
         }
         currentKey.set(e.getKey());
      });
      for (IntermediateOperation intermediateOperation : intermediateOperations) {
         stream = intermediateOperation.perform(stream);
      }

      S convertedStream = ((S) stream);
      // We rely on the fact that iterator processes 1 entry at a time
      handleStreamForEach(convertedStream, currentList);
      if (!currentList.isEmpty()) {
         handleList(currentList);
         collectedValues.add(new ImmortalCacheEntry(currentKey.get(), null));
      }
      return collectedValues;
   }

   public int getBatchSize() {
      return batchSize;
   }
}