io.atomix.protocols.raft.protocol.QueryResponse 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);
}
 

/**
 * Registers server handlers on the configured protocol.
 */
private void registerHandlers(RaftServerProtocol protocol) {
  protocol.registerOpenSessionHandler(request -> runOnContextIfReady(() -> role.onOpenSession(request), OpenSessionResponse::builder));
  protocol.registerCloseSessionHandler(request -> runOnContextIfReady(() -> role.onCloseSession(request), CloseSessionResponse::builder));
  protocol.registerKeepAliveHandler(request -> runOnContextIfReady(() -> role.onKeepAlive(request), KeepAliveResponse::builder));
  protocol.registerMetadataHandler(request -> runOnContextIfReady(() -> role.onMetadata(request), MetadataResponse::builder));
  protocol.registerConfigureHandler(request -> runOnContext(() -> role.onConfigure(request)));
  protocol.registerInstallHandler(request -> runOnContext(() -> role.onInstall(request)));
  protocol.registerJoinHandler(request -> runOnContext(() -> role.onJoin(request)));
  protocol.registerReconfigureHandler(request -> runOnContext(() -> role.onReconfigure(request)));
  protocol.registerLeaveHandler(request -> runOnContext(() -> role.onLeave(request)));
  protocol.registerTransferHandler(request -> runOnContext(() -> role.onTransfer(request)));
  protocol.registerAppendHandler(request -> runOnContext(() -> role.onAppend(request)));
  protocol.registerPollHandler(request -> runOnContext(() -> role.onPoll(request)));
  protocol.registerVoteHandler(request -> runOnContext(() -> role.onVote(request)));
  protocol.registerCommandHandler(request -> runOnContextIfReady(() -> role.onCommand(request), CommandResponse::builder));
  protocol.registerQueryHandler(request -> runOnContextIfReady(() -> role.onQuery(request), QueryResponse::builder));
}
 

/**
 * Tests sequencing callbacks with the sequencer.
 */
@Test
public void testSequenceResponses() throws Throwable {
  RaftSessionSequencer sequencer = new RaftSessionSequencer(new RaftSessionState("test", SessionId.from(1), UUID.randomUUID().toString(), TestPrimitiveType.instance(), 1000));
  long sequence1 = sequencer.nextRequest();
  long sequence2 = sequencer.nextRequest();
  assertTrue(sequence2 == sequence1 + 1);

  CommandResponse commandResponse = CommandResponse.builder()
      .withStatus(RaftResponse.Status.OK)
      .withIndex(2)
      .withEventIndex(0)
      .build();

  QueryResponse queryResponse = QueryResponse.builder()
      .withStatus(RaftResponse.Status.OK)
      .withIndex(2)
      .withEventIndex(0)
      .build();

  AtomicBoolean run = new AtomicBoolean();
  sequencer.sequenceResponse(sequence2, queryResponse, () -> run.set(true));
  sequencer.sequenceResponse(sequence1, commandResponse, () -> assertFalse(run.get()));
  assertTrue(run.get());
}
 

@Override
public void accept(QueryResponse response, Throwable error) {
  if (error == null) {
    if (response.status() == RaftResponse.Status.OK) {
      complete(response);
    } else if (response.error().type() == RaftError.Type.UNKNOWN_CLIENT
        || response.error().type() == RaftError.Type.UNKNOWN_SESSION) {
      complete(response.error().createException());
      state.setState(PrimitiveState.EXPIRED);
    } else if (response.error().type() == RaftError.Type.UNKNOWN_SERVICE
        || response.error().type() == RaftError.Type.CLOSED_SESSION) {
      complete(response.error().createException());
      state.setState(PrimitiveState.CLOSED);
    } else {
      complete(response.error().createException());
    }
  } else {
    fail(error);
  }
}
 

/**
 * 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());
}
 

/**
 * Applies a query to the state machine.
 */
protected CompletableFuture<QueryResponse> applyQuery(Indexed<QueryEntry> entry) {
  // In the case of the leader, the state machine is always up to date, so no queries will be queued and all query
  // indexes will be the last applied index.
  CompletableFuture<QueryResponse> future = new CompletableFuture<>();
  raft.getServiceManager().<OperationResult>apply(entry).whenComplete((result, error) -> {
    completeOperation(result, QueryResponse.builder(), error, future);
  });
  return future;
}
 

/**
 * Executes a linearizable query.
 * <p>
 * Linearizable queries are first sequenced with commands and then applied to the state machine. Once
 * applied, we verify the node's leadership prior to responding successfully to the query.
 */
private CompletableFuture<QueryResponse> queryLinearizable(Indexed<QueryEntry> entry) {
  return applyQuery(entry)
      .thenComposeAsync(response -> appender.appendEntries()
          .thenApply(index -> response)
          .exceptionally(error -> QueryResponse.builder()
              .withStatus(RaftResponse.Status.ERROR)
              .withError(RaftError.Type.QUERY_FAILURE, error.getMessage())
              .build()), raft.getThreadContext());
}
 

/**
 * 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;
}
 

@Override
@SuppressWarnings("unchecked")
protected void complete(QueryResponse response) {
  sequence(response, () -> {
    state.setResponseIndex(response.index());
    future.complete(response.result());
  });
}
 

/**
 * 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());
}
 

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

/**
 * 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());
}
 

/**
 * Performs a local query.
 */
protected CompletableFuture<QueryResponse> queryLocal(Indexed<QueryEntry> entry) {
  return applyQuery(entry);
}
 

@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
public void registerQueryHandler(Function<QueryRequest,
    CompletableFuture<QueryResponse>> handler) {
  registerHandler("query", handler);
}
 

@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 void registerQueryHandler(Function<QueryRequest,
    CompletableFuture<QueryResponse>> handler) {
  this.queryHandler = 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);
}