org.apache.ratis.protocol.RaftClientReply Java Examples

CompletableFuture<RaftClientReply> onNext(RaftClientRequest request) {
  final Map<Long, CompletableFuture<RaftClientReply>> map = replies.get();
  if (map == null) {
    return JavaUtils.completeExceptionally(new AlreadyClosedException(getName() + " is closed."));
  }
  final CompletableFuture<RaftClientReply> f = new CompletableFuture<>();
  CollectionUtils.putNew(request.getCallId(), f, map,
      () -> getName() + ":" + getClass().getSimpleName());
  try {
    requestStreamObserver.onNext(ClientProtoUtils.toRaftClientRequestProto(request));
    scheduler.onTimeout(requestTimeoutDuration, () -> timeoutCheck(request), LOG,
        () -> "Timeout check failed for client request: " + request);
  } catch(Throwable t) {
    handleReplyFuture(request.getCallId(), future -> future.completeExceptionally(t));
  }
  return f;
}
 

static <REPLY extends RaftClientReply, REPLY_PROTO> void asyncCall(
    StreamObserver<REPLY_PROTO> responseObserver,
    CheckedSupplier<CompletableFuture<REPLY>, IOException> supplier,
    Function<REPLY, REPLY_PROTO> toProto) {
  try {
    supplier.get().whenCompleteAsync((reply, exception) -> {
      if (exception != null) {
        responseObserver.onError(GrpcUtil.wrapException(exception));
      } else {
        responseObserver.onNext(toProto.apply(reply));
        responseObserver.onCompleted();
      }
    });
  } catch (Exception e) {
    responseObserver.onError(GrpcUtil.wrapException(e));
  }
}
 

private void runTestTimeoutException(CLUSTER cluster) throws Exception {
  // send a message to make sure the cluster is working
  try(RaftClient client = cluster.createClient()) {
    final RaftClientReply reply = client.send(new SimpleMessage("m0"));
    Assert.assertTrue(reply.isSuccess());

    RaftClientConfigKeys.Rpc.setRequestTimeout(properties.get(), ONE_SECOND);
    // Block StartTransaction
    cluster.getServers().stream()
        .map(cluster::getRaftServerImpl)
        .map(SimpleStateMachine4Testing::get)
        .forEach(SimpleStateMachine4Testing::blockStartTransaction);
    final CompletableFuture<RaftClientReply> replyFuture = client.sendAsync(new SimpleMessage("m1"));
    FIVE_SECONDS.sleep();
    // Unblock StartTransaction
    cluster.getServers().stream()
        .map(cluster::getRaftServerImpl)
        .map(SimpleStateMachine4Testing::get)
        .forEach(SimpleStateMachine4Testing::unblockStartTransaction);
    // The request should succeed after start transaction is unblocked
    Assert.assertTrue(replyFuture.get(FIVE_SECONDS.getDuration(), FIVE_SECONDS.getUnit()).isSuccess());
  }
}
 

static <REQUEST extends RaftClientRequest, REPLY extends RaftClientReply,
    PROTO_REQ, PROTO_REP> REPLY handleRequest(
    REQUEST request,
    Function<REQUEST, PROTO_REQ> reqToProto,
    Function<PROTO_REP, REPLY> repToProto,
    CheckedFunction<PROTO_REQ, PROTO_REP, ServiceException> handler)
    throws IOException {
  final PROTO_REQ proto = reqToProto.apply(request);
  try {
    final PROTO_REP reply = handler.apply(proto);
    return repToProto.apply(reply);
  } catch (ServiceException se) {
    LOG.trace("Failed to handle " + request, se);
    throw ProtoUtils.toIOException(se);
  }
}
 

static void checkAll(List<CompletableFuture<WatchReplies>> watches, Logger LOG) throws Exception {
  for(int i = 0; i < watches.size(); i++) {
    final WatchReplies watchReplies = watches.get(i).get(GET_TIMEOUT_SECOND, TimeUnit.SECONDS);
    final long logIndex = watchReplies.logIndex;
    LOG.info("checkAll {}: logIndex={}", i, logIndex);
    final RaftClientReply watchAllReply = watchReplies.getAll();
    Assert.assertTrue(watchAllReply.isSuccess());

    final RaftClientReply watchAllCommittedReply = watchReplies.getAllCommitted();
    Assert.assertTrue(watchAllCommittedReply.isSuccess());
    { // check commit infos
      final Collection<CommitInfoProto> commitInfos = watchAllCommittedReply.getCommitInfos();
      final String message = "logIndex=" + logIndex + ", " + ProtoUtils.toString(commitInfos);
      Assert.assertEquals(NUM_SERVERS, commitInfos.size());
      commitInfos.forEach(info -> Assert.assertTrue(message, logIndex <= info.getCommitIndex()));
    }
  }
}
 

protected void testDisconnectLeader() throws Exception {
  try(final MiniRaftCluster cluster = newCluster(3)) {
    cluster.start();

    final RaftServerImpl leader = waitForLeader(cluster);
    try (RaftClient client = cluster.createClient(leader.getId())) {
      client.send(new RaftTestUtil.SimpleMessage("message"));
      Thread.sleep(1000);
      isolate(cluster, leader.getId());
      RaftClientReply reply = client.send(new RaftTestUtil.SimpleMessage("message"));
      Assert.assertNotEquals(reply.getReplierId(), leader.getId());
      Assert.assertTrue(reply.isSuccess());
    } finally {
      deIsolate(cluster, leader.getId());
    }

    cluster.shutdown();
  }
}
 

@Override
public void submitRequest(ContainerCommandRequestProto request,
    HddsProtos.PipelineID pipelineID) throws IOException {
  RaftClientReply reply;
  Span span = TracingUtil
      .importAndCreateSpan(
          "XceiverServerRatis." + request.getCmdType().name(),
          request.getTraceID());
  try (Scope scope = GlobalTracer.get().activateSpan(span)) {

    RaftClientRequest raftClientRequest =
        createRaftClientRequest(request, pipelineID,
            RaftClientRequest.writeRequestType());
    try {
      reply = server.submitClientRequestAsync(raftClientRequest)
          .get(requestTimeout, TimeUnit.MILLISECONDS);
    } catch (Exception e) {
      throw new IOException(e.getMessage(), e);
    }
    processReply(reply);
  } finally {
    span.finish();
  }
}
 

@Override
public void addGroup(HddsProtos.PipelineID pipelineId,
    Collection<DatanodeDetails> peers) throws IOException {
  final PipelineID pipelineID = PipelineID.getFromProtobuf(pipelineId);
  final RaftGroupId groupId = RaftGroupId.valueOf(pipelineID.getId());
  final RaftGroup group = RatisHelper.newRaftGroup(groupId, peers);
  GroupManagementRequest request = GroupManagementRequest.newAdd(
      clientId, server.getId(), nextCallId(), group);

  RaftClientReply reply;
  try {
    reply = server.groupManagement(request);
  } catch (Exception e) {
    throw new IOException(e.getMessage(), e);
  }
  processReply(reply);
}
 

public void setConfiguration(RaftPeer... peers) throws IOException {
  try(RaftClient client = createClient()) {
    LOG.info("Start changing the configuration: {}", Arrays.asList(peers));
    final RaftClientReply reply = client.setConfiguration(peers);
    Preconditions.assertTrue(reply.isSuccess());
  }
}
 

static CompletableFuture<ByteString> sendAsync(
    ByteString request, Function<Message, CompletableFuture<RaftClientReply>> sendFunction) {
  return sendFunction.apply(() -> request
  ).thenApply(reply -> {
    final StateMachineException sme = reply.getStateMachineException();
    if (sme != null) {
      throw new CompletionException("Failed to send request " + request, sme);
    }
    Preconditions.assertTrue(reply.isSuccess(), () -> "Failed " + request + ", reply=" + reply);
    return reply.getMessage().getContent();
  });
}
 

private void completeReplyExceptionally(Throwable t, String event) {
  final Map<Long, CompletableFuture<RaftClientReply>> map = replies.getAndSetNull();
  if (map == null) {
    return;
  }
  for (Map.Entry<Long, CompletableFuture<RaftClientReply>> entry : map.entrySet()) {
    final CompletableFuture<RaftClientReply> f = entry.getValue();
    if (!f.isDone()) {
      f.completeExceptionally(t != null? t
          : new AlreadyClosedException(getName() + ": Stream " + event
              + ": no reply for async request cid=" + entry.getKey()));
    }
  }
}
 

static void assertNotReplicatedException(long logIndex, ReplicationLevel replication, RaftClientReply reply) {
  Assert.assertFalse(reply.isSuccess());
  final NotReplicatedException nre = reply.getNotReplicatedException();
  Assert.assertNotNull(nre);
  Assert.assertEquals(logIndex, nre.getLogIndex());
  Assert.assertEquals(replication, nre.getRequiredReplication());
}
 

@Override
protected void operation(RaftClient client) throws IOException {
  RaftClientReply getValue =
      client.sendReadOnly(Expression.Utils.toMessage(new Variable(name)));
  Expression response =
      Expression.Utils.bytes2Expression(getValue.getMessage().getContent().toByteArray(), 0);
  System.out.println(String.format("%s=%s", name, (DoubleValue) response).toString());
}
 

void runTestAppendEntriesTimeout(CLUSTER cluster) throws Exception {
  LOG.info("Running testAppendEntriesTimeout");
  final TimeDuration oldExpiryTime = RaftServerConfigKeys.RetryCache.expiryTime(getProperties());
  RaftServerConfigKeys.RetryCache.setExpiryTime(getProperties(), TimeDuration.valueOf(20, TimeUnit.SECONDS));
  waitForLeader(cluster);
  long time = System.currentTimeMillis();
  long waitTime = 5000;
  try (final RaftClient client = cluster.createClient()) {
    // block append requests
    cluster.getServerAliveStream()
        .filter(impl -> !impl.isLeader())
        .map(SimpleStateMachine4Testing::get)
        .forEach(SimpleStateMachine4Testing::blockWriteStateMachineData);

    CompletableFuture<RaftClientReply> replyFuture = client.sendAsync(new SimpleMessage("abc"));
    Thread.sleep(waitTime);
    // replyFuture should not be completed until append request is unblocked.
    Assert.assertFalse(replyFuture.isDone());
    // unblock append request.
    cluster.getServerAliveStream()
        .filter(impl -> !impl.isLeader())
        .map(SimpleStateMachine4Testing::get)
        .forEach(SimpleStateMachine4Testing::unblockWriteStateMachineData);

    Assert.assertTrue(replyFuture.get().isSuccess());
    Assert.assertTrue(System.currentTimeMillis() - time > waitTime);
  }

  //reset for the other tests
  RaftServerConfigKeys.RetryCache.setExpiryTime(getProperties(), oldExpiryTime);
}
 

static ByteString send(
    ByteString request, CheckedFunction<Message, RaftClientReply, IOException> sendFunction)
    throws IOException {
  final RaftClientReply reply = sendFunction.apply(Message.valueOf(request));
  final StateMachineException sme = reply.getStateMachineException();
  if (sme != null) {
    throw new IOException("Failed to send request " + request, sme);
  }
  Preconditions.assertTrue(reply.isSuccess(), () -> "Failed " + request + ", reply=" + reply);
  return reply.getMessage().getContent();
}
 

@Override
public RaftClientReply submitClientRequest(RaftClientRequest request)
    throws IOException {
  return handleRequest(request,
      ClientProtoUtils::toRaftClientRequestProto,
      ClientProtoUtils::toRaftClientReply,
      p -> getProtocol().submitClientRequest(null, p));
}
 

/**
 * Keep generating writing traffic and make sure snapshots are taken.
 * We then restart the whole raft peer and check if it can correctly load
 * snapshots + raft log.
 */
@Test
public void testRestartPeer() throws Exception {
  RaftTestUtil.waitForLeader(cluster);
  final RaftPeerId leaderId = cluster.getLeader().getId();
  int i = 0;
  try(final RaftClient client = cluster.createClient(leaderId)) {
    for (; i < SNAPSHOT_TRIGGER_THRESHOLD * 2 - 1; i++) {
      RaftClientReply reply = client.send(new SimpleMessage("m" + i));
      Assert.assertTrue(reply.isSuccess());
    }
  }

  final long nextIndex = cluster.getLeader().getState().getLog().getNextIndex();
  LOG.info("nextIndex = {}", nextIndex);
  // wait for the snapshot to be done
  final List<File> snapshotFiles = getSnapshotFiles(cluster, nextIndex - SNAPSHOT_TRIGGER_THRESHOLD, nextIndex);
  JavaUtils.attemptRepeatedly(() -> {
    Assert.assertTrue(snapshotFiles.stream().anyMatch(RaftSnapshotBaseTest::exists));
    return null;
  }, 10, ONE_SECOND, "snapshotFile.exist", LOG);

  // restart the peer and check if it can correctly load snapshot
  cluster.restart(false);
  try {
    // 200 messages + two leader elections --> last committed = 201
    assertLeaderContent(cluster);
  } finally {
    cluster.shutdown();
  }
}
 

@Override
protected void operation(RaftClient client) throws IOException {
  RaftClientReply send = client.send(
      new AssignmentMessage(new Variable(name), createExpression(value)));
  System.out.println("Success: " + send.isSuccess());
  System.out.println("Response: " + send.getMessage().getClass());
}
 

void runTestStream(CLUSTER cluster) throws Exception {
  RaftTestUtil.waitForLeader(cluster);

  // stream multiple parts
  final int numParts = 9;
  final int endOfRequest = 6;
  final StringBuilder key = new StringBuilder();
  try(RaftClient client = cluster.createClient();
      MessageOutputStream out = client.getStreamApi().stream()) {
    for (int i = 1; i <= numParts; i++) {
      key.append(i);
      out.sendAsync(new SimpleMessage(i + ""), i == endOfRequest);
    }
  }

  // check if all the parts are streamed as a single message.
  final String k = key.toString();
  try(RaftClient client = cluster.createClient()) {
    final String k1 = k.substring(0, endOfRequest);
    final RaftClientReply r1= client.sendReadOnly(new SimpleMessage(k1));
    Assert.assertTrue(r1.isSuccess());

    final String k2 = k.substring(endOfRequest);
    final RaftClientReply r2 = client.sendReadOnly(new SimpleMessage(k2));
    Assert.assertTrue(r2.isSuccess());
  }
}
 

static CompletableFuture<RaftClientReply> failWithException(
    Throwable t, CacheEntry entry) {
  if (entry != null) {
    entry.failWithException(t);
    return entry.getReplyFuture();
  } else {
    return JavaUtils.completeExceptionally(t);
  }
}
 

static CompletableFuture<RaftClientReply> failWithReply(
    RaftClientReply reply, CacheEntry entry) {
  if (entry != null) {
    entry.failWithReply(reply);
    return entry.getReplyFuture();
  } else {
    return CompletableFuture.completedFuture(reply);
  }
}
 

static CompletableFuture<RaftClientReply> failWithException(
    Throwable t, CacheEntry entry) {
  if (entry != null) {
    entry.failWithException(t);
    return entry.getReplyFuture();
  } else {
    return JavaUtils.completeExceptionally(t);
  }
}
 

@Override
public void onNext(RaftClientReplyProto proto) {
  final long callId = proto.getRpcReply().getCallId();
  try {
    final RaftClientReply reply = ClientProtoUtils.toRaftClientReply(proto);
    final NotLeaderException nle = reply.getNotLeaderException();
    if (nle != null) {
      completeReplyExceptionally(nle, NotLeaderException.class.getName());
      return;
    }
    handleReplyFuture(callId, f -> f.complete(reply));
  } catch (Throwable t) {
    handleReplyFuture(callId, f -> f.completeExceptionally(t));
  }
}
 

void replyPendingRequest(long index, RaftClientReply reply) {
  final PendingRequest pending = pendingRequests.remove(index);
  if (pending != null) {
    Preconditions.assertTrue(pending.getIndex() == index);
    pending.setReply(reply);
  }
}
 

void replySetConfiguration(Supplier<Collection<CommitInfoProto>> getCommitInfos) {
  // we allow the pendingRequest to be null in case that the new leader
  // commits the new configuration while it has not received the retry
  // request from the client
  if (pendingSetConf != null) {
    final RaftClientRequest request = pendingSetConf.getRequest();
    LOG.debug("{}: sends success for {}", name, request);
    // for setConfiguration we do not need to wait for statemachine. send back
    // reply after it's committed.
    pendingSetConf.setReply(new RaftClientReply(request, getCommitInfos.get()));
    pendingSetConf = null;
  }
}
 

RaftClientReply get(CompletableFuture<RaftClientReply> f, String name) throws Exception {
  final RaftClientReply reply;
  try {
    reply = f.get(GET_TIMEOUT_SECOND, TimeUnit.SECONDS);
  } catch (Exception e) {
    log.error("Failed to get {}({})", name, logIndex);
    throw e;
  }
  log.info("{}-Watch({}) returns {}", name, logIndex, reply);
  return reply;
}
 

private void sendRequestWithRetry(PendingOrderedRequest pending) {
  final CompletableFuture<RaftClientReply> f = pending.getReplyFuture();
  if (f.isDone()) {
    return;
  }

  final RaftClientRequest request = pending.newRequestImpl();
  if (request == null) { // already done
    LOG.debug("{} newRequestImpl returns null", pending);
    return;
  }

  final RetryPolicy retryPolicy = client.getRetryPolicy();
  sendRequest(pending).thenAccept(reply -> {
    if (f.isDone()) {
      return;
    }
    if (reply == null) {
      scheduleWithTimeout(pending, request, retryPolicy, null);
    } else {
      f.complete(reply);
    }
  }).exceptionally(e -> {
    if (e instanceof CompletionException) {
      e = JavaUtils.unwrapCompletionException(e);
      scheduleWithTimeout(pending, request, retryPolicy, e);
      return null;
    }
    f.completeExceptionally(e);
    return null;
  });
}
 

public void setConfiguration(RaftPeer... peers) throws IOException {
  try(RaftClient client = createClient()) {
    LOG.info("Start changing the configuration: {}", Arrays.asList(peers));
    final RaftClientReply reply = client.setConfiguration(peers);
    Preconditions.assertTrue(reply.isSuccess());
  }
}
 

@Override
public CompletableFuture<RaftClientReply> streamAsync(Message message, SizeInBytes subSize) {
  final int n = subSize.getSizeInt();
  final MessageOutputStream out = stream();
  final ByteString bytes = message.getContent();
  for(int i = 0; i < bytes.size(); ) {
    final int j = Math.min(i + n, bytes.size());
    final ByteString sub = bytes.substring(i, j);
    out.sendAsync(Message.valueOf(sub));
    i = j;
  }
  return out.closeAsync();
}
 

public static OMResponse getOMResponseFromRaftClientReply(
    RaftClientReply reply) throws InvalidProtocolBufferException {
  byte[] bytes = reply.getMessage().getContent().toByteArray();
  return OMResponse.newBuilder(OMResponse.parseFrom(bytes))
      .setLeaderOMNodeId(reply.getReplierId())
      .build();
}