org.infinispan.client.hotrod.impl.iteration.RemoteInnerPublisherHandler Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of infinispan-client-hotrod Show documentation
Infinispan Hot Rod Client
There is a newer version: 15.1.0.Dev04
Show newest version
package org.infinispan.client.hotrod.impl.iteration;

import java.lang.invoke.MethodHandles;
import java.net.ConnectException;
import java.net.SocketAddress;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Supplier;

import org.infinispan.client.hotrod.exceptions.RemoteIllegalLifecycleStateException;
import org.infinispan.client.hotrod.exceptions.TransportException;
import org.infinispan.client.hotrod.impl.operations.IterationNextResponse;
import org.infinispan.client.hotrod.impl.operations.IterationStartResponse;
import org.infinispan.client.hotrod.logging.Log;
import org.infinispan.client.hotrod.logging.LogFactory;
import org.infinispan.commons.reactive.AbstractAsyncPublisherHandler;
import org.infinispan.commons.util.IntSet;
import org.infinispan.commons.util.logging.TraceException;

import io.netty.channel.Channel;

class RemoteInnerPublisherHandler extends AbstractAsyncPublisherHandler,
      Map.Entry, IterationStartResponse, IterationNextResponse> {
   private static final Log log = LogFactory.getLog(MethodHandles.lookup().lookupClass());

   protected final RemotePublisher publisher;

   // Need to be volatile since cancel can come on a different thread
   protected volatile Channel channel;
   private volatile byte[] iterationId;
   private final AtomicBoolean cancelled = new AtomicBoolean();

   protected RemoteInnerPublisherHandler(RemotePublisher parent, int batchSize,
         Supplier> supplier, Map.Entry firstTarget) {
      super(batchSize, supplier, firstTarget);
      this.publisher = parent;
   }

   private String iterationId() {
      return publisher.iterationId(iterationId);
   }

   @Override
   protected void sendCancel(Map.Entry target) {
      if (!cancelled.getAndSet(true)) {
         actualCancel();
      }
   }

   private void actualCancel() {
      if (iterationId != null && channel != null) {
         // Just let cancel complete asynchronously
         publisher.sendCancel(iterationId, channel);
      }
   }

   @Override
   protected CompletionStage sendInitialCommand(
         Map.Entry target, int batchSize) {
      SocketAddress address = target.getKey();
      IntSet segments = target.getValue();
      log.tracef("Starting iteration with segments %s", segments);
      return publisher.newIteratorStartOperation(address, segments, batchSize);
   }

   @Override
   protected CompletionStage> sendNextCommand(Map.Entry target, int batchSize) {
      return publisher.newIteratorNextOperation(iterationId, channel);
   }

   @Override
   protected long handleInitialResponse(IterationStartResponse startResponse, Map.Entry target) {
      this.channel = startResponse.getChannel();
      this.iterationId = startResponse.getIterationId();
      if (log.isDebugEnabled()) {
         log.iterationTransportObtained(channel.remoteAddress(), iterationId());
         log.startedIteration(iterationId());
      }

      // We could have been cancelled while the initial response was sent
      if (cancelled.get()) {
         actualCancel();
      }
      return 0;
   }

   @Override
   protected long handleNextResponse(IterationNextResponse nextResponse, Map.Entry target) {
      if (!nextResponse.hasMore()) {
         // server doesn't clean up when complete
         sendCancel(target);
         publisher.completeSegments(target.getValue());
         targetComplete();
      }
      IntSet completedSegments = nextResponse.getCompletedSegments();
      if (completedSegments != null && log.isTraceEnabled()) {
         IntSet targetSegments = target.getValue();
         if (targetSegments != null) {
            targetSegments.removeAll(completedSegments);
         }
      }
      publisher.completeSegments(completedSegments);
      List> entries = nextResponse.getEntries();
      for (Map.Entry entry : entries) {
         if (!onNext(entry)) {
            break;
         }
      }
      return entries.size();
   }

   @Override
   protected void handleThrowableInResponse(Throwable t, Map.Entry target) {
      if (t instanceof TransportException || t instanceof RemoteIllegalLifecycleStateException || t instanceof ConnectException) {
         log.throwableDuringPublisher(t);
         if (log.isTraceEnabled()) {
            IntSet targetSegments = target.getValue();
            if (targetSegments != null) {
               log.tracef("There are still outstanding segments %s that will need to be retried", targetSegments);
            }
         }
         publisher.erroredServer(target.getKey());
         // Try next target if possible
         targetComplete();

         accept(0);
      } else {
         t.addSuppressed(new TraceException());
         super.handleThrowableInResponse(t, target);
      }

   }
}