org.infinispan.client.hotrod.impl.operations.GetAllBulkOperation Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of infinispan-client-hotrod-new Show documentation

Infinispan Hot Rod Client New

The newest version!

package org.infinispan.client.hotrod.impl.operations;

import java.net.SocketAddress;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletionStage;
import java.util.function.Function;
import java.util.stream.Collectors;

import org.infinispan.client.hotrod.DataFormat;

public class GetAllBulkOperation extends HotRodBulkOperation, Map, HotRodOperation>> {
   private final Set keys;

   public GetAllBulkOperation(Set keys, DataFormat dataFormat,
                              Function, HotRodOperation>> setGetAllOperationFunction) {
      super(dataFormat, setGetAllOperationFunction);
      this.keys = keys;
   }

   @Override
   protected Map>> gatherOperations(
         Function routingFunction) {
      Map> serializedKeys = new HashMap<>();
      for (Object key : keys) {
         byte[] bytes = dataFormat.keyToBytes(key);
         SocketAddress socketAddress = getAddressForKey(key, bytes, routingFunction);
         Set keyBytes = serializedKeys.computeIfAbsent(socketAddress, ___ -> new HashSet<>());
         keyBytes.add(bytes);
      }
      return serializedKeys.entrySet().stream()
            .collect(Collectors.toMap(Map.Entry::getKey, e ->
               opFunction.apply(e.getValue())));
   }

   @Override
   public CompletionStage> reduce(CompletionStage>> stage) {
      return stage
            .thenApply(c -> c.stream()
                  .reduce((m1, m2) -> {
                     m1.putAll(m2);
                     return m1;
                  }).orElse(Map.of()));
   }
}