use of org.apache.ratis.protocol.DataStreamReply in project incubator-ratis by apache.
the class NettyClientStreamRpc method streamAsync.
@Override
public CompletableFuture<DataStreamReply> streamAsync(DataStreamRequest request) {
final CompletableFuture<DataStreamReply> f = new CompletableFuture<>();
ClientInvocationId clientInvocationId = ClientInvocationId.valueOf(request.getClientId(), request.getStreamId());
final ReplyQueue q = replies.computeIfAbsent(clientInvocationId, key -> new ReplyQueue());
if (!q.offer(f)) {
f.completeExceptionally(new IllegalStateException(this + ": Failed to offer a future for " + request));
return f;
}
LOG.debug("{}: write {}", this, request);
getChannel().writeAndFlush(request);
return f;
}
use of org.apache.ratis.protocol.DataStreamReply in project incubator-ratis by apache.
the class DataStreamManagement method readImpl.
private void readImpl(DataStreamRequestByteBuf request, ChannelHandlerContext ctx, ByteBuf buf, CheckedBiFunction<RaftClientRequest, Set<RaftPeer>, Set<DataStreamOutputRpc>, IOException> getStreams) {
boolean close = WriteOption.containsOption(request.getWriteOptions(), StandardWriteOption.CLOSE);
ClientInvocationId key = ClientInvocationId.valueOf(request.getClientId(), request.getStreamId());
final StreamInfo info;
if (request.getType() == Type.STREAM_HEADER) {
final MemoizedSupplier<StreamInfo> supplier = JavaUtils.memoize(() -> newStreamInfo(buf, getStreams));
info = streams.computeIfAbsent(key, id -> supplier.get());
if (!supplier.isInitialized()) {
streams.remove(key);
throw new IllegalStateException("Failed to create a new stream for " + request + " since a stream already exists Key: " + key + " StreamInfo:" + info);
}
getMetrics().onRequestCreate(RequestType.HEADER);
} else if (close) {
info = Optional.ofNullable(streams.remove(key)).orElseThrow(() -> new IllegalStateException("Failed to remove StreamInfo for " + request));
} else {
info = Optional.ofNullable(streams.get(key)).orElseThrow(() -> {
streams.remove(key);
return new IllegalStateException("Failed to get StreamInfo for " + request);
});
}
final CompletableFuture<Long> localWrite;
final List<CompletableFuture<DataStreamReply>> remoteWrites;
if (request.getType() == Type.STREAM_HEADER) {
localWrite = CompletableFuture.completedFuture(0L);
remoteWrites = Collections.emptyList();
} else if (request.getType() == Type.STREAM_DATA) {
localWrite = info.getLocal().write(buf, request.getWriteOptions(), writeExecutor);
remoteWrites = info.applyToRemotes(out -> out.write(request, requestExecutor));
} else {
throw new IllegalStateException(this + ": Unexpected type " + request.getType() + ", request=" + request);
}
composeAsync(info.getPrevious(), requestExecutor, n -> JavaUtils.allOf(remoteWrites).thenCombineAsync(localWrite, (v, bytesWritten) -> {
if (request.getType() == Type.STREAM_HEADER || (request.getType() == Type.STREAM_DATA && !close)) {
sendReply(remoteWrites, request, bytesWritten, info.getCommitInfos(), ctx);
} else if (close) {
if (info.isPrimary()) {
// after all server close stream, primary server start transaction
startTransaction(info, request, bytesWritten, ctx);
} else {
sendReply(remoteWrites, request, bytesWritten, info.getCommitInfos(), ctx);
}
} else {
throw new IllegalStateException(this + ": Unexpected type " + request.getType() + ", request=" + request);
}
return null;
}, requestExecutor)).whenComplete((v, exception) -> {
try {
if (exception != null) {
streams.remove(key);
replyDataStreamException(server, exception, info.getRequest(), request, ctx);
}
} finally {
buf.release();
}
});
}
use of org.apache.ratis.protocol.DataStreamReply in project incubator-ratis by apache.
the class DataStreamManagement method checkSuccessRemoteWrite.
static boolean checkSuccessRemoteWrite(List<CompletableFuture<DataStreamReply>> replyFutures, long bytesWritten, final DataStreamRequestByteBuf request) {
for (CompletableFuture<DataStreamReply> replyFuture : replyFutures) {
final DataStreamReply reply = replyFuture.join();
assertReplyCorrespondingToRequest(request, reply);
if (!reply.isSuccess()) {
LOG.warn("reply is not success, request: {}", request);
return false;
}
if (reply.getBytesWritten() != bytesWritten) {
LOG.warn("reply written bytes not match, local size: {} remote size: {} request: {}", bytesWritten, reply.getBytesWritten(), request);
return false;
}
}
return true;
}
use of org.apache.ratis.protocol.DataStreamReply in project incubator-ratis by apache.
the class FileStoreWriter method streamWriteAndVerify.
public FileStoreWriter streamWriteAndVerify(RoutingTable routingTable) {
final int size = fileSize.getSizeInt();
final DataStreamOutput dataStreamOutput = client.getStreamOutput(fileName, size, routingTable);
final List<CompletableFuture<DataStreamReply>> futures = new ArrayList<>();
final List<Integer> sizes = new ArrayList<>();
for (int offset = 0; offset < size; ) {
final int remaining = size - offset;
final int length = Math.min(remaining, bufferSize);
final boolean close = length == remaining;
LOG.trace("write {}, offset={}, length={}, close? {}", fileName, offset, length, close);
final ByteBuffer bf = DataStreamTestUtils.initBuffer(0, length);
futures.add(close ? dataStreamOutput.writeAsync(bf, StandardWriteOption.CLOSE) : dataStreamOutput.writeAsync(bf));
sizes.add(length);
offset += length;
}
DataStreamReply reply = dataStreamOutput.closeAsync().join();
Assert.assertTrue(reply.isSuccess());
// check writeAsync requests
for (int i = 0; i < futures.size(); i++) {
reply = futures.get(i).join();
Assert.assertTrue(reply.isSuccess());
Assert.assertEquals(sizes.get(i).longValue(), reply.getBytesWritten());
Assert.assertEquals(reply.getType(), RaftProtos.DataStreamPacketHeaderProto.Type.STREAM_DATA);
}
return this;
}
use of org.apache.ratis.protocol.DataStreamReply in project incubator-ratis by apache.
the class DataStreamClusterTests method writeAsyncDefaultFileRegion.
static CheckedConsumer<DataStreamOutputImpl, Exception> writeAsyncDefaultFileRegion(File f, int size) {
return new CheckedConsumer<DataStreamOutputImpl, Exception>() {
@Override
public void accept(DataStreamOutputImpl out) {
final DataStreamReply dataStreamReply = out.writeAsync(f).join();
DataStreamTestUtils.assertSuccessReply(Type.STREAM_DATA, size, dataStreamReply);
}
@Override
public String toString() {
return "writeAsyncDefaultFileRegion";
}
};
}
Aggregations