org.infinispan.hotrod.impl.operations.GetAllParallelOperation Maven / Gradle / Ivy

Go to download
package org.infinispan.hotrod.impl.operations;

import java.net.SocketAddress;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import org.infinispan.api.common.CacheOptions;
import org.infinispan.hotrod.impl.DataFormat;

/**
 *
 */
public class GetAllParallelOperation extends ParallelHotRodOperation, GetAllOperation> {

   private final Set keys;

   protected GetAllParallelOperation(OperationContext operationContext, Set keys, CacheOptions options, DataFormat dataFormat) {
      super(operationContext, options, dataFormat);
      this.keys = keys;
   }

   @Override
   protected List> mapOperations() {
      Map> splittedKeys = new HashMap<>();

      for (byte[] key : keys) {
         SocketAddress socketAddress = operationContext.getChannelFactory().getHashAwareServer(key, operationContext.getCacheNameBytes());
         Set keys = splittedKeys.computeIfAbsent(socketAddress, k -> new HashSet<>());
         keys.add(key);
      }

      return splittedKeys.values().stream().map(
            keysSubset -> new GetAllOperation(operationContext, keysSubset, options, dataFormat())).collect(Collectors.toList());
   }

   @Override
   protected Map createCollector() {
      return new HashMap<>();
   }

   @Override
   protected void combine(Map collector, Map result) {
      collector.putAll(result);
   }
}