io.atomix.protocols.raft.protocol.QueryRequest Java Examples

/**
 * Forwards the query to the leader.
 */
private CompletableFuture<QueryResponse> queryForward(QueryRequest request) {
  if (raft.getLeader() == null) {
    return CompletableFuture.completedFuture(logResponse(QueryResponse.builder()
        .withStatus(RaftResponse.Status.ERROR)
        .withError(RaftError.Type.NO_LEADER)
        .build()));
  }

  log.trace("Forwarding {}", request);
  return forward(request, raft.getProtocol()::query)
      .exceptionally(error -> QueryResponse.builder()
          .withStatus(RaftResponse.Status.ERROR)
          .withError(RaftError.Type.NO_LEADER)
          .build())
      .thenApply(this::logResponse);
}
 

/**
 * Tests submitting a query to the cluster.
 */
@Test
public void testSubmitQuery() throws Throwable {
  RaftSessionConnection connection = mock(RaftSessionConnection.class);
  when(connection.query(any(QueryRequest.class)))
      .thenReturn(CompletableFuture.completedFuture(QueryResponse.builder()
          .withStatus(RaftResponse.Status.OK)
          .withIndex(10)
          .withResult("Hello world!".getBytes())
          .build()));

  RaftSessionState state = new RaftSessionState("test", SessionId.from(1), UUID.randomUUID().toString(), TestPrimitiveType.instance(), 1000);
  RaftSessionManager manager = mock(RaftSessionManager.class);
  ThreadContext threadContext = new TestContext();

  RaftSessionInvoker submitter = new RaftSessionInvoker(mock(RaftSessionConnection.class), connection, state, new RaftSessionSequencer(state), manager, threadContext);
  assertTrue(Arrays.equals(submitter.invoke(operation(QUERY)).get(), "Hello world!".getBytes()));
  assertEquals(10, state.getResponseIndex());
}
 

/**
 * Sends a query request to the given node.
 *
 * @param request the request to send
 * @return a future to be completed with the response
 */
public CompletableFuture<QueryResponse> query(QueryRequest request) {
  CompletableFuture<QueryResponse> future = new CompletableFuture<>();
  if (context.isCurrentContext()) {
    sendRequest(request, protocol::query, future);
  } else {
    context.execute(() -> sendRequest(request, protocol::query, future));
  }
  return future;
}
 

/**
 * Submits a query to the cluster.
 */
private void invokeQuery(PrimitiveOperation operation, CompletableFuture<byte[]> future) {
  QueryRequest request = QueryRequest.builder()
      .withSession(state.getSessionId().id())
      .withSequence(state.getCommandRequest())
      .withOperation(operation)
      .withIndex(Math.max(state.getResponseIndex(), state.getEventIndex()))
      .build();
  invokeQuery(request, future);
}
 

@Override
public CompletableFuture<QueryResponse> onQuery(QueryRequest request) {
  logRequest(request);
  return Futures.completedFuture(logResponse(QueryResponse.builder()
      .withStatus(Status.ERROR)
      .withError(RaftError.Type.UNAVAILABLE)
      .build()));
}
 

/**
 * Tests skipping over a failed query attempt.
 */
@Test
public void testSkippingOverFailedQuery() throws Throwable {
  CompletableFuture<QueryResponse> future1 = new CompletableFuture<>();
  CompletableFuture<QueryResponse> future2 = new CompletableFuture<>();

  RaftSessionConnection connection = mock(RaftSessionConnection.class);
  Mockito.when(connection.query(any(QueryRequest.class)))
      .thenReturn(future1)
      .thenReturn(future2);

  RaftSessionState state = new RaftSessionState("test", SessionId.from(1), UUID.randomUUID().toString(), TestPrimitiveType.instance(), 1000);
  RaftSessionManager manager = mock(RaftSessionManager.class);
  ThreadContext threadContext = new TestContext();

  RaftSessionInvoker submitter = new RaftSessionInvoker(mock(RaftSessionConnection.class), connection, state, new RaftSessionSequencer(state), manager, threadContext);

  CompletableFuture<byte[]> result1 = submitter.invoke(operation(QUERY));
  CompletableFuture<byte[]> result2 = submitter.invoke(operation(QUERY));

  assertEquals(1, state.getResponseIndex());

  assertFalse(result1.isDone());
  assertFalse(result2.isDone());

  future1.completeExceptionally(new RaftException.QueryFailure("failure"));
  future2.complete(QueryResponse.builder()
      .withStatus(RaftResponse.Status.OK)
      .withIndex(10)
      .withResult("Hello world!".getBytes())
      .build());

  assertTrue(result1.isCompletedExceptionally());
  assertTrue(result2.isDone());
  assertTrue(Arrays.equals(result2.get(), "Hello world!".getBytes()));

  assertEquals(10, state.getResponseIndex());
}
 

/**
 * Tests resequencing a query response.
 */
@Test
public void testResequenceQuery() throws Throwable {
  CompletableFuture<QueryResponse> future1 = new CompletableFuture<>();
  CompletableFuture<QueryResponse> future2 = new CompletableFuture<>();

  RaftSessionConnection connection = mock(RaftSessionConnection.class);
  Mockito.when(connection.query(any(QueryRequest.class)))
      .thenReturn(future1)
      .thenReturn(future2);

  RaftSessionState state = new RaftSessionState("test", SessionId.from(1), UUID.randomUUID().toString(), TestPrimitiveType.instance(), 1000);
  RaftSessionManager manager = mock(RaftSessionManager.class);
  ThreadContext threadContext = new TestContext();

  RaftSessionInvoker submitter = new RaftSessionInvoker(mock(RaftSessionConnection.class), connection, state, new RaftSessionSequencer(state), manager, threadContext);

  CompletableFuture<byte[]> result1 = submitter.invoke(operation(QUERY));
  CompletableFuture<byte[]> result2 = submitter.invoke(operation(QUERY));

  future2.complete(QueryResponse.builder()
      .withStatus(RaftResponse.Status.OK)
      .withIndex(10)
      .withResult("Hello world again!".getBytes())
      .build());

  assertEquals(1, state.getResponseIndex());

  assertFalse(result1.isDone());
  assertFalse(result2.isDone());

  future1.complete(QueryResponse.builder()
      .withStatus(RaftResponse.Status.OK)
      .withIndex(9)
      .withResult("Hello world!".getBytes())
      .build());

  assertTrue(result1.isDone());
  assertTrue(Arrays.equals(result1.get(), "Hello world!".getBytes()));
  assertTrue(result2.isDone());
  assertTrue(Arrays.equals(result2.get(), "Hello world again!".getBytes()));

  assertEquals(10, state.getResponseIndex());
}
 

@Override
public CompletableFuture<QueryResponse> query(MemberId memberId, QueryRequest request) {
  return sendAndReceive(context.querySubject, request, memberId);
}
 

@Override
public void registerQueryHandler(Function<QueryRequest, CompletableFuture<QueryResponse>> handler) {
  clusterCommunicator.subscribe(context.querySubject, serializer::decode, handler, serializer::encode);
}
 

@Override
public CompletableFuture<QueryResponse> query(MemberId memberId, QueryRequest request) {
  return sendAndReceive(context.querySubject, request, memberId);
}
 

@Override
public void registerQueryHandler(Function<QueryRequest, CompletableFuture<QueryResponse>> handler) {
  this.queryHandler = handler;
}
 

/**
 * Tests that the client's session is expired when an UnknownSessionException is received from the cluster.
 */
@Test
public void testExpireSessionOnQueryFailure() throws Throwable {
  CompletableFuture<QueryResponse> future = new CompletableFuture<>();

  RaftSessionConnection connection = mock(RaftSessionConnection.class);
  Mockito.when(connection.query(any(QueryRequest.class)))
      .thenReturn(future);

  RaftSessionState state = new RaftSessionState("test", SessionId.from(1), UUID.randomUUID().toString(), TestPrimitiveType.instance(), 1000);
  RaftSessionManager manager = mock(RaftSessionManager.class);
  ThreadContext threadContext = new TestContext();

  RaftSessionInvoker submitter = new RaftSessionInvoker(mock(RaftSessionConnection.class), connection, state, new RaftSessionSequencer(state), manager, threadContext);

  CompletableFuture<byte[]> result = submitter.invoke(operation(QUERY));

  assertEquals(1, state.getResponseIndex());

  assertFalse(result.isDone());

  future.completeExceptionally(new RaftException.UnknownSession("unknown session"));

  assertTrue(result.isCompletedExceptionally());
}
 

@Override
public CompletableFuture<QueryResponse> query(MemberId memberId, QueryRequest request) {
  return sendAndReceive(memberId, "query", request);
}
 

@Override
public CompletableFuture<QueryResponse> query(MemberId memberId, QueryRequest request) {
  return getServer(memberId).thenCompose(protocol -> protocol.query(encode(request))).thenApply(this::decode);
}
 

@Override
public CompletableFuture<QueryResponse> query(MemberId memberId, QueryRequest request) {
  return sendAndReceive(memberId, "query", request);
}
 

@Override
public void registerQueryHandler(Function<QueryRequest, CompletableFuture<QueryResponse>> handler) {
  registerHandler("query", handler);
}
 

@Override
public CompletableFuture<QueryResponse> query(MemberId memberId, QueryRequest request) {
  return getServer(memberId).thenCompose(listener -> listener.query(encode(request))).thenApply(this::decode);
}
 

@Override
public CompletableFuture<QueryResponse> query(MemberId memberId, QueryRequest request) {
  return sendAndReceive(memberId, "query", request);
}
 

@Override
protected OperationAttempt<QueryRequest, QueryResponse> next() {
  return new QueryAttempt(sequence, this.attempt + 1, request, future);
}
 

QueryAttempt(long sequence, int attempt, QueryRequest request, CompletableFuture<byte[]> future) {
  super(sequence, attempt, request, future);
}
 

QueryAttempt(long sequence, QueryRequest request, CompletableFuture<byte[]> future) {
  super(sequence, 1, request, future);
}
 

/**
 * Submits a query request to the cluster.
 */
private void invokeQuery(QueryRequest request, CompletableFuture<byte[]> future) {
  invoke(new QueryAttempt(sequencer.nextRequest(), request, future));
}
 

@Override
public void registerQueryHandler(Function<QueryRequest,
    CompletableFuture<QueryResponse>> handler) {
  this.queryHandler = handler;
}
 

@Override
public CompletableFuture<QueryResponse> query(MemberId memberId, QueryRequest request) {
  return getServer(memberId).thenCompose(listener ->
      listener.query(encode(request))).thenApply(this::decode);
}
 

@Override
public void registerQueryHandler(Function<QueryRequest,
    CompletableFuture<QueryResponse>> handler) {
  registerHandler("query", handler);
}
 

@Override
public CompletableFuture<QueryResponse> query(MemberId memberId, QueryRequest request) {
  return sendAndReceive(memberId, "query", request);
}
 

@Override
public CompletableFuture<QueryResponse> query(MemberId memberId, QueryRequest request) {
  return getServer(memberId).thenCompose(protocol ->
      protocol.query(encode(request))).thenApply(this::decode);
}
 

@Override
public CompletableFuture<QueryResponse> query(MemberId memberId, QueryRequest request) {
  return sendAndReceive(memberId, "query", request);
}
 

@Override
public void registerQueryHandler(Function<QueryRequest,
    CompletableFuture<QueryResponse>> handler) {
  this.queryHandler = handler;
}
 

@Override
public CompletableFuture<QueryResponse> query(MemberId memberId, QueryRequest request) {
  return getServer(memberId).thenCompose(listener ->
      listener.query(encode(request))).thenApply(this::decode);
}