org.infinispan.stream.impl.tx.TxClusterStreamManager Maven / Gradle / Ivy

Go to download
package org.infinispan.stream.impl.tx;

import org.infinispan.context.impl.LocalTxInvocationContext;
import org.infinispan.distribution.ch.ConsistentHash;
import org.infinispan.remoting.transport.Address;
import org.infinispan.stream.impl.ClusterStreamManager;
import org.infinispan.stream.impl.KeyTrackingTerminalOperation;
import org.infinispan.stream.impl.TerminalOperation;
import org.infinispan.util.AbstractDelegatingMap;

import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import java.util.function.Predicate;

/**
 * This is a delegating cluster stream manager that sends all calls to the underlying cluster stream manager.  However
 * in the case of performing an operation it adds all entries that are in the provided tx context to the map of
 * keys to exclude so those values are not processed in the remote nodes.
 * @param  the key type
 */
public class TxClusterStreamManager implements ClusterStreamManager {
   private final ClusterStreamManager manager;
   private final LocalTxInvocationContext ctx;
   private final ConsistentHash hash;

   public TxClusterStreamManager(ClusterStreamManager manager, LocalTxInvocationContext ctx, ConsistentHash hash) {
      this.manager = manager;
      this.ctx = ctx;
      this.hash = hash;
   }

   @Override
   public  Object remoteStreamOperation(boolean parallelDistribution, boolean parallelStream, ConsistentHash ch,
           Set segments, Set keysToInclude, Map> keysToExclude, boolean includeLoader,
           TerminalOperation operation, ResultsCallback callback, Predicate earlyTerminatePredicate) {
      TxExcludedKeys txExcludedKeys = new TxExcludedKeys<>(keysToExclude, ctx, hash);
      return manager.remoteStreamOperation(parallelDistribution, parallelStream, ch, segments, keysToInclude,
              txExcludedKeys, includeLoader, operation, callback, earlyTerminatePredicate);
   }

   @Override
   public  Object remoteStreamOperationRehashAware(boolean parallelDistribution, boolean parallelStream,
           ConsistentHash ch, Set segments, Set keysToInclude, Map> keysToExclude,
           boolean includeLoader, TerminalOperation operation, ResultsCallback callback,
           Predicate earlyTerminatePredicate) {
      TxExcludedKeys txExcludedKeys = new TxExcludedKeys<>(keysToExclude, ctx, hash);
      return manager.remoteStreamOperationRehashAware(parallelDistribution, parallelStream, ch, segments, keysToInclude,
              txExcludedKeys, includeLoader, operation, callback, earlyTerminatePredicate);
   }

   @Override
   public  Object remoteStreamOperation(boolean parallelDistribution, boolean parallelStream, ConsistentHash ch,
           Set segments, Set keysToInclude, Map> keysToExclude, boolean includeLoader,
           KeyTrackingTerminalOperation operation, ResultsCallback> callback) {
      TxExcludedKeys txExcludedKeys = new TxExcludedKeys<>(keysToExclude, ctx, hash);
      return manager.remoteStreamOperation(parallelDistribution, parallelStream, ch, segments, keysToInclude,
              txExcludedKeys, includeLoader, operation, callback);
   }

   @Override
   public  Object remoteStreamOperationRehashAware(boolean parallelDistribution, boolean parallelStream,
           ConsistentHash ch, Set segments, Set keysToInclude, Map> keysToExclude,
           boolean includeLoader, KeyTrackingTerminalOperation operation, ResultsCallback> callback) {
      TxExcludedKeys txExcludedKeys = new TxExcludedKeys<>(keysToExclude, ctx, hash);
      return manager.remoteStreamOperationRehashAware(parallelDistribution, parallelStream, ch, segments, keysToInclude,
              txExcludedKeys, includeLoader, operation, callback);
   }

   @Override
   public boolean isComplete(Object id) {
      return manager.isComplete(id);
   }

   @Override
   public boolean awaitCompletion(Object id, long time, TimeUnit unit) throws InterruptedException {
      return manager.awaitCompletion(id, time, unit);
   }

   @Override
   public void forgetOperation(Object id) {
      manager.forgetOperation(id);
   }

   @Override
   public  boolean receiveResponse(Object id, Address origin, boolean complete, Set segments, R1 response) {
      return manager.receiveResponse(id, origin, complete, segments, response);
   }

   private static class TxExcludedKeys extends AbstractDelegatingMap> {
      private final Map> map;
      private final Map> ctxMap;

      private TxExcludedKeys(Map> map, LocalTxInvocationContext ctx, ConsistentHash hash) {
         this.map = map;
         this.ctxMap = contextToMap(ctx, hash);
      }

      Map> contextToMap(LocalTxInvocationContext ctx, ConsistentHash hash) {
         Map> contextMap = new HashMap<>();
         ctx.getLookedUpEntries().forEach((k, v) -> {
            Integer segment = hash.getSegment(k);
            Set innerSet = contextMap.get(segment);
            if (innerSet == null) {
               innerSet = new HashSet();
               contextMap.put(segment, innerSet);
            }
            innerSet.add((K) k);
         });
         return contextMap;
      }

      @Override
      protected Map> delegate() {
         return map;
      }

      @Override
      public Set get(Object key) {
         if (!(key instanceof Integer)) {
            return null;
         }
         Set ctxSet = ctxMap.get(key);
         Set excludedSet = super.get(key);
         if (ctxSet != null) {
            if (excludedSet != null) {
               ctxSet.addAll(excludedSet);
            }
            return ctxSet;
         }
         return excludedSet;
      }

      @Override
      public boolean isEmpty() {
         return ctxMap.isEmpty() && super.isEmpty();
      }
   }
}