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

@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);
  }
}
 

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

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

/**
 * Submits a command to the cluster.
 */
private void invokeCommand(PrimitiveOperation operation, CompletableFuture<byte[]> future) {
  CommandRequest request = CommandRequest.builder()
      .withSession(state.getSessionId().id())
      .withSequence(state.nextCommandRequest())
      .withOperation(operation)
      .build();
  invokeCommand(request, future);
}
 

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

@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) {
  clusterCommunicator.subscribe(context.commandSubject, serializer::decode, handler, serializer::encode);
}
 

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

public PendingCommand(CommandRequest request, CompletableFuture<CommandResponse> future) {
  this.request = request;
  this.future = future;
}
 

/**
 * 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 CompletableFuture<CommandResponse> command(MemberId memberId, CommandRequest request) {
  return sendAndReceive(context.commandSubject, request, memberId);
}
 

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

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

@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 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 void registerCommandHandler(Function<CommandRequest,
    CompletableFuture<CommandResponse>> handler) {
  registerHandler("command", handler);
}
 

@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 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);
}
 

@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 sendAndReceive(memberId, "command", request);
}
 

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