All Downloads are FREE. Search and download functionalities are using the official Maven repository.

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

There is a newer version: 9.1.7.Final
Show newest version
package org.infinispan.stream.impl.tx;

import org.infinispan.container.entries.CacheEntry;
import org.infinispan.context.impl.LocalTxInvocationContext;
import org.infinispan.distribution.ch.ConsistentHash;
import org.infinispan.remoting.transport.Address;
import org.infinispan.stream.impl.AbstractCacheStream;
import org.infinispan.stream.impl.DistributedCacheStream;
import org.infinispan.stream.impl.DistributedDoubleCacheStream;
import org.infinispan.stream.impl.DistributedIntCacheStream;
import org.infinispan.stream.impl.DistributedLongCacheStream;

import java.util.Set;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * Int variant of tx cache stream
 * @see TxDistributedCacheStream
 */
public class TxDistributedIntCacheStream extends DistributedIntCacheStream {
   private final Address localAddress;
   private final LocalTxInvocationContext ctx;
   private final ConsistentHash hash;

   TxDistributedIntCacheStream(AbstractCacheStream stream, Address localAddress, ConsistentHash hash,
           LocalTxInvocationContext ctx) {
      super(stream);
      this.localAddress = localAddress;
      this.ctx = ctx;
      this.hash = hash;
   }

   @Override
   protected Supplier> supplierForSegments(ConsistentHash ch, Set targetSegments,
           Set excludedKeys, boolean primaryOnly) {
      return () -> {
         Supplier> supplier = super.supplierForSegments(ch, targetSegments, excludedKeys, primaryOnly);
         Set set = ctx.getLookedUpEntries().values().stream().filter(
                 e -> !localAddress.equals(ch.locatePrimaryOwner(e.getKey()))).collect(Collectors.toSet());
         Stream suppliedStream = supplier.get();
         if (!set.isEmpty()) {
            return Stream.concat(set.stream(), suppliedStream);
         }
         return suppliedStream;
      };
   }

   @Override
   protected  DistributedCacheStream cacheStream() {
      return new TxDistributedCacheStream<>(this, localAddress, hash, ctx);
   }

   @Override
   protected DistributedLongCacheStream longCacheStream() {
      return new TxDistributedLongCacheStream(this, localAddress, hash, ctx);
   }

   @Override
   protected DistributedDoubleCacheStream doubleCacheStream() {
      return new TxDistributedDoubleCacheStream(this, localAddress, hash, ctx);
   }
}