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

com.datastax.driver.core.ContinuousPagingQueue Maven / Gradle / Ivy

/*
 * Copyright DataStax, Inc.
 *
 * This software can be used solely with DataStax Enterprise. Please consult the license at
 * http://www.datastax.com/terms/datastax-dse-driver-license-terms
 */
package com.datastax.driver.core;

import static com.datastax.driver.core.Message.Response.Type.ERROR;
import static com.datastax.driver.core.Message.Response.Type.RESULT;
import static com.datastax.driver.core.Responses.Result.Kind.ROWS;

import com.datastax.driver.core.Message.Request;
import com.datastax.driver.core.Message.Response;
import com.datastax.driver.core.Requests.ReviseRequest;
import com.datastax.driver.core.Responses.Result;
import com.datastax.driver.core.Responses.Result.Rows;
import com.datastax.driver.core.exceptions.DriverInternalError;
import com.datastax.driver.core.exceptions.OperationTimedOutException;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture;
import io.netty.channel.EventLoop;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantLock;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Buffers the stream of responses to a continuous query between the network handler and the client.
 */
class ContinuousPagingQueue implements MultiResponseRequestHandler.Callback {
  private static final Logger logger = LoggerFactory.getLogger(ContinuousPagingQueue.class);

  private final Request request;

  // the continuous paging options specified by the user
  private final ContinuousPagingOptions continuousPagingOptions;

  // Coordinates access to shared state. This is acquired from the I/O thread, but in practice there
  // is little
  // contention.
  private final ReentrantLock lock;
  // Responses that we have received and have not been consumed by the client yet.
  // Only accessed while holding the lock.
  private final Queue queue;
  // If the client requested a page while the queue was empty, then it's waiting on that future.
  // Only accessed while holding the lock.
  private SettableFuture pendingResult;

  private volatile MultiResponseRequestHandler handler;
  // How long the client waits between each page
  private volatile long timeoutMillis;

  // An integer that represents the state of the continuous paging request:
  // - if positive, it is the sequence number of the next expected page;
  // - if negative, it is a terminal state, identified by the constants below.
  // This is only mutated from the connection's event loop, so no synchronization is needed.
  private volatile int state;
  private static final int STATE_FINISHED = -1;
  private static final int STATE_FAILED = -2;

  // These are set by the first response, and are constant for the rest of the execution
  private volatile Connection connection;
  private volatile ColumnDefinitions columnDefinitions;

  // How many pages were requested. This is the total number of pages requested
  // from the beginning. It will be zero if the protocol does not support numPagesRequested (DSE_V1)
  private volatile int numPagesRequested;

  ContinuousPagingQueue(
      Request request,
      ProtocolVersion protocolVersion,
      SettableFuture firstResult) {
    this.request = request;
    this.continuousPagingOptions = request.options().continuousPagingOptions;

    this.lock = new ReentrantLock();
    this.pendingResult = firstResult;
    this.queue = new ConcurrentLinkedQueue();

    this.state = 1;
    this.numPagesRequested =
        protocolVersion.compareTo(ProtocolVersion.DSE_V2) >= 0
            ? continuousPagingOptions.getMaxEnqueuedPages()
            : 0;
  }

  @Override
  public void register(MultiResponseRequestHandler handler) {
    this.handler = handler;

    // Same timeout as the initial request
    this.timeoutMillis = handler.timeoutMillis;
  }

  @Override
  public Request getRequest() {
    return request;
  }

  @Override
  public Request getCancelRequest(int streamId) {
    return ReviseRequest.continuousPagingCancel(streamId);
  }

  @Override
  public Request getBackpressureRequest(int streamId, int nextPages) {
    assert numPagesRequested > 0;
    return ReviseRequest.continuousPagingBackpressure(streamId, nextPages);
  }

  @Override
  public void onResponse(
      Connection connection, Response response, ExecutionInfo info, Statement statement) {
    assert connection.channel.eventLoop().inEventLoop();
    if (state < 0) {
      logger.debug(
          "Discarding {} response because the request has already completed", response.type);
      return;
    }
    this.connection = connection;
    if (response.type == RESULT && ((Result) response).kind == ROWS) {
      Rows rows = (Rows) response;
      if (rows.metadata.continuousPage.seqNo != state) {
        fail(
            new DriverInternalError(
                String.format(
                    "Received page number %d but was expecting %d",
                    rows.metadata.continuousPage.seqNo, state)),
            false);
      } else {
        if (rows.metadata.continuousPage.last) {
          logger.debug("Received last page ({})", rows.metadata.continuousPage.seqNo);
          state = STATE_FINISHED;
          // Make sure we don't leave it stuck
          connection.channel.config().setAutoRead(true);
          handler.release();
        } else {
          logger.debug("Received page {}", rows.metadata.continuousPage.seqNo);
          state = state + 1;
        }
        enqueueOrCompletePending(newResult(rows, info));
      }
    } else if (response.type == ERROR) {
      fail(((Responses.Error) response).asException(connection.address), true);
    } else {
      fail(new DriverInternalError("Unexpected response " + response.type), false);
    }
  }

  @Override
  public void onException(
      final Connection connection, final Exception exception, final boolean fromServer) {
    if (connection == null) {
      // This only happens when sending the initial request, if no host was available or if the
      // iterator returned
      // by the LBP threw an exception. In either case the write was not even attempted, so we're
      // sure we're not
      // going to race with responses or timeouts and we can complete without checking the state.
      logger.debug("Fail {} ({})", exception.getClass().getSimpleName(), exception.getMessage());
      enqueueOrCompletePending(exception);
    } else {
      EventLoop eventLoop = connection.channel.eventLoop();
      if (!eventLoop.inEventLoop()) {
        // reschedule so that the state is accessed from the right thread
        eventLoop.execute(
            new Runnable() {
              @Override
              public void run() {
                onException(connection, exception, fromServer);
              }
            });
      } else if (state > 0) {
        fail(exception, fromServer);
      }
    }
  }

  private void fail(Exception exception, boolean fromServer) {
    logger.debug(
        "Got failure {} ({})", exception.getClass().getSimpleName(), exception.getMessage());
    if (state >= 0) {
      if (fromServer) {
        // We can safely assume the server won't send any more responses, so release the streamId
        handler.release();
      } else {
        handler.cancel(); // notify server to stop sending responses
      }
      if (connection != null) {
        // Make sure we don't leave it stuck
        connection.channel.config().setAutoRead(true);
      }
      // Enqueue the exception *before* setting the state to failed to avoid a race
      // where the client tries to dequeue the exception before it has been enqueued,
      // but the sate has already been set to failed.
      enqueueOrCompletePending(exception);
      state = STATE_FAILED;
    }
  }

  // Enqueue a response or, if the client was already waiting for it, complete the pending future.
  private void enqueueOrCompletePending(Object pageOrError) {
    lock.lock();
    try {
      if (pendingResult != null) {
        if (logger.isDebugEnabled()) {
          logger.debug(
              "Client was waiting on empty queue, completing with {}", asDebugString(pageOrError));
        }
        SettableFuture tmp = pendingResult;
        pendingResult = null;
        complete(tmp, pageOrError);
      } else {
        if (logger.isDebugEnabled()) {
          logger.debug("Enqueuing {}", asDebugString(pageOrError));
        }
        enqueue(pageOrError);
      }
    } finally {
      lock.unlock();
    }
  }

  // Dequeue a response or, if the queue is empty, create the future that will get notified of the
  // next response.
  ListenableFuture dequeueOrCreatePending() {
    lock.lock();
    try {
      // Precondition: the client will not call this method until the previous call has completed
      // (this is guaranteed
      // by our public API because in order to ask for the next page, you need the reference to the
      // previous page --
      // see AsyncContinuousPagingResult#nextPage())
      assert pendingResult == null;

      Object head = dequeue();
      maybeRequestMore();

      if (head != null) {
        if (logger.isDebugEnabled()) {
          logger.debug(
              "Client queries on non-empty queue, returning immediate future of {}",
              asDebugString(head));
        }
        return immediateFuture(head);
      } else if (state == STATE_FAILED) {
        logger.debug("Client queries on failed empty queue, returning failed future");
        return immediateFuture(
            new IllegalStateException(
                "Can't get more results because the continuous query has failed already. "
                    + "Most likely this is because the query was cancelled"));
      } else {
        logger.debug("Client queries on empty queue, installing future");
        final SettableFuture future = SettableFuture.create();
        future.addListener(
            new Runnable() {
              @Override
              public void run() {
                if (future.isCancelled()) {
                  ContinuousPagingQueue.this.cancel();
                }
              }
            },
            GuavaCompatibility.INSTANCE.sameThreadExecutor());
        pendingResult = future;
        startTimeout();
        return future;
      }
    } finally {
      lock.unlock();
    }
  }

  private void enqueue(Object pageOrError) {
    assert lock.isHeldByCurrentThread();
    queue.add(pageOrError);
    // Backpressure without protocol support: if the queue grows too large, disable auto-read so
    // that the channel eventually becomes
    // non-writable on the server side (causing it to back off for a while)
    if (numPagesRequested == 0
        && queue.size() == continuousPagingOptions.getMaxEnqueuedPages()
        && state > 0) {
      if (logger.isDebugEnabled()) {
        logger.debug("Exceeded {} queued response pages, disabling auto-read", queue.size());
      }
      connection.channel.config().setAutoRead(false);
    }
  }

  private Object dequeue() {
    assert lock.isHeldByCurrentThread();
    Object head = queue.poll();
    if (numPagesRequested == 0
        && head != null
        && queue.size() == continuousPagingOptions.getMaxEnqueuedPages() - 1) {
      if (logger.isDebugEnabled()) {
        logger.debug("Back to {} queued response pages, re-enabling auto-read", queue.size());
      }
      connection.channel.config().setAutoRead(true);
    }
    return head;
  }

  /**
   * If the total number of results in the queue and in-flight (requested - received) is less than
   * half the queue size, then request more pages, unless the {@link #state} is failed or we don't
   * support backpressure at the protocol level, that is {@link #numPagesRequested} is zero.
   */
  private void maybeRequestMore() {
    if (state < 0 || numPagesRequested == 0) return;

    assert lock.isHeldByCurrentThread();

    // the pages received so far, which is the state minus one
    int numPagesReceived = state - 1;

    int maxEnqueuedPages = continuousPagingOptions.getMaxEnqueuedPages();

    // the pages that fit in the queue, which is the queue free space minus the requests in flight
    int numPagesFittingInQueue =
        maxEnqueuedPages - queue.size() - (numPagesRequested - numPagesReceived);

    if (numPagesFittingInQueue >= maxEnqueuedPages / 2) {
      // if we have already requested more than the client needs, then no need to request some more
      if (continuousPagingOptions.getMaxPages() > 0
          && numPagesRequested >= continuousPagingOptions.getMaxPages()) return;

      numPagesRequested += numPagesFittingInQueue;
      logger.debug("Requesting pages ({}/{})", numPagesRequested, numPagesReceived);

      handler.requestMore(numPagesFittingInQueue);
    }
  }

  private void complete(SettableFuture future, Object pageOrError) {
    if (pageOrError instanceof AsyncContinuousPagingResult) {
      future.set((AsyncContinuousPagingResult) pageOrError);
    } else {
      future.setException((Throwable) pageOrError);
    }
  }

  private ListenableFuture immediateFuture(Object pageOrError) {
    return (pageOrError instanceof AsyncContinuousPagingResult)
        ? Futures.immediateFuture((AsyncContinuousPagingResult) pageOrError)
        : Futures.immediateFailedFuture((Throwable) pageOrError);
  }

  private AsyncContinuousPagingResult newResult(Rows rows, ExecutionInfo info) {

    if (columnDefinitions == null) {
      // Contrary to ROWS responses from regular queries, the first page always includes metadata so
      // we use
      // this regardless of whether or not the query was from a prepared statement.
      columnDefinitions = rows.metadata.columns;
    }

    Token.Factory tokenFactory = handler.manager.cluster.getMetadata().tokenFactory();
    ProtocolVersion protocolVersion = handler.manager.cluster.manager.protocolVersion();
    CodecRegistry codecRegistry = handler.manager.configuration().getCodecRegistry();

    info =
        info.with(
            null, // Don't handle query trace, it's unlikely to be used with continuous paging
            rows.warnings,
            rows.metadata.pagingState,
            handler.statement,
            protocolVersion,
            codecRegistry);

    return new DefaultAsyncContinuousPagingResult(
        rows.data,
        columnDefinitions,
        rows.metadata.continuousPage.seqNo,
        rows.metadata.continuousPage.last,
        info,
        tokenFactory,
        protocolVersion,
        this);
  }

  void cancel() {
    if (logger.isTraceEnabled()) {
      logger.trace(
          "Cancelling cont. paging session with state {} and connection {}", state, connection);
    }
    if (state >= 0) {
      state = STATE_FAILED;
      handler.cancel();
      cancelPendingResult(); // if another thread is waiting on an empty queue, unblock it
      if (connection != null) {
        // Make sure we don't leave it stuck
        connection.channel.config().setAutoRead(true);
      }
    }
  }

  private void cancelPendingResult() {
    lock.lock();
    try {
      if (pendingResult != null) {
        pendingResult.cancel(true);
      }
    } finally {
      lock.unlock();
    }
  }

  private void startTimeout() {
    // We don't set a timeout for the initial query (because MultiResponseRequestHandler handles
    // it). We set
    // connection on the initial response so it will be set.
    assert connection != null;
    final int expectedPage = state;
    if (expectedPage < 0) {
      return;
    }
    assert expectedPage > 1 : expectedPage;
    if (timeoutMillis > 0) {
      connection
          .channel
          .eventLoop()
          .schedule(
              new Runnable() {
                @Override
                public void run() {
                  if (state == expectedPage) {
                    fail(
                        new OperationTimedOutException(
                            connection.address,
                            String.format("Timed out waiting for page %d", expectedPage)),
                        false);
                  } else {
                    // Ignore if the request has moved on. This is simpler than trying to cancel the
                    // timeout.
                    logger.trace(
                        "Timeout fired for page {} but query already at state {}, skipping",
                        expectedPage,
                        state);
                  }
                }
              },
              timeoutMillis,
              TimeUnit.MILLISECONDS);
    }
  }

  private String asDebugString(Object pageOrError) {
    return (pageOrError instanceof AsyncContinuousPagingResult)
        ? "page " + ((AsyncContinuousPagingResult) pageOrError).pageNumber()
        : ((Exception) pageOrError).getClass().getSimpleName();
  }
}