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

/**
 * Tests sequencing an event that arrives before a command response.
 */
@Test
public void testSequenceEventBeforeCommand() throws Throwable {
  RaftSessionSequencer sequencer = new RaftSessionSequencer(new RaftSessionState("test", SessionId.from(1), UUID.randomUUID().toString(), TestPrimitiveType.instance(), 1000));
  long sequence = sequencer.nextRequest();

  PublishRequest request = PublishRequest.builder()
      .withSession(1)
      .withEventIndex(1)
      .withPreviousIndex(0)
      .withEvents(Collections.emptyList())
      .build();

  CommandResponse response = CommandResponse.builder()
      .withStatus(RaftResponse.Status.OK)
      .withIndex(2)
      .withEventIndex(1)
      .build();

  AtomicInteger run = new AtomicInteger();
  sequencer.sequenceEvent(request, () -> assertEquals(0, run.getAndIncrement()));
  sequencer.sequenceResponse(sequence, response, () -> assertEquals(1, run.getAndIncrement()));
  assertEquals(2, run.get());
}
 

@Override
public CompletableFuture<CommandResponse> onCommand(CommandRequest request) {
  raft.checkThread();
  logRequest(request);

  if (raft.getLeader() == null) {
    return CompletableFuture.completedFuture(logResponse(CommandResponse.builder()
        .withStatus(RaftResponse.Status.ERROR)
        .withError(RaftError.Type.NO_LEADER)
        .build()));
  } else {
    return forward(request, raft.getProtocol()::command)
        .exceptionally(error -> CommandResponse.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 submitting a command to the cluster.
 */
@Test
public void testSubmitCommand() throws Throwable {
  RaftSessionConnection connection = mock(RaftSessionConnection.class);
  when(connection.command(any(CommandRequest.class)))
      .thenReturn(CompletableFuture.completedFuture(CommandResponse.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(connection, mock(RaftSessionConnection.class), state, new RaftSessionSequencer(state), manager, threadContext);
  assertArrayEquals(submitter.invoke(operation(COMMAND, HeapBytes.EMPTY)).get(), "Hello world!".getBytes());
  assertEquals(1, state.getCommandRequest());
  assertEquals(1, state.getCommandResponse());
  assertEquals(10, state.getResponseIndex());
}
 

/**
 * Tests sequencing an event that arrives before a command response.
 */
@Test
public void testSequenceEventAfterCommand() throws Throwable {
  RaftSessionSequencer sequencer = new RaftSessionSequencer(new RaftSessionState("test", SessionId.from(1), UUID.randomUUID().toString(), TestPrimitiveType.instance(), 1000));
  long sequence = sequencer.nextRequest();

  PublishRequest request = PublishRequest.builder()
      .withSession(1)
      .withEventIndex(1)
      .withPreviousIndex(0)
      .withEvents(Collections.emptyList())
      .build();

  CommandResponse response = CommandResponse.builder()
      .withStatus(RaftResponse.Status.OK)
      .withIndex(2)
      .withEventIndex(1)
      .build();

  AtomicInteger run = new AtomicInteger();
  sequencer.sequenceResponse(sequence, response, () -> assertEquals(0, run.getAndIncrement()));
  sequencer.sequenceEvent(request, () -> assertEquals(1, run.getAndIncrement()));
  assertEquals(2, run.get());
}
 

/**
 * Tests sequencing an event that arrives before a command response.
 */
@Test
public void testSequenceEventAtCommand() throws Throwable {
  RaftSessionSequencer sequencer = new RaftSessionSequencer(new RaftSessionState("test", SessionId.from(1), UUID.randomUUID().toString(), TestPrimitiveType.instance(), 1000));
  long sequence = sequencer.nextRequest();

  PublishRequest request = PublishRequest.builder()
      .withSession(1)
      .withEventIndex(2)
      .withPreviousIndex(0)
      .withEvents(Collections.emptyList())
      .build();

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

  AtomicInteger run = new AtomicInteger();
  sequencer.sequenceResponse(sequence, response, () -> assertEquals(1, run.getAndIncrement()));
  sequencer.sequenceEvent(request, () -> assertEquals(0, run.getAndIncrement()));
  assertEquals(2, run.get());
}
 

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

/**
 * Tests sequencing an event that arrives before a command response.
 */
@Test
public void testSequenceEventAfterAllCommands() throws Throwable {
  RaftSessionSequencer sequencer = new RaftSessionSequencer(new RaftSessionState("test", SessionId.from(1), UUID.randomUUID().toString(), TestPrimitiveType.instance(), 1000));
  long sequence = sequencer.nextRequest();

  PublishRequest request1 = PublishRequest.builder()
      .withSession(1)
      .withEventIndex(2)
      .withPreviousIndex(0)
      .withEvents(Collections.emptyList())
      .build();

  PublishRequest request2 = PublishRequest.builder()
      .withSession(1)
      .withEventIndex(3)
      .withPreviousIndex(2)
      .withEvents(Collections.emptyList())
      .build();

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

  AtomicInteger run = new AtomicInteger();
  sequencer.sequenceEvent(request1, () -> assertEquals(0, run.getAndIncrement()));
  sequencer.sequenceEvent(request2, () -> assertEquals(2, run.getAndIncrement()));
  sequencer.sequenceResponse(sequence, response, () -> assertEquals(1, run.getAndIncrement()));
  assertEquals(3, run.get());
}
 

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

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

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

/**
 * Fails pending commands.
 */
private void failPendingCommands() {
  for (RaftSession session : raft.getSessions().getSessions()) {
    for (PendingCommand command : session.clearCommands()) {
      command.future().complete(logResponse(CommandResponse.builder()
          .withStatus(RaftResponse.Status.ERROR)
          .withError(RaftError.Type.COMMAND_FAILURE, "Request sequence number " + command.request().sequenceNumber() + " out of sequence")
          .withLastSequence(session.getRequestSequence())
          .build()));
    }
  }
}
 

@Override
public CompletableFuture<CommandResponse> command(MemberId memberId, CommandRequest request) {
  return sendAndReceive(memberId, "command", request);
}
 

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

@Override
public CompletableFuture<CommandResponse> command(MemberId memberId, CommandRequest request) {
  return sendAndReceive(memberId, "command", request);
}
 

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

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

  RaftSessionConnection connection = mock(RaftSessionConnection.class);
  Mockito.when(connection.command(any(CommandRequest.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(connection, mock(RaftSessionConnection.class), state, new RaftSessionSequencer(state), manager, threadContext);

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

  assertEquals(1, state.getResponseIndex());

  assertFalse(result.isDone());

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

  assertTrue(result.isCompletedExceptionally());
}
 

/**
 * Test resequencing a command response.
 */
@Test
public void testResequenceCommand() throws Throwable {
  CompletableFuture<CommandResponse> future1 = new CompletableFuture<>();
  CompletableFuture<CommandResponse> future2 = new CompletableFuture<>();

  RaftSessionConnection connection = mock(RaftSessionConnection.class);
  Mockito.when(connection.command(any(CommandRequest.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(connection, mock(RaftSessionConnection.class), state, new RaftSessionSequencer(state), manager, threadContext);

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

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

  assertEquals(2, state.getCommandRequest());
  assertEquals(0, state.getCommandResponse());
  assertEquals(1, state.getResponseIndex());

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

  future1.complete(CommandResponse.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(2, state.getCommandRequest());
  assertEquals(2, state.getCommandResponse());
  assertEquals(10, state.getResponseIndex());
}
 

@Test
public void testReSubmitCommand() throws Throwable {
  // given
  // setup  thread context
  final Logger log = LoggerFactory.getLogger(getClass());
  final int threadPoolSize = Math.max(Math.min(Runtime.getRuntime().availableProcessors() * 2, 16), 4);
  final ThreadContext context = new BlockingAwareThreadPoolContextFactory(
      "raft-partition-group-data-%d", threadPoolSize, log).createContext();

  // collecting request futures
  final List<CompletableFuture<CommandResponse>> futures = new CopyOnWriteArrayList<>();
  final List<Long> sequences = new CopyOnWriteArrayList<>();
  final RaftSessionConnection connection = mock(RaftSessionConnection.class);
  when(connection.command(any(CommandRequest.class)))
      .thenAnswer(invocationOnMock -> {
        final CommandRequest commandRequest = (CommandRequest) invocationOnMock.getArguments()[0];
        sequences.add(commandRequest.sequenceNumber());
        final CompletableFuture<CommandResponse> future = new CompletableFuture<>();
        futures.add(future);
        return future;
      });

  // raft session invoker
  final RaftSessionState state = new RaftSessionState("test", SessionId.from(1), UUID.randomUUID().toString(), TestPrimitiveType.instance(), 1000);
  final RaftSessionManager manager = mock(RaftSessionManager.class);
  final RaftSessionInvoker submitter =
      new RaftSessionInvoker(connection,
          mock(RaftSessionConnection.class),
          state,
          new RaftSessionSequencer(state),
          manager,
          context);

  // send five commands
  submitter.invoke(operation(COMMAND, HeapBytes.EMPTY));
  submitter.invoke(operation(COMMAND, HeapBytes.EMPTY));
  submitter.invoke(operation(COMMAND, HeapBytes.EMPTY));
  submitter.invoke(operation(COMMAND, HeapBytes.EMPTY));
  final CompletableFuture<byte[]> lastFuture = submitter.invoke(operation(COMMAND, HeapBytes.EMPTY));

  verify(connection, timeout(2000).times(5)).command(any());

  // when we missed second command
  futures.get(0).complete(CommandResponse.builder()
      .withStatus(RaftResponse.Status.OK)
      .withIndex(10)
      .withResult("Hello world!".getBytes())
      .build());
  final ArrayList<CompletableFuture<CommandResponse>> copiedFutures = new ArrayList<>(futures.subList(1, futures.size()));
  futures.clear();
  sequences.clear();
  copiedFutures.forEach(f ->
      f.complete(CommandResponse.builder()
        .withStatus(RaftResponse.Status.ERROR)
        .withError(RaftError.Type.COMMAND_FAILURE)
        .withLastSequence(1)
        .build()));

  // then session invoker should resubmit 4 commands
  verify(connection, timeout(2000).times(9)).command(any());
  final CountDownLatch latch = new CountDownLatch(1);
  // context thread pool ensures execution of task in order
  context.execute(latch::countDown);
  latch.await();

  assertEquals(4, futures.size());
  assertArrayEquals(new Long[]{2L, 3L, 4L, 5L}, sequences.toArray(new Long[0]));

  futures.forEach(f -> f.complete(CommandResponse.builder()
          .withStatus(RaftResponse.Status.OK)
          .withIndex(10)
          .withResult("Hello world!".getBytes())
          .build()));

  assertArrayEquals(lastFuture.get(), "Hello world!".getBytes());
  assertEquals(5, state.getCommandRequest());
  assertEquals(5, state.getCommandResponse());
  assertEquals(10, state.getResponseIndex());
}
 

@Override
public void registerCommandHandler(Function<CommandRequest, CompletableFuture<CommandResponse>> handler) {
  registerHandler("command", handler);
}
 

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

/**
 * Commits a command.
 *
 * @param request the command request
 * @param future  the command response future
 */
private void commitCommand(CommandRequest request, CompletableFuture<CommandResponse> future) {
  final long term = raft.getTerm();
  final long timestamp = System.currentTimeMillis();

  CommandEntry command = new CommandEntry(term, timestamp, request.session(), request.sequenceNumber(), request.operation());
  appendAndCompact(command)
      .whenCompleteAsync((entry, error) -> {
        if (error != null) {
          Throwable cause = Throwables.getRootCause(error);
          if (Throwables.getRootCause(error) instanceof StorageException.TooLarge) {
            log.warn("Failed to append command {}", command, cause);
            future.complete(CommandResponse.builder()
                .withStatus(RaftResponse.Status.ERROR)
                .withError(RaftError.Type.PROTOCOL_ERROR)
                .build());
          } else {
            future.complete(CommandResponse.builder()
                .withStatus(RaftResponse.Status.ERROR)
                .withError(RaftError.Type.COMMAND_FAILURE)
                .build());
          }
          return;
        }

        // Replicate the command to followers.
        appender.appendEntries(entry.index()).whenComplete((commitIndex, commitError) -> {
          raft.checkThread();
          if (isRunning()) {
            // If the command was successfully committed, apply it to the state machine.
            if (commitError == null) {
              raft.getServiceManager().<OperationResult>apply(entry.index()).whenComplete((r, e) -> {
                completeOperation(r, CommandResponse.builder(), e, future);
              });
            } else {
              future.complete(CommandResponse.builder()
                  .withStatus(RaftResponse.Status.ERROR)
                  .withError(RaftError.Type.COMMAND_FAILURE)
                  .build());
            }
          } else {
            future.complete(CommandResponse.builder()
                .withStatus(RaftResponse.Status.ERROR)
                .withError(RaftError.Type.COMMAND_FAILURE)
                .build());
          }
        });
      }, raft.getThreadContext());
}
 

@Override
public void registerCommandHandler(Function<CommandRequest, CompletableFuture<CommandResponse>> handler) {
  this.commandHandler = handler;
}
 

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

@Override
public CompletableFuture<CommandResponse> command(MemberId memberId,
                                                  CommandRequest request) {
  return sendAndReceive(memberId, "command", request);
}
 

@Override
public CompletableFuture<CommandResponse> command(MemberId memberId,
                                                  CommandRequest request) {
  return sendAndReceive(memberId, "command", request);
}
 

@Override
public void registerCommandHandler(Function<CommandRequest,
    CompletableFuture<CommandResponse>> handler) {
  registerHandler("command", handler);
}
 

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

@Override
public void registerCommandHandler(Function<CommandRequest,
    CompletableFuture<CommandResponse>> handler) {
  this.commandHandler = handler;
}
 

@Override
public CompletableFuture<CommandResponse> command(MemberId memberId,
                                                  CommandRequest request) {
  return sendAndReceive(memberId, "command", request);
}